diff --git a/.devcontainer/build_cxx.sh b/.devcontainer/build_cxx.sh index 007c99b9e5..33a9f890ef 100755 --- a/.devcontainer/build_cxx.sh +++ b/.devcontainer/build_cxx.sh @@ -5,7 +5,7 @@ NPROC=$(nproc --all) SCRIPT_PATH=$(dirname $(realpath -s $0)) export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch -TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') +TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') mkdir -p ${SCRIPT_PATH}/../buildcxx/ cd ${SCRIPT_PATH}/../buildcxx/ @@ -13,7 +13,7 @@ cmake -D ENABLE_TENSORFLOW=ON \ -D ENABLE_PYTORCH=ON \ -D ENABLE_PADDLE=ON \ -D CMAKE_INSTALL_PREFIX=${SCRIPT_PATH}/../dp/ \ - -D LAMMPS_VERSION=stable_29Aug2024_update1 \ + -D LAMMPS_VERSION=stable_22Jul2025_update2 \ -D CMAKE_BUILD_TYPE=Debug \ -D BUILD_TESTING:BOOL=TRUE \ -D TENSORFLOW_ROOT=${TENSORFLOW_ROOT} \ diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 27c40bbe6a..85d67db2a9 100644 --- a/.devcontainer/devcontainer.json +++ b/.devcontainer/devcontainer.json @@ -11,7 +11,6 @@ "PATH": "${containerEnv:PATH}:${containerWorkspaceFolder}/.venv/bin", "DP_ENABLE_PYTORCH": "1", "DP_VARIANT": "cpu", - "LMP_CXX11_ABI_0": "1", "UV_EXTRA_INDEX_URL": "https://download.pytorch.org/whl/cpu" } } diff --git a/.devcontainer/download_libtorch.sh b/.devcontainer/download_libtorch.sh index 43a46805c5..8c1e480b7c 100755 --- a/.devcontainer/download_libtorch.sh +++ b/.devcontainer/download_libtorch.sh @@ -4,5 +4,5 @@ set -ev SCRIPT_PATH=$(dirname $(realpath -s $0)) cd ${SCRIPT_PATH}/.. -wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.7.0%2Bcpu.zip -O ~/libtorch.zip +wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.8.0%2Bcpu.zip -O ~/libtorch.zip unzip ~/libtorch.zip diff --git a/.devcontainer/gdb_lmp b/.devcontainer/gdb_lmp index 33e883780b..fc1c8b90fe 100755 --- a/.devcontainer/gdb_lmp +++ b/.devcontainer/gdb_lmp @@ -2,7 +2,7 @@ SCRIPT_PATH=$(dirname $(realpath -s $0)) export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch -TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') +TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \ LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \ diff --git a/.devcontainer/gdb_pytest_lmp b/.devcontainer/gdb_pytest_lmp index e27e40d4b0..d27587ec43 100755 --- a/.devcontainer/gdb_pytest_lmp +++ b/.devcontainer/gdb_pytest_lmp @@ -2,7 +2,7 @@ SCRIPT_PATH=$(dirname $(realpath -s $0))/../.. export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch -TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') +TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \ LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \ diff --git a/.devcontainer/lmp b/.devcontainer/lmp index c8e781aa57..524f99b326 100755 --- a/.devcontainer/lmp +++ b/.devcontainer/lmp @@ -2,7 +2,7 @@ SCRIPT_PATH=$(dirname $(realpath -s $0)) export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch -TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') +TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \ LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \ diff --git a/.devcontainer/pytest_lmp b/.devcontainer/pytest_lmp index 9371ba72d5..bb88da883f 100755 --- a/.devcontainer/pytest_lmp +++ b/.devcontainer/pytest_lmp @@ -2,7 +2,7 @@ SCRIPT_PATH=$(dirname $(realpath -s $0))/../.. export CMAKE_PREFIX_PATH=${SCRIPT_PATH}/../libtorch -TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') +TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') env LAMMPS_PLUGIN_PATH=${SCRIPT_PATH}/../dp/lib/deepmd_lmp \ LD_LIBRARY_PATH=${SCRIPT_PATH}/../dp/lib:${CMAKE_PREFIX_PATH}/lib:${TENSORFLOW_ROOT} \ diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 5855aef6c5..cbd920f6b3 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -4,4 +4,7 @@ updates: directory: "/" schedule: interval: "weekly" - target-branch: "devel" + - package-ecosystem: "pip" + directory: "/" + schedule: + interval: "weekly" diff --git a/.github/workflows/build_cc.yml b/.github/workflows/build_cc.yml index 59873c5101..79bacec204 100644 --- a/.github/workflows/build_cc.yml +++ b/.github/workflows/build_cc.yml @@ -2,6 +2,9 @@ on: push: branches-ignore: - "gh-readonly-queue/**" + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" pull_request: merge_group: concurrency: @@ -17,8 +20,6 @@ jobs: include: - variant: cpu dp_variant: cpu - - variant: cuda - dp_variant: cuda - variant: cuda120 dp_variant: cuda - variant: rocm @@ -26,20 +27,13 @@ jobs: - variant: clang dp_variant: clang steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' - uses: lukka/get-cmake@latest - run: python -m pip install uv - - run: source/install/uv_with_retry.sh pip install --system tensorflow - - run: source/install/uv_with_retry.sh pip install --system 'torch==2.7' --index-url https://download.pytorch.org/whl/cpu - - run: | - wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \ - && sudo dpkg -i cuda-keyring_1.0-1_all.deb \ - && sudo apt-get update \ - && sudo apt-get -y install cuda-cudart-dev-11-8 cuda-nvcc-11-8 - if: matrix.variant == 'cuda' + - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu - run: | wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \ && sudo dpkg -i cuda-keyring_1.0-1_all.deb \ diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index c520b2af59..5385855e2e 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -4,6 +4,9 @@ on: push: branches-ignore: - "gh-readonly-queue/**" + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" tags: - "v*" pull_request: @@ -26,15 +29,9 @@ jobs: python: 311 platform_id: manylinux_x86_64 dp_variant: cuda - cuda_version: 12.2 - - os: ubuntu-latest - python: 311 - platform_id: manylinux_x86_64 - dp_variant: cuda - cuda_version: 11.8 - dp_pkg_name: deepmd-kit-cu11 + cuda_version: 12.8 # macos-x86-64 - - os: macos-13 + - os: macos-15-intel python: 311 platform_id: macosx_x86_64 dp_variant: cpu @@ -54,23 +51,15 @@ jobs: platform_id: manylinux_aarch64 dp_variant: cpu steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 with: # https://github.com/pypa/setuptools_scm/issues/480 fetch-depth: 0 - name: Install uv run: curl --proto '=https' --tlsv1.2 -LsSf https://github.com/astral-sh/uv/releases/download/0.2.24/uv-installer.sh | sh if: runner.os != 'Linux' - - uses: docker/setup-qemu-action@v3 - name: Setup QEMU - if: matrix.platform_id == 'manylinux_aarch64' && matrix.os == 'ubuntu-latest' - # detect version in advance. See #3168 - - run: | - echo "SETUPTOOLS_SCM_PRETEND_VERSION=$(pipx run uv tool run --from setuptools_scm python -m setuptools_scm)" >> $GITHUB_ENV - rm -rf .git - if: matrix.dp_pkg_name == 'deepmd-kit-cu11' - name: Build wheels - uses: pypa/cibuildwheel@v2.23 + uses: pypa/cibuildwheel@v3.3 env: CIBW_BUILD_VERBOSITY: 1 CIBW_ARCHS: all @@ -79,7 +68,7 @@ jobs: CUDA_VERSION: ${{ matrix.cuda_version }} DP_PKG_NAME: ${{ matrix.dp_pkg_name }} CIBW_BUILD_FRONTEND: 'build[uv]' - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v6 with: name: cibw-cp${{ matrix.python }}-${{ matrix.platform_id }}-cu${{ matrix.cuda_version }}-${{ strategy.job-index }} path: ./wheelhouse/*.whl @@ -87,13 +76,13 @@ jobs: name: Build source distribution runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 with: fetch-depth: 0 - name: Build sdist run: pipx run uv tool run --with build[uv] --from build python -m build --installer uv --sdist - - uses: actions/upload-artifact@v4 + - uses: actions/upload-artifact@v6 with: name: cibw-sdist path: dist/*.tar.gz @@ -106,7 +95,7 @@ jobs: id-token: write if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v') steps: - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@v7 with: pattern: cibw-* path: dist @@ -123,13 +112,11 @@ jobs: include: - variant: "" cuda_version: "12" - - variant: "_cu11" - cuda_version: "11" steps: - name: Delete huge unnecessary tools folder run: rm -rf /opt/hostedtoolcache - - uses: actions/checkout@v4 - - uses: actions/download-artifact@v4 + - uses: actions/checkout@v6 + - uses: actions/download-artifact@v7 with: path: source/install/docker/dist pattern: cibw-*-manylinux_x86_64-cu${{ matrix.cuda_version }}* @@ -162,12 +149,12 @@ jobs: needs: [build_wheels, build_sdist] runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@v4 + - uses: actions/download-artifact@v7 with: path: dist/packages pattern: cibw-* merge-multiple: true - - uses: actions/setup-python@v5 + - uses: actions/setup-python@v6 name: Install Python with: python-version: '3.11' @@ -176,7 +163,7 @@ jobs: ls dist/packages > package_list.txt dumb-pypi --output-dir dist --packages-url ../../packages --package-list package_list.txt --title "DeePMD-kit Developed Packages" - name: Upload Pages artifact - uses: actions/upload-pages-artifact@v3 + uses: actions/upload-pages-artifact@v4 with: path: dist deploy_pypi_index: @@ -188,7 +175,7 @@ jobs: name: github-pages url: ${{ steps.deployment.outputs.page_url }} runs-on: ubuntu-latest - if: github.event_name == 'push' && github.ref == 'refs/heads/devel' && github.repository_owner == 'deepmodeling' + if: github.event_name == 'push' && github.ref == 'refs/heads/master' && github.repository_owner == 'deepmodeling' steps: - name: Deploy to GitHub Pages id: deployment diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index f8005ab831..7460fd7b22 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -4,6 +4,9 @@ on: push: branches-ignore: - "gh-readonly-queue/**" + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" pull_request: schedule: - cron: '45 2 * * 2' @@ -27,8 +30,8 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' cache: 'pip' @@ -40,13 +43,13 @@ jobs: && sudo dpkg -i cuda-keyring_1.0-1_all.deb \ && sudo apt-get update \ && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2 - python -m pip install tensorflow - python -m pip install 'torch==2.7' --index-url https://download.pytorch.org/whl/cpu + python -m pip install uv + uv pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu env: DEBIAN_FRONTEND: noninteractive # Initializes the CodeQL tools for scanning. - name: Initialize CodeQL - uses: github/codeql-action/init@v3 + uses: github/codeql-action/init@v4 with: languages: ${{ matrix.language }} queries: security-extended,security-and-quality @@ -59,6 +62,6 @@ jobs: if: matrix.language == 'c-cpp' - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@v3 + uses: github/codeql-action/analyze@v4 with: category: "/language:${{matrix.language}}" diff --git a/.github/workflows/copilot-setup-steps.yml b/.github/workflows/copilot-setup-steps.yml new file mode 100644 index 0000000000..51db243c1c --- /dev/null +++ b/.github/workflows/copilot-setup-steps.yml @@ -0,0 +1,66 @@ +name: "Copilot Setup Steps" + +# Automatically run the setup steps when they are changed to allow for easy validation, and +# allow manual testing through the repository's "Actions" tab +on: + workflow_dispatch: + push: + branches-ignore: + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" + paths: + - .github/workflows/copilot-setup-steps.yml + pull_request: + paths: + - .github/workflows/copilot-setup-steps.yml + +jobs: + # The job MUST be called `copilot-setup-steps` or it will not be picked up by Copilot. + copilot-setup-steps: + runs-on: ubuntu-latest + + # Set the permissions to the lowest permissions possible needed for your steps. + # Copilot will be given its own token for its operations. + permissions: + # If you want to clone the repository as part of your setup steps, for example to install dependencies, you'll need the `contents: read` permission. If you don't clone the repository in your setup steps, Copilot will do this for you automatically after the steps complete. + contents: read + + # You can define any steps you want, and they will run before the agent starts. + # If you do not check out your code, Copilot will do this for you. + steps: + - name: Checkout code + uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.10" + + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + + - name: Create virtual environment + run: uv venv venv + + - name: Activate virtual environment + run: echo "VIRTUAL_ENV=$PWD/venv" >> $GITHUB_ENV && echo "$PWD/venv/bin" >> $GITHUB_PATH + + - name: Install base dependencies + run: uv pip install --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu + + - name: Build Python package + run: uv pip install -e .[cpu,test] + + - name: Install pre-commit tools + run: uv tool install pre-commit + + - name: Install pre-commit hooks + run: pre-commit install --install-hooks + + - name: Verify installation + run: | + dp --version + python -c "import deepmd; import deepmd.tf; print('DeePMD-kit installation verified')" diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml index be43c5cff2..77f06528fe 100644 --- a/.github/workflows/labeler.yml +++ b/.github/workflows/labeler.yml @@ -9,6 +9,6 @@ jobs: pull-requests: write runs-on: ubuntu-latest steps: - - uses: actions/labeler@v5 + - uses: actions/labeler@v6 with: repo-token: "${{ secrets.GITHUB_TOKEN }}" diff --git a/.github/workflows/mirror_gitee.yml b/.github/workflows/mirror_gitee.yml index 2d090c0539..79d1ce11a5 100644 --- a/.github/workflows/mirror_gitee.yml +++ b/.github/workflows/mirror_gitee.yml @@ -1,6 +1,13 @@ name: Mirror to Gitee Repo -on: [ push, delete, create ] +on: + push: + branches-ignore: + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" + delete: + create: # Ensures that only one mirror task will run at a time. concurrency: diff --git a/.github/workflows/package_c.yml b/.github/workflows/package_c.yml index f5e4a97d56..655dfb4eca 100644 --- a/.github/workflows/package_c.yml +++ b/.github/workflows/package_c.yml @@ -4,6 +4,9 @@ on: push: branches-ignore: - "gh-readonly-queue/**" + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" tags: - "v*" pull_request: @@ -21,11 +24,18 @@ jobs: - tensorflow_build_version: "2.18" tensorflow_version: "" filename: libdeepmd_c.tar.gz - - tensorflow_build_version: "2.14" - tensorflow_version: ">=2.5.0rc0,<2.15" - filename: libdeepmd_c_cu11.tar.gz steps: - - uses: actions/checkout@v4 + - name: Free Disk Space (Ubuntu) + uses: insightsengineering/disk-space-reclaimer@v1 + with: + tools-cache: true + android: true + dotnet: true + haskell: true + large-packages: true + swap-storage: true + docker-images: true + - uses: actions/checkout@v6 with: fetch-depth: 0 - name: Package C library @@ -37,7 +47,7 @@ jobs: if: matrix.filename != 'libdeepmd_c.tar.gz' # for download and debug - name: Upload artifact - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: libdeepmd_c-${{ strategy.job-index }}-${{ matrix.filename }} path: ${{ matrix.filename }} @@ -53,9 +63,9 @@ jobs: needs: [build_c] runs-on: ubuntu-22.04 steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Download artifact - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: pattern: libdeepmd_c-* merge-multiple: true diff --git a/.github/workflows/suppr.txt b/.github/workflows/suppr.txt index 17e8c743a0..2a43b31eb3 100644 --- a/.github/workflows/suppr.txt +++ b/.github/workflows/suppr.txt @@ -1 +1,5 @@ leak:libpaddle_inference +# TensorFlow 2.20 +leak:xla:: +leak:mlir:: +leak:llvm:: diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index 3407a58d1e..32155cc309 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -2,6 +2,9 @@ on: push: branches-ignore: - "gh-readonly-queue/**" + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" pull_request: merge_group: concurrency: @@ -14,25 +17,36 @@ jobs: runs-on: ubuntu-22.04 strategy: matrix: - check_memleak: [true, false] + include: + - check_memleak: true + enable_tensorflow: true + enable_pytorch: true + enable_paddle: false + - check_memleak: true + enable_tensorflow: false + enable_pytorch: false + enable_paddle: true + - check_memleak: false + enable_tensorflow: true + enable_pytorch: true + enable_paddle: false + - check_memleak: false + enable_tensorflow: false + enable_pytorch: false + enable_paddle: true steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' cache: 'pip' - - name: Setup MPI - uses: mpi4py/setup-mpi@v1 - with: - mpi: mpich - uses: lukka/get-cmake@latest - run: python -m pip install uv - name: Install Python dependencies run: | - source/install/uv_with_retry.sh pip install --system tensorflow-cpu~=2.18.0 jax==0.5.0 - export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') - source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax] mpi4py - source/install/uv_with_retry.sh pip install --system 'torch==2.7' --index-url https://download.pytorch.org/whl/cpu + source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_cpu --group pin_pytorch_cpu --group pin_jax --torch-backend cpu + export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') + source/install/uv_with_retry.sh pip install --system -e .[cpu,test,lmp,jax] mpi4py mpich - name: Convert models run: source/tests/infer/convert-models.sh # https://github.com/actions/runner-images/issues/9491 @@ -45,35 +59,44 @@ jobs: OMP_NUM_THREADS: 1 TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 - LMP_CXX11_ABI_0: 1 CMAKE_GENERATOR: Ninja CXXFLAGS: ${{ matrix.check_memleak && '-fsanitize=leak' || '' }} LSAN_OPTIONS: suppressions=${{ github.workspace }}/.github/workflows/suppr.txt + ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && 'TRUE' || 'FALSE' }} + ENABLE_PYTORCH: ${{ matrix.enable_pytorch && 'TRUE' || 'FALSE' }} + ENABLE_PADDLE: ${{ matrix.enable_paddle && 'TRUE' || 'FALSE' }} # test lammps - run: | cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/paddle/lib/*.so ${{ github.workspace }}/dp_test/lib/ cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/third_party/install/onednn/lib/* ${{ github.workspace }}/dp_test/lib/ cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/third_party/install/mklml/lib/* ${{ github.workspace }}/dp_test/lib/ + if: matrix.enable_paddle + - run: | + export LD_LIBRARY_PATH=${{ github.workspace }}/dp_test/lib:$LD_LIBRARY_PATH pytest --cov=deepmd source/lmp/tests env: OMP_NUM_THREADS: 1 TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp - LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib + ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && '1' || '0' }} + ENABLE_PYTORCH: ${{ matrix.enable_pytorch && '1' || '0' }} + ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }} + ENABLE_PADDLE: ${{ matrix.enable_paddle && '1' || '0' }} if: ${{ !matrix.check_memleak }} # test ipi - run: | export PATH=${{ github.workspace }}/dp_test/bin:$PATH - cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/paddle/lib/*.so ${{ github.workspace }}/dp_test/lib/ - cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/third_party/install/onednn/lib/* ${{ github.workspace }}/dp_test/lib/ - cp ${{ github.workspace }}/source/build_tests/paddle_inference_install_dir/third_party/install/mklml/lib/* ${{ github.workspace }}/dp_test/lib/ pytest --cov=deepmd source/ipi/tests env: OMP_NUM_THREADS: 1 TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 LD_LIBRARY_PATH: ${{ github.workspace }}/dp_test/lib + ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && '1' || '0' }} + ENABLE_PYTORCH: ${{ matrix.enable_pytorch && '1' || '0' }} + ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }} + ENABLE_PADDLE: ${{ matrix.enable_paddle && '1' || '0' }} if: ${{ !matrix.check_memleak }} - uses: codecov/codecov-action@v5 env: diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 5f6b0e73ab..29c4b5263f 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -19,21 +19,17 @@ jobs: runs-on: nvidia # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845 container: - image: nvidia/cuda:12.6.2-cudnn-devel-ubuntu22.04 + image: nvidia/cuda:12.9.1-cudnn-devel-ubuntu22.04 options: --gpus all if: github.repository_owner == 'deepmodeling' && (github.event_name == 'pull_request' && github.event.label && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' || github.event_name == 'merge_group') steps: - name: Make sudo and git work run: apt-get update && apt-get install -y sudo git - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: '3.11' # cache: 'pip' - - name: Setup MPI - uses: mpi4py/setup-mpi@v1 - with: - mpi: mpich - name: Install wget and unzip run: apt-get update && apt-get install -y wget unzip - uses: lukka/get-cmake@latest @@ -47,12 +43,14 @@ jobs: && sudo apt-get -y install cuda-12-3 libcudnn8=8.9.5.*-1+cuda12.3 if: false # skip as we use nvidia image - run: python -m pip install -U uv - - run: source/install/uv_with_retry.sh pip install --system "tensorflow~=2.18.0rc2" "torch~=2.7.0" "jax[cuda12]==0.5.0" + - run: source/install/uv_with_retry.sh pip install --system --group pin_tensorflow_gpu --group pin_pytorch_gpu --group pin_jax "jax[cuda12]" - run: | export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])') - export TENSORFLOW_ROOT=$(python -c 'import importlib,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') - pip install "paddlepaddle-gpu==3.0.0" -i https://www.paddlepaddle.org.cn/packages/stable/cu126/ + export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') + pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple "paddlepaddle-gpu==3.3.0.dev20251204" source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit + # See https://github.com/jax-ml/jax/issues/29042 + source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13' env: DP_VARIANT: cuda DP_ENABLE_NATIVE_OPTIMIZATION: 1 @@ -65,6 +63,7 @@ jobs: # See https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html XLA_PYTHON_CLIENT_PREALLOCATE: false XLA_PYTHON_CLIENT_ALLOCATOR: platform + FLAGS_use_stride_compute_kernel: 0 - name: Convert models run: source/tests/infer/convert-models.sh - run: | @@ -74,7 +73,6 @@ jobs: OMP_NUM_THREADS: 1 TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 - LMP_CXX11_ABI_0: 1 CMAKE_GENERATOR: Ninja DP_VARIANT: cuda DP_USE_MPICH2: 1 diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 0190dc002c..4e8d1617b8 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -2,6 +2,9 @@ on: push: branches-ignore: - "gh-readonly-queue/**" + - "copilot/**" + - "dependabot/**" + - "pre-commit-ci-update-config" pull_request: merge_group: concurrency: @@ -16,23 +19,20 @@ jobs: fail-fast: false matrix: group: [1, 2, 3, 4, 5, 6] - python: ["3.9", "3.12"] + python: ["3.10", "3.13"] steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 with: python-version: ${{ matrix.python }} - run: python -m pip install -U uv - run: | - source/install/uv_with_retry.sh pip install --system openmpi tensorflow-cpu~=2.18.0 - source/install/uv_with_retry.sh pip install --system torch -i https://download.pytorch.org/whl/cpu - export TENSORFLOW_ROOT=$(python -c 'import tensorflow;print(tensorflow.__path__[0])') + source/install/uv_with_retry.sh pip install --system openmpi --group pin_tensorflow_cpu --group pin_pytorch_cpu --torch-backend cpu + export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)') export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])') - source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py "jax==0.5.0;python_version>='3.10'" - source/install/uv_with_retry.sh pip install --system -U setuptools - source/install/uv_with_retry.sh pip install --system horovod --no-build-isolation - source/install/uv_with_retry.sh pip install --system --pre "paddlepaddle==3.0.0" -i https://www.paddlepaddle.org.cn/packages/stable/cpu/ + source/install/uv_with_retry.sh pip install --system -e .[test,jax] mpi4py --group pin_jax + source/install/uv_with_retry.sh pip install --system --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cpu/paddlepaddle/" --index-url https://pypi.org/simple paddlepaddle==3.3.0.dev20251204 env: # Please note that uv has some issues with finding # existing TensorFlow package. Currently, it uses @@ -40,7 +40,6 @@ jobs: # changes, setting `TENSORFLOW_ROOT`. DP_ENABLE_PYTORCH: 1 DP_BUILD_TESTING: 1 - UV_EXTRA_INDEX_URL: "https://pypi.anaconda.org/mpi4py/simple" HOROVOD_WITH_TENSORFLOW: 1 HOROVOD_WITHOUT_PYTORCH: 1 HOROVOD_WITH_MPI: 1 @@ -48,7 +47,7 @@ jobs: CMAKE_POLICY_VERSION_MINIMUM: 3.5 - run: dp --version - name: Get durations from cache - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: .test_durations # the key must never match, even when restarting workflows, as that @@ -61,6 +60,8 @@ jobs: - run: pytest --cov=deepmd source/tests --splits 6 --group ${{ matrix.group }} --store-durations --clean-durations --durations-path=.test_durations --splitting-algorithm least_duration env: NUM_WORKERS: 0 + DP_CI_IMPORT_PADDLE_BEFORE_TF: 1 + FLAGS_use_stride_compute_kernel: 0 - name: Test TF2 eager mode run: pytest --cov=deepmd --cov-append source/tests/consistent/io/test_io.py source/jax2tf_tests env: @@ -70,7 +71,7 @@ jobs: if: matrix.group == 1 - run: mv .test_durations .test_durations_${{ matrix.group }} - name: Upload partial durations - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v6 with: name: split-${{ matrix.python }}-${{ matrix.group }} path: .test_durations_${{ matrix.group }} @@ -84,11 +85,11 @@ jobs: strategy: fail-fast: false matrix: - python: ["3.9", "3.12"] + python: ["3.10", "3.13"] needs: testpython steps: - name: Get durations from cache - uses: actions/cache@v4 + uses: actions/cache@v5 with: path: .test_durations # key won't match during the first run for the given commit, but @@ -97,7 +98,7 @@ jobs: key: test2-durations-combined-${{ matrix.python }}-${{ github.sha }} restore-keys: test2-durations-combined-${{ matrix.python }} - name: Download artifacts - uses: actions/download-artifact@v4 + uses: actions/download-artifact@v7 with: pattern: split-${{ matrix.python }}-* merge-multiple: true diff --git a/.github/workflows/todo.yml b/.github/workflows/todo.yml index 2608bb1071..25cb02aba2 100644 --- a/.github/workflows/todo.yml +++ b/.github/workflows/todo.yml @@ -2,13 +2,13 @@ name: TODO workflow on: push: branches: - - devel + - master jobs: build: if: github.repository_owner == 'deepmodeling' runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Run tdg-github-action uses: ribtoks/tdg-github-action@master with: diff --git a/.gitignore b/.gitignore index c574da757a..6382ecedd2 100644 --- a/.gitignore +++ b/.gitignore @@ -28,7 +28,7 @@ venv* .vscode/** _build _templates -API_CC +doc/API_CC/ doc/api_py/ doc/api_core/ doc/api_c/ @@ -50,3 +50,24 @@ uv.lock buildcxx/ node_modules/ *.bib.original + +# Coverage files +.coverage +.coverage.* + +# Test output files (temporary) +test_dp_test/ +test_dp_test_*.out +*_detail.out + +# Training and model output files +*.pth +*.ckpt* +checkpoint +lcurve.out +out.json +input_v2_compat.json +frozen_model.* + +# Test system directories +system/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index cee3d7f2ce..599b005309 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,7 +2,7 @@ # See https://pre-commit.com/hooks.html for more hooks repos: - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v5.0.0 + rev: v6.0.0 hooks: - id: trailing-whitespace exclude: "^.+\\.pbtxt$" @@ -22,14 +22,14 @@ repos: - id: check-toml # Python - repo: https://github.com/PyCQA/isort - rev: 6.0.1 + rev: 7.0.0 hooks: - id: isort files: \.py$ exclude: ^source/3rdparty - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.11.13 + rev: v0.14.10 hooks: - id: ruff args: ["--fix"] @@ -40,7 +40,7 @@ repos: types_or: [python, pyi, jupyter] - repo: https://github.com/pycqa/flake8 # flake8 cannot autofix - rev: "7.2.0" + rev: "7.3.0" hooks: - id: flake8 additional_dependencies: @@ -55,12 +55,12 @@ repos: exclude: ^source/3rdparty # Python inside docs - repo: https://github.com/asottile/blacken-docs - rev: 1.19.1 + rev: 1.20.0 hooks: - id: blacken-docs # C++ - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v20.1.5 + rev: v21.1.8 hooks: - id: clang-format exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|.+\.json$) @@ -74,7 +74,7 @@ repos: exclude: ^(source/3rdparty|\.github/workflows|\.clang-format) # Shell - repo: https://github.com/scop/pre-commit-shfmt - rev: v3.11.0-1 + rev: v3.12.0-2 hooks: - id: shfmt # CMake @@ -84,7 +84,7 @@ repos: - id: cmake-format #- id: cmake-lint - repo: https://github.com/njzjz/mirrors-bibtex-tidy - rev: v1.13.0 + rev: v1.14.0 hooks: - id: bibtex-tidy args: @@ -154,10 +154,8 @@ repos: exclude: .pre-commit-config.yaml|source/lmp # customized pylint rules - repo: https://github.com/pylint-dev/pylint/ - rev: v3.3.7 + rev: v4.0.4 hooks: - id: pylint entry: env PYTHONPATH=source/checker pylint files: ^deepmd/ -ci: - autoupdate_branch: devel diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000..c629a08def --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,191 @@ +# DeePMD-kit + +DeePMD-kit is a deep learning package for many-body potential energy representation and molecular dynamics. It supports multiple backends (TensorFlow, PyTorch, JAX, Paddle) and integrates with MD packages like LAMMPS, GROMACS, and i-PI. + +**Always reference these instructions first and fallback to search or bash commands only when you encounter unexpected information that does not match the info here.** + +## Working Effectively + +### Bootstrap and Build Repository + +- Create virtual environment: `uv venv venv && source venv/bin/activate` +- Install base dependencies: `uv pip install tensorflow-cpu` (takes ~8 seconds) +- Install PyTorch: `uv pip install torch --index-url https://download.pytorch.org/whl/cpu` (takes ~5 seconds) +- Build Python package: `uv pip install -e .[cpu,test]` -- takes 67 seconds. **NEVER CANCEL. Set timeout to 120+ seconds.** +- Build C++ components: `export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')` then `export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')` then `./source/install/build_cc.sh` -- takes 164 seconds. **NEVER CANCEL. Set timeout to 300+ seconds.** + +### Test Repository + +- Run single test: `pytest source/tests/tf/test_dp_test.py::TestDPTestEner::test_1frame -v` -- takes 8-13 seconds +- Run test subset: `pytest source/tests/tf/test_dp_test.py -v` -- takes 15 seconds. **NEVER CANCEL. Set timeout to 60+ seconds.** +- **Recommended: Use single test cases for validation instead of full test suite** -- full suite has 314 test files and takes 60+ minutes + +### Lint and Format Code + +- Install linter: `uv pip install ruff` +- Run linting: `ruff check .` -- takes <1 second +- Format code: `ruff format .` -- takes <1 second +- **Always run `ruff check .` and `ruff format .` before committing changes or the CI will fail.** + +### Training and Validation + +- Test TensorFlow training: `cd examples/water/se_e2_a && dp train input.json --skip-neighbor-stat` -- training proceeds but is slow on CPU +- Test PyTorch training: `cd examples/water/se_e2_a && dp --pt train input_torch.json --skip-neighbor-stat` -- training proceeds but is slow on CPU +- **Training examples are for validation only. Real training takes hours/days. Timeout training tests after 60 seconds for validation.** + +## Validation Scenarios + +**ALWAYS manually validate any new code through at least one complete scenario:** + +### Basic Functionality Validation + +1. **CLI Interface**: Run `dp --version` and `dp -h` to verify installation +2. **Python Interface**: Run `python -c "import deepmd; import deepmd.tf; print('Both interfaces work')"` +3. **Backend Selection**: Test `dp --tf -h`, `dp --pt -h`, `dp --jax -h`, `dp --pd -h` + +### Training Workflow Validation + +1. **TensorFlow Training**: `cd examples/water/se_e2_a && timeout 60 dp train input.json --skip-neighbor-stat` -- should start training and show decreasing loss +2. **PyTorch Training**: `cd examples/water/se_e2_a && timeout 60 dp --pt train input_torch.json --skip-neighbor-stat` -- should start training and show decreasing loss +3. **Verify training output**: Look for "batch X: trn: rmse" messages showing decreasing error values + +### Test-Based Validation + +1. **Core Tests**: `pytest source/tests/tf/test_dp_test.py::TestDPTestEner::test_1frame -v` -- should pass in ~10 seconds +2. **Multi-backend**: Test both TensorFlow and PyTorch components work + +## Common Commands and Timing + +### Repository Structure + +``` +ls -la [repo-root] +.github/ # GitHub workflows and templates +CONTRIBUTING.md # Contributing guide +README.md # Project overview +deepmd/ # Python package source +doc/ # Documentation +examples/ # Training examples and configurations +pyproject.toml # Python build configuration +source/ # C++ source code and tests +``` + +### Key Directories and Files + +- `deepmd/` - Main Python package with backend implementations +- `source/lib/` - Core C++ library +- `source/op/` - Backend-specific operators (TF, PyTorch, etc.) +- `source/api_cc/` - C++ API +- `source/api_c/` - C API +- `source/tests/` - Test suite (314 test files) +- `examples/water/se_e2_a/` - Basic water training example +- `examples/` - Various model examples for different scenarios + +### Common CLI Commands + +- `dp --version` - Show version information +- `dp -h` - Show help and available commands +- `dp train input.json` - Train a model (TensorFlow backend) +- `dp --pt train input.json` - Train with PyTorch backend +- `dp --jax train input.json` - Train with JAX backend +- `dp --pd train input.json` - Train with Paddle backend +- `dp test -m model.pb -s system/` - Test a trained model +- `dp freeze -o model.pb` - Freeze/save a model + +### Build Dependencies and Setup + +- **Python 3.10+** required +- **Virtual environment** strongly recommended: `uv venv venv && source venv/bin/activate` +- **Backend dependencies**: TensorFlow, PyTorch, JAX, or Paddle (install before building) +- **Build tools**: CMake, C++ compiler, scikit-build-core +- **C++ build requires**: Both TensorFlow and PyTorch installed, set TENSORFLOW_ROOT and PYTORCH_ROOT environment variables + +### Key Configuration Files + +- `pyproject.toml` - Python build configuration and dependencies +- `source/CMakeLists.txt` - C++ build configuration +- `examples/water/se_e2_a/input.json` - Basic TensorFlow training config +- `examples/water/se_e2_a/input_torch.json` - Basic PyTorch training config + +## Frequent Patterns and Time Expectations + +### Installation and Build Times + +- **Virtual environment setup**: ~5 seconds +- **TensorFlow CPU install**: ~8 seconds +- **PyTorch CPU install**: ~5 seconds +- **Python package build**: ~67 seconds. **NEVER CANCEL.** +- **C++ components build**: ~164 seconds. **NEVER CANCEL.** +- **Full fresh setup**: ~3-4 minutes total + +### Testing Times + +- **Single test**: 8-13 seconds +- **Test file (~5 tests)**: ~15 seconds +- **Backend-specific test subset**: 15-30 minutes. **Use sparingly.** +- **Full test suite (314 files)**: 60+ minutes. **Avoid in development - use single tests instead.** + +### Linting and Formatting + +- **Ruff check**: <1 second +- **Ruff format**: <1 second +- **Pre-commit hooks**: May have network issues, use individual tools + +### Commit Messages and PR Titles + +**All commit messages and PR titles must follow [conventional commit specification](https://www.conventionalcommits.org/):** + +- **Format**: `type(scope): description` +- **Common types**: `feat`, `fix`, `docs`, `style`, `refactor`, `test`, `chore`, `ci` +- **Examples**: + - `feat(core): add new descriptor type` + - `fix(tf): resolve memory leak in training` + - `docs: update installation guide` + - `ci: add workflow for testing` + +### Training and Model Operations + +- **Training initialization**: 10-30 seconds +- **Training per batch**: 0.1-1 second (CPU), much faster on GPU +- **Model freezing**: 5-15 seconds +- **Model testing**: 10-30 seconds + +## Backend-Specific Notes + +### TensorFlow Backend + +- **Default backend** when no flag specified +- **Configuration**: Use `input.json` format +- **Training**: `dp train input.json` +- **Requirements**: `tensorflow` or `tensorflow-cpu` package + +### PyTorch Backend + +- **Activation**: Use `--pt` flag or `export DP_BACKEND=pytorch` +- **Configuration**: Use `input_torch.json` format typically +- **Training**: `dp --pt train input_torch.json` +- **Requirements**: `torch` package + +### JAX Backend + +- **Activation**: Use `--jax` flag +- **Training**: `dp --jax train input.json` +- **Requirements**: `jax` and related packages +- **Note**: Experimental backend, may have limitations + +### Paddle Backend + +- **Activation**: Use `--pd` flag +- **Training**: `dp --pd train input.json` +- **Requirements**: `paddlepaddle` package +- **Note**: Less commonly used + +## Critical Warnings + +- **NEVER CANCEL BUILD OPERATIONS**: Python build takes 67 seconds, C++ build takes 164 seconds +- **USE SINGLE TESTS FOR VALIDATION**: Run individual tests instead of full test suite for faster feedback +- **ALWAYS activate virtual environment**: Build and runtime failures occur without proper environment +- **ALWAYS install backend dependencies first**: TensorFlow/PyTorch required before building C++ components +- **ALWAYS run linting before commits**: `ruff check . && ruff format .` or CI will fail +- **ALWAYS test both Python and C++ components**: Some features require both to be built +- **ALWAYS follow conventional commit format**: All commit messages and PR titles must use conventional commit specification (`type(scope): description`) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cb08609c2b..67491672e8 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -36,13 +36,6 @@ You can help other users of deepmd-kit in the following way ## Before you contribute -### Overview of DeePMD-kit - -Currently, we maintain two main branch: - -- master: stable branch with version tag -- devel : branch for developers - ### Developer guide See [documentation](https://deepmd.readthedocs.io/) for coding conventions, API and other needs-to-know of the code. @@ -58,13 +51,11 @@ Please perform the following steps to create your Pull Request to this repositor ### Step 2: Clone the forked repository to local storage and set configurations -1. Clone your own repo, not the public repo (from deepmodeling) ! And change the branch to devel. +1. Clone your own repo, not the public repo (from deepmodeling) ! ```bash git clone https://github.com/$username/deepmd-kit.git # Replace `$username` with your GitHub ID - - git checkout devel ``` 2. Add deepmodeling's repo as your remote repo, we can name it "upstream". And fetch upstream's latest codes to your workstation. @@ -76,15 +67,15 @@ Please perform the following steps to create your Pull Request to this repositor git fetch upstream # If your current codes are behind the latest codes, you should merge latest codes first. - # Notice you should merge from "devel"! - git merge upstream/devel + git merge upstream/master ``` 3. Modify your codes and design unit tests. -4. Commit your changes +4. Commit your changes to a new branch ```bash + git checkout -b branch1 git status # Checks the local status git add ... # Adds the file(s) you want to commit. If you want to commit all changes, you can directly use `git add.` git commit -m "commit-message: update the xx" @@ -92,7 +83,7 @@ Please perform the following steps to create your Pull Request to this repositor 5. Push the changed codes to your original repo on github. ```bash - git push origin devel + git push origin branch1 ``` ### Alternatively: Create a new branch @@ -145,4 +136,4 @@ Now, your PR is successfully submitted! After this PR is merged, you will automa ## Contact us -E-mail: contact@deepmodeling.org +E-mail: deepmodeling@deepmodeling.com diff --git a/backend/dp_backend.py b/backend/dp_backend.py index 81c3f20f19..e32d5db38b 100644 --- a/backend/dp_backend.py +++ b/backend/dp_backend.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """A PEP-517 backend to find TensorFlow.""" +import os + from scikit_build_core import build as _orig from .find_pytorch import ( @@ -39,10 +41,15 @@ def __dir__() -> list[str]: def get_requires_for_build_wheel( config_settings: dict, ) -> list[str]: + if os.environ.get("CIBUILDWHEEL", "0") == "1": + cibw_deps = ["mpich"] + else: + cibw_deps = [] return ( _orig.get_requires_for_build_wheel(config_settings) + find_tensorflow()[1] + find_pytorch()[1] + + cibw_deps ) diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index a66e9a2759..e7763cac84 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -3,9 +3,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) from .find_pytorch import ( get_pt_requirement, @@ -31,7 +28,7 @@ def __dir__() -> list[str]: def dynamic_metadata( field: str, - settings: Optional[dict[str, object]] = None, + settings: dict[str, object] | None = None, ): assert field in ["optional-dependencies", "entry-points", "scripts"] _, _, find_libpython_requires, extra_scripts, tf_version, pt_version = ( diff --git a/backend/find_paddle.py b/backend/find_paddle.py index 247526c717..c9b0319504 100644 --- a/backend/find_paddle.py +++ b/backend/find_paddle.py @@ -17,14 +17,10 @@ from sysconfig import ( get_path, ) -from typing import ( - Optional, - Union, -) @lru_cache -def find_paddle() -> tuple[Optional[str], list[str]]: +def find_paddle() -> tuple[str | None, list[str]]: """Find PaddlePadle library. Tries to find PaddlePadle in the order of: @@ -111,7 +107,7 @@ def get_pd_requirement(pd_version: str = "") -> dict: @lru_cache -def get_pd_version(pd_path: Optional[Union[str, Path]]) -> str: +def get_pd_version(pd_path: str | Path | None) -> str: """Get Paddle version from a Paddle Python library path. Parameters diff --git a/backend/find_pytorch.py b/backend/find_pytorch.py index c51c8905b8..d50f57bf5e 100644 --- a/backend/find_pytorch.py +++ b/backend/find_pytorch.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import importlib import os -import platform import site from functools import ( lru_cache, @@ -18,10 +17,6 @@ from sysconfig import ( get_path, ) -from typing import ( - Optional, - Union, -) from packaging.specifiers import ( SpecifierSet, @@ -30,9 +25,13 @@ Version, ) +from .utils import ( + read_dependencies_from_dependency_group, +) + @lru_cache -def find_pytorch() -> tuple[Optional[str], list[str]]: +def find_pytorch() -> tuple[str | None, list[str]]: """Find PyTorch library. Tries to find PyTorch in the order of: @@ -108,22 +107,24 @@ def get_pt_requirement(pt_version: str = "") -> dict: """ if pt_version is None: return {"torch": []} - if ( - os.environ.get("CIBUILDWHEEL", "0") == "1" - and platform.system() == "Linux" - and platform.machine() == "x86_64" - ): + cibw_requirement = [] + if os.environ.get("CIBUILDWHEEL", "0") == "1": cuda_version = os.environ.get("CUDA_VERSION", "12.2") if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"): # CUDA 12.2, cudnn 9 - pt_version = "2.7.0" - elif cuda_version in SpecifierSet(">=11,<12"): - # CUDA 11.8, cudnn 8 - pt_version = "2.3.1" + # or CPU builds + cibw_requirement = read_dependencies_from_dependency_group( + "pin_pytorch_cpu" + ) else: raise RuntimeError("Unsupported CUDA version") from None if pt_version == "": pt_version = os.environ.get("PYTORCH_VERSION", "") + if os.environ.get("CIBUILDWHEEL", "0") == "1": + # PyTorch OP library is built against mpich + mpi_requirement = ["mpich"] + else: + mpi_requirement = [] return { "torch": [ @@ -134,13 +135,15 @@ def get_pt_requirement(pt_version: str = "") -> dict: f"torch=={Version(pt_version).base_version}.*" if pt_version != "" # https://github.com/pytorch/pytorch/commit/7e0c26d4d80d6602aed95cb680dfc09c9ce533bc - else "torch>=2.1.0" + else "torch>=2.1.0", + *mpi_requirement, + *cibw_requirement, ], } @lru_cache -def get_pt_version(pt_path: Optional[Union[str, Path]]) -> str: +def get_pt_version(pt_path: str | Path | None) -> str: """Get TF version from a TF Python library path. Parameters diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py index 1fc3a8a6d9..b1528f1a47 100644 --- a/backend/find_tensorflow.py +++ b/backend/find_tensorflow.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import os +import re import site from functools import ( lru_cache, @@ -16,18 +17,18 @@ from sysconfig import ( get_path, ) -from typing import ( - Optional, - Union, -) from packaging.specifiers import ( SpecifierSet, ) +from .utils import ( + read_dependencies_from_dependency_group, +) + @lru_cache -def find_tensorflow() -> tuple[Optional[str], list[str]]: +def find_tensorflow() -> tuple[str | None, list[str]]: """Find TensorFlow library. Tries to find TensorFlow in the order of: @@ -56,6 +57,10 @@ def find_tensorflow() -> tuple[Optional[str], list[str]]: ) is not None: site_packages = Path(os.environ.get("TENSORFLOW_ROOT")).parent.absolute() tf_spec = FileFinder(str(site_packages)).find_spec("tensorflow") + if tf_spec is None: + raise RuntimeError( + f"cannot find TensorFlow under TENSORFLOW_ROOT {os.environ.get('TENSORFLOW_ROOT')}" + ) # get tensorflow spec # note: isolated build will not work for backend @@ -86,19 +91,10 @@ def find_tensorflow() -> tuple[Optional[str], list[str]]: cuda_version = os.environ.get("CUDA_VERSION", "12.2") if cuda_version == "" or cuda_version in SpecifierSet(">=12,<13"): # CUDA 12.2, cudnn 9 + # or CPU builds requires.extend( - [ - "tensorflow-cpu>=2.18.0rc0; platform_machine=='x86_64' and platform_system == 'Linux'", - ] + read_dependencies_from_dependency_group("pin_tensorflow_cpu") ) - elif cuda_version in SpecifierSet(">=11,<12"): - # CUDA 11.8, cudnn 8 - requires.extend( - [ - "tensorflow-cpu>=2.5.0rc0,<2.15; platform_machine=='x86_64' and platform_system == 'Linux'", - ] - ) - tf_version = "2.14.1" else: raise RuntimeError("Unsupported CUDA version") from None requires.extend(get_tf_requirement(tf_version)["cpu"]) @@ -153,7 +149,8 @@ def get_tf_requirement(tf_version: str = "") -> dict: "tensorflow-cpu; platform_machine!='aarch64' and (platform_machine!='arm64' or platform_system != 'Darwin')", "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')", # https://github.com/tensorflow/tensorflow/issues/61830 - "tensorflow-cpu!=2.15.*; platform_system=='Windows'", + # Since TF 2.20, not all symbols are exported to the public API. + "tensorflow-cpu!=2.15.*,<2.20; platform_system=='Windows'", # https://github.com/h5py/h5py/issues/2408 "h5py>=3.6.0,!=3.11.0; platform_system=='Linux' and platform_machine=='aarch64'", *extra_requires, @@ -199,7 +196,7 @@ def get_tf_requirement(tf_version: str = "") -> dict: @lru_cache -def get_tf_version(tf_path: Optional[Union[str, Path]]) -> str: +def get_tf_version(tf_path: str | Path | None) -> str: """Get TF version from a TF Python library path. Parameters @@ -228,6 +225,22 @@ def get_tf_version(tf_path: Optional[Union[str, Path]]) -> str: patch = line.split()[-1] elif line.startswith("#define TF_VERSION_SUFFIX"): suffix = line.split()[-1].strip('"') + if None in (major, minor, patch): + # since TF 2.20.0, version information is no more contained in version.h + # try to read version from tools/pip_package/setup.py + # _VERSION = '2.20.0' + setup_file = Path(tf_path) / "tools" / "pip_package" / "setup.py" + if setup_file.exists(): + with open(setup_file) as f: + for line in f: + # parse with regex + match = re.search( + r"_VERSION[ \t]*=[ \t]*'(\d+)\.(\d+)\.(\d+)([a-zA-Z0-9]*)?'", + line, + ) + if match: + major, minor, patch, suffix = match.groups() + break if None in (major, minor, patch): raise RuntimeError("Failed to read TF version") return ".".join((major, minor, patch)) + suffix diff --git a/backend/read_env.py b/backend/read_env.py index f28e2917f3..8a173513f9 100644 --- a/backend/read_env.py +++ b/backend/read_env.py @@ -46,15 +46,15 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]: # get variant option from the environment variables, available: cpu, cuda, rocm dp_variant = os.environ.get("DP_VARIANT", "cpu").lower() if dp_variant == "cpu" or dp_variant == "": - cmake_minimum_required_version = "3.16" + cmake_minimum_required_version = "3.25.2" elif dp_variant == "cuda": - cmake_minimum_required_version = "3.23" + cmake_minimum_required_version = "3.25.2" cmake_args.append("-DUSE_CUDA_TOOLKIT:BOOL=TRUE") cuda_root = os.environ.get("CUDAToolkit_ROOT") if cuda_root: cmake_args.append(f"-DCUDAToolkit_ROOT:STRING={cuda_root}") elif dp_variant == "rocm": - cmake_minimum_required_version = "3.21" + cmake_minimum_required_version = "3.25.2" cmake_args.append("-DUSE_ROCM_TOOLKIT:BOOL=TRUE") rocm_root = os.environ.get("ROCM_ROOT") if not rocm_root: @@ -119,6 +119,7 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]: cmake_args = [ "-DBUILD_PY_IF:BOOL=TRUE", + f"-DCIBUILDWHEEL={os.environ.get('CIBUILDWHEEL', '0')}", *cmake_args, ] return ( diff --git a/backend/utils.py b/backend/utils.py new file mode 100644 index 0000000000..0769879d24 --- /dev/null +++ b/backend/utils.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import sys +from pathlib import ( + Path, +) + +from dependency_groups import ( + resolve, +) + +if sys.version_info >= (3, 11): + import tomllib +else: + import tomli as tomllib + + +def read_dependencies_from_dependency_group(group: str) -> tuple[str, ...]: + """ + Reads dependencies from a dependency group. + + Parameters + ---------- + group : str + The name of the dependency group. + + Returns + ------- + tuple[str, ...] + A tuple of dependencies in the specified group. + """ + with Path("pyproject.toml").open("rb") as f: + pyproject = tomllib.load(f) + + groups = pyproject["dependency-groups"] + + return resolve(groups, group) diff --git a/deepmd/__init__.py b/deepmd/__init__.py index 14d933da11..bc351ee59b 100644 --- a/deepmd/__init__.py +++ b/deepmd/__init__.py @@ -8,6 +8,14 @@ modules for performance. """ +from typing import ( + TYPE_CHECKING, + Any, +) + +if TYPE_CHECKING: + from deepmd.infer import DeepPotential as DeepPotentialType + try: from deepmd._version import version as __version__ except ImportError: @@ -16,7 +24,7 @@ ) -def DeepPotential(*args, **kwargs): +def DeepPotential(*args: Any, **kwargs: Any) -> "DeepPotentialType": """Factory function that forwards to DeepEval (for compatibility and performance). diff --git a/deepmd/backend/backend.py b/deepmd/backend/backend.py index 3263169f6f..58dcfe427d 100644 --- a/deepmd/backend/backend.py +++ b/deepmd/backend/backend.py @@ -2,13 +2,15 @@ from abc import ( abstractmethod, ) +from collections.abc import ( + Callable, +) from enum import ( Flag, auto, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/dpmodel.py b/deepmd/backend/dpmodel.py index 7c21b256ae..31585aa7a6 100644 --- a/deepmd/backend/dpmodel.py +++ b/deepmd/backend/dpmodel.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/jax.py b/deepmd/backend/jax.py index 7a714c2090..9c0055b4f2 100644 --- a/deepmd/backend/jax.py +++ b/deepmd/backend/jax.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/paddle.py b/deepmd/backend/paddle.py index b1f664e76a..670130e86a 100644 --- a/deepmd/backend/paddle.py +++ b/deepmd/backend/paddle.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/pytorch.py b/deepmd/backend/pytorch.py index f5b0dd92b2..d155ef1f41 100644 --- a/deepmd/backend/pytorch.py +++ b/deepmd/backend/pytorch.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/backend/suffix.py b/deepmd/backend/suffix.py index e77aecb5d9..3d1602da0e 100644 --- a/deepmd/backend/suffix.py +++ b/deepmd/backend/suffix.py @@ -4,10 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) from deepmd.backend.backend import ( Backend, @@ -16,9 +12,9 @@ def format_model_suffix( filename: str, - feature: Optional[Backend.Feature] = None, - preferred_backend: Optional[Union[str, type["Backend"]]] = None, - strict_prefer: Optional[bool] = None, + feature: Backend.Feature | None = None, + preferred_backend: str | type["Backend"] | None = None, + strict_prefer: bool | None = None, ) -> str: """Check and format the suffixes of a filename. diff --git a/deepmd/backend/tensorflow.py b/deepmd/backend/tensorflow.py index 6b73d7c469..244b4d9980 100644 --- a/deepmd/backend/tensorflow.py +++ b/deepmd/backend/tensorflow.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from importlib.util import ( find_spec, ) from typing import ( TYPE_CHECKING, - Callable, ClassVar, ) diff --git a/deepmd/calculator.py b/deepmd/calculator.py index cb386d091a..dca2343ce0 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -6,6 +6,7 @@ ) from typing import ( TYPE_CHECKING, + Any, ClassVar, Optional, Union, @@ -25,6 +26,9 @@ from ase import ( Atoms, ) + from ase.neighborlist import ( + NeighborList, + ) __all__ = ["DP"] @@ -84,10 +88,10 @@ def __init__( self, model: Union[str, "Path"], label: str = "DP", - type_dict: Optional[dict[str, int]] = None, - neighbor_list=None, - head=None, - **kwargs, + type_dict: dict[str, int] | None = None, + neighbor_list: Optional["NeighborList"] = None, + head: str | None = None, + **kwargs: Any, ) -> None: Calculator.__init__(self, label=label, **kwargs) self.dp = DeepPot( @@ -130,7 +134,12 @@ def calculate( cell = None symbols = self.atoms.get_chemical_symbols() atype = [self.type_dict[k] for k in symbols] - e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype)[:3] + + fparam = self.atoms.info.get("fparam", None) + aparam = self.atoms.info.get("aparam", None) + e, f, v = self.dp.eval( + coords=coord, cells=cell, atom_types=atype, fparam=fparam, aparam=aparam + )[:3] self.results["energy"] = e[0][0] # see https://gitlab.com/ase/ase/-/merge_requests/2485 self.results["free_energy"] = e[0][0] diff --git a/deepmd/common.py b/deepmd/common.py index 3ab936db67..98cf2461bd 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -15,14 +15,13 @@ TYPE_CHECKING, Any, TypeVar, - Union, get_args, ) try: from typing import Literal # python >=3.8 except ImportError: - from typing_extensions import Literal # type: ignore + from typing import Literal # type: ignore import numpy as np import yaml @@ -35,6 +34,7 @@ ) __all__ = [ + "GLOBAL_NP_FLOAT_PRECISION", "VALID_ACTIVATION", "VALID_PRECISION", "expand_sys_str", @@ -44,7 +44,7 @@ "select_idx_map", ] -_PRECISION = Literal["default", "float16", "float32", "float64"] +_PRECISION = Literal["default", "float16", "bfloat16", "float32", "float64"] _ACTIVATION = Literal[ "relu", "relu6", @@ -158,7 +158,7 @@ def j_deprecated( return jdata[key] -def j_loader(filename: Union[str, Path]) -> dict[str, Any]: +def j_loader(filename: str | Path) -> dict[str, Any]: """Load yaml or json settings file. Parameters @@ -187,7 +187,7 @@ def j_loader(filename: Union[str, Path]) -> dict[str, Any]: raise TypeError("config file must be json, or yaml/yml") -def expand_sys_str(root_dir: Union[str, Path]) -> list[str]: +def expand_sys_str(root_dir: str | Path) -> list[str]: """Recursively iterate over directories taking those that contain `type.raw` file. Parameters @@ -249,16 +249,11 @@ def get_np_precision(precision: "_PRECISION") -> np.dtype: RuntimeError if string is invalid """ - if precision == "default": - return GLOBAL_NP_FLOAT_PRECISION - elif precision == "float16": - return np.float16 - elif precision == "float32": - return np.float32 - elif precision == "float64": - return np.float64 - else: - raise RuntimeError(f"{precision} is not a valid precision") + from deepmd.dpmodel.common import ( + get_xp_precision, + ) + + return get_xp_precision(np, precision) def symlink_prefix_files(old_prefix: str, new_prefix: str) -> None: @@ -288,7 +283,7 @@ def symlink_prefix_files(old_prefix: str, new_prefix: str) -> None: shutil.copyfile(ori_ff, new_ff) -def get_hash(obj) -> str: +def get_hash(obj: Any) -> str: """Get hash of object. Parameters diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index 723718529c..2db7550eee 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -1,14 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Utilities for the array API.""" +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + import array_api_compat import numpy as np from packaging.version import ( Version, ) +# Type alias for array_api compatible arrays +Array = np.ndarray | Any # Any to support JAX, PyTorch, etc. arrays + -def support_array_api(version: str) -> callable: +def support_array_api(version: str) -> Callable: """Mark a function as supporting the specific version of the array API. Parameters @@ -18,7 +28,7 @@ def support_array_api(version: str) -> callable: Returns ------- - callable + Callable The decorated function Examples @@ -28,7 +38,7 @@ def support_array_api(version: str) -> callable: ... pass """ - def set_version(func: callable) -> callable: + def set_version(func: Callable) -> Callable: func.array_api_version = version return func @@ -39,7 +49,7 @@ def set_version(func: callable) -> callable: # but it hasn't been released yet # below is a pure Python implementation of take_along_axis # https://github.com/data-apis/array-api/issues/177#issuecomment-2093630595 -def xp_swapaxes(a, axis1, axis2): +def xp_swapaxes(a: Array, axis1: int, axis2: int) -> Array: xp = array_api_compat.array_namespace(a) axes = list(range(a.ndim)) axes[axis1], axes[axis2] = axes[axis2], axes[axis1] @@ -47,7 +57,7 @@ def xp_swapaxes(a, axis1, axis2): return a -def xp_take_along_axis(arr, indices, axis): +def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array: xp = array_api_compat.array_namespace(arr) if Version(xp.__array_api_version__) >= Version("2024.12"): # see: https://github.com/data-apis/array-api-strict/blob/d086c619a58f35c38240592ef994aa19ca7beebc/array_api_strict/_indexing_functions.py#L30-L39 @@ -60,7 +70,7 @@ def xp_take_along_axis(arr, indices, axis): shape = list(arr.shape) shape.pop(-1) - shape = [*shape, n] + shape = (*shape, n) arr = xp.reshape(arr, (-1,)) if n != 0: @@ -76,7 +86,7 @@ def xp_take_along_axis(arr, indices, axis): return xp_swapaxes(out, axis, -1) -def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray: +def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array: """Reduces all values from the src tensor to the indices specified in the index tensor.""" # jax only if array_api_compat.is_jax_array(input): @@ -94,7 +104,7 @@ def xp_scatter_sum(input, dim, index: np.ndarray, src: np.ndarray) -> np.ndarray raise NotImplementedError("Only JAX arrays are supported.") -def xp_add_at(x, indices, values): +def xp_add_at(x: Array, indices: Array, values: Array) -> Array: """Adds values to the specified indices of x in place or returns new x (for JAX).""" xp = array_api_compat.array_namespace(x, indices, values) if array_api_compat.is_numpy_array(x): @@ -115,7 +125,7 @@ def xp_add_at(x, indices, values): return x -def xp_bincount(x, weights=None, minlength=0): +def xp_bincount(x: Array, weights: Array | None = None, minlength: int = 0) -> Array: """Counts the number of occurrences of each value in x.""" xp = array_api_compat.array_namespace(x) if array_api_compat.is_numpy_array(x) or array_api_compat.is_jax_array(x): diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py index eb95886598..6fb48eb659 100644 --- a/deepmd/dpmodel/atomic_model/base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/base_atomic_model.py @@ -1,12 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import math from typing import ( - Optional, + Any, ) import array_api_compat import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( NativeOP, to_numpy_array, @@ -41,8 +44,8 @@ def __init__( type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], - rcond: Optional[float] = None, - preset_out_bias: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_out_bias: dict[str, Array] | None = None, ) -> None: super().__init__() self.type_map = type_map @@ -68,7 +71,7 @@ def init_out_stat(self) -> None: self.out_bias = out_bias_data self.out_std = out_std_data - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ["out_bias"]: self.out_bias = value elif key in ["out_std"]: @@ -76,7 +79,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ["out_bias"]: return self.out_bias elif key in ["out_std"]: @@ -88,6 +91,10 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + def reinit_atom_exclude( self, exclude_types: list[int] = [], @@ -125,7 +132,7 @@ def atomic_output_def(self) -> FittingOutputDef: ) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -143,13 +150,13 @@ def change_type_map( def forward_common_atomic( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Common interface for atomic inference. This method accept extended coordinates, extended atom typs, neighbor list, @@ -219,13 +226,13 @@ def forward_common_atomic( def call( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: return self.forward_common_atomic( extended_coord, extended_atype, @@ -260,9 +267,9 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel": def apply_out_stat( self, - ret: dict[str, np.ndarray], - atype: np.ndarray, - ): + ret: dict[str, Array], + atype: Array, + ) -> dict[str, Array]: """Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. @@ -305,7 +312,7 @@ def _get_bias_index( def _fetch_out_stat( self, keys: list[str], - ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + ) -> tuple[dict[str, Array], dict[str, Array]]: ret_bias = {} ret_std = {} ntypes = self.get_ntypes() diff --git a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py index 00428f4e95..7cfa24526a 100644 --- a/deepmd/dpmodel/atomic_model/dipole_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dipole_atomic_model.py @@ -1,6 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import numpy as np +from typing import ( + Any, +) +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) from deepmd.dpmodel.fitting.dipole_fitting import ( DipoleFitting, ) @@ -11,7 +22,13 @@ class DPDipoleAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, + descriptor: BaseDescriptor, + fitting: BaseFitting, + type_map: list[str], + **kwargs: Any, + ) -> None: if not isinstance(fitting, DipoleFitting): raise TypeError( "fitting must be an instance of DipoleFitting for DPDipoleAtomicModel" @@ -20,8 +37,8 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): def apply_out_stat( self, - ret: dict[str, np.ndarray], - atype: np.ndarray, - ): + ret: dict[str, Array], + atype: Array, + ) -> dict[str, Array]: # dipole not applying bias return ret diff --git a/deepmd/dpmodel/atomic_model/dos_atomic_model.py b/deepmd/dpmodel/atomic_model/dos_atomic_model.py index 7ef6d10ebf..ce457cb472 100644 --- a/deepmd/dpmodel/atomic_model/dos_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dos_atomic_model.py @@ -1,4 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +from deepmd.dpmodel.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) from deepmd.dpmodel.fitting.dos_fitting import ( DOSFittingNet, ) @@ -9,7 +19,13 @@ class DPDOSAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, + descriptor: BaseDescriptor, + fitting: BaseFitting, + type_map: list[str], + **kwargs: Any, + ) -> None: if not isinstance(fitting, DOSFittingNet): raise TypeError( "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel" diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index 2fa072cc78..07a02ad56b 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) -import numpy as np - +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) @@ -41,15 +42,17 @@ class DPAtomicModel(BaseAtomicModel): def __init__( self, - descriptor, - fitting, + descriptor: BaseDescriptor, + fitting: BaseFitting, type_map: list[str], - **kwargs, + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) self.type_map = type_map self.descriptor = descriptor self.fitting = fitting + if hasattr(self.fitting, "reinit_exclude"): + self.fitting.reinit_exclude(self.atom_exclude_types) self.type_map = type_map super().init_out_stat() @@ -65,7 +68,7 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.descriptor.get_sel() - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -125,13 +128,13 @@ def enable_compression( def forward_atomic( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Models' atomic predictions. Parameters @@ -175,7 +178,7 @@ def forward_atomic( return ret def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -190,7 +193,7 @@ def change_type_map( if model_with_new_type_stat is not None else None, ) - self.fitting_net.change_type_map(type_map=type_map) + self.fitting.change_type_map(type_map=type_map) def serialize(self) -> dict: dd = super().serialize() @@ -213,7 +216,7 @@ def serialize(self) -> dict: """The base fitting class.""" @classmethod - def deserialize(cls, data) -> "DPAtomicModel": + def deserialize(cls, data: dict[str, Any]) -> "DPAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 2) data.pop("@class") @@ -233,6 +236,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.fitting.get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.fitting.has_default_fparam() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/dpmodel/atomic_model/energy_atomic_model.py b/deepmd/dpmodel/atomic_model/energy_atomic_model.py index 4f9f8ec005..6deb87662d 100644 --- a/deepmd/dpmodel/atomic_model/energy_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/energy_atomic_model.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + from deepmd.dpmodel.fitting.ener_fitting import ( EnergyFittingNet, InvarFitting, @@ -10,7 +14,9 @@ class DPEnergyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any + ) -> None: if not ( isinstance(fitting, EnergyFittingNet) or isinstance(fitting, InvarFitting) ): diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index ce0f1d0cb9..3fa1119418 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -1,12 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, ) import array_api_compat import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.utils.nlist import ( build_multiple_neighbor_list, get_multiple_nlist_key, @@ -51,7 +53,7 @@ def __init__( self, models: list[BaseAtomicModel], type_map: list[str], - **kwargs, + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() @@ -111,7 +113,7 @@ def get_type_map(self) -> list[str]: return self.type_map def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -134,7 +136,7 @@ def get_model_rcuts(self) -> list[float]: def get_sel(self) -> list[int]: return [max([model.get_nsel() for model in self.models])] - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -146,11 +148,11 @@ def get_model_nsels(self) -> list[int]: """Get the processed sels for each individual models. Not distinguishing types.""" return [model.get_nsel() for model in self.models] - def get_model_sels(self) -> list[Union[int, list[int]]]: + def get_model_sels(self) -> list[int | list[int]]: """Get the sels for each individual models.""" return [model.get_sel() for model in self.models] - def _sort_rcuts_sels(self) -> tuple[list[float], list[int]]: + def _sort_rcuts_sels(self) -> tuple[tuple[Array, Array], list[int]]: # sort the pair of rcut and sels in ascending order, first based on sel, then on rcut. zipped = sorted( zip(self.get_model_rcuts(), self.get_model_nsels()), @@ -192,13 +194,13 @@ def enable_compression( def forward_atomic( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Return atomic prediction. Parameters @@ -262,7 +264,7 @@ def forward_atomic( return fit_ret @staticmethod - def remap_atype(ori_map: list[str], new_map: list[str]) -> np.ndarray: + def remap_atype(ori_map: list[str], new_map: list[str]) -> Array: """ This method is used to map the atype from the common type_map to the original type_map of indivial AtomicModels. @@ -325,10 +327,10 @@ def deserialize(cls, data: dict) -> "LinearEnergyAtomicModel": def _compute_weight( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlists_: list[np.ndarray], - ) -> list[np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlists_: list[Array], + ) -> list[Array]: """This should be a list of user defined weights that matches the number of models to be combined.""" xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlists_) nmodels = len(self.models) @@ -397,8 +399,8 @@ def __init__( sw_rmin: float, sw_rmax: float, type_map: list[str], - smin_alpha: Optional[float] = 0.1, - **kwargs, + smin_alpha: float | None = 0.1, + **kwargs: Any, ) -> None: models = [dp_model, zbl_model] kwargs["models"] = models @@ -424,7 +426,7 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": + def deserialize(cls, data: Any) -> "DPZBLLinearEnergyAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 2) models = [ @@ -436,7 +438,7 @@ def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": data.pop("type", None) return super().deserialize(data) - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -446,15 +448,15 @@ def set_case_embd(self, case_idx: int): def _compute_weight( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlists_: list[np.ndarray], - ) -> list[np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlists_: list[Array], + ) -> list[Array]: """ZBL weight. Returns ------- - list[np.ndarray] + list[Array] the atomic ZBL weight for interpolation. (nframes, nloc, 1) """ assert self.sw_rmax > self.sw_rmin, ( diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index 01caa7cd64..3e48e88c87 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -4,7 +4,7 @@ abstractmethod, ) from typing import ( - Optional, + Any, ) from deepmd.dpmodel.output_def import ( @@ -17,9 +17,9 @@ def make_base_atomic_model( - t_tensor, + t_tensor: type, fwd_method_name: str = "forward_atomic", -): +) -> type: """Make the base class for the atomic model. Parameters @@ -135,9 +135,9 @@ def fwd( extended_coord: t_tensor, extended_atype: t_tensor, nlist: t_tensor, - mapping: Optional[t_tensor] = None, - fparam: Optional[t_tensor] = None, - aparam: Optional[t_tensor] = None, + mapping: t_tensor | None = None, + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, ) -> dict[str, t_tensor]: pass @@ -147,12 +147,12 @@ def serialize(self) -> dict: @classmethod @abstractmethod - def deserialize(cls, data: dict): + def deserialize(cls, data: dict) -> Any: pass @abstractmethod def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: pass @@ -205,7 +205,7 @@ def make_atom_mask( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -222,7 +222,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 9d7739d5c8..5385d4c56c 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, + NoReturn, ) import array_api_compat import numpy as np from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) from deepmd.dpmodel.output_def import ( @@ -61,11 +62,11 @@ def __init__( self, tab_file: str, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], type_map: list[str], - rcond: Optional[float] = None, - atom_ener: Optional[list[float]] = None, - **kwargs, + rcond: float | None = None, + atom_ener: list[float] | None = None, + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() @@ -120,7 +121,7 @@ def get_type_map(self) -> list[str]: def get_sel(self) -> list[int]: return [self.sel] - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> NoReturn: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -154,7 +155,7 @@ def need_sorted_nlist_for_lower(self) -> bool: return False def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -181,7 +182,7 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "PairTabAtomicModel": + def deserialize(cls, data: dict) -> "PairTabAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 2) data.pop("@class") @@ -198,13 +199,13 @@ def deserialize(cls, data) -> "PairTabAtomicModel": def forward_atomic( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: xp = array_api_compat.array_namespace(extended_coord, extended_atype, nlist) nframes, nloc, nnei = nlist.shape extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) @@ -237,22 +238,22 @@ def forward_atomic( def _pair_tabulated_inter( self, - nlist: np.ndarray, - i_type: np.ndarray, - j_type: np.ndarray, - rr: np.ndarray, - ) -> np.ndarray: + nlist: Array, + i_type: Array, + j_type: Array, + rr: Array, + ) -> Array: """Pairwise tabulated energy. Parameters ---------- - nlist : np.ndarray + nlist : Array The unmasked neighbour list. (nframes, nloc) - i_type : np.ndarray + i_type : Array The integer representation of atom type for all local atoms for all frames. (nframes, nloc) - j_type : np.ndarray + j_type : Array The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) - rr : np.ndarray + rr : Array The salar distance vector between two atoms. (nframes, nloc, nnei) Returns @@ -310,12 +311,12 @@ def _pair_tabulated_inter( return ener @staticmethod - def _get_pairwise_dist(coords: np.ndarray, nlist: np.ndarray) -> np.ndarray: + def _get_pairwise_dist(coords: Array, nlist: Array) -> Array: """Get pairwise distance `dr`. Parameters ---------- - coords : np.ndarray + coords : Array The coordinate of the atoms, shape of (nframes, nall, 3). nlist The masked nlist, shape of (nframes, nloc, nnei). @@ -337,23 +338,23 @@ def _get_pairwise_dist(coords: np.ndarray, nlist: np.ndarray) -> np.ndarray: @staticmethod def _extract_spline_coefficient( - i_type: np.ndarray, - j_type: np.ndarray, - idx: np.ndarray, - tab_data: np.ndarray, + i_type: Array, + j_type: Array, + idx: Array, + tab_data: Array, nspline: np.int64, - ) -> np.ndarray: + ) -> Array: """Extract the spline coefficient from the table. Parameters ---------- - i_type : np.ndarray + i_type : Array The integer representation of atom type for all local atoms for all frames. (nframes, nloc) - j_type : np.ndarray + j_type : Array The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) - idx : np.ndarray + idx : Array The index of the spline coefficient. (nframes, nloc, nnei) - tab_data : np.ndarray + tab_data : Array The table storing all the spline coefficient. (ntype, ntype, nspline, 4) nspline : int The number of splines in the table. @@ -391,14 +392,14 @@ def _extract_spline_coefficient( return final_coef @staticmethod - def _calculate_ener(coef: np.ndarray, uu: np.ndarray) -> np.ndarray: + def _calculate_ener(coef: Array, uu: Array) -> Array: """Calculate energy using spline coeeficients. Parameters ---------- - coef : np.ndarray + coef : Array The spline coefficients. (nframes, nloc, nnei, 4) - uu : np.ndarray + uu : Array The atom displancemnt used in interpolation and extrapolation (nframes, nloc, nnei) Returns diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index bc7860491c..2180e48265 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -1,8 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import array_api_compat -import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.fitting.polarizability_fitting import ( PolarFitting, ) @@ -13,7 +18,9 @@ class DPPolarAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any + ) -> None: if not isinstance(fitting, PolarFitting): raise TypeError( "fitting must be an instance of PolarFitting for DPPolarAtomicModel" @@ -22,9 +29,9 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): def apply_out_stat( self, - ret: dict[str, np.ndarray], - atype: np.ndarray, - ): + ret: dict[str, Array], + atype: Array, + ) -> dict[str, Array]: """Apply the stat to each atomic output. Parameters diff --git a/deepmd/dpmodel/atomic_model/property_atomic_model.py b/deepmd/dpmodel/atomic_model/property_atomic_model.py index e3c038e695..ec65f949e0 100644 --- a/deepmd/dpmodel/atomic_model/property_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/property_atomic_model.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import numpy as np +from typing import ( + Any, +) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.fitting.property_fitting import ( PropertyFittingNet, ) @@ -11,7 +16,9 @@ class DPPropertyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: list[str], **kwargs: Any + ) -> None: if not isinstance(fitting, PropertyFittingNet): raise TypeError( "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel" @@ -20,9 +27,9 @@ def __init__(self, descriptor, fitting, type_map, **kwargs): def apply_out_stat( self, - ret: dict[str, np.ndarray], - atype: np.ndarray, - ): + ret: dict[str, Array], + atype: Array, + ) -> dict[str, Array]: """Apply the stat to each atomic output. In property fitting, each output will be multiplied by label std and then plus the label average value. diff --git a/deepmd/dpmodel/common.py b/deepmd/dpmodel/common.py index 1f9d4817a2..bd6f7dac49 100644 --- a/deepmd/dpmodel/common.py +++ b/deepmd/dpmodel/common.py @@ -3,12 +3,15 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from functools import ( wraps, ) from typing import ( + TYPE_CHECKING, Any, - Callable, Optional, overload, ) @@ -20,6 +23,10 @@ from deepmd.common import ( VALID_PRECISION, ) + +if TYPE_CHECKING: + from deepmd.dpmodel.array_api import Array + from deepmd.env import ( GLOBAL_ENER_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION, @@ -59,7 +66,7 @@ def get_xp_precision( xp: Any, precision: str, -): +) -> Any: """Get the precision from the API compatible namespace.""" if precision == "float16" or precision == "half": return xp.float16 @@ -87,16 +94,16 @@ class NativeOP(ABC): """The unit operation of a native model.""" @abstractmethod - def call(self, *args, **kwargs): + def call(self, *args: Any, **kwargs: Any) -> "Array": """Forward pass in NumPy implementation.""" pass - def __call__(self, *args, **kwargs): + def __call__(self, *args: Any, **kwargs: Any) -> "Array": """Forward pass in NumPy implementation.""" return self.call(*args, **kwargs) -def to_numpy_array(x: Any) -> Optional[np.ndarray]: +def to_numpy_array(x: Optional["Array"]) -> np.ndarray | None: """Convert an array to a NumPy array. Parameters @@ -158,7 +165,7 @@ def cast_precision(func: Callable[..., Any]) -> Callable[..., Any]: """ @wraps(func) - def wrapper(self, *args, **kwargs): + def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: # only convert tensors returned_tensor = func( self, @@ -185,13 +192,13 @@ def wrapper(self, *args, **kwargs): @overload def safe_cast_array( - input: np.ndarray, from_precision: str, to_precision: str -) -> np.ndarray: ... + input: "Array", from_precision: str, to_precision: str +) -> "Array": ... @overload def safe_cast_array(input: None, from_precision: str, to_precision: str) -> None: ... def safe_cast_array( - input: Optional[np.ndarray], from_precision: str, to_precision: str -) -> Optional[np.ndarray]: + input: Optional["Array"], from_precision: str, to_precision: str +) -> Optional["Array"]: """Convert an array from a precision to another precision. If input is not an array or without the specific precision, the method will not @@ -201,7 +208,7 @@ def safe_cast_array( Parameters ---------- - input : np.ndarray or None + input : Array or None Input array from_precision : str Array data type that is casted from diff --git a/deepmd/dpmodel/descriptor/descriptor.py b/deepmd/dpmodel/descriptor/descriptor.py index 443a2a66f1..9b0e067972 100644 --- a/deepmd/dpmodel/descriptor/descriptor.py +++ b/deepmd/dpmodel/descriptor/descriptor.py @@ -4,15 +4,19 @@ ABC, abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, NoReturn, - Optional, - Union, ) import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -34,7 +38,7 @@ class DescriptorBlock(ABC, make_plugin_registry("DescriptorBlock")): local_cluster = False - def __new__(cls, *args, **kwargs): + def __new__(cls, *args: Any, **kwargs: Any) -> Any: if cls is DescriptorBlock: try: descrpt_type = kwargs["type"] @@ -82,8 +86,8 @@ def get_dim_emb(self) -> int: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -107,7 +111,9 @@ def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" raise NotImplementedError - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: Any, shared_level: Any, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -118,13 +124,13 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: @abstractmethod def call( self, - nlist: np.ndarray, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - extended_atype_embd: Optional[np.ndarray] = None, - mapping: Optional[np.ndarray] = None, - type_embedding: Optional[np.ndarray] = None, - ): + nlist: Array, + extended_coord: Array, + extended_atype: Array, + extended_atype_embd: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, + ) -> Any: """Calculate DescriptorBlock.""" pass @@ -137,7 +143,9 @@ def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" -def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None: +def extend_descrpt_stat( + des: Any, type_map: list[str], des_with_stat: Any = None +) -> None: r""" Extend the statistics of a descriptor block with types from newly provided `type_map`. diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 51c56e9681..5228ba55b2 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import math -from typing import ( - Any, +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, Union, @@ -17,6 +18,7 @@ NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) from deepmd.dpmodel.common import ( @@ -74,7 +76,7 @@ ) -def np_softmax(x, axis=-1): +def np_softmax(x: Array, axis: int = -1) -> Array: xp = array_api_compat.array_namespace(x) # x = xp.nan_to_num(x) # to avoid value warning x = xp.where(xp.isnan(x), xp.zeros_like(x), x) @@ -82,7 +84,7 @@ def np_softmax(x, axis=-1): return e_x / xp.sum(e_x, axis=axis, keepdims=True) -def np_normalize(x, axis=-1): +def np_normalize(x: Array, axis: int = -1) -> Array: xp = array_api_compat.array_namespace(x) return x / xp.linalg.vector_norm(x, axis=axis, keepdims=True) @@ -244,7 +246,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -262,20 +264,20 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - scaling_factor=1.0, + scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, - spin: Optional[Any] = None, - stripped_type_embedding: Optional[bool] = None, + spin: None = None, + stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: ## seed, uniform_seed, not included. # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -319,6 +321,7 @@ def __init__( trainable_ln=trainable_ln, ln_eps=ln_eps, seed=child_seed(seed, 0), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd self.use_tebd_bias = use_tebd_bias @@ -333,6 +336,7 @@ def __init__( use_tebd_bias=use_tebd_bias, type_map=type_map, seed=child_seed(seed, 1), + trainable=trainable, ) self.tebd_dim = tebd_dim self.concat_output_tebd = concat_output_tebd @@ -397,7 +401,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_atten.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: "DescrptDPA1", shared_level: int, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -406,18 +412,18 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: raise NotImplementedError @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -438,19 +444,21 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: np.ndarray, - stddev: np.ndarray, + mean: Array, + stddev: Array, ) -> None: """Update mean and stddev for descriptor.""" self.se_atten.mean = mean self.se_atten.stddev = stddev - def get_stat_mean_and_stddev(self) -> tuple[np.ndarray, np.ndarray]: + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: """Get mean and stddev for descriptor.""" return self.se_atten.mean, self.se_atten.stddev def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["DescrptDPA1"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -479,11 +487,11 @@ def change_type_map( @cast_precision def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> Array: """Compute the descriptor. Parameters @@ -520,7 +528,7 @@ def call( type_embedding = self.type_embedding.call() # nf x nall x tebd_dim atype_embd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0), + xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), (nf, nall, self.tebd_dim), ) # nfnl x tebd_dim @@ -632,9 +640,9 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters @@ -667,7 +675,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -684,13 +692,14 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - scaling_factor=1.0, + scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: self.rcut = rcut self.rcut_smth = rcut_smth @@ -741,6 +750,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.embeddings = embeddings if self.tebd_input_mode in ["strip"]: @@ -756,6 +766,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.embeddings_strip = embeddings_strip else: @@ -774,6 +785,7 @@ def __init__( smooth=self.smooth, precision=self.precision, seed=child_seed(seed, 2), + trainable=trainable, ) wanted_shape = (self.ntypes, self.nnei, 4) @@ -814,7 +826,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -822,7 +834,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -847,24 +859,24 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] * self.axis_neuron @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.tebd_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -918,9 +930,9 @@ def reinit_exclude( def cal_g( self, - ss, - embedding_idx, - ): + ss: Array, + embedding_idx: int, + ) -> Array: xp = array_api_compat.array_namespace(ss) nfnl, nnei = ss.shape[0:2] shape2 = math.prod(ss.shape[2:]) @@ -931,9 +943,9 @@ def cal_g( def cal_g_strip( self, - ss, - embedding_idx, - ): + ss: Array, + embedding_idx: int, + ) -> Array: assert self.embeddings_strip is not None # nfnl x nnei x ng gg = self.embeddings_strip[embedding_idx].call(ss) @@ -941,13 +953,13 @@ def cal_g_strip( def call( self, - nlist: np.ndarray, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - atype_embd_ext: Optional[np.ndarray] = None, - mapping: Optional[np.ndarray] = None, - type_embedding: Optional[np.ndarray] = None, - ): + nlist: Array, + coord_ext: Array, + atype_ext: Array, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, + ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) # nf x nloc x nnei x 4 dmatrix, diff, sw = self.env_mat.call( @@ -1027,7 +1039,7 @@ def call( xp.tile( (xp.reshape(atype, (-1, 1)) * ntypes_with_padding), (1, nnei) ), - (-1), + (-1,), ) idx_j = xp.reshape(nei_type, (-1,)) # (nf x nl x nnei) x ng @@ -1180,12 +1192,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a neighbor-wise attention net.""" super().__init__() @@ -1219,29 +1232,32 @@ def __init__( smooth=smooth, precision=precision, seed=child_seed(seed, ii), + trainable=trainable, ) for ii in range(layer_num) ] def call( self, - input_G, - nei_mask, - input_r: Optional[np.ndarray] = None, - sw: Optional[np.ndarray] = None, - ): + input_G: Array, + nei_mask: Array, + input_r: Array | None = None, + sw: Array | None = None, + ) -> Array: out = input_G for layer in self.attention_layers: out = layer(out, nei_mask, input_r=input_r, sw=sw) return out - def __getitem__(self, key): + def __getitem__(self, key: int) -> "NeighborGatedAttentionLayer": if isinstance(key, int): return self.attention_layers[key] else: raise TypeError(key) - def __setitem__(self, key, value) -> None: + def __setitem__( + self, key: int, value: Union["NeighborGatedAttentionLayer", dict] + ) -> None: if not isinstance(key, int): raise TypeError(key) if isinstance(value, self.network_type): @@ -1252,7 +1268,7 @@ def __setitem__(self, key, value) -> None: raise TypeError(value) self.attention_layers[key] = value - def serialize(self): + def serialize(self) -> dict: """Serialize the networks to a dict. Returns @@ -1308,12 +1324,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a neighbor-wise attention layer.""" super().__init__() @@ -1340,6 +1357,7 @@ def __init__( smooth=smooth, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.attn_layer_norm = LayerNorm( self.embed_dim, @@ -1351,11 +1369,11 @@ def __init__( def call( self, - x, - nei_mask, - input_r: Optional[np.ndarray] = None, - sw: Optional[np.ndarray] = None, - ): + x: Array, + nei_mask: Array, + input_r: Array | None = None, + sw: Array | None = None, + ) -> Array: residual = x x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) x = residual + x @@ -1387,7 +1405,7 @@ def serialize(self) -> dict: } @classmethod - def deserialize(cls, data) -> "NeighborGatedAttentionLayer": + def deserialize(cls, data: dict) -> "NeighborGatedAttentionLayer": """Deserialize the networks from a dict. Parameters @@ -1415,11 +1433,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a multi-head neighbor-wise attention net.""" super().__init__() @@ -1449,6 +1468,7 @@ def __init__( use_timestep=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.out_proj = NativeLayer( hidden_dim, @@ -1457,9 +1477,17 @@ def __init__( use_timestep=False, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) - def call(self, query, nei_mask, input_r=None, sw=None, attnw_shift=20.0): + def call( + self, + query: Array, + nei_mask: Array, + input_r: Array | None = None, + sw: Array | None = None, + attnw_shift: float = 20.0, + ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(query, nei_mask) # Linear projection # q, k, v = xp.split(self.in_proj(query), 3, axis=-1) @@ -1520,7 +1548,7 @@ def call(self, query, nei_mask, input_r=None, sw=None, attnw_shift=20.0): output = self.out_proj(o) return output, attn_weights - def serialize(self): + def serialize(self) -> dict: return { "nnei": self.nnei, "embed_dim": self.embed_dim, @@ -1539,7 +1567,7 @@ def serialize(self): } @classmethod - def deserialize(cls, data): + def deserialize(cls, data: dict) -> "GatedAttentionLayer": data = data.copy() in_proj = data.pop("in_proj") out_proj = data.pop("out_proj") diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 70accefa30..e5e02d312c 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -1,18 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, NoReturn, - Optional, - Union, ) import array_api_compat -import numpy as np from deepmd.dpmodel import ( NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) from deepmd.dpmodel.common import ( @@ -83,7 +84,7 @@ def __init__( tebd_dim: int = 8, tebd_input_mode: str = "concat", set_davg_zero: bool = True, - activation_function="tanh", + activation_function: str = "tanh", resnet_dt: bool = False, type_one_side: bool = False, use_three_body: bool = False, @@ -151,7 +152,7 @@ def __init__( self.three_body_rcut = three_body_rcut self.three_body_rcut_smth = three_body_rcut_smth - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if hasattr(self, key): return getattr(self, key) else: @@ -214,7 +215,7 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, ) -> None: r"""The constructor for the RepformerArgs class which defines the parameters of the repformer block in DPA2 descriptor. @@ -321,7 +322,7 @@ def __init__( ln_eps = 1e-5 self.ln_eps = ln_eps - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if hasattr(self, key): return getattr(self, key) else: @@ -372,9 +373,9 @@ def __init__( self, ntypes: int, # args for repinit - repinit: Union[RepinitArgs, dict], + repinit: RepinitArgs | dict, # args for repformer - repformer: Union[RepformerArgs, dict], + repformer: RepformerArgs | dict, # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -382,11 +383,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: r"""The DPA-2 descriptor[1]_. @@ -442,7 +443,7 @@ def __init__( Comput Mater 10, 293 (2024). https://doi.org/10.1038/s41524-024-01493-2 """ - def init_subclass_params(sub_data, sub_class): + def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: if isinstance(sub_data, dict): return sub_class(**sub_data) elif isinstance(sub_data, sub_class): @@ -474,6 +475,7 @@ def init_subclass_params(sub_data, sub_class): smooth=smooth, type_one_side=self.repinit_args.type_one_side, seed=child_seed(seed, 0), + trainable=trainable, ) self.use_three_body = self.repinit_args.use_three_body if self.use_three_body: @@ -493,6 +495,7 @@ def init_subclass_params(sub_data, sub_class): resnet_dt=self.repinit_args.resnet_dt, smooth=smooth, seed=child_seed(seed, 5), + trainable=trainable, ) else: self.repinit_three_body = None @@ -533,6 +536,7 @@ def init_subclass_params(sub_data, sub_class): g1_out_mlp=self.repformer_args.g1_out_mlp, ln_eps=self.repformer_args.ln_eps, seed=child_seed(seed, 1), + trainable=trainable, ) self.rcsl_list = [ (self.repformers.get_rcut(), self.repformers.get_nsel()), @@ -562,6 +566,7 @@ def init_subclass_params(sub_data, sub_class): use_tebd_bias=use_tebd_bias, type_map=type_map, seed=child_seed(seed, 2), + trainable=trainable, ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -585,6 +590,7 @@ def init_subclass_params(sub_data, sub_class): bias=False, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) self.tebd_transform = None if self.add_tebd_to_repinit_out: @@ -594,6 +600,7 @@ def init_subclass_params(sub_data, sub_class): bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) assert self.repinit.rcut > self.repformers.rcut assert self.repinit.sel[0] > self.repformers.sel[0] @@ -665,7 +672,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -674,7 +683,7 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -729,19 +738,19 @@ def change_type_map( repinit_three_body["dstd"] = repinit_three_body["dstd"][remap_index] @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -766,8 +775,8 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: list[np.ndarray], - stddev: list[np.ndarray], + mean: list[Array], + stddev: list[Array], ) -> None: """Update mean and stddev for descriptor.""" descrpt_list = [self.repinit, self.repformers] @@ -777,7 +786,9 @@ def set_stat_mean_and_stddev( descrpt.mean = mean[ii] descrpt.stddev = stddev[ii] - def get_stat_mean_and_stddev(self) -> tuple[list[np.ndarray], list[np.ndarray]]: + def get_stat_mean_and_stddev( + self, + ) -> tuple[list[Array], list[Array]]: """Get mean and stddev for descriptor.""" mean_list = [self.repinit.mean, self.repformers.mean] stddev_list = [ @@ -792,11 +803,11 @@ def get_stat_mean_and_stddev(self) -> tuple[list[np.ndarray], list[np.ndarray]]: @cast_precision def call( self, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> tuple[Array, Array]: """Compute the descriptor. Parameters @@ -841,7 +852,7 @@ def call( type_embedding = self.type_embedding.call() # repinit g1_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0), + xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), (nframes, nall, self.tebd_dim), ) g1_inp = g1_ext[:, :nloc, :] @@ -1062,9 +1073,9 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/dpa3.py b/deepmd/dpmodel/descriptor/dpa3.py index 25550be926..47a4fb1478 100644 --- a/deepmd/dpmodel/descriptor/dpa3.py +++ b/deepmd/dpmodel/descriptor/dpa3.py @@ -1,15 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, ) import array_api_compat -import numpy as np from deepmd.dpmodel import ( NativeOP, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, to_numpy_array, @@ -208,7 +209,7 @@ def __init__( self.use_dynamic_sel = use_dynamic_sel self.sel_reduce_factor = sel_reduce_factor - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if hasattr(self, key): return getattr(self, key) else: @@ -294,7 +295,7 @@ def __init__( self, ntypes: int, # args for repflow - repflow: Union[RepFlowArgs, dict], + repflow: RepFlowArgs | dict, # kwargs for descriptor concat_output_tebd: bool = False, activation_function: str = "silu", @@ -302,15 +303,15 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__() - def init_subclass_params(sub_data, sub_class): + def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any: if isinstance(sub_data, dict): return sub_class(**sub_data) elif isinstance(sub_data, sub_class): @@ -357,6 +358,7 @@ def init_subclass_params(sub_data, sub_class): env_protection=env_protection, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd @@ -374,6 +376,7 @@ def init_subclass_params(sub_data, sub_class): use_tebd_bias=use_tebd_bias, type_map=type_map, seed=child_seed(seed, 2), + trainable=trainable, ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -448,7 +451,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.repflows.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -457,7 +462,7 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -486,15 +491,17 @@ def change_type_map( repflow["dstd"] = repflow["dstd"][remap_index] @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.get_dim_emb() - def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None): + def compute_input_stats( + self, merged: list[dict], path: DPPath | None = None + ) -> None: """Update mean and stddev for descriptor elements.""" descrpt_list = [self.repflows] for ii, descrpt in enumerate(descrpt_list): @@ -502,8 +509,8 @@ def compute_input_stats(self, merged: list[dict], path: Optional[DPPath] = None) def set_stat_mean_and_stddev( self, - mean: list[np.ndarray], - stddev: list[np.ndarray], + mean: list[Array], + stddev: list[Array], ) -> None: """Update mean and stddev for descriptor.""" descrpt_list = [self.repflows] @@ -511,7 +518,7 @@ def set_stat_mean_and_stddev( descrpt.mean = mean[ii] descrpt.stddev = stddev[ii] - def get_stat_mean_and_stddev(self) -> tuple[list[np.ndarray], list[np.ndarray]]: + def get_stat_mean_and_stddev(self) -> tuple[list[Array], list[Array]]: """Get mean and stddev for descriptor.""" mean_list = [self.repflows.mean] stddev_list = [self.repflows.stddev] @@ -520,11 +527,11 @@ def get_stat_mean_and_stddev(self) -> tuple[list[np.ndarray], list[np.ndarray]]: @cast_precision def call( self, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> tuple[Array, Array]: """Compute the descriptor. Parameters @@ -562,12 +569,12 @@ def call( type_embedding = self.type_embedding.call() if self.use_loc_mapping: node_ebd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], [-1]), axis=0), + xp.take(type_embedding, xp.reshape(atype_ext[:, :nloc], (-1,)), axis=0), (nframes, nloc, self.tebd_dim), ) else: node_ebd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0), + xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), (nframes, nall, self.tebd_dim), ) node_ebd_inp = node_ebd_ext[:, :nloc, :] @@ -654,9 +661,9 @@ def deserialize(cls, data: dict) -> "DescrptDPA3": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index f050bb6222..3682e17c4d 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -3,13 +3,14 @@ from typing import ( Any, NoReturn, - Optional, - Union, ) import array_api_compat import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( NativeOP, ) @@ -43,9 +44,9 @@ class DescrptHybrid(BaseDescriptor, NativeOP): def __init__( self, - list: list[Union[BaseDescriptor, dict[str, Any]]], - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + list: list[BaseDescriptor | dict[str, Any]], + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input ) -> None: super().__init__() # warning: list is conflict with built-in list @@ -76,7 +77,7 @@ def __init__( ) # if hybrid sel is larger than sub sel, the nlist needs to be cut for each type hybrid_sel = self.get_sel() - nlist_cut_idx: list[np.ndarray] = [] + nlist_cut_idx: list[Array] = [] if self.mixed_types() and not all( descrpt.mixed_types() for descrpt in self.descrpt_list ): @@ -144,7 +145,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension.""" return np.sum([descrpt.get_dim_emb() for descrpt in self.descrpt_list]).item() - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -168,7 +169,9 @@ def get_env_protection(self) -> float: ) return all_protection[0] - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: Any, shared_level: Any, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -177,7 +180,7 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -191,7 +194,7 @@ def change_type_map( ) def compute_input_stats( - self, merged: list[dict], path: Optional[DPPath] = None + self, merged: list[dict], path: DPPath | None = None ) -> None: """Update mean and stddev for descriptor elements.""" for descrpt in self.descrpt_list: @@ -199,8 +202,8 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: list[Union[np.ndarray, list[np.ndarray]]], - stddev: list[Union[np.ndarray, list[np.ndarray]]], + mean: list[np.ndarray | list[Array]], + stddev: list[np.ndarray | list[Array]], ) -> None: """Update mean and stddev for descriptor.""" for ii, descrpt in enumerate(self.descrpt_list): @@ -209,8 +212,8 @@ def set_stat_mean_and_stddev( def get_stat_mean_and_stddev( self, ) -> tuple[ - list[Union[np.ndarray, list[np.ndarray]]], - list[Union[np.ndarray, list[np.ndarray]]], + list[Array | list[Array]], + list[Array | list[Array]], ]: """Get mean and stddev for descriptor.""" mean_list = [] @@ -255,11 +258,17 @@ def enable_compression( def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> tuple[ + Array, + Array | None, + Array | None, + Array | None, + Array | None, + ]: """Compute the descriptor. Parameters @@ -322,9 +331,9 @@ def call( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index f45e85e516..f87ca2c5b6 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -3,16 +3,20 @@ ABC, abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, NoReturn, - Optional, - Union, ) from deepmd.common import ( j_get_type, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -26,9 +30,9 @@ def make_base_descriptor( - t_tensor, + t_tensor: type, fwd_method_name: str = "forward", -): +) -> type: """Make the base class for the descriptor. Parameters @@ -44,7 +48,7 @@ def make_base_descriptor( class BD(ABC, PluginVariant, make_plugin_registry("descriptor")): """Base descriptor provides the interfaces of descriptor.""" - def __new__(cls, *args, **kwargs): + def __new__(cls, *args: Any, **kwargs: Any) -> Any: if cls is BD: cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__)) return super().__new__(cls) @@ -113,7 +117,9 @@ def get_env_protection(self) -> float: pass @abstractmethod - def share_params(self, base_class, shared_level, resume=False): + def share_params( + self, base_class: Any, shared_level: Any, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -123,7 +129,7 @@ def share_params(self, base_class, shared_level, resume=False): @abstractmethod def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -131,19 +137,19 @@ def change_type_map( pass @abstractmethod - def set_stat_mean_and_stddev(self, mean, stddev) -> None: + def set_stat_mean_and_stddev(self, mean: Any, stddev: Any) -> None: """Update mean and stddev for descriptor.""" pass @abstractmethod - def get_stat_mean_and_stddev(self): + def get_stat_mean_and_stddev(self) -> Any: """Get mean and stddev for descriptor.""" pass def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """Update mean and stddev for descriptor elements.""" raise NotImplementedError @@ -176,11 +182,11 @@ def enable_compression( @abstractmethod def fwd( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[t_tensor] = None, - ): + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + ) -> Array: """Calculate descriptor.""" pass @@ -212,9 +218,9 @@ def deserialize(cls, data: dict) -> "BD": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 0a39e4c596..706fc690e4 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import array_api_compat @@ -13,6 +11,7 @@ NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) from deepmd.dpmodel.common import ( @@ -167,15 +166,17 @@ class DescrptBlockRepflows(NativeOP, DescriptorBlock): For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection. seed : int, optional Random seed for parameter initialization. + trainable : bool, default: True + Whether the block is trainable """ def __init__( self, - e_rcut, - e_rcut_smth, + e_rcut: float, + e_rcut_smth: float, e_sel: int, - a_rcut, - a_rcut_smth, + a_rcut: float, + a_rcut_smth: float, a_sel: int, ntypes: int, nlayers: int = 6, @@ -204,7 +205,8 @@ def __init__( use_dynamic_sel: bool = False, sel_reduce_factor: float = 10.0, use_loc_mapping: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.e_rcut = float(e_rcut) @@ -269,10 +271,19 @@ def __init__( self.seed = seed self.edge_embd = NativeLayer( - 1, self.e_dim, precision=precision, seed=child_seed(seed, 0) + 1, + self.e_dim, + precision=precision, + seed=child_seed(seed, 0), + trainable=trainable, ) self.angle_embd = NativeLayer( - 1, self.a_dim, precision=precision, bias=False, seed=child_seed(seed, 1) + 1, + self.a_dim, + precision=precision, + bias=False, + seed=child_seed(seed, 1), + trainable=trainable, ) layers = [] for ii in range(nlayers): @@ -304,6 +315,7 @@ def __init__( sel_reduce_factor=self.sel_reduce_factor, smooth_edge_update=self.smooth_edge_update, seed=child_seed(child_seed(seed, 1), ii), + trainable=trainable, ) ) self.layers = layers @@ -358,7 +370,7 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension e_dim.""" return self.e_dim - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -366,7 +378,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -391,24 +403,24 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.n_dim @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.n_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension e_dim.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -462,12 +474,12 @@ def reinit_exclude( def call( self, - nlist: np.ndarray, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - atype_embd_ext: Optional[np.ndarray] = None, - mapping: Optional[np.ndarray] = None, - ): + nlist: Array, + coord_ext: Array, + atype_ext: Array, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) nframes, nloc, nnei = nlist.shape nall = xp.reshape(coord_ext, (nframes, -1)).shape[1] // 3 @@ -490,7 +502,7 @@ def call( sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) # get angle nlist (maybe smaller) - a_dist_mask = (xp.linalg.vector_norm(diff, axis=-1) < self.a_rcut)[ + a_dist_mask = (safe_for_vector_norm(diff, axis=-1) < self.a_rcut)[ :, :, : self.a_sel ] a_nlist = nlist[:, :, : self.a_sel] @@ -578,7 +590,8 @@ def call( # n_angle x 1 a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask] else: - edge_index = angle_index = xp.zeros([1, 3], dtype=nlist.dtype) + edge_index = xp.zeros([2, 1], dtype=nlist.dtype) + angle_index = xp.zeros([3, 1], dtype=nlist.dtype) # get edge and angle embedding # nb x nloc x nnei x e_dim [OR] n_edge x e_dim @@ -622,7 +635,7 @@ def call( edge_ebd, h2, sw, - owner=edge_index[:, 0], + owner=edge_index[0, :], num_owner=nframes * nloc, nb=nframes, nloc=nloc, @@ -649,7 +662,7 @@ def need_sorted_nlist_for_lower(self) -> bool: return True @classmethod - def deserialize(cls, data): + def deserialize(cls, data: dict) -> "DescrptBlockRepflows": """Deserialize the descriptor block.""" data = data.copy() edge_embd = NativeLayer.deserialize(data.pop("edge_embd")) @@ -670,7 +683,7 @@ def deserialize(cls, data): obj.stddev = dstd return obj - def serialize(self): + def serialize(self) -> dict: """Serialize the descriptor block.""" return { "e_rcut": self.e_rcut, @@ -720,15 +733,15 @@ def serialize(self): def _cal_hg_dynamic( - flat_edge_ebd: np.ndarray, - flat_h2: np.ndarray, - flat_sw: np.ndarray, - owner: np.ndarray, + flat_edge_ebd: Array, + flat_h2: Array, + flat_sw: Array, + owner: Array, num_owner: int, nb: int, nloc: int, scale_factor: float, -) -> np.ndarray: +) -> Array: """ Calculate the transposed rotation matrix. @@ -775,16 +788,16 @@ def _cal_hg_dynamic( def symmetrization_op_dynamic( - flat_edge_ebd: np.ndarray, - flat_h2: np.ndarray, - flat_sw: np.ndarray, - owner: np.ndarray, + flat_edge_ebd: Array, + flat_h2: Array, + flat_sw: Array, + owner: Array, num_owner: int, nb: int, nloc: int, scale_factor: float, axis_neuron: int, -) -> np.ndarray: +) -> Array: """ Symmetrization operator to obtain atomic invariant rep. @@ -859,7 +872,8 @@ def __init__( update_residual: float = 0.1, update_residual_init: str = "const", precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei @@ -922,6 +936,7 @@ def __init__( n_dim, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) if self.update_style == "res_residual": self.n_residual.append( @@ -931,6 +946,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) ) @@ -941,6 +957,7 @@ def __init__( n_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) if self.update_style == "res_residual": self.n_residual.append( @@ -950,6 +967,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) ) @@ -959,6 +977,7 @@ def __init__( self.n_multi_edge_message * n_dim, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) if self.update_style == "res_residual": for head_index in range(self.n_multi_edge_message): @@ -969,6 +988,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(child_seed(seed, 5), head_index), + trainable=trainable, ) ) @@ -978,6 +998,7 @@ def __init__( e_dim, precision=precision, seed=child_seed(seed, 6), + trainable=trainable, ) if self.update_style == "res_residual": self.e_residual.append( @@ -987,6 +1008,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 7), + trainable=trainable, ) ) @@ -1015,6 +1037,7 @@ def __init__( precision=precision, bias=False, seed=child_seed(seed, 8), + trainable=trainable, ) self.a_compress_e_linear = NativeLayer( self.e_dim, @@ -1022,6 +1045,7 @@ def __init__( precision=precision, bias=False, seed=child_seed(seed, 9), + trainable=trainable, ) else: self.a_compress_n_linear = None @@ -1033,12 +1057,14 @@ def __init__( self.e_dim, precision=precision, seed=child_seed(seed, 10), + trainable=trainable, ) self.edge_angle_linear2 = NativeLayer( self.e_dim, self.e_dim, precision=precision, seed=child_seed(seed, 11), + trainable=trainable, ) if self.update_style == "res_residual": self.e_residual.append( @@ -1048,6 +1074,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 12), + trainable=trainable, ) ) @@ -1057,6 +1084,7 @@ def __init__( self.a_dim, precision=precision, seed=child_seed(seed, 13), + trainable=trainable, ) if self.update_style == "res_residual": self.a_residual.append( @@ -1066,6 +1094,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 14), + trainable=trainable, ) ) else: @@ -1078,11 +1107,11 @@ def __init__( def optim_angle_update( self, - angle_ebd: np.ndarray, - node_ebd: np.ndarray, - edge_ebd: np.ndarray, + angle_ebd: Array, + node_ebd: Array, + edge_ebd: Array, feat: str = "edge", - ) -> np.ndarray: + ) -> Array: xp = array_api_compat.array_namespace(angle_ebd, node_ebd, edge_ebd) if feat == "edge": @@ -1126,14 +1155,14 @@ def optim_angle_update( def optim_angle_update_dynamic( self, - flat_angle_ebd: np.ndarray, - node_ebd: np.ndarray, - flat_edge_ebd: np.ndarray, - n2a_index: np.ndarray, - eij2a_index: np.ndarray, - eik2a_index: np.ndarray, - feat="edge", - ): + flat_angle_ebd: Array, + node_ebd: Array, + flat_edge_ebd: Array, + n2a_index: Array, + eij2a_index: Array, + eik2a_index: Array, + feat: str = "edge", + ) -> Array: xp = array_api_compat.array_namespace( flat_angle_ebd, node_ebd, flat_edge_ebd, n2a_index, eij2a_index, eik2a_index ) @@ -1185,12 +1214,12 @@ def optim_angle_update_dynamic( def optim_edge_update( self, - node_ebd: np.ndarray, - node_ebd_ext: np.ndarray, - edge_ebd: np.ndarray, - nlist: np.ndarray, + node_ebd: Array, + node_ebd_ext: Array, + edge_ebd: Array, + nlist: Array, feat: str = "node", - ) -> np.ndarray: + ) -> Array: xp = array_api_compat.array_namespace(node_ebd, node_ebd_ext, edge_ebd, nlist) if feat == "node": @@ -1228,13 +1257,13 @@ def optim_edge_update( def optim_edge_update_dynamic( self, - node_ebd: np.ndarray, - node_ebd_ext: np.ndarray, - flat_edge_ebd: np.ndarray, - n2e_index: np.ndarray, - n_ext2e_index: np.ndarray, + node_ebd: Array, + node_ebd_ext: Array, + flat_edge_ebd: Array, + n2e_index: Array, + n_ext2e_index: Array, feat: str = "node", - ): + ) -> Array: xp = array_api_compat.array_namespace( node_ebd, node_ebd_ext, flat_edge_ebd, n2e_index, n_ext2e_index ) @@ -1276,19 +1305,19 @@ def optim_edge_update_dynamic( def call( self, - node_ebd_ext: np.ndarray, # nf x nall x n_dim - edge_ebd: np.ndarray, # nf x nloc x nnei x e_dim - h2: np.ndarray, # nf x nloc x nnei x 3 - angle_ebd: np.ndarray, # nf x nloc x a_nnei x a_nnei x a_dim - nlist: np.ndarray, # nf x nloc x nnei - nlist_mask: np.ndarray, # nf x nloc x nnei - sw: np.ndarray, # switch func, nf x nloc x nnei - a_nlist: np.ndarray, # nf x nloc x a_nnei - a_nlist_mask: np.ndarray, # nf x nloc x a_nnei - a_sw: np.ndarray, # switch func, nf x nloc x a_nnei - edge_index: np.ndarray, # n_edge x 2 - angle_index: np.ndarray, # n_angle x 3 - ): + node_ebd_ext: Array, # nf x nall x n_dim + edge_ebd: Array, # nf x nloc x nnei x e_dim + h2: Array, # nf x nloc x nnei x 3 + angle_ebd: Array, # nf x nloc x a_nnei x a_nnei x a_dim + nlist: Array, # nf x nloc x nnei + nlist_mask: Array, # nf x nloc x nnei + sw: Array, # switch func, nf x nloc x nnei + a_nlist: Array, # nf x nloc x a_nnei + a_nlist_mask: Array, # nf x nloc x a_nnei + a_sw: Array, # switch func, nf x nloc x a_nnei + edge_index: Array, # 2 x n_edge + angle_index: Array, # 3 x n_angle + ) -> tuple[Array, Array]: """ Parameters ---------- @@ -1312,12 +1341,12 @@ def call( Masks of the neighbor list for angle. real nei 1 otherwise 0 a_sw : nf x nloc x a_nnei Switch function for angle. - edge_index : Optional for dynamic sel, n_edge x 2 + edge_index : Optional for dynamic sel, 2 x n_edge n2e_index : n_edge Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i). n_ext2e_index : n_edge Broadcast indices from extended node(j) to edge(ij). - angle_index : Optional for dynamic sel, n_angle x 3 + angle_index : Optional for dynamic sel, 3 x n_angle n2a_index : n_angle Broadcast indices from extended node(j) to angle(ijk). eij2a_index : n_angle @@ -1362,11 +1391,11 @@ def call( assert (n_edge, 3) == h2.shape del a_nlist # may be used in the future - n2e_index, n_ext2e_index = edge_index[:, 0], edge_index[:, 1] + n2e_index, n_ext2e_index = edge_index[0, :], edge_index[1, :] n2a_index, eij2a_index, eik2a_index = ( - angle_index[:, 0], - angle_index[:, 1], - angle_index[:, 2], + angle_index[0, :], + angle_index[1, :], + angle_index[2, :], ) # nb x nloc x nnei x n_dim [OR] n_edge x n_dim @@ -1378,16 +1407,16 @@ def call( ) ) - n_update_list: list[np.ndarray] = [node_ebd] - e_update_list: list[np.ndarray] = [edge_ebd] - a_update_list: list[np.ndarray] = [angle_ebd] + n_update_list: list[Array] = [node_ebd] + e_update_list: list[Array] = [edge_ebd] + a_update_list: list[Array] = [angle_ebd] # node self mlp node_self_mlp = self.act(self.node_self_mlp(node_ebd)) n_update_list.append(node_self_mlp) # node sym (grrg + drrd) - node_sym_list: list[np.ndarray] = [] + node_sym_list: list[Array] = [] node_sym_list.append( symmetrization_op( edge_ebd, @@ -1757,15 +1786,15 @@ def call( def list_update_res_avg( self, - update_list: list[np.ndarray], - ) -> np.ndarray: + update_list: list[Array], + ) -> Array: nitem = len(update_list) uu = update_list[0] for ii in range(1, nitem): uu = uu + update_list[ii] return uu / (float(nitem) ** 0.5) - def list_update_res_incr(self, update_list: list[np.ndarray]) -> np.ndarray: + def list_update_res_incr(self, update_list: list[Array]) -> Array: nitem = len(update_list) uu = update_list[0] scale = 1.0 / (float(nitem - 1) ** 0.5) if nitem > 1 else 0.0 @@ -1774,8 +1803,8 @@ def list_update_res_incr(self, update_list: list[np.ndarray]) -> np.ndarray: return uu def list_update_res_residual( - self, update_list: list[np.ndarray], update_name: str = "node" - ) -> np.ndarray: + self, update_list: list[Array], update_name: str = "node" + ) -> Array: nitem = len(update_list) uu = update_list[0] if update_name == "node": @@ -1791,9 +1820,7 @@ def list_update_res_residual( raise NotImplementedError return uu - def list_update( - self, update_list: list[np.ndarray], update_name: str = "node" - ) -> np.ndarray: + def list_update(self, update_list: list[Array], update_name: str = "node") -> Array: if self.update_style == "res_avg": return self.list_update_res_avg(update_list) elif self.update_style == "res_incr": diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index 3d02054350..79d4f9228f 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import array_api_compat @@ -13,6 +14,7 @@ NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) from deepmd.dpmodel.common import ( @@ -54,7 +56,7 @@ ) -def xp_transpose_01423(x): +def xp_transpose_01423(x: Array) -> Array: xp = array_api_compat.array_namespace(x) x_shape2 = x.shape[2] x_shape3 = x.shape[3] @@ -65,7 +67,7 @@ def xp_transpose_01423(x): return x -def xp_transpose_01342(x): +def xp_transpose_01342(x: Array) -> Array: xp = array_api_compat.array_namespace(x) x_shape2 = x.shape[2] x_shape3 = x.shape[3] @@ -164,17 +166,19 @@ class DescrptBlockRepformers(NativeOP, DescriptorBlock): The epsilon value for layer normalization. seed : int, optional The random seed for initialization. + trainable : bool, default: True + Whether the block is trainable """ def __init__( self, - rcut, - rcut_smth, + rcut: float, + rcut_smth: float, sel: int, ntypes: int, nlayers: int = 3, - g1_dim=128, - g2_dim=16, + g1_dim: int = 128, + g2_dim: int = 16, axis_neuron: int = 4, direct_dist: bool = False, update_g1_has_conv: bool = True, @@ -202,8 +206,9 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.rcut = rcut @@ -252,7 +257,11 @@ def __init__( self.epsilon = 1e-4 self.g2_embd = NativeLayer( - 1, self.g2_dim, precision=precision, seed=child_seed(seed, 0) + 1, + self.g2_dim, + precision=precision, + seed=child_seed(seed, 0), + trainable=trainable, ) layers = [] for ii in range(nlayers): @@ -290,6 +299,7 @@ def __init__( g1_out_conv=self.g1_out_conv, g1_out_mlp=self.g1_out_mlp, seed=child_seed(child_seed(seed, 1), ii), + trainable=trainable, ) ) self.layers = layers @@ -328,7 +338,7 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.g2_dim - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -336,7 +346,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -357,24 +367,24 @@ def mixed_types(self) -> bool: return True @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.g1_dim @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.g1_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -428,13 +438,13 @@ def reinit_exclude( def call( self, - nlist: np.ndarray, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - atype_embd_ext: Optional[np.ndarray] = None, - mapping: Optional[np.ndarray] = None, - type_embedding: Optional[np.ndarray] = None, - ): + nlist: Array, + coord_ext: Array, + atype_ext: Array, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, + ) -> Array: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext) exclude_mask = xp.astype(exclude_mask, xp.bool) @@ -509,7 +519,7 @@ def need_sorted_nlist_for_lower(self) -> bool: return False @classmethod - def deserialize(cls, data): + def deserialize(cls, data: dict[str, Any]) -> "DescrptBlockRepformers": """Deserialize the descriptor block.""" data = data.copy() g2_embd = NativeLayer.deserialize(data.pop("g2_embd")) @@ -526,7 +536,7 @@ def deserialize(cls, data): obj.stddev = dstd return obj - def serialize(self): + def serialize(self) -> dict[str, Any]: """Serialize the descriptor block.""" return { "rcut": self.rcut, @@ -582,8 +592,8 @@ def get_residual( _mode: str = "norm", trainable: bool = True, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, -) -> np.ndarray: + seed: int | list[int] | None = None, +) -> Array: """ Get residual tensor for one update vector. @@ -617,9 +627,9 @@ def get_residual( def _make_nei_g1( - g1_ext: np.ndarray, - nlist: np.ndarray, -) -> np.ndarray: + g1_ext: Array, + nlist: Array, +) -> Array: """ Make neighbor-wise atomic invariant rep. @@ -632,7 +642,7 @@ def _make_nei_g1( Returns ------- - gg1: np.ndarray + gg1: Array Neighbor-wise atomic invariant rep, with shape [nf, nloc, nnei, ng1]. """ xp = array_api_compat.array_namespace(g1_ext, nlist) @@ -650,9 +660,9 @@ def _make_nei_g1( def _apply_nlist_mask( - gg: np.ndarray, - nlist_mask: np.ndarray, -) -> np.ndarray: + gg: Array, + nlist_mask: Array, +) -> Array: """ Apply nlist mask to neighbor-wise rep tensors. @@ -668,7 +678,7 @@ def _apply_nlist_mask( return masked_gg -def _apply_switch(gg: np.ndarray, sw: np.ndarray) -> np.ndarray: +def _apply_switch(gg: Array, sw: Array) -> Array: """ Apply switch function to neighbor-wise rep tensors. @@ -686,14 +696,14 @@ def _apply_switch(gg: np.ndarray, sw: np.ndarray) -> np.ndarray: def _cal_hg( - g: np.ndarray, - h: np.ndarray, - nlist_mask: np.ndarray, - sw: np.ndarray, + g: Array, + h: Array, + nlist_mask: Array, + sw: Array, smooth: bool = True, epsilon: float = 1e-4, use_sqrt_nnei: bool = True, -) -> np.ndarray: +) -> Array: """ Calculate the transposed rotation matrix. @@ -751,7 +761,7 @@ def _cal_hg( return hg -def _cal_grrg(hg: np.ndarray, axis_neuron: int) -> np.ndarray: +def _cal_grrg(hg: Array, axis_neuron: int) -> Array: """ Calculate the atomic invariant rep. @@ -780,15 +790,15 @@ def _cal_grrg(hg: np.ndarray, axis_neuron: int) -> np.ndarray: def symmetrization_op( - g: np.ndarray, - h: np.ndarray, - nlist_mask: np.ndarray, - sw: np.ndarray, + g: Array, + h: Array, + nlist_mask: Array, + sw: Array, axis_neuron: int, smooth: bool = True, epsilon: float = 1e-4, use_sqrt_nnei: bool = True, -) -> np.ndarray: +) -> Array: """ Symmetrization operator to obtain atomic invariant rep. @@ -846,7 +856,8 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" super().__init__() @@ -859,6 +870,7 @@ def __init__( bias=False, precision=precision, seed=seed, + trainable=trainable, ) self.has_gate = has_gate self.smooth = smooth @@ -867,11 +879,11 @@ def __init__( def call( self, - g2: np.ndarray, # nf x nloc x nnei x ng2 - h2: np.ndarray, # nf x nloc x nnei x 3 - nlist_mask: np.ndarray, # nf x nloc x nnei - sw: np.ndarray, # nf x nloc x nnei - ) -> np.ndarray: + g2: Array, # nf x nloc x nnei x ng2 + h2: Array, # nf x nloc x nnei x 3 + nlist_mask: Array, # nf x nloc x nnei + sw: Array, # nf x nloc x nnei + ) -> Array: xp = array_api_compat.array_namespace(g2, h2, nlist_mask, sw) ( nf, @@ -969,7 +981,8 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim @@ -980,20 +993,22 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.head_map = NativeLayer( input_dim * head_num, input_dim, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.precision = precision def call( self, - AA: np.ndarray, # nf x nloc x nnei x nnei x nh - g2: np.ndarray, # nf x nloc x nnei x ng2 - ) -> np.ndarray: + AA: Array, # nf x nloc x nnei x nnei x nh + g2: Array, # nf x nloc x nnei x ng2 + ) -> Array: xp = array_api_compat.array_namespace(AA, g2) nf, nloc, nnei, ng2 = g2.shape nh = self.head_num @@ -1057,21 +1072,27 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num self.head_map = NativeLayer( - head_num, 1, bias=False, precision=precision, seed=seed + head_num, + 1, + bias=False, + precision=precision, + seed=seed, + trainable=trainable, ) self.precision = precision def call( self, - AA: np.ndarray, # nf x nloc x nnei x nnei x nh - h2: np.ndarray, # nf x nloc x nnei x 3 - ) -> np.ndarray: + AA: Array, # nf x nloc x nnei x nnei x nh + h2: Array, # nf x nloc x nnei x 3 + ) -> Array: xp = array_api_compat.array_namespace(AA, h2) nf, nloc, nnei, _ = h2.shape nh = self.head_num @@ -1132,7 +1153,8 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim @@ -1144,6 +1166,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.mapkv = NativeLayer( input_dim, @@ -1151,12 +1174,14 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.head_map = NativeLayer( input_dim * head_num, input_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) self.smooth = smooth self.attnw_shift = attnw_shift @@ -1164,11 +1189,11 @@ def __init__( def call( self, - g1: np.ndarray, # nf x nloc x ng1 - gg1: np.ndarray, # nf x nloc x nnei x ng1 - nlist_mask: np.ndarray, # nf x nloc x nnei - sw: np.ndarray, # nf x nloc x nnei - ) -> np.ndarray: + g1: Array, # nf x nloc x ng1 + gg1: Array, # nf x nloc x nnei x ng1 + nlist_mask: Array, # nf x nloc x nnei + sw: Array, # nf x nloc x nnei + ) -> Array: xp = array_api_compat.array_namespace(g1, gg1, nlist_mask, sw) nf, nloc, nnei = nlist_mask.shape ni, nd, nh = self.input_dim, self.hidden_dim, self.head_num @@ -1263,12 +1288,12 @@ def deserialize(cls, data: dict) -> "LocalAtten": class RepformerLayer(NativeOP): def __init__( self, - rcut, - rcut_smth, + rcut: float, + rcut_smth: float, sel: int, ntypes: int, - g1_dim=128, - g2_dim=16, + g1_dim: int = 128, + g2_dim: int = 16, axis_neuron: int = 4, update_chnnl_2: bool = True, update_g1_has_conv: bool = True, @@ -1293,8 +1318,9 @@ def __init__( use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei @@ -1354,6 +1380,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) ) @@ -1363,6 +1390,7 @@ def __init__( g1_dim, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.linear2 = None self.proj_g1g2 = None @@ -1379,6 +1407,7 @@ def __init__( g2_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) if self.update_style == "res_residual": g2_residual.append( @@ -1388,6 +1417,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) ) if self.g1_out_mlp: @@ -1396,6 +1426,7 @@ def __init__( g1_dim, precision=precision, seed=child_seed(seed, 15), + trainable=trainable, ) if self.update_style == "res_residual": g1_residual.append( @@ -1405,6 +1436,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 16), + trainable=trainable, ) ) else: @@ -1417,6 +1449,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) else: self.proj_g1g2 = NativeLayer( @@ -1425,6 +1458,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) if self.update_style == "res_residual": g1_residual.append( @@ -1434,6 +1468,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 17), + trainable=trainable, ) ) if self.update_g2_has_g1g1: @@ -1443,6 +1478,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 5), + trainable=trainable, ) if self.update_style == "res_residual": g2_residual.append( @@ -1452,6 +1488,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 6), + trainable=trainable, ) ) if self.update_g2_has_attn or self.update_h2: @@ -1463,10 +1500,15 @@ def __init__( self.smooth, precision=precision, seed=child_seed(seed, 7), + trainable=trainable, ) if self.update_g2_has_attn: self.attn2_mh_apply = Atten2MultiHeadApply( - g2_dim, attn2_nhead, precision=precision, seed=child_seed(seed, 8) + g2_dim, + attn2_nhead, + precision=precision, + seed=child_seed(seed, 8), + trainable=trainable, ) self.attn2_lm = LayerNorm( g2_dim, @@ -1483,12 +1525,17 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 10), + trainable=trainable, ) ) if self.update_h2: self.attn2_ev_apply = Atten2EquiVarApply( - g2_dim, attn2_nhead, precision=precision, seed=child_seed(seed, 11) + g2_dim, + attn2_nhead, + precision=precision, + seed=child_seed(seed, 11), + trainable=trainable, ) if self.update_style == "res_residual": h2_residual.append( @@ -1498,6 +1545,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 12), + trainable=trainable, ) ) if self.update_g1_has_attn: @@ -1508,6 +1556,7 @@ def __init__( self.smooth, precision=precision, seed=child_seed(seed, 13), + trainable=trainable, ) if self.update_style == "res_residual": g1_residual.append( @@ -1517,6 +1566,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 14), + trainable=trainable, ) ) @@ -1536,9 +1586,9 @@ def cal_1_dim(self, g1d: int, g2d: int, ax: int) -> int: def _update_h2( self, - h2: np.ndarray, - attn: np.ndarray, - ) -> np.ndarray: + h2: Array, + attn: Array, + ) -> Array: """ Calculate the attention weights update for pair-wise equivariant rep. @@ -1556,11 +1606,11 @@ def _update_h2( def _update_g1_conv( self, - gg1: np.ndarray, - g2: np.ndarray, - nlist_mask: np.ndarray, - sw: np.ndarray, - ) -> np.ndarray: + gg1: Array, + g2: Array, + nlist_mask: Array, + sw: Array, + ) -> Array: """ Calculate the convolution update for atomic invariant rep. @@ -1614,11 +1664,11 @@ def _update_g1_conv( def _update_g2_g1g1( self, - g1: np.ndarray, # nf x nloc x ng1 - gg1: np.ndarray, # nf x nloc x nnei x ng1 - nlist_mask: np.ndarray, # nf x nloc x nnei - sw: np.ndarray, # nf x nloc x nnei - ) -> np.ndarray: + g1: Array, # nf x nloc x ng1 + gg1: Array, # nf x nloc x nnei x ng1 + nlist_mask: Array, # nf x nloc x nnei + sw: Array, # nf x nloc x nnei + ) -> Array: """ Update the g2 using element-wise dot g1_i * g1_j. @@ -1644,13 +1694,13 @@ def _update_g2_g1g1( def call( self, - g1_ext: np.ndarray, # nf x nall x ng1 - g2: np.ndarray, # nf x nloc x nnei x ng2 - h2: np.ndarray, # nf x nloc x nnei x 3 - nlist: np.ndarray, # nf x nloc x nnei - nlist_mask: np.ndarray, # nf x nloc x nnei - sw: np.ndarray, # switch func, nf x nloc x nnei - ): + g1_ext: Array, # nf x nall x ng1 + g2: Array, # nf x nloc x nnei x ng2 + h2: Array, # nf x nloc x nnei x 3 + nlist: Array, # nf x nloc x nnei + nlist_mask: Array, # nf x nloc x nnei + sw: Array, # switch func, nf x nloc x nnei + ) -> tuple[Array, Array]: """ Parameters ---------- @@ -1682,10 +1732,10 @@ def call( assert (nf, nloc) == g1.shape[:2] assert (nf, nloc, nnei) == h2.shape[:3] - g2_update: list[np.ndarray] = [g2] - h2_update: list[np.ndarray] = [h2] - g1_update: list[np.ndarray] = [g1] - g1_mlp: list[np.ndarray] = [g1] if not self.g1_out_mlp else [] + g2_update: list[Array] = [g2] + h2_update: list[Array] = [h2] + g1_update: list[Array] = [g1] + g1_mlp: list[Array] = [g1] if not self.g1_out_mlp else [] if self.g1_out_mlp: assert self.g1_self_mlp is not None g1_self_mlp = self.act(self.g1_self_mlp(g1)) @@ -1787,15 +1837,15 @@ def call( def list_update_res_avg( self, - update_list: list[np.ndarray], - ) -> np.ndarray: + update_list: list[Array], + ) -> Array: nitem = len(update_list) uu = update_list[0] for ii in range(1, nitem): uu = uu + update_list[ii] return uu / (float(nitem) ** 0.5) - def list_update_res_incr(self, update_list: list[np.ndarray]) -> np.ndarray: + def list_update_res_incr(self, update_list: list[Array]) -> Array: nitem = len(update_list) uu = update_list[0] scale = 1.0 / (float(nitem - 1) ** 0.5) if nitem > 1 else 0.0 @@ -1804,8 +1854,8 @@ def list_update_res_incr(self, update_list: list[np.ndarray]) -> np.ndarray: return uu def list_update_res_residual( - self, update_list: list[np.ndarray], update_name: str = "g1" - ) -> np.ndarray: + self, update_list: list[Array], update_name: str = "g1" + ) -> Array: nitem = len(update_list) uu = update_list[0] if update_name == "g1": @@ -1821,9 +1871,7 @@ def list_update_res_residual( raise NotImplementedError return uu - def list_update( - self, update_list: list[np.ndarray], update_name: str = "g1" - ) -> np.ndarray: + def list_update(self, update_list: list[Array], update_name: str = "g1") -> Array: if self.update_style == "res_avg": return self.list_update_res_avg(update_list) elif self.update_style == "res_incr": diff --git a/deepmd/dpmodel/descriptor/se_atten_v2.py b/deepmd/dpmodel/descriptor/se_atten_v2.py index 897863ec0f..99074fb652 100644 --- a/deepmd/dpmodel/descriptor/se_atten_v2.py +++ b/deepmd/dpmodel/descriptor/se_atten_v2.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -39,7 +37,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -56,19 +54,19 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - scaling_factor=1.0, + scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, concat_output_tebd: bool = True, - spin: Optional[Any] = None, - stripped_type_embedding: Optional[bool] = None, + spin: Any | None = None, + stripped_type_embedding: bool | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: DescrptDPA1.__init__( self, diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index bd72d936e3..c09a6cbdc3 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import array_api_compat @@ -16,6 +16,9 @@ PRECISION_DICT, NativeOP, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, to_numpy_array, @@ -162,11 +165,11 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - spin: Optional[Any] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + spin: Any | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: del ntypes ## seed, uniform_seed, not included. @@ -207,6 +210,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(seed, ii), + trainable=trainable, ) self.embeddings = embeddings self.env_mat = EnvMat(self.rcut, self.rcut_smth, protection=self.env_protection) @@ -221,7 +225,7 @@ def __init__( self.sel_cumsum = [0, *np.cumsum(self.sel).tolist()] self.ndescrpt = self.nnei * 4 - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.davg = value elif key in ("std", "data_std", "dstd"): @@ -229,7 +233,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.davg elif key in ("std", "data_std", "dstd"): @@ -238,19 +242,19 @@ def __getitem__(self, key): raise KeyError(key) @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.get_dim_out() - def get_dim_out(self): + def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.neuron[-1] * self.axis_neuron - def get_dim_emb(self): + def get_dim_emb(self) -> int: """Returns the embedding (g2) dimension of this descriptor.""" return self.neuron[-1] - def get_rcut(self): + def get_rcut(self) -> float: """Returns cutoff radius.""" return self.rcut @@ -258,7 +262,7 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth - def get_sel(self): + def get_sel(self) -> list[int]: """Returns cutoff radius.""" return self.sel @@ -280,7 +284,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: Any, shared_level: Any, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -289,7 +295,7 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -310,8 +316,8 @@ def get_type_map(self) -> list[str]: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -350,22 +356,22 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: np.ndarray, - stddev: np.ndarray, + mean: Array, + stddev: Array, ) -> None: """Update mean and stddev for descriptor.""" self.davg = mean self.dstd = stddev - def get_stat_mean_and_stddev(self) -> tuple[np.ndarray, np.ndarray]: + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: """Get mean and stddev for descriptor.""" return self.davg, self.dstd def cal_g( self, - ss, - embedding_idx, - ): + ss: Array, + embedding_idx: int, + ) -> Array: xp = array_api_compat.array_namespace(ss) nf_times_nloc, nnei = ss.shape[0:2] ss = xp.reshape(ss, (nf_times_nloc, nnei, 1)) @@ -383,11 +389,11 @@ def reinit_exclude( @cast_precision def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> Array: """Compute the descriptor. Parameters @@ -516,9 +522,9 @@ def deserialize(cls, data: dict) -> "DescrptSeA": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters @@ -548,11 +554,11 @@ class DescrptSeAArrayAPI(DescrptSeA): @cast_precision def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> Array: """Compute the descriptor. Parameters diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 5b2931b23f..6decd91a23 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, NoReturn, - Optional, - Union, ) import array_api_compat @@ -15,6 +15,9 @@ PRECISION_DICT, NativeOP, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, get_xp_precision, @@ -120,11 +123,11 @@ def __init__( set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - spin: Optional[Any] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + spin: Any | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input # consistent with argcheck, not used though - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: del ntypes ## seed, uniform_seed, not included. @@ -166,6 +169,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(seed, ii), + trainable=trainable, ) self.embeddings = embeddings self.env_mat = EnvMat(self.rcut, self.rcut_smth, protection=self.env_protection) @@ -180,7 +184,7 @@ def __init__( self.sel_cumsum = [0, *np.cumsum(self.sel).tolist()] self.ndescrpt = self.nnei - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.davg = value elif key in ("std", "data_std", "dstd"): @@ -188,7 +192,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.davg elif key in ("std", "data_std", "dstd"): @@ -197,11 +201,11 @@ def __getitem__(self, key): raise KeyError(key) @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.get_dim_out() - def get_dim_out(self): + def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.neuron[-1] @@ -209,7 +213,7 @@ def get_dim_emb(self) -> NoReturn: """Returns the embedding (g2) dimension of this descriptor.""" raise NotImplementedError - def get_rcut(self): + def get_rcut(self) -> float: """Returns cutoff radius.""" return self.rcut @@ -217,7 +221,7 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth - def get_sel(self): + def get_sel(self) -> list[int]: """Returns cutoff radius.""" return self.sel @@ -239,7 +243,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: Any, shared_level: Any, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -248,7 +254,7 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -269,8 +275,8 @@ def get_type_map(self) -> list[str]: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -309,22 +315,22 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: np.ndarray, - stddev: np.ndarray, + mean: Array, + stddev: Array, ) -> None: """Update mean and stddev for descriptor.""" self.davg = mean self.dstd = stddev - def get_stat_mean_and_stddev(self) -> tuple[np.ndarray, np.ndarray]: + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: """Get mean and stddev for descriptor.""" return self.davg, self.dstd def cal_g( self, - ss, - ll, - ): + ss: Array, + ll: int, + ) -> Array: xp = array_api_compat.array_namespace(ss) nf, nloc, nnei = ss.shape[0:3] ss = xp.reshape(ss, (nf, nloc, nnei, 1)) @@ -335,11 +341,11 @@ def cal_g( @cast_precision def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> Array: """Compute the descriptor. Parameters @@ -453,9 +459,9 @@ def deserialize(cls, data: dict) -> "DescrptSeR": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index fb30f04961..863187dd4c 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, NoReturn, - Optional, - Union, ) import array_api_compat @@ -15,6 +16,9 @@ PRECISION_DICT, NativeOP, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, get_xp_precision, @@ -107,9 +111,9 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = DEFAULT_PRECISION, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input ) -> None: del ntypes self.rcut = rcut @@ -147,6 +151,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(self.seed, ii), + trainable=trainable, ) self.embeddings = embeddings self.env_mat = EnvMat(self.rcut, self.rcut_smth, protection=self.env_protection) @@ -160,7 +165,7 @@ def __init__( self.orig_sel = self.sel self.ndescrpt = self.nnei * 4 - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.davg = value elif key in ("std", "data_std", "dstd"): @@ -168,7 +173,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.davg elif key in ("std", "data_std", "dstd"): @@ -177,12 +182,12 @@ def __getitem__(self, key): raise KeyError(key) @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.get_dim_out() def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -193,15 +198,15 @@ def change_type_map( "We may consider adding this support in the future if there is a clear demand for it." ) - def get_dim_out(self): + def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.neuron[-1] - def get_dim_emb(self): + def get_dim_emb(self) -> int: """Returns the embedding (g2) dimension of this descriptor.""" return self.neuron[-1] - def get_rcut(self): + def get_rcut(self) -> float: """Returns cutoff radius.""" return self.rcut @@ -209,7 +214,7 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth - def get_sel(self): + def get_sel(self) -> list: """Returns cutoff radius.""" return self.sel @@ -231,7 +236,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -249,8 +256,8 @@ def get_type_map(self) -> list[str]: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -289,14 +296,14 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: np.ndarray, - stddev: np.ndarray, + mean: Array, + stddev: Array, ) -> None: """Update mean and stddev for descriptor.""" self.davg = mean self.dstd = stddev - def get_stat_mean_and_stddev(self) -> tuple[np.ndarray, np.ndarray]: + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: """Get mean and stddev for descriptor.""" return self.davg, self.dstd @@ -310,11 +317,11 @@ def reinit_exclude( @cast_precision def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> tuple[Array, Array]: """Compute the descriptor. Parameters @@ -451,9 +458,9 @@ def deserialize(cls, data: dict) -> "DescrptSeT": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/descriptor/se_t_tebd.py b/deepmd/dpmodel/descriptor/se_t_tebd.py index ff26024aad..e118d5abd4 100644 --- a/deepmd/dpmodel/descriptor/se_t_tebd.py +++ b/deepmd/dpmodel/descriptor/se_t_tebd.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, - Union, ) import array_api_compat @@ -14,6 +15,7 @@ NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) from deepmd.dpmodel.common import ( @@ -122,7 +124,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, @@ -134,11 +136,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, - use_tebd_bias=False, + use_tebd_bias: bool = False, smooth: bool = True, ) -> None: self.se_ttebd = DescrptBlockSeTTebd( @@ -157,6 +159,7 @@ def __init__( env_protection=env_protection, smooth=smooth, seed=child_seed(seed, 0), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd self.type_map = type_map @@ -171,6 +174,7 @@ def __init__( use_tebd_bias=use_tebd_bias, type_map=type_map, seed=child_seed(seed, 1), + trainable=trainable, ) self.tebd_dim = tebd_dim self.concat_output_tebd = concat_output_tebd @@ -235,7 +239,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_ttebd.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> NoReturn: + def share_params( + self, base_class: "DescrptSeTTebd", shared_level: int, resume: bool = False + ) -> NoReturn: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -244,18 +250,18 @@ def share_params(self, base_class, shared_level, resume=False) -> NoReturn: raise NotImplementedError @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -276,19 +282,21 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: np.ndarray, - stddev: np.ndarray, + mean: Array, + stddev: Array, ) -> None: """Update mean and stddev for descriptor.""" self.se_ttebd.mean = mean self.se_ttebd.stddev = stddev - def get_stat_mean_and_stddev(self) -> tuple[np.ndarray, np.ndarray]: + def get_stat_mean_and_stddev(self) -> tuple[Array, Array]: """Get mean and stddev for descriptor.""" return self.se_ttebd.mean, self.se_ttebd.stddev def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["DescrptSeTTebd"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -317,11 +325,11 @@ def change_type_map( @cast_precision def call( self, - coord_ext, - atype_ext, - nlist, - mapping: Optional[np.ndarray] = None, - ): + coord_ext: Array, + atype_ext: Array, + nlist: Array, + mapping: Array | None = None, + ) -> tuple[Array, Array]: """Compute the descriptor. Parameters @@ -358,7 +366,7 @@ def call( type_embedding = self.type_embedding.call() # nf x nall x tebd_dim atype_embd_ext = xp.reshape( - xp.take(type_embedding, xp.reshape(atype_ext, [-1]), axis=0), + xp.take(type_embedding, xp.reshape(atype_ext, (-1,)), axis=0), (nf, nall, self.tebd_dim), ) # nfnl x tebd_dim @@ -449,9 +457,9 @@ def deserialize(cls, data: dict) -> "DescrptSeTTebd": def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[Array, Array]: """Update the selection and perform neighbor statistics. Parameters @@ -484,19 +492,20 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, tebd_input_mode: str = "concat", set_davg_zero: bool = True, - activation_function="tanh", + activation_function: str = "tanh", precision: str = "float64", resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: self.rcut = rcut self.rcut_smth = rcut_smth @@ -542,6 +551,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.embeddings = embeddings if self.tebd_input_mode in ["strip"]: @@ -557,6 +567,7 @@ def __init__( self.resnet_dt, self.precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.embeddings_strip = embeddings_strip else: @@ -600,7 +611,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -608,7 +619,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -633,24 +644,24 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.tebd_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -704,18 +715,18 @@ def reinit_exclude( def cal_g( self, - ss, - embedding_idx, - ): + ss: Array, + embedding_idx: int, + ) -> Array: # nfnl x nt_i x nt_j x ng gg = self.embeddings[embedding_idx].call(ss) return gg def cal_g_strip( self, - ss, - embedding_idx, - ): + ss: Array, + embedding_idx: int, + ) -> Array: assert self.embeddings_strip is not None # nfnl x nt_i x nt_j x ng gg = self.embeddings_strip[embedding_idx].call(ss) @@ -723,13 +734,13 @@ def cal_g_strip( def call( self, - nlist: np.ndarray, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - atype_embd_ext: Optional[np.ndarray] = None, - mapping: Optional[np.ndarray] = None, - type_embedding: Optional[np.ndarray] = None, - ): + nlist: Array, + coord_ext: Array, + atype_ext: Array, + atype_embd_ext: Array | None = None, + mapping: Array | None = None, + type_embedding: Array | None = None, + ) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(nlist, coord_ext, atype_ext) # nf x nloc x nnei x 4 dmatrix, diff, sw = self.env_mat.call( diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index fcaea43338..2f10d7b487 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -1,16 +1,16 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat -import numpy as np from deepmd.dpmodel import ( DEFAULT_PRECISION, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, ) @@ -84,6 +84,9 @@ class DipoleFitting(GeneralFitting): Only reducible variable are differentiable. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -96,20 +99,21 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, + trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], r_differentiable: bool = True, c_differentiable: bool = True, - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -144,9 +148,10 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" return self.embedding_width @@ -161,12 +166,12 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) var_name = data.pop("var_name", None) assert var_name == "dipole" return super().deserialize(data) - def output_def(self): + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ OutputVariableDef( @@ -182,14 +187,14 @@ def output_def(self): @cast_precision def call( self, - descriptor: np.ndarray, - atype: np.ndarray, - gr: Optional[np.ndarray] = None, - g2: Optional[np.ndarray] = None, - h2: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + descriptor: Array, + atype: Array, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Calculate the fitting. Parameters diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index 2f6df77eac..803f31b30f 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( TYPE_CHECKING, - Optional, - Union, ) import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( DEFAULT_PRECISION, to_numpy_array, @@ -37,15 +38,16 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - bias_dos: Optional[np.ndarray] = None, - rcond: Optional[float] = None, - trainable: Union[bool, list[bool]] = True, + bias_dos: Array | None = None, + rcond: float | None = None, + trainable: bool | list[bool] = True, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = False, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list | None = None, ) -> None: if bias_dos is not None: self.bias_dos = bias_dos @@ -70,12 +72,13 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data["numb_dos"] = data.pop("dim_out") data.pop("tot_ener_zero", None) data.pop("var_name", None) diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 6435b6468f..713aef5117 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -2,8 +2,6 @@ from typing import ( TYPE_CHECKING, Any, - Optional, - Union, ) from deepmd.dpmodel.common import ( @@ -33,19 +31,20 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, - atom_ener: Optional[list[float]] = None, + trainable: list[bool] | None = None, + atom_ener: list[float] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list | None = None, ) -> None: super().__init__( var_name="energy", @@ -70,12 +69,13 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index cd0d4e72d4..c61ab234b1 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -2,10 +2,11 @@ from abc import ( abstractmethod, ) +from collections.abc import ( + Callable, +) from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -16,6 +17,9 @@ PRECISION_DICT, NativeOP, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( get_xp_precision, to_numpy_array, @@ -94,6 +98,9 @@ class GeneralFitting(NativeOP, BaseFitting): A list of strings. Give the name to each type of atoms. seed: Optional[Union[int, list[int]]] Random seed for initializing the network parameters. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -106,20 +113,21 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - bias_atom_e: Optional[np.ndarray] = None, - rcond: Optional[float] = None, + bias_atom_e: Array | None = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, + trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = True, exclude_types: list[int] = [], - remove_vaccum_contribution: Optional[list[bool]] = None, - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, + remove_vaccum_contribution: list[bool] | None = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: self.var_name = var_name self.ntypes = ntypes @@ -129,6 +137,7 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam self.rcond = rcond self.tot_ener_zero = tot_ener_zero self.trainable = trainable @@ -177,6 +186,15 @@ def __init__( self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec) else: self.case_embd = None + + if self.default_fparam is not None: + if self.numb_fparam > 0: + assert len(self.default_fparam) == self.numb_fparam, ( + "default_fparam length mismatch!" + ) + self.default_fparam_tensor = np.array(self.default_fparam, dtype=self.prec) + else: + self.default_fparam_tensor = None # init networks in_dim = ( self.dim_descrpt @@ -198,13 +216,79 @@ def __init__( self.precision, bias_out=True, seed=child_seed(seed, ii), + trainable=trainable, ) for ii in range(self.ntypes if not self.mixed_types else 1) ], ) + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + ) -> None: + """ + Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + + Parameters + ---------- + merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `numpy.ndarray` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + protection : float + Divided-by-zero protection + """ + if self.numb_fparam == 0 and self.numb_aparam == 0: + # skip data statistics + return + if callable(merged): + sampled = merged() + else: + sampled = merged + # stat fparam + if self.numb_fparam > 0: + cat_data = np.concatenate([frame["fparam"] for frame in sampled], axis=0) + cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) + fparam_avg = np.mean(cat_data, axis=0) + fparam_std = np.std(cat_data, axis=0, ddof=0) # ddof=0 for population std + fparam_std = np.where( + fparam_std < protection, + np.array(protection, dtype=fparam_std.dtype), + fparam_std, + ) + fparam_inv_std = 1.0 / fparam_std + self.fparam_avg = fparam_avg.astype(self.fparam_avg.dtype) + self.fparam_inv_std = fparam_inv_std.astype(self.fparam_inv_std.dtype) + # stat aparam + if self.numb_aparam > 0: + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in [frame["aparam"] for frame in sampled]: + ss = np.reshape(ss_, [-1, self.numb_aparam]) + sys_sumv.append(np.sum(ss, axis=0)) + sys_sumv2.append(np.sum(ss * ss, axis=0)) + sys_sumn.append(ss.shape[0]) + sumv = np.sum(np.stack(sys_sumv), axis=0) + sumv2 = np.sum(np.stack(sys_sumv2), axis=0) + sumn = sum(sys_sumn) + aparam_avg = sumv / sumn + aparam_std = np.sqrt(sumv2 / sumn - (sumv / sumn) ** 2) + aparam_std = np.where( + aparam_std < protection, + np.array(protection, dtype=aparam_std.dtype), + aparam_std, + ) + aparam_inv_std = 1.0 / aparam_std + self.aparam_avg = aparam_avg.astype(self.aparam_avg.dtype) + self.aparam_inv_std = aparam_inv_std.astype(self.aparam_inv_std.dtype) + @abstractmethod - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" pass @@ -216,6 +300,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam + def has_default_fparam(self) -> bool: + """Check if the fitting has default frame parameters.""" + return self.default_fparam is not None + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. @@ -229,7 +317,7 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this fitting net by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -237,7 +325,7 @@ def set_case_embd(self, case_idx: int): self.case_embd = np.eye(self.dim_case_embd, dtype=self.prec)[case_idx] def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -258,7 +346,7 @@ def change_type_map( ) self.bias_atom_e = self.bias_atom_e[remap_index] - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ["bias_atom_e"]: self.bias_atom_e = value elif key in ["fparam_avg"]: @@ -273,10 +361,12 @@ def __setitem__(self, key, value) -> None: self.case_embd = value elif key in ["scale"]: self.scale = value + elif key in ["default_fparam_tensor"]: + self.default_fparam_tensor = value else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ["bias_atom_e"]: return self.bias_atom_e elif key in ["fparam_avg"]: @@ -291,6 +381,8 @@ def __getitem__(self, key): return self.case_embd elif key in ["scale"]: return self.scale + elif key in ["default_fparam_tensor"]: + return self.default_fparam_tensor else: raise KeyError(key) @@ -305,7 +397,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 3, + "@version": 4, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -314,6 +406,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "rcond": self.rcond, "activation_function": self.activation_function, "precision": self.precision, @@ -352,14 +445,14 @@ def deserialize(cls, data: dict) -> "GeneralFitting": def _call_common( self, - descriptor: np.ndarray, - atype: np.ndarray, - gr: Optional[np.ndarray] = None, - g2: Optional[np.ndarray] = None, - h2: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + descriptor: Array, + atype: Array, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Calculate the fitting. Parameters @@ -402,6 +495,14 @@ def _call_common( xx_zeros = xp.zeros_like(xx) else: xx_zeros = None + + if self.numb_fparam > 0 and fparam is None: + # use default fparam + assert self.default_fparam_tensor is not None + fparam = xp.tile( + xp.reshape(self.default_fparam_tensor, (1, self.numb_fparam)), (nf, 1) + ) + # check fparam dim, concate to input descriptor if self.numb_fparam > 0: assert fparam is not None, "fparam should not be None" @@ -412,7 +513,7 @@ def _call_common( ) fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...] fparam = xp.tile( - xp.reshape(fparam, [nf, 1, self.numb_fparam]), (1, nloc, 1) + xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1) ) xx = xp.concat( [xx, fparam], @@ -431,7 +532,7 @@ def _call_common( f"get an input aparam of dim {aparam.shape[-1]}, " f"which is not consistent with {self.numb_aparam}." ) - aparam = xp.reshape(aparam, [nf, nloc, self.numb_aparam]) + aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam)) aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...] xx = xp.concat( [xx, aparam], @@ -446,7 +547,7 @@ def _call_common( if self.dim_case_embd > 0: assert self.case_embd is not None case_embd = xp.tile( - xp.reshape(self.case_embd[...], [1, 1, -1]), [nf, nloc, 1] + xp.reshape(self.case_embd[...], (1, 1, -1)), (nf, nloc, 1) ) xx = xp.concat( [xx, case_embd], @@ -465,7 +566,7 @@ def _call_common( ) for type_i in range(self.ntypes): mask = xp.tile( - xp.reshape((atype == type_i), [nf, nloc, 1]), (1, 1, net_dim_out) + xp.reshape((atype == type_i), (nf, nloc, 1)), (1, 1, net_dim_out) ) atom_property = self.nets[(type_i,)](xx) if self.remove_vaccum_contribution is not None and not ( @@ -485,10 +586,10 @@ def _call_common( outs += xp.reshape( xp.take( xp.astype(self.bias_atom_e[...], outs.dtype), - xp.reshape(atype, [-1]), + xp.reshape(atype, (-1,)), axis=0, ), - [nf, nloc, net_dim_out], + (nf, nloc, net_dim_out), ) # nf x nloc exclude_mask = self.emask.build_type_exclude_mask(atype) diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index b5d3a02d86..f771f927cd 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -2,15 +2,14 @@ from typing import ( Any, NoReturn, - Optional, - Union, ) -import numpy as np - from deepmd.dpmodel import ( DEFAULT_PRECISION, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, ) @@ -110,6 +109,9 @@ class InvarFitting(GeneralFitting): Atomic contributions of the excluded atom types are set zero. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ @@ -124,20 +126,21 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - bias_atom: Optional[np.ndarray] = None, - rcond: Optional[float] = None, + bias_atom: Array | None = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, - atom_ener: Optional[list[float]] = None, + trainable: list[bool] | None = None, + atom_ener: list[float] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = True, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -173,6 +176,7 @@ def __init__( else [x is not None for x in atom_ener], type_map=type_map, seed=seed, + default_fparam=default_fparam, ) def serialize(self) -> dict: @@ -185,18 +189,18 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) return super().deserialize(data) - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" return self.dim_out - def compute_output_stats(self, merged) -> NoReturn: + def compute_output_stats(self, merged: Any) -> NoReturn: """Update the output bias for fitting net.""" raise NotImplementedError - def output_def(self): + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ OutputVariableDef( @@ -212,14 +216,14 @@ def output_def(self): @cast_precision def call( self, - descriptor: np.ndarray, - atype: np.ndarray, - gr: Optional[np.ndarray] = None, - g2: Optional[np.ndarray] = None, - h2: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + descriptor: Array, + atype: Array, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Calculate the fitting. Parameters diff --git a/deepmd/dpmodel/fitting/make_base_fitting.py b/deepmd/dpmodel/fitting/make_base_fitting.py index 201b5e27d1..7b65a150b2 100644 --- a/deepmd/dpmodel/fitting/make_base_fitting.py +++ b/deepmd/dpmodel/fitting/make_base_fitting.py @@ -4,8 +4,8 @@ abstractmethod, ) from typing import ( + Any, NoReturn, - Optional, ) from deepmd.common import ( @@ -21,9 +21,9 @@ def make_base_fitting( - t_tensor, + t_tensor: Any, fwd_method_name: str = "forward", -): +) -> type: """Make the base class for the fitting. Parameters @@ -39,7 +39,7 @@ def make_base_fitting( class BF(ABC, PluginVariant, make_plugin_registry("fitting")): """Base fitting provides the interfaces of fitting net.""" - def __new__(cls, *args, **kwargs): + def __new__(cls: type, *args: Any, **kwargs: Any) -> Any: if cls is BF: cls = cls.get_class_by_type(j_get_type(kwargs, cls.__name__)) return super().__new__(cls) @@ -54,16 +54,16 @@ def fwd( self, descriptor: t_tensor, atype: t_tensor, - gr: Optional[t_tensor] = None, - g2: Optional[t_tensor] = None, - h2: Optional[t_tensor] = None, - fparam: Optional[t_tensor] = None, - aparam: Optional[t_tensor] = None, + gr: t_tensor | None = None, + g2: t_tensor | None = None, + h2: t_tensor | None = None, + fparam: t_tensor | None = None, + aparam: t_tensor | None = None, ) -> dict[str, t_tensor]: """Calculate fitting.""" pass - def compute_output_stats(self, merged) -> NoReturn: + def compute_output_stats(self, merged: Any) -> NoReturn: """Update the output bias for fitting net.""" raise NotImplementedError @@ -74,7 +74,7 @@ def get_type_map(self) -> list[str]: @abstractmethod def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index 8acb818a46..f3e6318ba5 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import array_api_compat @@ -14,6 +12,9 @@ from deepmd.dpmodel import ( DEFAULT_PRECISION, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( cast_precision, to_numpy_array, @@ -90,6 +91,9 @@ class PolarFitting(GeneralFitting): Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -102,21 +106,22 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, + trainable: list[bool] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, spin: Any = None, mixed_types: bool = False, exclude_types: list[int] = [], fit_diag: bool = True, - scale: Optional[list[float]] = None, + scale: list[float] | None = None, shift_diag: bool = True, - type_map: Optional[list[str]] = None, - seed: Optional[Union[int, list[int]]] = None, + type_map: list[str] | None = None, + seed: int | list[int] | None = None, + default_fparam: list[float] | None = None, ) -> None: if tot_ener_zero: raise NotImplementedError("tot_ener_zero is not implemented") @@ -164,9 +169,10 @@ def __init__( exclude_types=exclude_types, type_map=type_map, seed=seed, + default_fparam=default_fparam, ) - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" return ( self.embedding_width @@ -174,13 +180,13 @@ def _net_out_dim(self): else self.embedding_width * self.embedding_width ) - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Array) -> None: if key in ["constant_matrix"]: self.constant_matrix = value else: super().__setitem__(key, value) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Array: if key in ["constant_matrix"]: return self.constant_matrix else: @@ -189,7 +195,7 @@ def __getitem__(self, key): def serialize(self) -> dict: data = super().serialize() data["type"] = "polar" - data["@version"] = 4 + data["@version"] = 5 data["embedding_width"] = self.embedding_width data["fit_diag"] = self.fit_diag data["shift_diag"] = self.shift_diag @@ -200,12 +206,12 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) var_name = data.pop("var_name", None) assert var_name == "polar" return super().deserialize(data) - def output_def(self): + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ OutputVariableDef( @@ -219,7 +225,7 @@ def output_def(self): ) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -247,14 +253,14 @@ def change_type_map( @cast_precision def call( self, - descriptor: np.ndarray, - atype: np.ndarray, - gr: Optional[np.ndarray] = None, - g2: Optional[np.ndarray] = None, - h2: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> dict[str, np.ndarray]: + descriptor: Array, + atype: Array, + gr: Array | None = None, + g2: Array | None = None, + h2: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> dict[str, Array]: """Calculate the fitting. Parameters @@ -289,7 +295,7 @@ def call( ] # out = out * self.scale[atype, ...] scale_atype = xp.reshape( - xp.take(xp.astype(self.scale, out.dtype), xp.reshape(atype, [-1]), axis=0), + xp.take(xp.astype(self.scale, out.dtype), xp.reshape(atype, (-1,)), axis=0), (*atype.shape, 1), ) out = out * scale_atype @@ -315,7 +321,7 @@ def call( bias = xp.reshape( xp.take( xp.astype(self.constant_matrix, out.dtype), - xp.reshape(atype, [-1]), + xp.reshape(atype, (-1,)), axis=0, ), (nframes, nloc), diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index 6d0aa3546f..4df91825b1 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -1,17 +1,18 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) - -import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( DEFAULT_PRECISION, ) from deepmd.dpmodel.fitting.invar_fitting import ( InvarFitting, ) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -61,6 +62,9 @@ class PropertyFittingNet(InvarFitting): Atomic contributions of the excluded atom types are set zero. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -69,9 +73,9 @@ def __init__( dim_descrpt: int, task_dim: int = 1, neuron: list[int] = [128, 128, 128], - bias_atom_p: Optional[np.ndarray] = None, - rcond: Optional[float] = None, - trainable: Union[bool, list[bool]] = True, + bias_atom_p: Array | None = None, + rcond: float | None = None, + trainable: bool | list[bool] = True, intensive: bool = False, property_name: str = "property", resnet_dt: bool = True, @@ -82,9 +86,10 @@ def __init__( precision: str = DEFAULT_PRECISION, mixed_types: bool = True, exclude_types: list[int] = [], - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, # not used - seed: Optional[int] = None, + seed: int | None = None, ) -> None: self.task_dim = task_dim self.intensive = intensive @@ -106,12 +111,27 @@ def __init__( mixed_types=mixed_types, exclude_types=exclude_types, type_map=type_map, + default_fparam=default_fparam, + ) + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reducible=True, + r_differentiable=False, + c_differentiable=False, + intensive=self.intensive, + ), + ] ) @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version"), 4, 1) + check_version_compatibility(data.pop("@version"), 5, 1) data.pop("dim_out") data["property_name"] = data.pop("var_name") data.pop("tot_ener_zero") @@ -131,6 +151,6 @@ def serialize(self) -> dict: "task_dim": self.task_dim, "intensive": self.intensive, } - dd["@version"] = 4 + dd["@version"] = 5 return dd diff --git a/deepmd/dpmodel/infer/deep_eval.py b/deepmd/dpmodel/infer/deep_eval.py index 91fa0ac2ac..3bd0f435e8 100644 --- a/deepmd/dpmodel/infer/deep_eval.py +++ b/deepmd/dpmodel/infer/deep_eval.py @@ -1,15 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.model.base_model import ( BaseModel, ) @@ -62,7 +66,7 @@ class DeepEval(DeepEvalBackend): The output definition of the model. *args : list Positional arguments. - auto_batch_size : bool or int or AutomaticBatchSize, default: False + auto_batch_size : bool or int or AutomaticBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional @@ -77,7 +81,7 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -86,6 +90,7 @@ def __init__( model_data = load_dp_model(model_file) self.dp = BaseModel.deserialize(model_data["model"]) + self.dp.model_def_script = json.dumps(model_data.get("model_def_script", {})) self.rcut = self.dp.get_rcut() self.type_map = self.dp.get_type_map() if isinstance(auto_batch_size, bool): @@ -120,6 +125,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.dp.get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.dp.has_default_fparam() + @property def model_type(self) -> type["DeepEvalWrapper"]: """The the evaluator of the model type.""" @@ -130,7 +139,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepDOS elif "dipole" in model_output_type: return DeepDipole - elif "polar" in model_output_type: + elif "polar" in model_output_type or "polarizability" in model_output_type: return DeepPolar elif "wfc" in model_output_type: return DeepWFC @@ -160,14 +169,14 @@ def get_ntypes_spin(self) -> int: def eval( self, - coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: np.ndarray, + coords: Array, + cells: Array | None, + atom_types: Array, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: Array | None = None, + aparam: Array | None = None, **kwargs: Any, - ) -> dict[str, np.ndarray]: + ) -> dict[str, Array]: """Evaluate the energy, force and virial by using this DP. Parameters @@ -273,7 +282,7 @@ def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Calla """ if self.auto_batch_size is not None: - def eval_func(*args, **kwargs): + def eval_func(*args: Any, **kwargs: Any) -> Any: return self.auto_batch_size.execute_all( inner_func, numb_test, natoms, *args, **kwargs ) @@ -284,8 +293,8 @@ def eval_func(*args, **kwargs): def _get_natoms_and_nframes( self, - coords: np.ndarray, - atom_types: np.ndarray, + coords: Array, + atom_types: Array, mixed_type: bool = False, ) -> tuple[int, int]: if mixed_type: @@ -301,13 +310,13 @@ def _get_natoms_and_nframes( def _eval_model( self, - coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + coords: Array, + cells: Array | None, + atom_types: Array, + fparam: Array | None, + aparam: Array | None, request_defs: list[OutputVariableDef], - ): + ) -> dict[str, Array]: model = self.dp nframes = coords.shape[0] @@ -365,7 +374,9 @@ def _eval_model( ) # this is kinda hacky return tuple(results) - def _get_output_shape(self, odef, nframes, natoms): + def _get_output_shape( + self, odef: OutputVariableDef, nframes: int, natoms: int + ) -> list[int]: if odef.category == OutputVariableCategory.DERV_C_REDU: # virial return [nframes, *odef.shape[:-1], 9] @@ -391,4 +402,14 @@ def _get_output_shape(self, odef, nframes, natoms): def get_model_def_script(self) -> dict: """Get model definition script.""" - return json.loads(self.model.get_model_def_script()) + return json.loads(self.dp.get_model_def_script()) + + def get_model(self) -> "BaseModel": + """Get the dpmodel BaseModel. + + Returns + ------- + BaseModel + The dpmodel BaseModel. + """ + return self.dp diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 7a17fcfcf0..7ac81d8a3d 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import array_api_compat -import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.loss.loss import ( Loss, ) @@ -17,7 +19,7 @@ ) -def custom_huber_loss(predictions, targets, delta=1.0): +def custom_huber_loss(predictions: Array, targets: Array, delta: float = 1.0) -> Array: xp = array_api_compat.array_namespace(predictions, targets) error = targets - predictions abs_error = xp.abs(error) @@ -41,14 +43,14 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, numb_generalized_coord: int = 0, - use_huber=False, - huber_delta=0.01, - **kwargs, + use_huber: bool = False, + huber_delta: float = 0.01, + **kwargs: Any, ) -> None: self.starter_learning_rate = starter_learning_rate self.start_pref_e = start_pref_e @@ -89,9 +91,9 @@ def call( self, learning_rate: float, natoms: int, - model_dict: dict[str, np.ndarray], - label_dict: dict[str, np.ndarray], - ) -> dict[str, np.ndarray]: + model_dict: dict[str, Array], + label_dict: dict[str, Array], + ) -> dict[str, Array]: """Calculate loss from model results and labeled results.""" energy = model_dict["energy_redu"] force = model_dict["energy_derv_r"] @@ -132,18 +134,18 @@ def call( atom_ener_coeff = xp.reshape(atom_ener_coeff, xp.shape(atom_ener)) energy = xp.sum(atom_ener_coeff * atom_ener, 1) if self.has_f or self.has_pf or self.relative_f or self.has_gf: - force_reshape = xp.reshape(force, [-1]) - force_hat_reshape = xp.reshape(force_hat, [-1]) + force_reshape = xp.reshape(force, (-1,)) + force_hat_reshape = xp.reshape(force_hat, (-1,)) diff_f = force_hat_reshape - force_reshape else: diff_f = None if self.relative_f is not None: - force_hat_3 = xp.reshape(force_hat, [-1, 3]) - norm_f = xp.reshape(xp.norm(force_hat_3, axis=1), [-1, 1]) + self.relative_f - diff_f_3 = xp.reshape(diff_f, [-1, 3]) + force_hat_3 = xp.reshape(force_hat, (-1, 3)) + norm_f = xp.reshape(xp.norm(force_hat_3, axis=1), (-1, 1)) + self.relative_f + diff_f_3 = xp.reshape(diff_f, (-1, 3)) diff_f_3 = diff_f_3 / norm_f - diff_f = xp.reshape(diff_f_3, [-1]) + diff_f = xp.reshape(diff_f_3, (-1,)) atom_norm = 1.0 / natoms atom_norm_ener = 1.0 / natoms @@ -177,22 +179,26 @@ def call( delta=self.huber_delta, ) loss += pref_e * l_huber_loss - more_loss["rmse_e"] = self.display_if_exist(l2_ener_loss, find_energy) + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(l2_ener_loss), find_energy + ) if self.has_f: l2_force_loss = xp.mean(xp.square(diff_f)) if not self.use_huber: loss += pref_f * l2_force_loss else: l_huber_loss = custom_huber_loss( - xp.reshape(force, [-1]), - xp.reshape(force_hat, [-1]), + xp.reshape(force, (-1,)), + xp.reshape(force_hat, (-1,)), delta=self.huber_delta, ) loss += pref_f * l_huber_loss - more_loss["rmse_f"] = self.display_if_exist(l2_force_loss, find_force) + more_loss["rmse_f"] = self.display_if_exist( + xp.sqrt(l2_force_loss), find_force + ) if self.has_v: - virial_reshape = xp.reshape(virial, [-1]) - virial_hat_reshape = xp.reshape(virial_hat, [-1]) + virial_reshape = xp.reshape(virial, (-1,)) + virial_hat_reshape = xp.reshape(virial_hat, (-1,)) l2_virial_loss = xp.mean( xp.square(virial_hat_reshape - virial_reshape), ) @@ -205,10 +211,12 @@ def call( delta=self.huber_delta, ) loss += pref_v * l_huber_loss - more_loss["rmse_v"] = self.display_if_exist(l2_virial_loss, find_virial) + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(l2_virial_loss), find_virial + ) if self.has_ae: - atom_ener_reshape = xp.reshape(atom_ener, [-1]) - atom_ener_hat_reshape = xp.reshape(atom_ener_hat, [-1]) + atom_ener_reshape = xp.reshape(atom_ener, (-1,)) + atom_ener_hat_reshape = xp.reshape(atom_ener_hat, (-1,)) l2_atom_ener_loss = xp.mean( xp.square(atom_ener_hat_reshape - atom_ener_reshape), ) @@ -222,24 +230,24 @@ def call( ) loss += pref_ae * l_huber_loss more_loss["rmse_ae"] = self.display_if_exist( - l2_atom_ener_loss, find_atom_ener + xp.sqrt(l2_atom_ener_loss), find_atom_ener ) if self.has_pf: - atom_pref_reshape = xp.reshape(atom_pref, [-1]) + atom_pref_reshape = xp.reshape(atom_pref, (-1,)) l2_pref_force_loss = xp.mean( xp.multiply(xp.square(diff_f), atom_pref_reshape), ) loss += pref_pf * l2_pref_force_loss more_loss["rmse_pf"] = self.display_if_exist( - l2_pref_force_loss, find_atom_pref + xp.sqrt(l2_pref_force_loss), find_atom_pref ) if self.has_gf: find_drdq = label_dict["find_drdq"] drdq = label_dict["drdq"] - force_reshape_nframes = xp.reshape(force, [-1, natoms[0] * 3]) - force_hat_reshape_nframes = xp.reshape(force_hat, [-1, natoms[0] * 3]) + force_reshape_nframes = xp.reshape(force, (-1, natoms[0] * 3)) + force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms[0] * 3)) drdq_reshape = xp.reshape( - drdq, [-1, natoms[0] * 3, self.numb_generalized_coord] + drdq, (-1, natoms[0] * 3, self.numb_generalized_coord) ) gen_force_hat = xp.einsum( "bij,bi->bj", drdq_reshape, force_hat_reshape_nframes @@ -252,7 +260,9 @@ def call( + (self.start_pref_gf - self.limit_pref_gf) * lr_ratio ) loss += pref_gf * l2_gen_force_loss - more_loss["rmse_gf"] = self.display_if_exist(l2_gen_force_loss, find_drdq) + more_loss["rmse_gf"] = self.display_if_exist( + xp.sqrt(l2_gen_force_loss), find_drdq + ) self.l2_l = loss more_loss["rmse"] = xp.sqrt(loss) diff --git a/deepmd/dpmodel/loss/loss.py b/deepmd/dpmodel/loss/loss.py index ff3a462cf1..6dc468582a 100644 --- a/deepmd/dpmodel/loss/loss.py +++ b/deepmd/dpmodel/loss/loss.py @@ -5,8 +5,10 @@ ) import array_api_compat -import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( NativeOP, ) @@ -24,9 +26,9 @@ def call( self, learning_rate: float, natoms: int, - model_dict: dict[str, np.ndarray], - label_dict: dict[str, np.ndarray], - ) -> dict[str, np.ndarray]: + model_dict: dict[str, Array], + label_dict: dict[str, Array], + ) -> dict[str, Array]: """Calculate loss from model results and labeled results.""" @property @@ -35,12 +37,12 @@ def label_requirement(self) -> list[DataRequirementItem]: """Return data label requirements needed for this loss calculation.""" @staticmethod - def display_if_exist(loss: np.ndarray, find_property: float) -> np.ndarray: + def display_if_exist(loss: Array, find_property: float) -> Array: """Display NaN if labeled property is not found. Parameters ---------- - loss : np.ndarray + loss : Array the loss scalar find_property : float whether the property is found diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 15c0bfc083..163cd62387 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -7,7 +7,6 @@ ) from typing import ( Any, - Optional, ) from deepmd.utils.data_system import ( @@ -36,7 +35,7 @@ class BaseBaseModel(ABC, PluginVariant, make_plugin_registry("model")): BaseModel class for DPModel backend. """ - def __new__(cls, *args, **kwargs): + def __new__(cls, *args: Any, **kwargs: Any) -> "BaseModel": if inspect.isabstract(cls): # getting model type based on fitting type model_type = kwargs.get("type", "standard") @@ -68,15 +67,15 @@ def get_type_map(self) -> list[str]: """Get the type map.""" @abstractmethod - def get_rcut(self): + def get_rcut(self) -> float: """Get the cut-off radius.""" @abstractmethod - def get_dim_fparam(self): + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" @abstractmethod - def get_dim_aparam(self): + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" @abstractmethod @@ -133,7 +132,7 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": model_def_script: str """The model definition script.""" - min_nbor_dist: Optional[float] + min_nbor_dist: float | None """The minimum distance between two atoms. Used for model compression. None when skipping neighbor statistics. """ @@ -143,7 +142,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" pass - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist @@ -163,9 +162,9 @@ def get_nsel(self) -> int: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/model/dipole_model.py b/deepmd/dpmodel/model/dipole_model.py index 4ca523f79b..d213514551 100644 --- a/deepmd/dpmodel/model/dipole_model.py +++ b/deepmd/dpmodel/model/dipole_model.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later - +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model import ( DPDipoleAtomicModel, @@ -24,8 +26,8 @@ class DipoleModel(DPModelCommon, DPDipoleModel_): def __init__( self, - *args, - **kwargs, - ): + *args: Any, + **kwargs: Any, + ) -> None: DPModelCommon.__init__(self) DPDipoleModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/dpmodel/model/dos_model.py b/deepmd/dpmodel/model/dos_model.py index 3df887b460..5c5d2a5e90 100644 --- a/deepmd/dpmodel/model/dos_model.py +++ b/deepmd/dpmodel/model/dos_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model import ( DPDOSAtomicModel, @@ -23,8 +26,8 @@ class DOSModel(DPModelCommon, DPDOSModel_): def __init__( self, - *args, - **kwargs, - ): + *args: Any, + **kwargs: Any, + ) -> None: DPModelCommon.__init__(self) DPDOSModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 769bba0b20..063533f2a7 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -1,13 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) - from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.dpmodel.fitting.base_fitting import ( + BaseFitting, +) from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -19,9 +18,9 @@ class DPModelCommon: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -45,3 +44,7 @@ def update_sel( train_data, type_map, local_jdata["descriptor"] ) return local_jdata_cpy, min_nbor_dist + + def get_fitting_net(self) -> BaseFitting: + """Get the fitting network.""" + return self.atomic_model.fitting diff --git a/deepmd/dpmodel/model/dp_zbl_model.py b/deepmd/dpmodel/model/dp_zbl_model.py index 7bf22dfc6b..b5940f4707 100644 --- a/deepmd/dpmodel/model/dp_zbl_model.py +++ b/deepmd/dpmodel/model/dp_zbl_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( @@ -29,8 +29,8 @@ class DPZBLModel(DPZBLModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) @@ -38,9 +38,9 @@ def __init__( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/dpmodel/model/ener_model.py b/deepmd/dpmodel/model/ener_model.py index 88e65a849a..9d38a17513 100644 --- a/deepmd/dpmodel/model/ener_model.py +++ b/deepmd/dpmodel/model/ener_model.py @@ -2,6 +2,9 @@ from copy import ( deepcopy, ) +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model import ( DPEnergyAtomicModel, @@ -27,15 +30,15 @@ class EnergyModel(DPModelCommon, DPEnergyModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPEnergyModel_.__init__(self, *args, **kwargs) self._enable_hessian = False self.hess_fitting_def = None - def enable_hessian(self): + def enable_hessian(self) -> None: self.hess_fitting_def = deepcopy(self.atomic_output_def()) self.hess_fitting_def["energy"].r_hessian = True self._enable_hessian = True diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index ec0b986394..3a88aac1e4 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, +) +from typing import ( + Any, ) import array_api_compat import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.atomic_model.base_atomic_model import ( BaseAtomicModel, ) @@ -47,23 +52,23 @@ def model_call_from_call_lower( np.ndarray, np.ndarray, np.ndarray, - Optional[np.ndarray], - Optional[np.ndarray], + np.ndarray | None, + np.ndarray | None, bool, ], - dict[str, np.ndarray], + dict[str, Array], ], rcut: float, sel: list[int], mixed_types: bool, model_output_def: ModelOutputDef, - coord: np.ndarray, - atype: np.ndarray, - box: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + coord: Array, + atype: Array, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, -): +) -> dict[str, Array]: """Return model prediction from lower interface. Parameters @@ -131,7 +136,7 @@ def model_call_from_call_lower( return model_predict -def make_model(T_AtomicModel: type[BaseAtomicModel]): +def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: """Make a model as a derived class of an atomic model. The model provide two interfaces. @@ -157,10 +162,10 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]): class CM(NativeOP, BaseModel): def __init__( self, - *args, + *args: Any, # underscore to prevent conflict with normal inputs - atomic_model_: Optional[T_AtomicModel] = None, - **kwargs, + atomic_model_: T_AtomicModel | None = None, + **kwargs: Any, ) -> None: BaseModel.__init__(self) if atomic_model_ is not None: @@ -173,7 +178,7 @@ def __init__( self.global_np_float_precision = GLOBAL_NP_FLOAT_PRECISION self.global_ener_float_precision = GLOBAL_ENER_FLOAT_PRECISION - def model_output_def(self): + def model_output_def(self) -> ModelOutputDef: """Get the output def for the model.""" return ModelOutputDef(self.atomic_output_def()) @@ -218,13 +223,13 @@ def enable_compression( def call( self, - coord, - atype, - box: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + coord: Array, + atype: Array, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, - ) -> dict[str, np.ndarray]: + ) -> dict[str, Array]: """Return model prediction. Parameters @@ -272,14 +277,14 @@ def call( def call_lower( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, Array]: """Return model prediction. Lower interface that takes extended atomic coordinates and types, nlist, and mapping as input, and returns the predictions on the extended region. @@ -334,14 +339,14 @@ def call_lower( def forward_common_atomic( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + extended_coord: Array, + extended_atype: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, Array]: atomic_ret = self.atomic_model.forward_common_atomic( extended_coord, extended_atype, @@ -355,29 +360,24 @@ def forward_common_atomic( self.atomic_output_def(), extended_coord, do_atomic_virial=do_atomic_virial, + mask=atomic_ret["mask"] if "mask" in atomic_ret else None, ) forward_lower = call_lower def input_type_cast( self, - coord: np.ndarray, - box: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - ) -> tuple[ - np.ndarray, - Optional[np.ndarray], - Optional[np.ndarray], - Optional[np.ndarray], - str, - ]: + coord: Array, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, + ) -> tuple[Array, Array, np.ndarray | None, np.ndarray | None, str]: """Cast the input data to global float type.""" input_prec = RESERVED_PRECISION_DICT[self.precision_dict[coord.dtype.name]] ### ### type checking would not pass jit, convert to coord prec anyway ### - _lst: list[Optional[np.ndarray]] = [ + _lst: list[np.ndarray | None] = [ vv.astype(coord.dtype) if vv is not None else None for vv in [box, fparam, aparam] ] @@ -396,9 +396,9 @@ def input_type_cast( def output_type_cast( self, - model_ret: dict[str, np.ndarray], + model_ret: dict[str, Array], input_prec: str, - ) -> dict[str, np.ndarray]: + ) -> dict[str, Array]: """Convert the model output to the input prec.""" do_cast = ( input_prec != RESERVED_PRECISION_DICT[self.global_np_float_precision] @@ -423,11 +423,11 @@ def output_type_cast( def format_nlist( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - nlist: np.ndarray, + extended_coord: Array, + extended_atype: Array, + nlist: Array, extra_nlist_sort: bool = False, - ): + ) -> Array: """Format the neighbor list. 1. If the number of neighbors in the `nlist` is equal to sum(self.sel), @@ -475,11 +475,11 @@ def format_nlist( def _format_nlist( self, - extended_coord: np.ndarray, - nlist: np.ndarray, + extended_coord: Array, + nlist: Array, nnei: int, extra_nlist_sort: bool = False, - ): + ) -> Array: xp = array_api_compat.array_namespace(extended_coord, nlist) n_nf, n_nloc, n_nnei = nlist.shape extended_coord = extended_coord.reshape([n_nf, -1, 3]) @@ -521,7 +521,7 @@ def _format_nlist( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -530,7 +530,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. @@ -538,7 +538,7 @@ def do_grad_c( return self.atomic_model.do_grad_c(var_name) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -549,10 +549,10 @@ def serialize(self) -> dict: return self.atomic_model.serialize() @classmethod - def deserialize(cls, data) -> "CM": + def deserialize(cls, data: dict) -> "CM": return cls(atomic_model_=T_AtomicModel.deserialize(data)) - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: self.atomic_model.set_case_embd(case_idx) def get_dim_fparam(self) -> int: @@ -563,6 +563,10 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.atomic_model.get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.atomic_model.has_default_fparam() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py index 1d18b70e8e..339998aa89 100644 --- a/deepmd/dpmodel/model/model.py +++ b/deepmd/dpmodel/model/model.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, @@ -45,7 +48,9 @@ ) -def _get_standard_model_components(data, ntypes): +def _get_standard_model_components( + data: dict[str, Any], ntypes: int +) -> tuple[BaseDescriptor, BaseFitting, str]: # descriptor data["descriptor"]["ntypes"] = ntypes data["descriptor"]["type_map"] = copy.deepcopy(data["type_map"]) @@ -181,7 +186,7 @@ def get_spin_model(data: dict) -> SpinModel: return SpinModel(backbone_model=backbone_model, spin=spin) -def get_model(data: dict): +def get_model(data: dict) -> BaseModel: """Get a model from a dictionary. Parameters diff --git a/deepmd/dpmodel/model/polar_model.py b/deepmd/dpmodel/model/polar_model.py index 994b3556c2..b898eababd 100644 --- a/deepmd/dpmodel/model/polar_model.py +++ b/deepmd/dpmodel/model/polar_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) from deepmd.dpmodel.atomic_model import ( DPPolarAtomicModel, @@ -23,8 +26,8 @@ class PolarModel(DPModelCommon, DPPolarModel_): def __init__( self, - *args, - **kwargs, - ): + *args: Any, + **kwargs: Any, + ) -> None: DPModelCommon.__init__(self) DPPolarModel_.__init__(self, *args, **kwargs) diff --git a/deepmd/dpmodel/model/property_model.py b/deepmd/dpmodel/model/property_model.py index 9bd07bd349..20c884cd20 100644 --- a/deepmd/dpmodel/model/property_model.py +++ b/deepmd/dpmodel/model/property_model.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + from deepmd.dpmodel.atomic_model import ( DPPropertyAtomicModel, ) @@ -20,8 +24,12 @@ class PropertyModel(DPModelCommon, DPPropertyModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPPropertyModel_.__init__(self, *args, **kwargs) + + def get_var_name(self) -> str: + """Get the name of the property.""" + return self.get_fitting_net().var_name diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index d149d427e0..521978bdde 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.atomic_model.dp_atomic_model import ( DPAtomicModel, ) @@ -27,7 +30,7 @@ class SpinModel(NativeOP): def __init__( self, - backbone_model, + backbone_model: DPAtomicModel, spin: Spin, ) -> None: super().__init__() @@ -37,7 +40,9 @@ def __init__( self.virtual_scale_mask = self.spin.get_virtual_scale_mask() self.spin_mask = self.spin.get_spin_mask() - def process_spin_input(self, coord, atype, spin): + def process_spin_input( + self, coord: Array, atype: Array, spin: Array + ) -> tuple[Array, Array]: """Generate virtual coordinates and types, concat into the input.""" nframes, nloc = coord.shape[:-1] atype_spin = np.concatenate([atype, atype + self.ntypes_real], axis=-1) @@ -49,12 +54,12 @@ def process_spin_input(self, coord, atype, spin): def process_spin_input_lower( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - extended_spin: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - ): + extended_coord: Array, + extended_atype: Array, + extended_spin: Array, + nlist: Array, + mapping: Array | None = None, + ) -> tuple[Array, Array]: """ Add `extended_spin` into `extended_coord` to generate virtual atoms, and extend `nlist` and `mapping`. Note that the final `extended_coord_updated` with shape [nframes, nall + nall, 3] has the following order: @@ -92,8 +97,12 @@ def process_spin_input_lower( ) def process_spin_output( - self, atype, out_tensor, add_mag: bool = True, virtual_scale: bool = True - ): + self, + atype: Array, + out_tensor: Array, + add_mag: bool = True, + virtual_scale: bool = True, + ) -> tuple[Array, Array]: """Split the output both real and virtual atoms, and scale the latter.""" nframes, nloc_double = out_tensor.shape[:2] nloc = nloc_double // 2 @@ -112,12 +121,12 @@ def process_spin_output( def process_spin_output_lower( self, - extended_atype, - extended_out_tensor, + extended_atype: Array, + extended_out_tensor: Array, nloc: int, add_mag: bool = True, virtual_scale: bool = True, - ): + ) -> tuple[Array, Array]: """Split the extended output of both real and virtual atoms with switch, and scale the latter.""" nframes, nall_double = extended_out_tensor.shape[:2] nall = nall_double // 2 @@ -148,7 +157,7 @@ def process_spin_output_lower( return extended_out_real, extended_out_mag, atomic_mask > 0.0 @staticmethod - def extend_nlist(extended_atype, nlist): + def extend_nlist(extended_atype: Array, nlist: Array) -> Array: nframes, nloc, nnei = nlist.shape nall = extended_atype.shape[1] nlist_mask = nlist != -1 @@ -178,7 +187,9 @@ def extend_nlist(extended_atype, nlist): return extended_nlist @staticmethod - def concat_switch_virtual(extended_tensor, extended_tensor_virtual, nloc: int): + def concat_switch_virtual( + extended_tensor: Array, extended_tensor_virtual: Array, nloc: int + ) -> Array: nframes, nall = extended_tensor.shape[:2] out_shape = list(extended_tensor.shape) out_shape[1] *= 2 @@ -197,7 +208,7 @@ def concat_switch_virtual(extended_tensor, extended_tensor_virtual, nloc: int): return extended_tensor_updated.reshape(out_shape) @staticmethod - def expand_aparam(aparam, nloc: int): + def expand_aparam(aparam: Array, nloc: int) -> Array: """Expand the atom parameters for virtual atoms if necessary.""" nframes, natom, numb_aparam = aparam.shape if natom == nloc: # good @@ -226,19 +237,19 @@ def get_type_map(self) -> list[str]: ntypes = len(tmap) // 2 # ignore the virtual type return tmap[:ntypes] - def get_ntypes(self): + def get_ntypes(self) -> int: """Returns the number of element types.""" return len(self.get_type_map()) - def get_rcut(self): + def get_rcut(self) -> float: """Get the cut-off radius.""" return self.backbone_model.get_rcut() - def get_dim_fparam(self): + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.backbone_model.get_dim_fparam() - def get_dim_aparam(self): + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.backbone_model.get_dim_aparam() @@ -264,7 +275,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.backbone_model.get_min_nbor_dist() @@ -288,7 +299,7 @@ def has_spin() -> bool: """Returns whether it has spin input and output.""" return True - def model_output_def(self): + def model_output_def(self) -> ModelOutputDef: """Get the output def for the model.""" model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: @@ -298,7 +309,7 @@ def model_output_def(self): backbone_model_atomic_output_def[var_name].magnetic = True return ModelOutputDef(backbone_model_atomic_output_def) - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: """Get attribute from the wrapped model.""" if name in self.__dict__: return self.__dict__[name] @@ -312,7 +323,7 @@ def serialize(self) -> dict: } @classmethod - def deserialize(cls, data) -> "SpinModel": + def deserialize(cls, data: dict) -> "SpinModel": backbone_model_obj = make_model(DPAtomicModel).deserialize( data["backbone_model"] ) @@ -324,14 +335,14 @@ def deserialize(cls, data) -> "SpinModel": def call( self, - coord, - atype, - spin, - box: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + coord: Array, + atype: Array, + spin: Array, + box: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, - ) -> dict[str, np.ndarray]: + ) -> dict[str, Array]: """Return model prediction. Parameters @@ -386,15 +397,15 @@ def call( def call_lower( self, - extended_coord: np.ndarray, - extended_atype: np.ndarray, - extended_spin: np.ndarray, - nlist: np.ndarray, - mapping: Optional[np.ndarray] = None, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + extended_coord: Array, + extended_atype: Array, + extended_spin: Array, + nlist: Array, + mapping: Array | None = None, + fparam: Array | None = None, + aparam: Array | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, Array]: """Return model prediction. Lower interface that takes extended atomic coordinates, types and spins, nlist, and mapping as input, and returns the predictions on the extended region. diff --git a/deepmd/dpmodel/model/transform_output.py b/deepmd/dpmodel/model/transform_output.py index 9d7873f081..d3315eda55 100644 --- a/deepmd/dpmodel/model/transform_output.py +++ b/deepmd/dpmodel/model/transform_output.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later + import array_api_compat import numpy as np from deepmd.dpmodel.array_api import ( + Array, xp_scatter_sum, ) from deepmd.dpmodel.common import ( @@ -20,11 +22,12 @@ def fit_output_to_model_output( - fit_ret: dict[str, np.ndarray], + fit_ret: dict[str, Array], fit_output_def: FittingOutputDef, - coord_ext: np.ndarray, + coord_ext: Array, do_atomic_virial: bool = False, -) -> dict[str, np.ndarray]: + mask: Array | None = None, +) -> dict[str, Array]: """Transform the output of the fitting network to the model output. @@ -38,9 +41,19 @@ def fit_output_to_model_output( if vdef.reducible: kk_redu = get_reduce_name(kk) # cast to energy prec before reduction - model_ret[kk_redu] = xp.sum( - vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis - ) + if vdef.intensive: + if mask is not None: + model_ret[kk_redu] = xp.sum( + vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis + ) / np.sum(mask, axis=-1, keepdims=True) + else: + model_ret[kk_redu] = xp.mean( + vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis + ) + else: + model_ret[kk_redu] = xp.sum( + vv.astype(GLOBAL_ENER_FLOAT_PRECISION), axis=atom_axis + ) if vdef.r_differentiable: kk_derv_r, kk_derv_c = get_deriv_name(kk) # name-holders @@ -53,14 +66,14 @@ def fit_output_to_model_output( def get_leading_dims( - vv: np.ndarray, + vv: Array, vdef: OutputVariableDef, -): +) -> list[int]: """Get the dimensions of nf x nloc. Parameters ---------- - vv : np.ndarray + vv : Array The input array from which to compute the leading dimensions. vdef : OutputVariableDef The output variable definition containing the shape to exclude from `vv`. @@ -75,11 +88,11 @@ def get_leading_dims( def communicate_extended_output( - model_ret: dict[str, np.ndarray], + model_ret: dict[str, Array], model_output_def: ModelOutputDef, - mapping: np.ndarray, # nf x nloc + mapping: Array, # nf x nloc do_atomic_virial: bool = False, -) -> dict[str, np.ndarray]: +) -> dict[str, Array]: """Transform the output of the model network defined on local and ghost (extended) atoms to local atoms. @@ -100,7 +113,9 @@ def communicate_extended_output( if vdef.r_differentiable: if model_ret[kk_derv_r] is not None: derv_r_ext_dims = list(vdef.shape) + [3] # noqa:RUF005 - mapping = xp.reshape(mapping, (mldims + [1] * len(derv_r_ext_dims))) + mapping = xp.reshape( + mapping, tuple(mldims + [1] * len(derv_r_ext_dims)) + ) mapping = xp.tile(mapping, [1] * len(mldims) + derv_r_ext_dims) force = xp.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype) force = xp_scatter_sum( diff --git a/deepmd/dpmodel/modifier/base_modifier.py b/deepmd/dpmodel/modifier/base_modifier.py index 9edc4722e1..febb9b75e8 100644 --- a/deepmd/dpmodel/modifier/base_modifier.py +++ b/deepmd/dpmodel/modifier/base_modifier.py @@ -4,6 +4,9 @@ ABC, abstractmethod, ) +from typing import ( + Any, +) from deepmd.utils.plugin import ( PluginVariant, @@ -15,7 +18,7 @@ def make_base_modifier() -> type[object]: class BaseModifier(ABC, PluginVariant, make_plugin_registry("modifier")): """Base class for data modifier.""" - def __new__(cls, *args, **kwargs): + def __new__(cls, *args: Any, **kwargs: Any) -> "BaseModifier": if cls is BaseModifier: cls = cls.get_class_by_type(kwargs["type"]) return super().__new__(cls) diff --git a/deepmd/dpmodel/output_def.py b/deepmd/dpmodel/output_def.py index c2a1147786..5028bc43a3 100644 --- a/deepmd/dpmodel/output_def.py +++ b/deepmd/dpmodel/output_def.py @@ -3,6 +3,9 @@ from enum import ( IntEnum, ) +from typing import ( + Any, +) def check_shape( @@ -19,7 +22,7 @@ def check_shape( raise ValueError(f"{shape} shape not matching def {def_shape}") -def check_var(var, var_def) -> None: +def check_var(var: Any, var_def: Any) -> None: if var_def.atomic: # var.shape == [nf, nloc, *var_def.shape] if len(var.shape) != len(var_def.shape) + 2: @@ -32,7 +35,7 @@ def check_var(var, var_def) -> None: check_shape(list(var.shape[1:]), var_def.shape) -def model_check_output(cls): +def model_check_output(cls: type) -> type: """Check if the output of the Model is consistent with the definition. Two methods are assumed to be provided by the Model: @@ -45,17 +48,17 @@ def model_check_output(cls): class wrapper(cls): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) self.md = self.output_def() def __call__( self, - *args, - **kwargs, - ): + *args: Any, + **kwargs: Any, + ) -> Any: ret = cls.__call__(self, *args, **kwargs) for kk in self.md.keys_outp(): dd = self.md[kk] @@ -74,7 +77,7 @@ def __call__( return wrapper -def fitting_check_output(cls): +def fitting_check_output(cls: type) -> type: """Check if the output of the Fitting is consistent with the definition. Two methods are assumed to be provided by the Fitting: @@ -87,17 +90,17 @@ def fitting_check_output(cls): class wrapper(cls): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) self.md = self.output_def() def __call__( self, - *args, - **kwargs, - ): + *args: Any, + **kwargs: Any, + ) -> Any: ret = cls.__call__(self, *args, **kwargs) for kk in self.md.keys(): dd = self.md[kk] @@ -227,10 +230,10 @@ def __init__( raise ValueError("only r_differentiable variable can calculate hessian") @property - def size(self): + def size(self) -> int: return self.output_size - def squeeze(self, dim) -> None: + def squeeze(self, dim: int) -> None: # squeeze the shape on given dimension if -len(self.shape) <= dim < len(self.shape) and self.shape[dim] == 1: self.shape.pop(dim) @@ -264,7 +267,7 @@ def __getitem__( def get_data(self) -> dict[str, OutputVariableDef]: return self.var_defs - def keys(self): + def keys(self): # noqa: ANN201 return self.var_defs.keys() @@ -316,25 +319,25 @@ def get_data( ) -> dict[str, OutputVariableDef]: return self.var_defs - def keys(self): + def keys(self): # noqa: ANN201 return self.var_defs.keys() - def keys_outp(self): + def keys_outp(self): # noqa: ANN201 return self.def_outp.keys() - def keys_redu(self): + def keys_redu(self): # noqa: ANN201 return self.def_redu.keys() - def keys_derv_r(self): + def keys_derv_r(self): # noqa: ANN201 return self.def_derv_r.keys() - def keys_hess_r(self): + def keys_hess_r(self): # noqa: ANN201 return self.def_hess_r.keys() - def keys_derv_c(self): + def keys_derv_c(self): # noqa: ANN201 return self.def_derv_c.keys() - def keys_derv_c_redu(self): + def keys_derv_c_redu(self): # noqa: ANN201 return self.def_derv_c_redu.keys() diff --git a/deepmd/dpmodel/utils/env_mat.py b/deepmd/dpmodel/utils/env_mat.py index ee11678d3a..f2943d5e44 100644 --- a/deepmd/dpmodel/utils/env_mat.py +++ b/deepmd/dpmodel/utils/env_mat.py @@ -1,15 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import array_api_compat -import numpy as np from deepmd.dpmodel import ( NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, support_array_api, xp_take_along_axis, ) @@ -20,10 +20,10 @@ @support_array_api(version="2023.12") def compute_smooth_weight( - distance: np.ndarray, + distance: Array, rmin: float, rmax: float, -): +) -> Array: """Compute smooth weight for descriptor elements.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") @@ -37,10 +37,10 @@ def compute_smooth_weight( @support_array_api(version="2023.12") def compute_exp_sw( - distance: np.ndarray, + distance: Array, rmin: float, rmax: float, -): +) -> Array: """Compute the exponential switch function for neighbor update.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") @@ -54,14 +54,14 @@ def compute_exp_sw( def _make_env_mat( - nlist, - coord, + nlist: Any, + coord: Any, rcut: float, ruct_smth: float, radial_only: bool = False, protection: float = 0.0, use_exp_switch: bool = False, -): +) -> tuple[Any, Any, Any]: """Make smooth environment matrix.""" xp = array_api_compat.array_namespace(nlist) nf, nloc, nnei = nlist.shape @@ -101,8 +101,8 @@ def _make_env_mat( class EnvMat(NativeOP): def __init__( self, - rcut, - rcut_smth, + rcut: float, + rcut_smth: float, protection: float = 0.0, use_exp_switch: bool = False, ) -> None: @@ -113,13 +113,13 @@ def __init__( def call( self, - coord_ext: np.ndarray, - atype_ext: np.ndarray, - nlist: np.ndarray, - davg: Optional[np.ndarray] = None, - dstd: Optional[np.ndarray] = None, + coord_ext: Array, + atype_ext: Array, + nlist: Array, + davg: Array | None = None, + dstd: Array | None = None, radial_only: bool = False, - ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + ) -> tuple[Array, Array, Array]: """Compute the environment matrix. Parameters @@ -159,7 +159,9 @@ def call( em /= xp.reshape(xp.take(dstd, xp.reshape(atype, (-1,)), axis=0), em.shape) return em, diff, sw - def _call(self, nlist, coord_ext, radial_only): + def _call( + self, nlist: Any, coord_ext: Any, radial_only: bool + ) -> tuple[Any, Any, Any]: em, diff, ww = _make_env_mat( nlist, coord_ext, diff --git a/deepmd/dpmodel/utils/env_mat_stat.py b/deepmd/dpmodel/utils/env_mat_stat.py index e25739fa56..b8befa0087 100644 --- a/deepmd/dpmodel/utils/env_mat_stat.py +++ b/deepmd/dpmodel/utils/env_mat_stat.py @@ -13,6 +13,9 @@ from deepmd.common import ( get_hash, ) +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( get_xp_precision, ) @@ -38,12 +41,12 @@ class EnvMatStat(BaseEnvMatStat): - def compute_stat(self, env_mat: dict[str, np.ndarray]) -> dict[str, StatItem]: + def compute_stat(self, env_mat: dict[str, Array]) -> dict[str, StatItem]: """Compute the statistics of the environment matrix for a single system. Parameters ---------- - env_mat : np.ndarray + env_mat : Array The environment matrix. Returns @@ -79,7 +82,7 @@ def __init__(self, descriptor: Union["Descriptor", "DescriptorBlock"]) -> None: ) # se_r=1, se_a=4 def iter( - self, data: list[dict[str, Union[np.ndarray, list[tuple[int, int]]]]] + self, data: list[dict[str, np.ndarray | list[tuple[int, int]]]] ) -> Iterator[dict[str, StatItem]]: """Get the iterator of the environment matrix. @@ -166,7 +169,7 @@ def iter( self.last_dim, ), ) - atype = xp.reshape(atype, (coord.shape[0] * coord.shape[1])) + atype = xp.reshape(atype, (coord.shape[0] * coord.shape[1],)) # (1, nloc) eq (ntypes, 1), so broadcast is possible # shape: (ntypes, nloc) type_idx = xp.equal( @@ -189,7 +192,7 @@ def iter( for type_i in range(self.descriptor.get_ntypes()): dd = env_mat[type_idx[type_i, ...]] dd = xp.reshape( - dd, [-1, self.last_dim] + dd, (-1, self.last_dim) ) # typen_atoms * unmasked_nnei, 4 env_mats = {} env_mats[f"r_{type_i}"] = dd[:, :1] @@ -218,7 +221,7 @@ def get_hash(self) -> str: } ) - def __call__(self): + def __call__(self) -> tuple[Array, Array]: avgs = self.get_avg() stds = self.get_std() diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py index f390bbc7c1..9d8f0c8572 100644 --- a/deepmd/dpmodel/utils/exclude_mask.py +++ b/deepmd/dpmodel/utils/exclude_mask.py @@ -4,6 +4,7 @@ import numpy as np from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) @@ -25,16 +26,16 @@ def __init__( # (ntypes) self.type_mask = type_mask.reshape([-1]) - def get_exclude_types(self): + def get_exclude_types(self) -> list[int]: return self.exclude_types - def get_type_mask(self): + def get_type_mask(self) -> Array: return self.type_mask def build_type_exclude_mask( self, - atype: np.ndarray, - ): + atype: Array, + ) -> Array: """Compute type exclusion mask for atoms. Parameters @@ -53,7 +54,7 @@ def build_type_exclude_mask( xp = array_api_compat.array_namespace(atype) nf, natom = atype.shape return xp.reshape( - xp.take(self.type_mask[...], xp.reshape(atype, [-1]), axis=0), + xp.take(self.type_mask[...], xp.reshape(atype, (-1,)), axis=0), (nf, natom), ) @@ -86,14 +87,14 @@ def __init__( # (ntypes+1 x ntypes+1) self.type_mask = type_mask.reshape([-1]) - def get_exclude_types(self): + def get_exclude_types(self) -> list[tuple[int, int]]: return self.exclude_types def build_type_exclude_mask( self, - nlist: np.ndarray, - atype_ext: np.ndarray, - ): + nlist: Array, + atype_ext: Array, + ) -> Array: """Compute type exclusion mask for atom pairs. Parameters @@ -137,5 +138,5 @@ def build_type_exclude_mask( ) return mask - def __contains__(self, item) -> bool: + def __contains__(self, item: tuple[int, int]) -> bool: return item in self.exclude_types diff --git a/deepmd/dpmodel/utils/learning_rate.py b/deepmd/dpmodel/utils/learning_rate.py index 90c18fca22..10f7ec8d04 100644 --- a/deepmd/dpmodel/utils/learning_rate.py +++ b/deepmd/dpmodel/utils/learning_rate.py @@ -1,16 +1,20 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + import numpy as np class LearningRateExp: def __init__( self, - start_lr, - stop_lr, - decay_steps, - stop_steps, - decay_rate=None, - **kwargs, + start_lr: float, + stop_lr: float, + decay_steps: int, + stop_steps: int, + decay_rate: float | None = None, + **kwargs: Any, ) -> None: """ Construct an exponential-decayed learning rate. @@ -45,7 +49,7 @@ def __init__( self.decay_rate = decay_rate self.min_lr = stop_lr - def value(self, step) -> np.float64: + def value(self, step: int) -> np.float64: """Get the learning rate at the given step.""" step_lr = self.start_lr * np.power(self.decay_rate, step // self.decay_steps) if step_lr < self.min_lr: diff --git a/deepmd/dpmodel/utils/neighbor_stat.py b/deepmd/dpmodel/utils/neighbor_stat.py index 3aea8ceeb9..1bcc894624 100644 --- a/deepmd/dpmodel/utils/neighbor_stat.py +++ b/deepmd/dpmodel/utils/neighbor_stat.py @@ -2,13 +2,13 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import array_api_compat import numpy as np +from deepmd.dpmodel.array_api import ( + Array, +) from deepmd.dpmodel.common import ( NativeOP, ) @@ -46,10 +46,10 @@ def __init__( def call( self, - coord: np.ndarray, - atype: np.ndarray, - cell: Optional[np.ndarray], - ) -> tuple[float, np.ndarray]: + coord: Array, + atype: Array, + cell: Array | None, + ) -> tuple[Array, Array]: """Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix. @@ -82,8 +82,8 @@ def call( nall = coord1.shape[1] // 3 coord0 = coord1[:, : nloc * 3] diff = ( - xp.reshape(coord1, [nframes, -1, 3])[:, None, :, :] - - xp.reshape(coord0, [nframes, -1, 3])[:, :, None, :] + xp.reshape(coord1, (nframes, -1, 3))[:, None, :, :] + - xp.reshape(coord0, (nframes, -1, 3))[:, :, None, :] ) assert list(diff.shape) == [nframes, nloc, nall, 3] # remove the diagonal elements diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index bf28b66b7b..e712adfdd8 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -5,11 +5,12 @@ """ import itertools -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, ClassVar, - Optional, - Union, ) import array_api_compat @@ -21,6 +22,7 @@ NativeOP, ) from deepmd.dpmodel.array_api import ( + Array, support_array_api, xp_add_at, xp_bincount, @@ -36,7 +38,7 @@ ) -def sigmoid_t(x: np.ndarray) -> np.ndarray: +def sigmoid_t(x): # noqa: ANN001, ANN201 """Sigmoid.""" if array_api_compat.is_jax_array(x): from deepmd.jax.env import ( @@ -53,7 +55,7 @@ class Identity(NativeOP): def __init__(self) -> None: super().__init__() - def call(self, x: np.ndarray) -> np.ndarray: + def call(self, x): # noqa: ANN001, ANN201 """The Identity operation layer.""" return x @@ -73,11 +75,11 @@ class NativeLayer(NativeOP): Parameters ---------- - w : np.ndarray, optional + w : Array, optional The weights of the layer. - b : np.ndarray, optional + b : Array, optional The biases of the layer. - idt : np.ndarray, optional + idt : Array, optional The identity matrix of the layer. activation_function : str, optional The activation function of the layer. @@ -87,19 +89,24 @@ class NativeLayer(NativeOP): The precision of the layer. seed : int, optional Random seed. + trainable : bool, default=True + Whether the layer is trainable. """ def __init__( self, - num_in, - num_out, + num_in: int, + num_out: int, bias: bool = True, use_timestep: bool = False, - activation_function: Optional[str] = None, + activation_function: str | None = None, resnet: bool = False, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: + # trainable must be set before any array attribute is set + self.trainable = trainable prec = PRECISION_DICT[precision.lower()] self.precision = precision # only use_timestep when skip connection is established. @@ -139,13 +146,14 @@ def serialize(self) -> dict: } return { "@class": "Layer", - "@version": 1, + "@version": 2, "bias": self.b is not None, "use_timestep": self.idt is not None, "activation_function": self.activation_function, "resnet": self.resnet, # make deterministic "precision": np.dtype(PRECISION_DICT[self.precision]).name, + "trainable": self.trainable, "@variables": data, } @@ -159,7 +167,7 @@ def deserialize(cls, data: dict) -> "NativeLayer": The dict to deserialize from. """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 1, 1) + check_version_compatibility(data.pop("@version", 1), 2, 1) data.pop("@class", None) variables = data.pop("@variables") assert variables["w"] is not None and len(variables["w"].shape) == 2 @@ -199,7 +207,7 @@ def check_shape_consistency(self) -> None: def check_type_consistency(self) -> None: precision = self.precision - def check_var(var) -> None: + def check_var(var: Array | None) -> None: if var is not None: # array api standard doesn't provide a API to get the dtype name # this is really hacked @@ -211,7 +219,7 @@ def check_var(var) -> None: check_var(self.b) check_var(self.idt) - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("w", "matrix"): self.w = value elif key in ("b", "bias"): @@ -227,7 +235,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("w", "matrix"): return self.w elif key in ("b", "bias"): @@ -240,6 +248,8 @@ def __getitem__(self, key): return self.resnet elif key == "precision": return self.precision + elif key == "trainable": + return self.trainable else: raise KeyError(key) @@ -250,12 +260,12 @@ def dim_out(self) -> int: return self.w.shape[1] @support_array_api(version="2022.12") - def call(self, x: np.ndarray) -> np.ndarray: + def call(self, x): # noqa: ANN001, ANN201 """Forward pass. Parameters ---------- - x : np.ndarray + x : Array The input. Returns @@ -291,14 +301,14 @@ def get_activation_fn(activation_function: str) -> Callable[[np.ndarray], np.nda activation_function = activation_function.lower() if activation_function == "tanh": - def fn(x): + def fn(x): # noqa: ANN001, ANN202 # noqa: ANN001, ANN202 xp = array_api_compat.array_namespace(x) return xp.tanh(x) return fn elif activation_function == "relu": - def fn(x): + def fn(x): # noqa: ANN001, ANN202 xp = array_api_compat.array_namespace(x) # https://stackoverflow.com/a/47936476/9567349 return x * xp.astype(x > 0, x.dtype) @@ -306,7 +316,7 @@ def fn(x): return fn elif activation_function in ("gelu", "gelu_tf"): - def fn(x): + def fn(x): # noqa: ANN001, ANN202 xp = array_api_compat.array_namespace(x) # generated by GitHub Copilot return ( @@ -318,7 +328,7 @@ def fn(x): return fn elif activation_function == "relu6": - def fn(x): + def fn(x): # noqa: ANN001, ANN202 xp = array_api_compat.array_namespace(x) # generated by GitHub Copilot return xp.where( @@ -328,7 +338,7 @@ def fn(x): return fn elif activation_function == "softplus": - def fn(x): + def fn(x): # noqa: ANN001, ANN202 xp = array_api_compat.array_namespace(x) # generated by GitHub Copilot return xp.log(1 + xp.exp(x)) @@ -336,14 +346,14 @@ def fn(x): return fn elif activation_function == "sigmoid": - def fn(x): + def fn(x): # noqa: ANN001, ANN202 # generated by GitHub Copilot return sigmoid_t(x) return fn elif activation_function == "silu": - def fn(x): + def fn(x): # noqa: ANN001, ANN202 # generated by GitHub Copilot return x * sigmoid_t(x) @@ -352,13 +362,13 @@ def fn(x): "custom_silu" ): - def sigmoid(x): + def sigmoid(x): # noqa: ANN001, ANN202 return 1 / (1 + np.exp(-x)) - def silu(x): + def silu(x): # noqa: ANN001, ANN202 return x * sigmoid(x) - def silu_grad(x): + def silu_grad(x): # noqa: ANN001, ANN202 sig = sigmoid(x) return sig + x * sig * (1 - sig) @@ -370,7 +380,7 @@ def silu_grad(x): slope = float(silu_grad(threshold)) const = float(silu(threshold)) - def fn(x): + def fn(x): # noqa: ANN001, ANN202 xp = array_api_compat.array_namespace(x) return xp.where( x < threshold, @@ -381,7 +391,7 @@ def fn(x): return fn elif activation_function.lower() in ("none", "linear"): - def fn(x): + def fn(x): # noqa: ANN001, ANN202 return x return fn @@ -415,7 +425,7 @@ def __init__( uni_init: bool = True, trainable: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: self.eps = eps self.uni_init = uni_init @@ -429,6 +439,7 @@ def __init__( resnet=False, precision=precision, seed=seed, + trainable=trainable, ) xp = array_api_compat.array_namespace(self.w, self.b) self.w = xp.squeeze(self.w, 0) # keep the weight shape to be [num_in] @@ -493,7 +504,7 @@ def _check_shape_consistency(self) -> None: f"of b {self.b.shape[0]}", ) - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("w", "matrix"): self.w = value elif key in ("b", "bias"): @@ -507,7 +518,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("w", "matrix"): return self.w elif key in ("b", "bias"): @@ -524,12 +535,12 @@ def __getitem__(self, key): def dim_out(self) -> int: return self.w.shape[0] - def call(self, x: np.ndarray) -> np.ndarray: + def call(self, x): # noqa: ANN001, ANN201 """Forward pass. Parameters ---------- - x : np.ndarray + x : Array The input. Returns @@ -541,7 +552,13 @@ def call(self, x: np.ndarray) -> np.ndarray: return y @staticmethod - def layer_norm_numpy(x, shape, weight=None, bias=None, eps=1e-5): + def layer_norm_numpy( # noqa: ANN205 + x, # noqa: ANN001 + shape: tuple[int, ...], + weight=None, # noqa: ANN001 + bias=None, # noqa: ANN001 + eps: float = 1e-5, + ): xp = array_api_compat.array_namespace(x) # mean and variance mean = xp.mean(x, axis=tuple(range(-len(shape), 0)), keepdims=True) @@ -554,7 +571,7 @@ def layer_norm_numpy(x, shape, weight=None, bias=None, eps=1e-5): return x_normalized -def make_multilayer_network(T_NetworkLayer, ModuleBase): +def make_multilayer_network(T_NetworkLayer: type, ModuleBase: type) -> type: class NN(ModuleBase): """Native representation of a neural network. @@ -564,7 +581,7 @@ class NN(ModuleBase): The layers of the network. """ - def __init__(self, layers: Optional[list[dict]] = None) -> None: + def __init__(self, layers: list[dict] | None = None) -> None: super().__init__() if layers is None: layers = [] @@ -599,11 +616,11 @@ def deserialize(cls, data: dict) -> "NN": data.pop("@class", None) return cls(data["layers"]) - def __getitem__(self, key): + def __getitem__(self, key: int) -> Any: assert isinstance(key, int) return self.layers[key] - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: int, value: Any) -> None: assert isinstance(key, int) self.layers[key] = value @@ -616,12 +633,12 @@ def check_shape_consistency(self) -> None: f"output {self.layers[ii].dim_out}", ) - def call(self, x): + def call(self, x): # noqa: ANN001, ANN202 """Forward pass. Parameters ---------- - x : np.ndarray + x : Array The input. Returns @@ -633,6 +650,25 @@ def call(self, x): x = layer(x) return x + def call_until_last(self, x): # noqa: ANN001, ANN202 + """Return the output before last layer. + + Parameters + ---------- + x : Array + The input. + + Returns + ------- + np.ndarray + The output before last layer. + """ + # avoid slice (self.layers[:-1]) for jit + for ii, layer in enumerate(self.layers): + if ii < len(self.layers) - 1: + x = layer(x) + return x + def clear(self) -> None: """Clear the network parameters to zero.""" for layer in self.layers: @@ -649,7 +685,7 @@ def clear(self) -> None: NativeNet = make_multilayer_network(NativeLayer, NativeOP) -def make_embedding_network(T_Network, T_NetworkLayer): +def make_embedding_network(T_Network: type, T_NetworkLayer: type) -> type: class EN(T_Network): """The embedding network. @@ -674,16 +710,19 @@ class EN(T_Network): def __init__( self, - in_dim, + in_dim: int, neuron: list[int] = [24, 48, 96], activation_function: str = "tanh", resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, bias: bool = True, + trainable: bool | list[bool] = True, ) -> None: layers = [] i_in = in_dim + if isinstance(trainable, bool): + trainable = [trainable] * len(neuron) for idx, ii in enumerate(neuron): i_ot = ii layers.append( @@ -696,6 +735,7 @@ def __init__( resnet=True, precision=precision, seed=child_seed(seed, idx), + trainable=trainable[idx], ).serialize() ) i_in = i_ot @@ -751,7 +791,9 @@ def deserialize(cls, data: dict) -> "EmbeddingNet": EmbeddingNet = make_embedding_network(NativeNet, NativeLayer) -def make_fitting_network(T_EmbeddingNet, T_Network, T_NetworkLayer): +def make_fitting_network( + T_EmbeddingNet: type, T_Network: type, T_NetworkLayer: type +) -> type: class FN(T_EmbeddingNet): """The fitting network. It may be implemented as an embedding net connected with a linear output layer. @@ -778,15 +820,22 @@ class FN(T_EmbeddingNet): def __init__( self, - in_dim, - out_dim, + in_dim: int, + out_dim: int, neuron: list[int] = [24, 48, 96], activation_function: str = "tanh", resnet_dt: bool = False, precision: str = DEFAULT_PRECISION, bias_out: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool | list[bool] = True, ) -> None: + if trainable is None: + trainable = [True] * (len(neuron) + 1) + elif isinstance(trainable, bool): + trainable = [trainable] * (len(neuron) + 1) + else: + pass super().__init__( in_dim, neuron=neuron, @@ -794,6 +843,7 @@ def __init__( resnet_dt=resnet_dt, precision=precision, seed=seed, + trainable=trainable[:-1], ) i_in = neuron[-1] if len(neuron) > 0 else in_dim i_ot = out_dim @@ -807,6 +857,7 @@ def __init__( resnet=False, precision=precision, seed=child_seed(seed, len(neuron)), + trainable=trainable[-1], ) ) self.out_dim = out_dim @@ -886,7 +937,7 @@ def __init__( ndim: int, ntypes: int, network_type: str = "network", - networks: list[Union[NativeNet, dict]] = [], + networks: list[NativeNet | dict] = [], ) -> None: self.ndim = ndim self.ntypes = ntypes @@ -894,7 +945,7 @@ def __init__( self._networks = [None for ii in range(ntypes**ndim)] for ii, network in enumerate(networks): self[ii] = network - if len(networks): + if len(networks) and all(net is not None for net in networks): self.check_completeness() def check_completeness(self) -> None: @@ -909,7 +960,7 @@ def check_completeness(self) -> None: if self[tuple(tt)] is None: raise RuntimeError(f"network for {tt} not found") - def _convert_key(self, key): + def _convert_key(self, key: int | tuple) -> int: if isinstance(key, int): idx = key else: @@ -924,11 +975,13 @@ def _convert_key(self, key): idx = sum([tt * self.ntypes**ii for ii, tt in enumerate(key)]) return idx - def __getitem__(self, key): + def __getitem__(self, key: int | tuple) -> Any: return self._networks[self._convert_key(key)] - def __setitem__(self, key, value) -> None: - if isinstance(value, self.network_type): + def __setitem__(self, key: int | tuple, value: Any) -> None: + if value is None: + pass + elif isinstance(value, self.network_type): pass elif isinstance(value, dict): value = self.network_type.deserialize(value) @@ -952,7 +1005,9 @@ def serialize(self) -> dict: "ndim": self.ndim, "ntypes": self.ntypes, "network_type": network_type_name, - "networks": [nn.serialize() for nn in self._networks], + "networks": [ + nn.serialize() if nn is not None else None for nn in self._networks + ], } @classmethod @@ -970,11 +1025,11 @@ def deserialize(cls, data: dict) -> "NetworkCollection": return cls(**data) -def aggregate( - data: np.ndarray, - owners: np.ndarray, - average=True, - num_owner=None, +def aggregate( # noqa: ANN201 + data, # noqa: ANN001 + owners, # noqa: ANN001 + average: bool = True, + num_owner: int | None = None, ): """ Aggregate rows in data by specifying the owners. @@ -1010,10 +1065,10 @@ def aggregate( return output -def get_graph_index( - nlist: np.ndarray, - nlist_mask: np.ndarray, - a_nlist_mask: np.ndarray, +def get_graph_index( # noqa: ANN201 + nlist, # noqa: ANN001 + nlist_mask, # noqa: ANN001 + a_nlist_mask, # noqa: ANN001 nall: int, use_loc_mapping: bool = True, ): @@ -1036,12 +1091,12 @@ def get_graph_index( Returns ------- - edge_index : n_edge x 2 + edge_index : 2 x n_edge n2e_index : n_edge Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i). n_ext2e_index : n_edge Broadcast indices from extended node(j) to edge(ij). - angle_index : n_angle x 3 + angle_index : 3 x n_angle n2a_index : n_angle Broadcast indices from extended node(j) to angle(ijk). eij2a_index : n_angle @@ -1111,7 +1166,7 @@ def get_graph_index( # n_angle eik2a_index = edge_index_ik[a_nlist_mask_3d] - edge_index_result = xp.stack([n2e_index, n_ext2e_index], axis=-1) - angle_index_result = xp.stack([n2a_index, eij2a_index, eik2a_index], axis=-1) + edge_index_result = xp.stack([n2e_index, n_ext2e_index], axis=0) + angle_index_result = xp.stack([n2a_index, eij2a_index, eik2a_index], axis=0) return edge_index_result, angle_index_result diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index 4115871f3b..a43cf46403 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -1,13 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import array_api_compat -import numpy as np from deepmd.dpmodel.array_api import ( + Array, xp_take_along_axis, ) @@ -18,13 +14,13 @@ def extend_input_and_build_neighbor_list( - coord, - atype, + coord: Array, + atype: Array, rcut: float, sel: list[int], mixed_types: bool = False, - box: Optional[np.ndarray] = None, -): + box: Array | None = None, +) -> tuple[Array, Array]: xp = array_api_compat.array_namespace(coord, atype) nframes, nloc = atype.shape[:2] if box is not None: @@ -51,20 +47,20 @@ def extend_input_and_build_neighbor_list( ## translated from torch implementation by chatgpt def build_neighbor_list( - coord: np.ndarray, - atype: np.ndarray, + coord: Array, + atype: Array, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, -) -> np.ndarray: +) -> Array: """Build neighbor list for a single frame. keeps nsel neighbors. Parameters ---------- - coord : np.ndarray + coord : Array exptended coordinates of shape [batch_size, nall x 3] - atype : np.ndarray + atype : Array extended atomic types of shape [batch_size, nall] type < 0 the atom is treat as virtual atoms. nloc : int @@ -81,7 +77,7 @@ def build_neighbor_list( Returns ------- - neighbor_list : np.ndarray + neighbor_list : Array Neighbor list of shape [batch_size, nloc, nsel], the neighbors are stored in an ascending order. If the number of neighbors is less than nsel, the positions are masked @@ -115,8 +111,8 @@ def build_neighbor_list( nsel = sum(sel) coord0 = coord1[:, : nloc * 3] diff = ( - xp.reshape(coord1, [batch_size, -1, 3])[:, None, :, :] - - xp.reshape(coord0, [batch_size, -1, 3])[:, :, None, :] + xp.reshape(coord1, (batch_size, -1, 3))[:, None, :, :] + - xp.reshape(coord0, (batch_size, -1, 3))[:, :, None, :] ) assert list(diff.shape) == [batch_size, nloc, nall, 3] rr = xp.linalg.vector_norm(diff, axis=-1) @@ -153,10 +149,10 @@ def build_neighbor_list( def nlist_distinguish_types( - nlist: np.ndarray, - atype: np.ndarray, + nlist: Array, + atype: Array, sel: list[int], -): +) -> Array: """Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types. @@ -188,20 +184,20 @@ def get_multiple_nlist_key(rcut: float, nsel: int) -> str: ## translated from torch implementation by chatgpt def build_multiple_neighbor_list( - coord: np.ndarray, - nlist: np.ndarray, + coord: Array, + nlist: Array, rcuts: list[float], nsels: list[int], -) -> dict[str, np.ndarray]: +) -> dict[str, Array]: """Input one neighbor list, and produce multiple neighbor lists with different cutoff radius and numbers of selection out of it. The required rcuts and nsels should be smaller or equal to the input nlist. Parameters ---------- - coord : np.ndarray + coord : Array exptended coordinates of shape [batch_size, nall x 3] - nlist : np.ndarray + nlist : Array Neighbor list of shape [batch_size, nloc, nsel], the neighbors should be stored in an ascending order. rcuts : list[float] @@ -211,7 +207,7 @@ def build_multiple_neighbor_list( Returns ------- - nlist_dict : dict[str, np.ndarray] + nlist_dict : dict[str, Array] A dict of nlists, key given by get_multiple_nlist_key(rc, nsel) value being the corresponding nlist. @@ -247,33 +243,33 @@ def build_multiple_neighbor_list( ## translated from torch implementation by chatgpt def extend_coord_with_ghosts( - coord: np.ndarray, - atype: np.ndarray, - cell: Optional[np.ndarray], + coord: Array, + atype: Array, + cell: Array | None, rcut: float, -): +) -> tuple[Array, Array]: """Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors within rcut are appended. Parameters ---------- - coord : np.ndarray + coord : Array original coordinates of shape [-1, nloc*3]. - atype : np.ndarray + atype : Array atom type of shape [-1, nloc]. - cell : np.ndarray + cell : Array simulation cell tensor of shape [-1, 9]. rcut : float the cutoff radius Returns ------- - extended_coord: np.ndarray + extended_coord: Array extended coordinates of shape [-1, nall*3]. - extended_atype: np.ndarray + extended_atype: Array extended atom type of shape [-1, nall]. - index_mapping: np.ndarray + index_mapping: Array mapping extended index to the local index """ diff --git a/deepmd/dpmodel/utils/region.py b/deepmd/dpmodel/utils/region.py index bc9b9479a0..6d8dfebf88 100644 --- a/deepmd/dpmodel/utils/region.py +++ b/deepmd/dpmodel/utils/region.py @@ -1,24 +1,27 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import array_api_compat -import numpy as np + +from deepmd.dpmodel.array_api import ( + Array, +) def phys2inter( - coord: np.ndarray, - cell: np.ndarray, -) -> np.ndarray: + coord: Array, + cell: Array, +) -> Array: """Convert physical coordinates to internal(direct) coordinates. Parameters ---------- - coord : np.ndarray + coord : Array physical coordinates of shape [*, na, 3]. - cell : np.ndarray + cell : Array simulation cell tensor of shape [*, 3, 3]. Returns ------- - inter_coord: np.ndarray + inter_coord: Array the internal coordinates """ @@ -28,21 +31,21 @@ def phys2inter( def inter2phys( - coord: np.ndarray, - cell: np.ndarray, -) -> np.ndarray: + coord: Array, + cell: Array, +) -> Array: """Convert internal(direct) coordinates to physical coordinates. Parameters ---------- - coord : np.ndarray + coord : Array internal coordinates of shape [*, na, 3]. - cell : np.ndarray + cell : Array simulation cell tensor of shape [*, 3, 3]. Returns ------- - phys_coord: np.ndarray + phys_coord: Array the physical coordinates """ @@ -51,21 +54,21 @@ def inter2phys( def normalize_coord( - coord: np.ndarray, - cell: np.ndarray, -) -> np.ndarray: + coord: Array, + cell: Array, +) -> Array: """Apply PBC according to the atomic coordinates. Parameters ---------- - coord : np.ndarray + coord : Array original coordinates of shape [*, na, 3]. - cell : np.ndarray + cell : Array simulation cell shape [*, 3, 3]. Returns ------- - wrapped_coord: np.ndarray + wrapped_coord: Array wrapped coordinates of shape [*, na, 3]. """ @@ -76,28 +79,28 @@ def normalize_coord( def to_face_distance( - cell: np.ndarray, -) -> np.ndarray: + cell: Array, +) -> Array: """Compute the to-face-distance of the simulation cell. Parameters ---------- - cell : np.ndarray + cell : Array simulation cell tensor of shape [*, 3, 3]. Returns ------- - dist: np.ndarray + dist: Array the to face distances of shape [*, 3] """ xp = array_api_compat.array_namespace(cell) cshape = cell.shape - dist = b_to_face_distance(xp.reshape(cell, [-1, 3, 3])) - return xp.reshape(dist, list(cshape[:-2]) + [3]) # noqa:RUF005 + dist = b_to_face_distance(xp.reshape(cell, (-1, 3, 3))) + return xp.reshape(dist, tuple(list(cshape[:-2]) + [3])) # noqa:RUF005 -def b_to_face_distance(cell): +def b_to_face_distance(cell: Array) -> Array: xp = array_api_compat.array_namespace(cell) volume = xp.linalg.det(cell) c_yz = xp.linalg.cross(cell[:, 1, ...], cell[:, 2, ...], axis=-1) diff --git a/deepmd/dpmodel/utils/safe_gradient.py b/deepmd/dpmodel/utils/safe_gradient.py index 2baf530c08..e58ff1d613 100644 --- a/deepmd/dpmodel/utils/safe_gradient.py +++ b/deepmd/dpmodel/utils/safe_gradient.py @@ -5,17 +5,23 @@ for more information. """ +from typing import ( + Any, +) + import array_api_compat -def safe_for_sqrt(x): +def safe_for_sqrt(x: Any) -> Any: """Safe version of sqrt that has a gradient of 0 at x = 0.""" xp = array_api_compat.array_namespace(x) mask = x > 0.0 return xp.where(mask, xp.sqrt(xp.where(mask, x, xp.ones_like(x))), xp.zeros_like(x)) -def safe_for_vector_norm(x, /, *, axis=None, keepdims=False, ord=2): +def safe_for_vector_norm( + x: Any, /, *, axis: Any | None = None, keepdims: bool = False, ord: Any = 2 +) -> Any: """Safe version of sqrt that has a gradient of 0 at x = 0.""" xp = array_api_compat.array_namespace(x) mask = xp.sum(xp.square(x), axis=axis, keepdims=True) > 0 diff --git a/deepmd/dpmodel/utils/seed.py b/deepmd/dpmodel/utils/seed.py index 165ff558b9..1603cc0164 100644 --- a/deepmd/dpmodel/utils/seed.py +++ b/deepmd/dpmodel/utils/seed.py @@ -1,7 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, overload, ) @@ -11,10 +9,10 @@ def child_seed(seed: None, idx: int) -> None: ... @overload -def child_seed(seed: Union[int, list[int]], idx: int) -> list[int]: ... +def child_seed(seed: int | list[int], idx: int) -> list[int]: ... -def child_seed(seed: Optional[Union[int, list[int]]], idx: int) -> Optional[list[int]]: +def child_seed(seed: int | list[int] | None, idx: int) -> list[int] | None: """Generate a child seed from a parent seed. Parameters diff --git a/deepmd/dpmodel/utils/serialization.py b/deepmd/dpmodel/utils/serialization.py index 5520933753..74d31ea589 100644 --- a/deepmd/dpmodel/utils/serialization.py +++ b/deepmd/dpmodel/utils/serialization.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import datetime import json +from collections.abc import ( + Callable, +) from pathlib import ( Path, ) from typing import ( - Callable, + Any, ) import h5py @@ -18,7 +21,9 @@ __version__ = "unknown" -def traverse_model_dict(model_obj, callback: Callable, is_variable: bool = False): +def traverse_model_dict( + model_obj: Any, callback: Callable, is_variable: bool = False +) -> Any: """Traverse a model dict and call callback on each variable. Parameters @@ -67,7 +72,7 @@ class Counter: def __init__(self) -> None: self.count = -1 - def __call__(self): + def __call__(self) -> int: self.count += 1 return self.count @@ -149,7 +154,7 @@ def load_dp_model(filename: str) -> dict: model_dict = traverse_model_dict(model_dict, lambda x: f[x][()].copy()) elif filename_extension in {".yaml", ".yml"}: - def convert_numpy_ndarray(x): + def convert_numpy_ndarray(x: Any) -> Any: if isinstance(x, dict) and x.get("@class") == "np.ndarray": dtype = np.dtype(x["dtype"]) value = np.asarray(x["value"], dtype=dtype) diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index 17e40f3592..a1b698b698 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, ) import array_api_compat import numpy as np from deepmd.dpmodel.array_api import ( + Array, support_array_api, ) from deepmd.dpmodel.common import ( @@ -63,11 +63,11 @@ def __init__( activation_function: str = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: self.ntypes = ntypes self.neuron = neuron @@ -93,10 +93,11 @@ def __init__( self.precision, seed=self.seed, bias=self.use_tebd_bias, + trainable=trainable, ) @support_array_api(version="2022.12") - def call(self) -> np.ndarray: + def call(self) -> Array: """Compute the type embedding network.""" sample_array = self.embedding_net[0]["w"] xp = array_api_compat.array_namespace(sample_array) @@ -110,7 +111,7 @@ def call(self) -> np.ndarray: return embed @classmethod - def deserialize(cls, data: dict): + def deserialize(cls, data: dict) -> "TypeEmbedNet": """Deserialize the model. Parameters @@ -161,7 +162,7 @@ def serialize(self) -> dict: } def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -218,7 +219,9 @@ def change_type_map( self.ntypes = len(type_map) -def get_econf_tebd(type_map, precision: str = "default"): +def get_econf_tebd( + type_map: list[str], precision: str = "default" +) -> tuple[Array, int]: from deepmd.utils.econf_embd import ( ECONF_DIM, ) diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py index 39967d565c..a8cf20c6b3 100644 --- a/deepmd/entrypoints/convert_backend.py +++ b/deepmd/entrypoints/convert_backend.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + from deepmd.backend.backend import ( Backend, ) @@ -8,7 +12,7 @@ def convert_backend( *, # Enforce keyword-only arguments INPUT: str, OUTPUT: str, - **kwargs, + **kwargs: Any, ) -> None: """Convert a model file from one backend to another. diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py index 5679d838ac..74fdd90ebd 100644 --- a/deepmd/entrypoints/doc.py +++ b/deepmd/entrypoints/doc.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Module that prints train input arguments docstrings.""" +from typing import ( + Any, +) + from deepmd.utils.argcheck import ( gen_doc, gen_json, @@ -11,7 +15,7 @@ def doc_train_input( - *, out_type: str = "rst", multi_task: bool = False, **kwargs + *, out_type: str = "rst", multi_task: bool = False, **kwargs: Any ) -> None: """Print out trining input arguments to console.""" if out_type == "rst": diff --git a/deepmd/entrypoints/eval_desc.py b/deepmd/entrypoints/eval_desc.py new file mode 100644 index 0000000000..a918b63edb --- /dev/null +++ b/deepmd/entrypoints/eval_desc.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Evaluate descriptors using trained DeePMD model.""" + +import logging +import os +from pathlib import ( + Path, +) +from typing import ( + Any, +) + +import numpy as np + +from deepmd.common import ( + expand_sys_str, +) +from deepmd.infer.deep_eval import ( + DeepEval, +) +from deepmd.utils.data import ( + DeepmdData, +) + +__all__ = ["eval_desc"] + +log = logging.getLogger(__name__) + + +def eval_desc( + *, + model: str, + system: str, + datafile: str, + output: str = "desc", + head: str | None = None, + **kwargs: Any, +) -> None: + """Evaluate descriptors for given systems. + + Parameters + ---------- + model : str + path where model is stored + system : str + system directory + datafile : str + the path to the list of systems to process + output : str + output directory for descriptor files + head : Optional[str], optional + (Supported backend: PyTorch) Task head if in multi-task mode. + **kwargs + additional arguments + + Notes + ----- + Descriptors are saved as 3D numpy arrays with shape (nframes, natoms, ndesc) + where each frame contains the descriptors for all atoms. + + Raises + ------ + RuntimeError + if no valid system was found + """ + if datafile is not None: + with open(datafile) as datalist: + all_sys = datalist.read().splitlines() + else: + all_sys = expand_sys_str(system) + + if len(all_sys) == 0: + raise RuntimeError("Did not find valid system") + + # init model + dp = DeepEval(model, head=head) + + # create output directory + output_dir = Path(output) + output_dir.mkdir(parents=True, exist_ok=True) + + for cc, system_path in enumerate(all_sys): + log.info("# -------output of dp eval_desc------- ") + log.info(f"# processing system : {system_path}") + + # create data class + tmap = dp.get_type_map() + data = DeepmdData( + system_path, + set_prefix="set", + shuffle_test=False, + type_map=tmap, + sort_atoms=False, + ) + + # get test data + test_data = data.get_test() + mixed_type = data.mixed_type + natoms = len(test_data["type"][0]) + nframes = test_data["box"].shape[0] + + # prepare input data + coord = test_data["coord"].reshape([nframes, -1]) + box = test_data["box"] + if not data.pbc: + box = None + if mixed_type: + atype = test_data["type"].reshape([nframes, -1]) + else: + atype = test_data["type"][0] + + # handle optional parameters + fparam = None + if dp.get_dim_fparam() > 0: + if "fparam" in test_data: + fparam = test_data["fparam"] + + aparam = None + if dp.get_dim_aparam() > 0: + if "aparam" in test_data: + aparam = test_data["aparam"] + + # evaluate descriptors + log.info(f"# evaluating descriptors for {nframes} frames") + descriptors = dp.eval_descriptor( + coord, + box, + atype, + fparam=fparam, + aparam=aparam, + ) + + # descriptors are kept in 3D format (nframes, natoms, ndesc) + + # save descriptors + system_name = os.path.basename(system_path.rstrip("/")) + desc_file = output_dir / f"{system_name}.npy" + np.save(desc_file, descriptors) + + log.info(f"# descriptors saved to {desc_file}") + log.info(f"# descriptor shape: {descriptors.shape}") + log.info("# ----------------------------------- ") + + log.info("# eval_desc completed successfully") diff --git a/deepmd/entrypoints/gui.py b/deepmd/entrypoints/gui.py index 7a61eb07ee..3d17810bec 100644 --- a/deepmd/entrypoints/gui.py +++ b/deepmd/entrypoints/gui.py @@ -1,8 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """DP-GUI entrypoint.""" +from typing import ( + Any, +) -def start_dpgui(*, port: int, bind_all: bool, **kwargs) -> None: + +def start_dpgui(*, port: int, bind_all: bool, **kwargs: Any) -> None: """Host DP-GUI server. Parameters diff --git a/deepmd/entrypoints/ipi.py b/deepmd/entrypoints/ipi.py index 2f7e1f5458..bd527defb1 100644 --- a/deepmd/entrypoints/ipi.py +++ b/deepmd/entrypoints/ipi.py @@ -15,7 +15,7 @@ ROOT_DIR = get_op_dir() -def _program(name: str, args: list[str]): +def _program(name: str, args: list[str]) -> None: """Execute a program. Parameters diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 2c91ca5f29..34ebe4d2e3 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -18,6 +18,9 @@ from deepmd.entrypoints.doc import ( doc_train_input, ) +from deepmd.entrypoints.eval_desc import ( + eval_desc, +) from deepmd.entrypoints.gui import ( start_dpgui, ) @@ -65,6 +68,14 @@ def main(args: argparse.Namespace) -> None: strict_prefer=False, ) test(**dict_args) + elif args.command == "eval-desc": + dict_args["model"] = format_model_suffix( + dict_args["model"], + feature=Backend.Feature.DEEP_EVAL, + preferred_backend=args.backend, + strict_prefer=False, + ) + eval_desc(**dict_args) elif args.command == "doc-train-input": doc_train_input(**dict_args) elif args.command == "model-devi": diff --git a/deepmd/entrypoints/neighbor_stat.py b/deepmd/entrypoints/neighbor_stat.py index 62dceb24fd..836f1e0174 100644 --- a/deepmd/entrypoints/neighbor_stat.py +++ b/deepmd/entrypoints/neighbor_stat.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, + Any, ) from deepmd.backend.backend import ( @@ -21,11 +21,11 @@ def neighbor_stat( *, system: str, rcut: float, - type_map: Optional[list[str]], + type_map: list[str] | None, mixed_type: bool = False, backend: str = "tensorflow", - **kwargs, -): + **kwargs: Any, +) -> None: """Calculate neighbor statistics. Parameters diff --git a/deepmd/entrypoints/show.py b/deepmd/entrypoints/show.py index e279fbe3e7..ddd097d22c 100644 --- a/deepmd/entrypoints/show.py +++ b/deepmd/entrypoints/show.py @@ -1,9 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from typing import ( + Any, +) from deepmd.infer.deep_eval import ( DeepEval, ) +from deepmd.utils.econf_embd import ( + sort_element_type, +) +from deepmd.utils.model_branch_dict import ( + OrderedDictTableWrapper, + get_model_dict, +) log = logging.getLogger(__name__) @@ -12,7 +22,7 @@ def show( *, INPUT: str, ATTRIBUTES: list[str], - **kwargs, + **kwargs: Any, ) -> None: model = DeepEval(INPUT, head=0) model_params = model.get_model_def_script() @@ -30,10 +40,15 @@ def show( ) model_branches = list(model_params["model_dict"].keys()) model_branches += ["RANDOM"] + _, model_branch_dict = get_model_dict(model_params["model_dict"]) log.info( f"Available model branches are {model_branches}, " f"where 'RANDOM' means using a randomly initialized fitting net." ) + log.info( + "Detailed information: \n" + + OrderedDictTableWrapper(model_branch_dict).as_table() + ) if "type-map" in ATTRIBUTES: if model_is_multi_task: model_branches = list(model_params["model_dict"].keys()) @@ -69,3 +84,48 @@ def show( log.info(f"Parameter counts{log_prefix}:") for k in sorted(size_dict): log.info(f"Parameters in {k}: {size_dict[k]:,}") + + if "observed-type" in ATTRIBUTES: + if model_is_multi_task: + log.info("The observed types for each branch: ") + total_observed_types_list = [] + model_branches = list(model_params["model_dict"].keys()) + for branch in model_branches: + if ( + model_params["model_dict"][branch] + .get("info", {}) + .get("observed_type", None) + is not None + ): + observed_type_list = model_params["model_dict"][branch]["info"][ + "observed_type" + ] + observed_types = { + "type_num": len(observed_type_list), + "observed_type": observed_type_list, + } + else: + tmp_model = DeepEval(INPUT, head=branch, no_jit=True) + observed_types = tmp_model.get_observed_types() + log.info( + f"{branch}: Number of observed types: {observed_types['type_num']} " + ) + log.info( + f"{branch}: Observed types: {observed_types['observed_type']} " + ) + total_observed_types_list += [ + tt + for tt in observed_types["observed_type"] + if tt not in total_observed_types_list + ] + log.info( + f"TOTAL number of observed types in the model: {len(total_observed_types_list)} " + ) + log.info( + f"TOTAL observed types in the model: {sort_element_type(total_observed_types_list)} " + ) + else: + log.info("The observed types for this model: ") + observed_types = model.get_observed_types() + log.info(f"Number of observed types: {observed_types['type_num']} ") + log.info(f"Observed types: {observed_types['observed_type']} ") diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py index 919d23f757..4a0cb27cb1 100644 --- a/deepmd/entrypoints/test.py +++ b/deepmd/entrypoints/test.py @@ -7,13 +7,14 @@ ) from typing import ( TYPE_CHECKING, - Optional, + Any, ) import numpy as np from deepmd.common import ( expand_sys_str, + j_loader, ) from deepmd.infer.deep_dipole import ( DeepDipole, @@ -38,9 +39,15 @@ DeepWFC, ) from deepmd.utils import random as dp_random +from deepmd.utils.compat import ( + update_deepmd_input, +) from deepmd.utils.data import ( DeepmdData, ) +from deepmd.utils.data_system import ( + process_systems, +) from deepmd.utils.weight_avg import ( weighted_average, ) @@ -58,15 +65,17 @@ def test( *, model: str, - system: str, - datafile: str, + system: str | None, + datafile: str | None, + train_json: str | None = None, + valid_json: str | None = None, numb_test: int, - rand_seed: Optional[int], + rand_seed: int | None, shuffle_test: bool, detail_file: str, atomic: bool, - head: Optional[str] = None, - **kwargs, + head: str | None = None, + **kwargs: Any, ) -> None: """Test model predictions. @@ -74,12 +83,16 @@ def test( ---------- model : str path where model is stored - system : str + system : str, optional system directory - datafile : str + datafile : str, optional the path to the list of systems to test + train_json : Optional[str] + Path to the input.json file provided via ``--train-data``. Training systems will be used for testing. + valid_json : Optional[str] + Path to the input.json file provided via ``--valid-data``. Validation systems will be used for testing. numb_test : int - munber of tests to do. 0 means all data. + number of tests to do. 0 means all data. rand_seed : Optional[int] seed for random generator shuffle_test : bool @@ -101,11 +114,41 @@ def test( if numb_test == 0: # only float has inf, but should work for min numb_test = float("inf") - if datafile is not None: + if train_json is not None: + jdata = j_loader(train_json) + jdata = update_deepmd_input(jdata) + data_params = jdata.get("training", {}).get("training_data", {}) + systems = data_params.get("systems") + if not systems: + raise RuntimeError("No training data found in input json") + root = Path(train_json).parent + if isinstance(systems, str): + systems = str((root / Path(systems)).resolve()) + else: + systems = [str((root / Path(ss)).resolve()) for ss in systems] + patterns = data_params.get("rglob_patterns", None) + all_sys = process_systems(systems, patterns=patterns) + elif valid_json is not None: + jdata = j_loader(valid_json) + jdata = update_deepmd_input(jdata) + data_params = jdata.get("training", {}).get("validation_data", {}) + systems = data_params.get("systems") + if not systems: + raise RuntimeError("No validation data found in input json") + root = Path(valid_json).parent + if isinstance(systems, str): + systems = str((root / Path(systems)).resolve()) + else: + systems = [str((root / Path(ss)).resolve()) for ss in systems] + patterns = data_params.get("rglob_patterns", None) + all_sys = process_systems(systems, patterns=patterns) + elif datafile is not None: with open(datafile) as datalist: all_sys = datalist.read().splitlines() - else: + elif system is not None: all_sys = expand_sys_str(system) + else: + raise RuntimeError("No data source specified for testing") if len(all_sys) == 0: raise RuntimeError("Did not find valid system") @@ -259,7 +302,7 @@ def test_ener( data: DeepmdData, system: str, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, has_atom_ener: bool, append_detail: bool = False, ) -> tuple[list[np.ndarray], list[int]]: @@ -291,6 +334,7 @@ def test_ener( data.add("energy", 1, atomic=False, must=False, high_prec=True) data.add("force", 3, atomic=True, must=False, high_prec=False) + data.add("atom_pref", 1, atomic=True, must=False, high_prec=False, repeat=3) data.add("virial", 9, atomic=False, must=False, high_prec=False) if dp.has_efield: data.add("efield", 3, atomic=True, must=True, high_prec=False) @@ -298,7 +342,11 @@ def test_ener( data.add("atom_ener", 1, atomic=True, must=True, high_prec=False) if dp.get_dim_fparam() > 0: data.add( - "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False + "fparam", + dp.get_dim_fparam(), + atomic=False, + must=not dp.has_default_fparam(), + high_prec=False, ) if dp.get_dim_aparam() > 0: data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) @@ -313,6 +361,7 @@ def test_ener( find_force = test_data.get("find_force") find_virial = test_data.get("find_virial") find_force_mag = test_data.get("find_force_mag") + find_atom_pref = test_data.get("find_atom_pref") mixed_type = data.mixed_type natoms = len(test_data["type"][0]) nframes = test_data["box"].shape[0] @@ -334,7 +383,7 @@ def test_ener( atype = test_data["type"][:numb_test].reshape([numb_test, -1]) else: atype = test_data["type"][0] - if dp.get_dim_fparam() > 0: + if dp.get_dim_fparam() > 0 and test_data["find_fparam"] != 0.0: fparam = test_data["fparam"][:numb_test] else: fparam = None @@ -419,6 +468,16 @@ def test_ener( diff_f = force - test_data["force"][:numb_test] mae_f = mae(diff_f) rmse_f = rmse(diff_f) + size_f = diff_f.size + if find_atom_pref == 1: + atom_weight = test_data["atom_pref"][:numb_test] + weight_sum = np.sum(atom_weight) + if weight_sum > 0: + mae_fw = np.sum(np.abs(diff_f) * atom_weight) / weight_sum + rmse_fw = np.sqrt(np.sum(diff_f * diff_f * atom_weight) / weight_sum) + else: + mae_fw = 0.0 + rmse_fw = 0.0 diff_v = virial - test_data["virial"][:numb_test] mae_v = mae(diff_v) rmse_v = rmse(diff_v) @@ -451,13 +510,18 @@ def test_ener( dict_to_return["rmse_e"] = (rmse_e, energy.size) dict_to_return["rmse_ea"] = (rmse_ea, energy.size) if not out_put_spin and find_force == 1: - log.info(f"Force MAE : {mae_f:e} eV/A") - log.info(f"Force RMSE : {rmse_f:e} eV/A") - dict_to_return["mae_f"] = (mae_f, force.size) - dict_to_return["rmse_f"] = (rmse_f, force.size) + log.info(f"Force MAE : {mae_f:e} eV/Å") + log.info(f"Force RMSE : {rmse_f:e} eV/Å") + dict_to_return["mae_f"] = (mae_f, size_f) + dict_to_return["rmse_f"] = (rmse_f, size_f) + if find_atom_pref == 1: + log.info(f"Force weighted MAE : {mae_fw:e} eV/Å") + log.info(f"Force weighted RMSE: {rmse_fw:e} eV/Å") + dict_to_return["mae_fw"] = (mae_fw, weight_sum) + dict_to_return["rmse_fw"] = (rmse_fw, weight_sum) if out_put_spin and find_force == 1: - log.info(f"Force atom MAE : {mae_fr:e} eV/A") - log.info(f"Force atom RMSE : {rmse_fr:e} eV/A") + log.info(f"Force atom MAE : {mae_fr:e} eV/Å") + log.info(f"Force atom RMSE : {rmse_fr:e} eV/Å") dict_to_return["mae_fr"] = (mae_fr, force_r.size) dict_to_return["rmse_fr"] = (rmse_fr, force_r.size) if out_put_spin and find_force_mag == 1: @@ -478,8 +542,8 @@ def test_ener( log.info(f"Atomic ener MAE : {mae_ae:e} eV") log.info(f"Atomic ener RMSE : {rmse_ae:e} eV") if dp.has_hessian: - log.info(f"Hessian MAE : {mae_h:e} eV/A^2") - log.info(f"Hessian RMSE : {rmse_h:e} eV/A^2") + log.info(f"Hessian MAE : {mae_h:e} eV/Å^2") + log.info(f"Hessian RMSE : {rmse_h:e} eV/Å^2") dict_to_return["mae_h"] = (mae_h, hessian.size) dict_to_return["rmse_h"] = (rmse_h, hessian.size) @@ -597,21 +661,25 @@ def print_ener_sys_avg(avg: dict[str, float]) -> None: log.info(f"Energy RMSE : {avg['rmse_e']:e} eV") log.info(f"Energy MAE/Natoms : {avg['mae_ea']:e} eV") log.info(f"Energy RMSE/Natoms : {avg['rmse_ea']:e} eV") - if "rmse_f" in avg.keys(): - log.info(f"Force MAE : {avg['mae_f']:e} eV/A") - log.info(f"Force RMSE : {avg['rmse_f']:e} eV/A") + if "rmse_f" in avg: + log.info(f"Force MAE : {avg['mae_f']:e} eV/Å") + log.info(f"Force RMSE : {avg['rmse_f']:e} eV/Å") + if "rmse_fw" in avg: + log.info(f"Force weighted MAE : {avg['mae_fw']:e} eV/Å") + log.info(f"Force weighted RMSE: {avg['rmse_fw']:e} eV/Å") else: - log.info(f"Force atom MAE : {avg['mae_fr']:e} eV/A") + log.info(f"Force atom MAE : {avg['mae_fr']:e} eV/Å") log.info(f"Force spin MAE : {avg['mae_fm']:e} eV/uB") - log.info(f"Force atom RMSE : {avg['rmse_fr']:e} eV/A") + log.info(f"Force atom RMSE : {avg['rmse_fr']:e} eV/Å") log.info(f"Force spin RMSE : {avg['rmse_fm']:e} eV/uB") - log.info(f"Virial MAE : {avg['mae_v']:e} eV") - log.info(f"Virial RMSE : {avg['rmse_v']:e} eV") - log.info(f"Virial MAE/Natoms : {avg['mae_va']:e} eV") - log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV") - if "rmse_h" in avg.keys(): - log.info(f"Hessian MAE : {avg['mae_h']:e} eV/A^2") - log.info(f"Hessian RMSE : {avg['rmse_h']:e} eV/A^2") + if "rmse_v" in avg: + log.info(f"Virial MAE : {avg['mae_v']:e} eV") + log.info(f"Virial RMSE : {avg['rmse_v']:e} eV") + log.info(f"Virial MAE/Natoms : {avg['mae_va']:e} eV") + log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV") + if "rmse_h" in avg: + log.info(f"Hessian MAE : {avg['mae_h']:e} eV/Å^2") + log.info(f"Hessian RMSE : {avg['rmse_h']:e} eV/Å^2") def test_dos( @@ -619,7 +687,7 @@ def test_dos( data: DeepmdData, system: str, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, has_atom_dos: bool, append_detail: bool = False, ) -> tuple[list[np.ndarray], list[int]]: @@ -779,7 +847,7 @@ def test_property( data: DeepmdData, system: str, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, has_atom_property: bool, append_detail: bool = False, ) -> tuple[list[np.ndarray], list[int]]: @@ -934,7 +1002,9 @@ def print_property_sys_avg(avg: dict[str, float]) -> None: log.info(f"PROPERTY RMSE : {avg['rmse_property']:e} units") -def run_test(dp: "DeepTensor", test_data: dict, numb_test: int, test_sys: DeepmdData): +def run_test( + dp: "DeepTensor", test_data: dict, numb_test: int, test_sys: DeepmdData +) -> dict: """Run tests. Parameters @@ -971,7 +1041,7 @@ def test_wfc( dp: "DeepWFC", data: DeepmdData, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, ) -> tuple[list[np.ndarray], list[int]]: """Test energy type model. @@ -998,8 +1068,8 @@ def test_wfc( wfc, numb_test, _ = run_test(dp, test_data, numb_test, data) rmse_f = rmse(wfc - test_data["wfc"][:numb_test]) - log.info("# number of test data : {numb_test:d} ") - log.info("WFC RMSE : {rmse_f:e} eV/A") + log.info(f"# number of test data : {numb_test:d} ") + log.info(f"WFC RMSE : {rmse_f:e}") if detail_file is not None: detail_path = Path(detail_file) @@ -1018,7 +1088,7 @@ def test_wfc( return {"rmse": (rmse_f, wfc.size)} -def print_wfc_sys_avg(avg) -> None: +def print_wfc_sys_avg(avg: dict) -> None: """Print errors summary for wfc type potential. Parameters @@ -1026,14 +1096,14 @@ def print_wfc_sys_avg(avg) -> None: avg : np.ndarray array with summaries """ - log.info(f"WFC RMSE : {avg['rmse']:e} eV/A") + log.info(f"WFC RMSE : {avg['rmse']:e}") def test_polar( dp: "DeepPolar", data: DeepmdData, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, *, atomic: bool, ) -> tuple[list[np.ndarray], list[int]]: @@ -1160,7 +1230,7 @@ def test_polar( return {"rmse": (rmse_f, polar.size)} -def print_polar_sys_avg(avg) -> None: +def print_polar_sys_avg(avg: dict) -> None: """Print errors summary for polar type potential. Parameters @@ -1168,14 +1238,14 @@ def print_polar_sys_avg(avg) -> None: avg : np.ndarray array with summaries """ - log.info(f"Polarizability RMSE : {avg['rmse']:e} eV/A") + log.info(f"Polarizability RMSE : {avg['rmse']:e}") def test_dipole( dp: "DeepDipole", data: DeepmdData, numb_test: int, - detail_file: Optional[str], + detail_file: str | None, atomic: bool, ) -> tuple[list[np.ndarray], list[int]]: """Test energy type model. @@ -1274,7 +1344,7 @@ def test_dipole( return {"rmse": (rmse_f, dipole.size)} -def print_dipole_sys_avg(avg) -> None: +def print_dipole_sys_avg(avg: dict) -> None: """Print errors summary for dipole type potential. Parameters @@ -1282,4 +1352,4 @@ def print_dipole_sys_avg(avg) -> None: avg : np.ndarray array with summaries """ - log.info(f"Dipole RMSE : {avg['rmse']:e} eV/A") + log.info(f"Dipole RMSE : {avg['rmse']:e}") diff --git a/deepmd/env.py b/deepmd/env.py index 2c1241a36b..7b29a338f1 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import os +import platform from configparser import ( ConfigParser, ) @@ -16,6 +17,7 @@ "GLOBAL_CONFIG", "GLOBAL_ENER_FLOAT_PRECISION", "GLOBAL_NP_FLOAT_PRECISION", + "LRU_CACHE_SIZE", "SHARED_LIB_DIR", "SHARED_LIB_MODULE", "global_float_prec", @@ -47,6 +49,20 @@ "DP_INTERFACE_PREC." ) +# Dynamic calculation of cache size +_default_lru_cache_size = 512 +LRU_CACHE_SIZE = _default_lru_cache_size + +if platform.system() != "Windows": + import resource + + soft_limit, hard_limit = resource.getrlimit(resource.RLIMIT_NOFILE) + safe_buffer = 128 + if soft_limit > safe_buffer + _default_lru_cache_size: + LRU_CACHE_SIZE = soft_limit - safe_buffer + else: + LRU_CACHE_SIZE = soft_limit // 2 + def set_env_if_empty(key: str, value: str, verbose: bool = True) -> None: """Set environment variable only if it is empty. diff --git a/deepmd/infer/__init__.py b/deepmd/infer/__init__.py index a0330c4ffb..77780d5922 100644 --- a/deepmd/infer/__init__.py +++ b/deepmd/infer/__init__.py @@ -17,7 +17,7 @@ ] -def DeepPotential(*args, **kwargs) -> "DeepEval": +def DeepPotential(*args: object, **kwargs: object) -> "DeepEval": """Factory function that forwards to DeepEval (for compatibility). Parameters diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index 0d7ccee2b6..4e63a156c1 100644 --- a/deepmd/infer/deep_dos.py +++ b/deepmd/infer/deep_dos.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -61,11 +59,11 @@ def numb_dos(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index ee61abe58c..0a8686a98b 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -8,7 +8,6 @@ Any, ClassVar, Optional, - Union, ) import numpy as np @@ -84,13 +83,15 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: pass - def __new__(cls, model_file: str, *args, **kwargs): + def __new__( + cls, model_file: str, *args: object, **kwargs: object + ) -> "DeepEvalBackend": if cls is DeepEvalBackend: backend = Backend.detect_backend_by_model(model_file) return super().__new__(backend().deep_eval) @@ -100,11 +101,11 @@ def __new__(cls, model_file: str, *args, **kwargs): def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -160,6 +161,10 @@ def get_type_map(self) -> list[str]: def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this DP.""" + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + @abstractmethod def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" @@ -167,11 +172,11 @@ def get_dim_aparam(self) -> int: def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> np.ndarray: @@ -215,6 +220,48 @@ def eval_descriptor( """ raise NotImplementedError + def eval_fitting_last_layer( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> np.ndarray: + """Evaluate fitting before last layer by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + + Returns + ------- + fitting + Fitting output before last layer. + """ + raise NotImplementedError + def eval_typeebd(self) -> np.ndarray: """Evaluate output of type embedding network by using this model. @@ -275,7 +322,7 @@ def get_has_spin(self) -> bool: """Check if the model has spin atom types.""" return False - def get_has_hessian(self): + def get_has_hessian(self) -> bool: """Check if the model has hessian.""" return False @@ -295,6 +342,24 @@ def get_model_size(self) -> dict: """Get model parameter count.""" raise NotImplementedError("Not implemented in this backend.") + def get_observed_types(self) -> dict: + """Get observed types (elements) of the model during data statistics.""" + raise NotImplementedError("Not implemented in this backend.") + + @abstractmethod + def get_model(self) -> Any: + """Get the model module implemented by the deep learning framework. + + For PyTorch, this returns the nn.Module. For Paddle, this returns + the paddle.nn.Layer. For TensorFlow, this returns the graph. + For dpmodel, this returns the BaseModel. + + Returns + ------- + model + The model module implemented by the deep learning framework. + """ + class DeepEval(ABC): """High-level Deep Evaluator interface. @@ -319,7 +384,7 @@ class DeepEval(ABC): Keyword arguments. """ - def __new__(cls, model_file: str, *args, **kwargs): + def __new__(cls, model_file: str, *args: object, **kwargs: object) -> "DeepEval": if cls is DeepEval: deep_eval = DeepEvalBackend( model_file, @@ -334,7 +399,7 @@ def __init__( self, model_file: str, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -370,6 +435,10 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this DP.""" return self.deep_eval.get_dim_fparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.deep_eval.has_default_fparam() + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.deep_eval.get_dim_aparam() @@ -391,7 +460,9 @@ def _get_natoms_and_nframes( nframes = coords.shape[0] return natoms, nframes - def _expande_atype(self, atype: np.ndarray, nframes: int, mixed_type: bool): + def _expande_atype( + self, atype: np.ndarray, nframes: int, mixed_type: bool + ) -> np.ndarray: if not mixed_type: atype = np.tile(atype.reshape(1, -1), (nframes, 1)) return atype @@ -399,10 +470,10 @@ def _expande_atype(self, atype: np.ndarray, nframes: int, mixed_type: bool): def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> np.ndarray: @@ -463,6 +534,73 @@ def eval_descriptor( ) return descriptor + def eval_fitting_last_layer( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + mixed_type: bool = False, + **kwargs: Any, + ) -> np.ndarray: + """Evaluate fitting before last layer by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + efield + The external field on atoms. + The array should be of size nframes x natoms x 3 + mixed_type + Whether to perform the mixed_type mode. + If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), + in which frames in a system may have different natoms_vec(s), with the same nloc. + + Returns + ------- + fitting + Fitting output before last layer. + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + fitting = self.deep_eval.eval_fitting_last_layer( + coords, + cells, + atom_types, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + return fitting + def eval_typeebd(self) -> np.ndarray: """Evaluate output of type embedding network by using this model. @@ -492,7 +630,21 @@ def eval_typeebd(self) -> np.ndarray: """ return self.deep_eval.eval_typeebd() - def _standard_input(self, coords, cells, atom_types, fparam, aparam, mixed_type): + def _standard_input( + self, + coords: np.ndarray | list, + cells: np.ndarray | list | None, + atom_types: np.ndarray | list, + fparam: np.ndarray | list | None, + aparam: np.ndarray | list | None, + mixed_type: bool, + ) -> tuple[ + np.ndarray, + np.ndarray | None, + np.ndarray, + np.ndarray | None, + np.ndarray | None, + ]: coords = np.array(coords) if cells is not None: cells = np.array(cells) @@ -539,7 +691,7 @@ def get_sel_type(self) -> list[int]: """ return self.deep_eval.get_sel_type() - def _get_sel_natoms(self, atype) -> int: + def _get_sel_natoms(self, atype: np.ndarray) -> int: return np.sum(np.isin(atype, self.get_sel_type()).astype(int)) @property @@ -568,3 +720,21 @@ def get_model_def_script(self) -> dict: def get_model_size(self) -> dict: """Get model parameter count.""" return self.deep_eval.get_model_size() + + def get_observed_types(self) -> dict: + """Get observed types (elements) of the model during data statistics.""" + return self.deep_eval.get_observed_types() + + def get_model(self) -> Any: + """Get the model module implemented by the deep learning framework. + + For PyTorch, this returns the nn.Module. For Paddle, this returns + the paddle.nn.Layer. For TensorFlow, this returns the graph. + For dpmodel, this returns the BaseModel. + + Returns + ------- + model + The model module implemented by the deep learning framework. + """ + return self.deep_eval.get_model() diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 7220e53637..9b74f43270 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, ) import numpy as np @@ -49,13 +48,13 @@ def output_tensor_name(self) -> str: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, - **kwargs, + **kwargs: Any, ) -> np.ndarray: """Evaluate the model. diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index 6e00a30f91..ce23c5981b 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -2,8 +2,6 @@ from typing import ( Any, Literal, - Optional, - Union, overload, ) @@ -93,45 +91,45 @@ def output_def_mag(self) -> ModelOutputDef: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: Literal[True], - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, mixed_type: bool, **kwargs: Any, - ) -> Union[ - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray], - ]: + ) -> ( + tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] + | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray] + ): pass @overload def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: Literal[False], - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, mixed_type: bool, **kwargs: Any, - ) -> Union[ - tuple[np.ndarray, np.ndarray, np.ndarray], - tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], - ]: + ) -> ( + tuple[np.ndarray, np.ndarray, np.ndarray] + | tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] + ): pass @overload def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, mixed_type: bool, **kwargs: Any, ) -> tuple[np.ndarray, ...]: @@ -140,11 +138,11 @@ def eval( def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: Any, ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index 5944491cc0..5e35dcd781 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -81,11 +79,11 @@ def task_dim(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict[str, Any], ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index bb5bc12697..877ade6e19 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -2,10 +2,6 @@ from abc import ( abstractmethod, ) -from typing import ( - Optional, - Union, -) import numpy as np @@ -41,11 +37,11 @@ class DeepTensor(DeepEval): def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: Union[list[int], np.ndarray], + cells: np.ndarray | None, + atom_types: list[int] | np.ndarray, atomic: bool = True, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict, ) -> np.ndarray: @@ -110,11 +106,11 @@ def eval( def eval_full( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict, ) -> tuple[np.ndarray, ...]: @@ -241,11 +237,11 @@ class OldDeepTensor(DeepTensor): def eval_full( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, mixed_type: bool = False, **kwargs: dict, ) -> tuple[np.ndarray, ...]: diff --git a/deepmd/infer/model_devi.py b/deepmd/infer/model_devi.py index 1828d40c49..31ee414cbd 100644 --- a/deepmd/infer/model_devi.py +++ b/deepmd/infer/model_devi.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, overload, ) @@ -19,14 +19,14 @@ try: from typing import Literal # python >=3.8 except ImportError: - from typing_extensions import Literal # type: ignore + from typing import Literal # type: ignore @overload def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: Literal[False] = ..., ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... @@ -34,8 +34,8 @@ def calc_model_devi_f( @overload def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: Literal[True] = ..., ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: ... @@ -43,16 +43,16 @@ def calc_model_devi_f( @overload def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: bool = False, ) -> tuple[np.ndarray, ...]: ... def calc_model_devi_f( fs: np.ndarray, - real_f: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_f: np.ndarray | None = None, + relative: float | None = None, atomic: bool = False, ) -> tuple[np.ndarray, ...]: """Calculate model deviation of force. @@ -106,9 +106,7 @@ def calc_model_devi_f( return max_devi_f, min_devi_f, avg_devi_f -def calc_model_devi_e( - es: np.ndarray, real_e: Optional[np.ndarray] = None -) -> np.ndarray: +def calc_model_devi_e(es: np.ndarray, real_e: np.ndarray | None = None) -> np.ndarray: """Calculate model deviation of total energy per atom. Here we don't use the atomic energy, as the decomposition @@ -138,8 +136,8 @@ def calc_model_devi_e( def calc_model_devi_v( vs: np.ndarray, - real_v: Optional[np.ndarray] = None, - relative: Optional[float] = None, + real_v: np.ndarray | None = None, + relative: float | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Calculate model deviation of virial. @@ -186,7 +184,7 @@ def calc_model_devi_v( def write_model_devi_out( devi: np.ndarray, fname: str, header: str = "", atomic: bool = False -): +) -> np.ndarray: """Write output of model deviation. Parameters @@ -225,7 +223,7 @@ def write_model_devi_out( return devi -def _check_tmaps(tmaps, ref_tmap=None): +def _check_tmaps(tmaps: list[list[str]], ref_tmap: list[str] | None = None) -> bool: """Check whether type maps are identical.""" assert isinstance(tmaps, list) if ref_tmap is None: @@ -241,20 +239,20 @@ def _check_tmaps(tmaps, ref_tmap=None): def calc_model_devi( - coord, - box, - atype, - models, - fname=None, - frequency=1, - mixed_type=False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - real_data: Optional[dict] = None, + coord: np.ndarray, + box: np.ndarray | None, + atype: np.ndarray, + models: list[DeepPot], + fname: str | None = None, + frequency: int = 1, + mixed_type: bool = False, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + real_data: dict | None = None, atomic: bool = False, - relative: Optional[float] = None, - relative_v: Optional[float] = None, -): + relative: float | None = None, + relative_v: float | None = None, +) -> np.ndarray: """Python interface to calculate model deviation. Parameters @@ -361,10 +359,10 @@ def make_model_devi( frequency: int, real_error: bool = False, atomic: bool = False, - relative: Optional[float] = None, - relative_v: Optional[float] = None, - **kwargs, -): + relative: float | None = None, + relative_v: float | None = None, + **kwargs: Any, +) -> None: """Make model deviation calculation. Parameters diff --git a/deepmd/jax/atomic_model/base_atomic_model.py b/deepmd/jax/atomic_model/base_atomic_model.py index ffd58daf5e..bed75077da 100644 --- a/deepmd/jax/atomic_model/base_atomic_model.py +++ b/deepmd/jax/atomic_model/base_atomic_model.py @@ -1,19 +1,33 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +from packaging.version import ( + Version, +) + from deepmd.jax.common import ( ArrayAPIVariable, to_jax_array, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( AtomExcludeMask, PairExcludeMask, ) -def base_atomic_model_set_attr(name, value): +def base_atomic_model_set_attr(name: str, value: Any) -> Any: if name in {"out_bias", "out_std"}: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name == "pair_excl" and value is not None: value = PairExcludeMask(value.ntypes, value.exclude_types) elif name == "atom_excl" and value is not None: diff --git a/deepmd/jax/atomic_model/dp_atomic_model.py b/deepmd/jax/atomic_model/dp_atomic_model.py index adfc22c6fa..7227839f1f 100644 --- a/deepmd/jax/atomic_model/dp_atomic_model.py +++ b/deepmd/jax/atomic_model/dp_atomic_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP @@ -55,9 +54,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, diff --git a/deepmd/jax/atomic_model/linear_atomic_model.py b/deepmd/jax/atomic_model/linear_atomic_model.py index 6ce82fa07c..1c183db7ac 100644 --- a/deepmd/jax/atomic_model/linear_atomic_model.py +++ b/deepmd/jax/atomic_model/linear_atomic_model.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, +) + +from packaging.version import ( + Version, ) from deepmd.dpmodel.atomic_model.linear_atomic_model import ( @@ -22,8 +25,10 @@ to_jax_array, ) from deepmd.jax.env import ( + flax_version, jax, jnp, + nnx, ) @@ -33,13 +38,19 @@ def __setattr__(self, name: str, value: Any) -> None: value = base_atomic_model_set_attr(name, value) if name == "mapping_list": value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) elif name == "zbl_weight": - value = ArrayAPIVariable(to_jax_array(value)) + # discard since it's only used in tests + # to fix flax.errors.TraceContextError: Cannot mutate 'FlaxModule' from different trace level + return elif name == "models": value = [ DPAtomicModel.deserialize(value[0].serialize()), PairTabAtomicModel.deserialize(value[1].serialize()), ] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) return super().__setattr__(name, value) def forward_common_atomic( @@ -47,9 +58,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, diff --git a/deepmd/jax/atomic_model/pairtab_atomic_model.py b/deepmd/jax/atomic_model/pairtab_atomic_model.py index 023f4e886a..7f18a6403c 100644 --- a/deepmd/jax/atomic_model/pairtab_atomic_model.py +++ b/deepmd/jax/atomic_model/pairtab_atomic_model.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, +) + +from packaging.version import ( + Version, ) from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( @@ -16,8 +19,10 @@ to_jax_array, ) from deepmd.jax.env import ( + flax_version, jax, jnp, + nnx, ) @@ -29,6 +34,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) return super().__setattr__(name, value) def forward_common_atomic( @@ -36,9 +43,9 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, ) -> dict[str, jnp.ndarray]: return super().forward_common_atomic( extended_coord, diff --git a/deepmd/jax/common.py b/deepmd/jax/common.py index 59f36d11ad..27c7f8883d 100644 --- a/deepmd/jax/common.py +++ b/deepmd/jax/common.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, overload, ) @@ -24,7 +23,7 @@ def to_jax_array(array: np.ndarray) -> jnp.ndarray: ... def to_jax_array(array: None) -> None: ... -def to_jax_array(array: Optional[np.ndarray]) -> Optional[jnp.ndarray]: +def to_jax_array(array: np.ndarray | None) -> jnp.ndarray | None: """Convert a numpy array to a JAX array. Parameters @@ -70,11 +69,11 @@ def flax_module( metas.add(type(nnx.Module)) class MixedMetaClass(*metas): - def __call__(self, *args, **kwargs): + def __call__(self, *args: Any, **kwargs: Any) -> Any: return type(nnx.Module).__call__(self, *args, **kwargs) class FlaxModule(module, nnx.Module, metaclass=MixedMetaClass): - def __init_subclass__(cls, **kwargs) -> None: + def __init_subclass__(cls, **kwargs: Any) -> None: return super().__init_subclass__(**kwargs) def __setattr__(self, name: str, value: Any) -> None: @@ -84,20 +83,22 @@ def __setattr__(self, name: str, value: Any) -> None: class ArrayAPIVariable(nnx.Variable): - def __array__(self, *args, **kwargs): + def __array__(self, *args: Any, **kwargs: Any) -> np.ndarray: return self.value.__array__(*args, **kwargs) - def __array_namespace__(self, *args, **kwargs): + def __array_namespace__(self, *args: Any, **kwargs: Any) -> Any: return self.value.__array_namespace__(*args, **kwargs) - def __dlpack__(self, *args, **kwargs): + def __dlpack__(self, *args: Any, **kwargs: Any) -> Any: return self.value.__dlpack__(*args, **kwargs) - def __dlpack_device__(self, *args, **kwargs): + def __dlpack_device__(self, *args: Any, **kwargs: Any) -> Any: return self.value.__dlpack_device__(*args, **kwargs) -def scatter_sum(input, dim, index: jnp.ndarray, src: jnp.ndarray) -> jnp.ndarray: +def scatter_sum( + input: jnp.ndarray, dim: int, index: jnp.ndarray, src: jnp.ndarray +) -> jnp.ndarray: """Reduces all values from the src tensor to the indices specified in the index tensor.""" idx = jnp.arange(input.size, dtype=jnp.int64).reshape(input.shape) new_idx = jnp.take_along_axis(idx, index, axis=dim).ravel() diff --git a/deepmd/jax/descriptor/dpa1.py b/deepmd/jax/descriptor/dpa1.py index fef9bd5448..07695b23ed 100644 --- a/deepmd/jax/descriptor/dpa1.py +++ b/deepmd/jax/descriptor/dpa1.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.dpa1 import DescrptBlockSeAtten as DescrptBlockSeAttenDP from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP from deepmd.dpmodel.descriptor.dpa1 import GatedAttentionLayer as GatedAttentionLayerDP @@ -20,6 +24,10 @@ from deepmd.jax.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -58,6 +66,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = [ NeighborGatedAttentionLayer.deserialize(ii.serialize()) for ii in value ] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) return super().__setattr__(name, value) @@ -68,9 +78,13 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"embeddings", "embeddings_strip"}: if value is not None: value = NetworkCollection.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name == "dpa1_attention": value = NeighborGatedAttention.deserialize(value.serialize()) elif name == "env_mat": diff --git a/deepmd/jax/descriptor/dpa2.py b/deepmd/jax/descriptor/dpa2.py index 8eea324b41..8da450d2ec 100644 --- a/deepmd/jax/descriptor/dpa2.py +++ b/deepmd/jax/descriptor/dpa2.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.dpa2 import DescrptDPA2 as DescrptDPA2DP from deepmd.dpmodel.utils.network import Identity as IdentityDP from deepmd.dpmodel.utils.network import NativeLayer as NativeLayerDP @@ -23,6 +27,10 @@ from deepmd.jax.descriptor.se_t_tebd import ( DescrptBlockSeTTebd, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.network import ( NativeLayer, ) @@ -39,18 +47,23 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"repinit"}: value = DescrptBlockSeAtten.deserialize(value.serialize()) elif name in {"repinit_three_body"}: if value is not None: value = DescrptBlockSeTTebd.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"repformers"}: value = DescrptBlockRepformers.deserialize(value.serialize()) elif name in {"type_embedding"}: value = TypeEmbedNet.deserialize(value.serialize()) elif name in {"g1_shape_tranform", "tebd_transform"}: if value is None: - pass + if Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif isinstance(value, NativeLayerDP): value = NativeLayer.deserialize(value.serialize()) elif isinstance(value, IdentityDP): diff --git a/deepmd/jax/descriptor/dpa3.py b/deepmd/jax/descriptor/dpa3.py index 299afaf01d..9f734bd553 100644 --- a/deepmd/jax/descriptor/dpa3.py +++ b/deepmd/jax/descriptor/dpa3.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.dpa3 import DescrptDPA3 as DescrptDPA3DP from deepmd.jax.common import ( ArrayAPIVariable, @@ -15,6 +19,10 @@ from deepmd.jax.descriptor.repflows import ( DescrptBlockRepflows, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.type_embed import ( TypeEmbedNet, ) @@ -28,6 +36,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"repflows"}: value = DescrptBlockRepflows.deserialize(value.serialize()) elif name in {"type_embedding"}: diff --git a/deepmd/jax/descriptor/hybrid.py b/deepmd/jax/descriptor/hybrid.py index 20fc5f838b..b76e515c54 100644 --- a/deepmd/jax/descriptor/hybrid.py +++ b/deepmd/jax/descriptor/hybrid.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.hybrid import DescrptHybrid as DescrptHybridDP from deepmd.jax.common import ( ArrayAPIVariable, @@ -12,6 +16,10 @@ from deepmd.jax.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) @BaseDescriptor.register("hybrid") @@ -20,7 +28,11 @@ class DescrptHybrid(DescrptHybridDP): def __setattr__(self, name: str, value: Any) -> None: if name in {"nlist_cut_idx"}: value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) elif name in {"descrpt_list"}: value = [BaseDescriptor.deserialize(vv.serialize()) for vv in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) return super().__setattr__(name, value) diff --git a/deepmd/jax/descriptor/repflows.py b/deepmd/jax/descriptor/repflows.py index df137639aa..be26012a52 100644 --- a/deepmd/jax/descriptor/repflows.py +++ b/deepmd/jax/descriptor/repflows.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.repflows import ( DescrptBlockRepflows as DescrptBlockRepflowsDP, ) @@ -12,6 +16,10 @@ flax_module, to_jax_array, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -27,8 +35,12 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"layers"}: value = [RepFlowLayer.deserialize(layer.serialize()) for layer in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) elif name in {"edge_embd", "angle_embd"}: value = NativeLayer.deserialize(value.serialize()) elif name in {"env_mat_edge", "env_mat_angle"}: @@ -58,8 +70,12 @@ def __setattr__(self, name: str, value: Any) -> None: }: if value is not None: value = NativeLayer.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"n_residual", "e_residual", "a_residual"}: value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) else: pass return super().__setattr__(name, value) diff --git a/deepmd/jax/descriptor/repformers.py b/deepmd/jax/descriptor/repformers.py index 77ca4a9a6b..5701677349 100644 --- a/deepmd/jax/descriptor/repformers.py +++ b/deepmd/jax/descriptor/repformers.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.repformers import ( Atten2EquiVarApply as Atten2EquiVarApplyDP, ) @@ -20,6 +24,10 @@ flax_module, to_jax_array, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -36,8 +44,12 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"layers"}: value = [RepformerLayer.deserialize(layer.serialize()) for layer in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) elif name == "g2_embd": value = NativeLayer.deserialize(value.serialize()) elif name == "env_mat": @@ -87,21 +99,35 @@ def __setattr__(self, name: str, value: Any) -> None: if name in {"linear1", "linear2", "g1_self_mlp", "proj_g1g2", "proj_g1g1g2"}: if value is not None: value = NativeLayer.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"g1_residual", "g2_residual", "h2_residual"}: value = [ArrayAPIVariable(to_jax_array(vv)) for vv in value] + if Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) elif name in {"attn2g_map"}: if value is not None: value = Atten2Map.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"attn2_mh_apply"}: if value is not None: value = Atten2MultiHeadApply.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"attn2_lm"}: if value is not None: value = LayerNorm.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"attn2_ev_apply"}: if value is not None: value = Atten2EquiVarApply.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"loc_attn"}: if value is not None: value = LocalAtten.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) return super().__setattr__(name, value) diff --git a/deepmd/jax/descriptor/se_e2_a.py b/deepmd/jax/descriptor/se_e2_a.py index 31c147ad9d..4d704a4b30 100644 --- a/deepmd/jax/descriptor/se_e2_a.py +++ b/deepmd/jax/descriptor/se_e2_a.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeAArrayAPI as DescrptSeADP from deepmd.jax.common import ( ArrayAPIVariable, @@ -12,6 +16,10 @@ from deepmd.jax.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -29,9 +37,13 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"embeddings"}: if value is not None: value = NetworkCollection.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name == "env_mat": # env_mat doesn't store any value pass diff --git a/deepmd/jax/descriptor/se_e2_r.py b/deepmd/jax/descriptor/se_e2_r.py index 96ff4103dd..e5827c42af 100644 --- a/deepmd/jax/descriptor/se_e2_r.py +++ b/deepmd/jax/descriptor/se_e2_r.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.se_r import DescrptSeR as DescrptSeRDP from deepmd.jax.common import ( ArrayAPIVariable, @@ -12,6 +16,10 @@ from deepmd.jax.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -29,9 +37,13 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"embeddings"}: if value is not None: value = NetworkCollection.deserialize(value.serialize()) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name == "env_mat": # env_mat doesn't store any value pass diff --git a/deepmd/jax/descriptor/se_t.py b/deepmd/jax/descriptor/se_t.py index 029f4231fe..6d0b026c94 100644 --- a/deepmd/jax/descriptor/se_t.py +++ b/deepmd/jax/descriptor/se_t.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.se_t import DescrptSeT as DescrptSeTDP from deepmd.jax.common import ( ArrayAPIVariable, @@ -12,6 +16,10 @@ from deepmd.jax.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -30,6 +38,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"embeddings"}: if value is not None: value = NetworkCollection.deserialize(value.serialize()) diff --git a/deepmd/jax/descriptor/se_t_tebd.py b/deepmd/jax/descriptor/se_t_tebd.py index 84e3d3f084..8e2dae782a 100644 --- a/deepmd/jax/descriptor/se_t_tebd.py +++ b/deepmd/jax/descriptor/se_t_tebd.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.descriptor.se_t_tebd import ( DescrptBlockSeTTebd as DescrptBlockSeTTebdDP, ) @@ -15,6 +19,10 @@ from deepmd.jax.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.exclude_mask import ( PairExcludeMask, ) @@ -33,6 +41,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name in {"embeddings", "embeddings_strip"}: if value is not None: value = NetworkCollection.deserialize(value.serialize()) diff --git a/deepmd/jax/env.py b/deepmd/jax/env.py index 422e4ea4ce..d425d5a0e8 100644 --- a/deepmd/jax/env.py +++ b/deepmd/jax/env.py @@ -5,6 +5,7 @@ import jax import jax.numpy as jnp +from flax import __version__ as flax_version from flax import ( nnx, ) @@ -19,6 +20,7 @@ jax.config.update("jax_numpy_dtype_promotion", "strict") __all__ = [ + "flax_version", "jax", "jax_export", "jnp", diff --git a/deepmd/jax/fitting/fitting.py b/deepmd/jax/fitting/fitting.py index d62681490c..8fea40cd57 100644 --- a/deepmd/jax/fitting/fitting.py +++ b/deepmd/jax/fitting/fitting.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.fitting.dipole_fitting import DipoleFitting as DipoleFittingNetDP from deepmd.dpmodel.fitting.dos_fitting import DOSFittingNet as DOSFittingNetDP from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP @@ -17,6 +21,10 @@ flax_module, to_jax_array, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.fitting.base_fitting import ( BaseFitting, ) @@ -35,10 +43,14 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any: "fparam_inv_std", "aparam_avg", "aparam_inv_std", + "case_embd", + "default_fparam_tensor", }: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) elif name == "emask": value = AtomExcludeMask(value.ntypes, value.exclude_types) elif name == "nets": @@ -90,4 +102,6 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) return super().__setattr__(name, value) diff --git a/deepmd/jax/infer/deep_eval.py b/deepmd/jax/infer/deep_eval.py index acfd42b66a..1e29ee1c78 100644 --- a/deepmd/jax/infer/deep_eval.py +++ b/deepmd/jax/infer/deep_eval.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np @@ -83,7 +84,7 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, **kwargs: Any, ) -> None: @@ -103,7 +104,7 @@ def __init__( stablehlo_atomic_virial_no_ghost=model_data["@variables"][ "stablehlo_atomic_virial_no_ghost" ].tobytes(), - model_def_script=model_data["model_def_script"], + model_def_script=json.dumps(model_data["model_def_script"]), **model_data["constants"], ) elif model_file.endswith(".savedmodel"): @@ -158,7 +159,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepDOS elif "dipole" in model_output_type: return DeepDipole - elif "polar" in model_output_type: + elif "polar" in model_output_type or "polarizability" in model_output_type: return DeepPolar elif "wfc" in model_output_type: return DeepWFC @@ -189,11 +190,11 @@ def get_ntypes_spin(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -301,7 +302,7 @@ def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Calla """ if self.auto_batch_size is not None: - def eval_func(*args, **kwargs): + def eval_func(*args: Any, **kwargs: Any) -> Any: return self.auto_batch_size.execute_all( inner_func, numb_test, natoms, *args, **kwargs ) @@ -330,12 +331,12 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], - ): + ) -> tuple[np.ndarray, ...]: model = self.dp nframes = coords.shape[0] @@ -395,7 +396,9 @@ def _eval_model( ) # this is kinda hacky return tuple(results) - def _get_output_shape(self, odef, nframes, natoms): + def _get_output_shape( + self, odef: OutputVariableDef, nframes: int, natoms: int + ) -> list[int]: if odef.category == OutputVariableCategory.DERV_C_REDU: # virial return [nframes, *odef.shape[:-1], 9] @@ -420,3 +423,13 @@ def _get_output_shape(self, odef, nframes, natoms): def get_model_def_script(self) -> dict: """Get model definition script.""" return json.loads(self.dp.get_model_def_script()) + + def get_model(self) -> Any: + """Get the JAX model as BaseModel. + + Returns + ------- + BaseModel + The JAX model as BaseModel instance. + """ + return self.dp diff --git a/deepmd/jax/jax2tf/format_nlist.py b/deepmd/jax/jax2tf/format_nlist.py index f0c630206f..5cf93610e7 100644 --- a/deepmd/jax/jax2tf/format_nlist.py +++ b/deepmd/jax/jax2tf/format_nlist.py @@ -9,7 +9,7 @@ def format_nlist( nlist: tnp.ndarray, nsel: int, rcut: float, -): +) -> tnp.ndarray: """Format neighbor list. If nnei == nsel, do nothing; diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index 29ed131f8e..3cd30de85a 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, ) @@ -44,7 +44,7 @@ def model_call_from_call_lower( fparam: tnp.ndarray, aparam: tnp.ndarray, do_atomic_virial: bool = False, -): +) -> dict[str, tnp.ndarray]: """Return model prediction from lower interface. Parameters diff --git a/deepmd/jax/jax2tf/nlist.py b/deepmd/jax/jax2tf/nlist.py index 5a0ed58b63..c44a1196c8 100644 --- a/deepmd/jax/jax2tf/nlist.py +++ b/deepmd/jax/jax2tf/nlist.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Union, -) import tensorflow as tf import tensorflow.experimental.numpy as tnp @@ -17,7 +14,7 @@ def build_neighbor_list( atype: tnp.ndarray, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> tnp.ndarray: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -115,7 +112,7 @@ def nlist_distinguish_types( nlist: tnp.ndarray, atype: tnp.ndarray, sel: list[int], -): +) -> tnp.ndarray: """Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types. @@ -140,7 +137,7 @@ def nlist_distinguish_types( return ret -def tf_outer(a, b): +def tf_outer(a: tnp.ndarray, b: tnp.ndarray) -> tnp.ndarray: return tf.einsum("i,j->ij", a, b) @@ -150,7 +147,7 @@ def extend_coord_with_ghosts( atype: tnp.ndarray, cell: tnp.ndarray, rcut: float, -): +) -> tuple[tnp.ndarray, tnp.ndarray, tnp.ndarray]: """Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors within rcut are appended. diff --git a/deepmd/jax/jax2tf/region.py b/deepmd/jax/jax2tf/region.py index 96024bd79a..a90e693478 100644 --- a/deepmd/jax/jax2tf/region.py +++ b/deepmd/jax/jax2tf/region.py @@ -93,7 +93,7 @@ def to_face_distance( return tnp.reshape(dist, tf.concat([cshape[:-2], [3]], axis=0)) -def b_to_face_distance(cell): +def b_to_face_distance(cell: tnp.ndarray) -> tnp.ndarray: volume = tf.linalg.det(cell) c_yz = tf.linalg.cross(cell[:, 1, ...], cell[:, 2, ...]) _h2yz = volume / tf.linalg.norm(c_yz, axis=-1) diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index aac022ace9..e819ebf65a 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json -from typing import ( - Optional, +from collections.abc import ( + Callable, ) import tensorflow as tf @@ -38,10 +38,17 @@ def deserialize_to_file(model_file: str, data: dict) -> None: tf_model = tf.Module() - def exported_whether_do_atomic_virial(do_atomic_virial, has_ghost_atoms): + def exported_whether_do_atomic_virial( + do_atomic_virial: bool, has_ghost_atoms: bool + ) -> Callable: def call_lower_with_fixed_do_atomic_virial( - coord, atype, nlist, mapping, fparam, aparam - ): + coord: tnp.ndarray, + atype: tnp.ndarray, + nlist: tnp.ndarray, + mapping: tnp.ndarray, + fparam: tnp.ndarray, + aparam: tnp.ndarray, + ) -> dict[str, tnp.ndarray]: return call_lower( coord, atype, @@ -86,8 +93,13 @@ def call_lower_with_fixed_do_atomic_virial( ], ) def call_lower_without_atomic_virial( - coord, atype, nlist, mapping, fparam, aparam - ): + coord: tnp.ndarray, + atype: tnp.ndarray, + nlist: tnp.ndarray, + mapping: tnp.ndarray, + fparam: tnp.ndarray, + aparam: tnp.ndarray, + ) -> dict[str, tnp.ndarray]: nlist = format_nlist(coord, nlist, model.get_nnei(), model.get_rcut()) return tf.cond( tf.shape(coord)[1] == tf.shape(nlist)[1], @@ -112,7 +124,14 @@ def call_lower_without_atomic_virial( tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), ], ) - def call_lower_with_atomic_virial(coord, atype, nlist, mapping, fparam, aparam): + def call_lower_with_atomic_virial( + coord: tnp.ndarray, + atype: tnp.ndarray, + nlist: tnp.ndarray, + mapping: tnp.ndarray, + fparam: tnp.ndarray, + aparam: tnp.ndarray, + ) -> dict[str, tnp.ndarray]: nlist = format_nlist(coord, nlist, model.get_nnei(), model.get_rcut()) return tf.cond( tf.shape(coord)[1] == tf.shape(nlist)[1], @@ -126,7 +145,7 @@ def call_lower_with_atomic_virial(coord, atype, nlist, mapping, fparam, aparam): tf_model.call_lower_atomic_virial = call_lower_with_atomic_virial - def make_call_whether_do_atomic_virial(do_atomic_virial: bool): + def make_call_whether_do_atomic_virial(do_atomic_virial: bool) -> Callable: if do_atomic_virial: call_lower = call_lower_with_atomic_virial else: @@ -135,10 +154,10 @@ def make_call_whether_do_atomic_virial(do_atomic_virial: bool): def call( coord: tnp.ndarray, atype: tnp.ndarray, - box: Optional[tnp.ndarray] = None, - fparam: Optional[tnp.ndarray] = None, - aparam: Optional[tnp.ndarray] = None, - ): + box: tnp.ndarray | None = None, + fparam: tnp.ndarray | None = None, + aparam: tnp.ndarray | None = None, + ) -> dict[str, tnp.ndarray]: """Return model prediction. Parameters @@ -194,7 +213,7 @@ def call_with_atomic_virial( box: tnp.ndarray, fparam: tnp.ndarray, aparam: tnp.ndarray, - ): + ) -> dict[str, tnp.ndarray]: return make_call_whether_do_atomic_virial(do_atomic_virial=True)( coord, atype, box, fparam, aparam ) @@ -217,7 +236,7 @@ def call_without_atomic_virial( box: tnp.ndarray, fparam: tnp.ndarray, aparam: tnp.ndarray, - ): + ) -> dict[str, tnp.ndarray]: return make_call_whether_do_atomic_virial(do_atomic_virial=False)( coord, atype, box, fparam, aparam ) @@ -226,49 +245,49 @@ def call_without_atomic_virial( # set functions to export other attributes @tf.function - def get_type_map(): + def get_type_map() -> tf.Tensor: return tf.constant(model.get_type_map(), dtype=tf.string) tf_model.get_type_map = get_type_map @tf.function - def get_rcut(): + def get_rcut() -> tf.Tensor: return tf.constant(model.get_rcut(), dtype=tf.double) tf_model.get_rcut = get_rcut @tf.function - def get_dim_fparam(): + def get_dim_fparam() -> tf.Tensor: return tf.constant(model.get_dim_fparam(), dtype=tf.int64) tf_model.get_dim_fparam = get_dim_fparam @tf.function - def get_dim_aparam(): + def get_dim_aparam() -> tf.Tensor: return tf.constant(model.get_dim_aparam(), dtype=tf.int64) tf_model.get_dim_aparam = get_dim_aparam @tf.function - def get_sel_type(): + def get_sel_type() -> tf.Tensor: return tf.constant(model.get_sel_type(), dtype=tf.int64) tf_model.get_sel_type = get_sel_type @tf.function - def is_aparam_nall(): + def is_aparam_nall() -> tf.Tensor: return tf.constant(model.is_aparam_nall(), dtype=tf.bool) tf_model.is_aparam_nall = is_aparam_nall @tf.function - def model_output_type(): + def model_output_type() -> tf.Tensor: return tf.constant(model.model_output_type(), dtype=tf.string) tf_model.model_output_type = model_output_type @tf.function - def mixed_types(): + def mixed_types() -> tf.Tensor: return tf.constant(model.mixed_types(), dtype=tf.bool) tf_model.mixed_types = mixed_types @@ -276,19 +295,19 @@ def mixed_types(): if model.get_min_nbor_dist() is not None: @tf.function - def get_min_nbor_dist(): + def get_min_nbor_dist() -> tf.Tensor: return tf.constant(model.get_min_nbor_dist(), dtype=tf.double) tf_model.get_min_nbor_dist = get_min_nbor_dist @tf.function - def get_sel(): + def get_sel() -> tf.Tensor: return tf.constant(model.get_sel(), dtype=tf.int64) tf_model.get_sel = get_sel @tf.function - def get_model_def_script(): + def get_model_def_script() -> tf.Tensor: return tf.constant( json.dumps(model_def_script, separators=(",", ":")), dtype=tf.string ) diff --git a/deepmd/jax/jax2tf/tfmodel.py b/deepmd/jax/jax2tf/tfmodel.py index 0d7b13ba1f..85115547d7 100644 --- a/deepmd/jax/jax2tf/tfmodel.py +++ b/deepmd/jax/jax2tf/tfmodel.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) import jax.experimental.jax2tf as jax2tf @@ -45,7 +44,7 @@ def decode_list_of_bytes(list_of_bytes: list[bytes]) -> list[str]: class TFModelWrapper(tf.Module): def __init__( self, - model, + model: str, ) -> None: self.model = tf.saved_model.load(model) self._call_lower = jax2tf.call_tf(self.model.call_lower) @@ -75,9 +74,9 @@ def __call__( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> Any: """Return model prediction. @@ -111,11 +110,11 @@ def call( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, jnp.ndarray]: """Return model prediction. Parameters @@ -165,7 +164,7 @@ def call( aparam, ) - def model_output_def(self): + def model_output_def(self) -> ModelOutputDef: return ModelOutputDef( FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()]) ) @@ -175,11 +174,11 @@ def call_lower( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, jnp.ndarray]: if do_atomic_virial: call_lower = self._call_lower_atomic_virial else: @@ -207,15 +206,15 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map - def get_rcut(self): + def get_rcut(self) -> float: """Get the cut-off radius.""" return self.rcut - def get_dim_fparam(self): + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.dim_fparam - def get_dim_aparam(self): + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.dim_aparam @@ -269,7 +268,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist @@ -291,9 +290,9 @@ def mixed_types(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/jax/model/base_model.py b/deepmd/jax/model/base_model.py index 7c97ff692f..533181e250 100644 --- a/deepmd/jax/model/base_model.py +++ b/deepmd/jax/model/base_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) from deepmd.dpmodel.model.base_model import ( make_base_model, @@ -20,15 +17,15 @@ def forward_common_atomic( - self, + self: "BaseModel", extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, -): +) -> dict[str, jnp.ndarray]: atomic_ret = self.atomic_model.forward_common_atomic( extended_coord, extended_atype, @@ -46,21 +43,30 @@ def forward_common_atomic( atom_axis = -(len(shap) + 1) if vdef.reducible: kk_redu = get_reduce_name(kk) - model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis) + if vdef.intensive: + mask = atomic_ret["mask"] if "mask" in atomic_ret else None + if mask is not None: + model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis) / jnp.sum( + mask, axis=-1, keepdims=True + ) + else: + model_predict[kk_redu] = jnp.mean(vv, axis=atom_axis) + else: + model_predict[kk_redu] = jnp.sum(vv, axis=atom_axis) kk_derv_r, kk_derv_c = get_deriv_name(kk) if vdef.r_differentiable: def eval_output( - cc_ext, - extended_atype, - nlist, - mapping, - fparam, - aparam, + cc_ext: jnp.ndarray, + extended_atype: jnp.ndarray, + nlist: jnp.ndarray, + mapping: jnp.ndarray | None, + fparam: jnp.ndarray | None, + aparam: jnp.ndarray | None, *, - _kk=kk, - _atom_axis=atom_axis, - ): + _kk: str = kk, + _atom_axis: int = atom_axis, + ) -> jnp.ndarray: atomic_ret = self.atomic_model.forward_common_atomic( cc_ext[None, ...], extended_atype[None, ...], @@ -108,16 +114,16 @@ def eval_output( if do_atomic_virial: def eval_ce( - cc_ext, - extended_atype, - nlist, - mapping, - fparam, - aparam, + cc_ext: jnp.ndarray, + extended_atype: jnp.ndarray, + nlist: jnp.ndarray, + mapping: jnp.ndarray | None, + fparam: jnp.ndarray | None, + aparam: jnp.ndarray | None, *, - _kk=kk, - _atom_axis=atom_axis - 1, - ): + _kk: str = kk, + _atom_axis: int = atom_axis - 1, + ) -> jnp.ndarray: # atomic_ret[_kk]: [nf, nloc, *def] atomic_ret = self.atomic_model.forward_common_atomic( cc_ext[None, ...], diff --git a/deepmd/jax/model/dp_model.py b/deepmd/jax/model/dp_model.py index 436582f22b..5545b5505b 100644 --- a/deepmd/jax/model/dp_model.py +++ b/deepmd/jax/model/dp_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.model import ( @@ -52,11 +51,11 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, jnp.ndarray]: return forward_common_atomic( self, extended_coord, @@ -74,7 +73,7 @@ def format_nlist( extended_atype: jnp.ndarray, nlist: jnp.ndarray, extra_nlist_sort: bool = False, - ): + ) -> jnp.ndarray: return dpmodel_model.format_nlist( self, jax.lax.stop_gradient(extended_coord), diff --git a/deepmd/jax/model/dp_zbl_model.py b/deepmd/jax/model/dp_zbl_model.py index babbc65233..d9be671a46 100644 --- a/deepmd/jax/model/dp_zbl_model.py +++ b/deepmd/jax/model/dp_zbl_model.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.model.dp_zbl_model import DPZBLModel as DPZBLModelDP @@ -34,11 +33,11 @@ def forward_common_atomic( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, jnp.ndarray]: return forward_common_atomic( self, extended_coord, @@ -56,7 +55,7 @@ def format_nlist( extended_atype: jnp.ndarray, nlist: jnp.ndarray, extra_nlist_sort: bool = False, - ): + ) -> jnp.ndarray: return DPZBLModelDP.format_nlist( self, jax.lax.stop_gradient(extended_coord), diff --git a/deepmd/jax/model/hlo.py b/deepmd/jax/model/hlo.py index 4d59957456..47959dd130 100644 --- a/deepmd/jax/model/hlo.py +++ b/deepmd/jax/model/hlo.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, ) from deepmd.dpmodel.model.make_model import ( @@ -44,21 +43,21 @@ class HLO(BaseModel): def __init__( self, - stablehlo, - stablehlo_atomic_virial, - stablehlo_no_ghost, - stablehlo_atomic_virial_no_ghost, - model_def_script, - type_map, - rcut, - dim_fparam, - dim_aparam, - sel_type, - is_aparam_nall, - model_output_type, - mixed_types, - min_nbor_dist, - sel, + stablehlo: bytearray, + stablehlo_atomic_virial: bytearray, + stablehlo_no_ghost: bytearray, + stablehlo_atomic_virial_no_ghost: bytearray, + model_def_script: str, + type_map: list[str], + rcut: float, + dim_fparam: int, + dim_aparam: int, + sel_type: list[int], + is_aparam_nall: bool, + model_output_type: str, + mixed_types: bool, + min_nbor_dist: float | None, + sel: list[int], ) -> None: self._call_lower = jax_export.deserialize(stablehlo).call self._call_lower_atomic_virial = jax_export.deserialize( @@ -85,9 +84,9 @@ def __call__( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, ) -> Any: """Return model prediction. @@ -121,11 +120,11 @@ def call( self, coord: jnp.ndarray, atype: jnp.ndarray, - box: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + box: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, jnp.ndarray]: """Return model prediction. Parameters @@ -165,7 +164,7 @@ def call( do_atomic_virial=do_atomic_virial, ) - def model_output_def(self): + def model_output_def(self) -> ModelOutputDef: return ModelOutputDef( FittingOutputDef([OUTPUT_DEFS[tt] for tt in self.model_output_type()]) ) @@ -175,11 +174,11 @@ def call_lower( extended_coord: jnp.ndarray, extended_atype: jnp.ndarray, nlist: jnp.ndarray, - mapping: Optional[jnp.ndarray] = None, - fparam: Optional[jnp.ndarray] = None, - aparam: Optional[jnp.ndarray] = None, + mapping: jnp.ndarray | None = None, + fparam: jnp.ndarray | None = None, + aparam: jnp.ndarray | None = None, do_atomic_virial: bool = False, - ): + ) -> dict[str, jnp.ndarray]: if extended_coord.shape[1] > nlist.shape[1]: if do_atomic_virial: call_lower = self._call_lower_atomic_virial @@ -203,15 +202,15 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map - def get_rcut(self): + def get_rcut(self) -> float: """Get the cut-off radius.""" return self.rcut - def get_dim_fparam(self): + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.dim_fparam - def get_dim_aparam(self): + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.dim_aparam @@ -265,7 +264,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist @@ -287,9 +286,9 @@ def mixed_types(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/jax/model/model.py b/deepmd/jax/model/model.py index dc350e968c..321f33b315 100644 --- a/deepmd/jax/model/model.py +++ b/deepmd/jax/model/model.py @@ -26,7 +26,7 @@ ) -def get_standard_model(data: dict): +def get_standard_model(data: dict) -> BaseModel: """Get a Model from a dictionary. Parameters @@ -103,7 +103,7 @@ def get_zbl_model(data: dict) -> DPZBLModel: ) -def get_model(data: dict): +def get_model(data: dict) -> BaseModel: """Get a model from a dictionary. Parameters diff --git a/deepmd/jax/utils/auto_batch_size.py b/deepmd/jax/utils/auto_batch_size.py index ef93d85d96..1ecf020086 100644 --- a/deepmd/jax/utils/auto_batch_size.py +++ b/deepmd/jax/utils/auto_batch_size.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import jaxlib from deepmd.jax.env import ( jax, @@ -52,7 +51,7 @@ def is_oom_error(self, e: Exception) -> bool: # several sources think CUSOLVER_STATUS_INTERNAL_ERROR is another out-of-memory error, # such as https://github.com/JuliaGPU/CUDA.jl/issues/1924 # (the meaningless error message should be considered as a bug in cusolver) - if isinstance(e, (jaxlib.xla_extension.XlaRuntimeError, ValueError)) and ( + if isinstance(e, (RuntimeError, ValueError)) and ( "RESOURCE_EXHAUSTED:" in e.args[0] ): return True diff --git a/deepmd/jax/utils/exclude_mask.py b/deepmd/jax/utils/exclude_mask.py index 18d13d9400..4ae230c8dc 100644 --- a/deepmd/jax/utils/exclude_mask.py +++ b/deepmd/jax/utils/exclude_mask.py @@ -3,6 +3,10 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.utils.exclude_mask import AtomExcludeMask as AtomExcludeMaskDP from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask as PairExcludeMaskDP from deepmd.jax.common import ( @@ -10,6 +14,10 @@ flax_module, to_jax_array, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) @flax_module @@ -19,6 +27,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) return super().__setattr__(name, value) @@ -29,4 +39,6 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) return super().__setattr__(name, value) diff --git a/deepmd/jax/utils/neighbor_stat.py b/deepmd/jax/utils/neighbor_stat.py index 6d9bc872e8..c462665626 100644 --- a/deepmd/jax/utils/neighbor_stat.py +++ b/deepmd/jax/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np @@ -81,8 +78,8 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - cell: Optional[np.ndarray], - ): + cell: np.ndarray | None, + ) -> tuple[np.ndarray, np.ndarray]: """Execute the operation. Parameters diff --git a/deepmd/jax/utils/network.py b/deepmd/jax/utils/network.py index 2c406095cd..72d9f760eb 100644 --- a/deepmd/jax/utils/network.py +++ b/deepmd/jax/utils/network.py @@ -4,6 +4,11 @@ ClassVar, ) +import numpy as np +from packaging.version import ( + Version, +) + from deepmd.dpmodel.common import ( NativeOP, ) @@ -16,25 +21,27 @@ make_multilayer_network, ) from deepmd.jax.common import ( + ArrayAPIVariable, flax_module, to_jax_array, ) from deepmd.jax.env import ( + flax_version, nnx, ) class ArrayAPIParam(nnx.Param): - def __array__(self, *args, **kwargs): + def __array__(self, *args: Any, **kwargs: Any) -> np.ndarray: return self.value.__array__(*args, **kwargs) - def __array_namespace__(self, *args, **kwargs): + def __array_namespace__(self, *args: Any, **kwargs: Any) -> Any: return self.value.__array_namespace__(*args, **kwargs) - def __dlpack__(self, *args, **kwargs): + def __dlpack__(self, *args: Any, **kwargs: Any) -> Any: return self.value.__dlpack__(*args, **kwargs) - def __dlpack_device__(self, *args, **kwargs): + def __dlpack_device__(self, *args: Any, **kwargs: Any) -> Any: return self.value.__dlpack_device__(*args, **kwargs) @@ -44,13 +51,19 @@ def __setattr__(self, name: str, value: Any) -> None: if name in {"w", "b", "idt"}: value = to_jax_array(value) if value is not None: - value = ArrayAPIParam(value) + if self.trainable: + value = ArrayAPIParam(value) + else: + value = ArrayAPIVariable(value) return super().__setattr__(name, value) @flax_module class NativeNet(make_multilayer_network(NativeLayer, NativeOP)): - pass + def __setattr__(self, name: str, value: Any) -> None: + if name in {"layers"} and Version(flax_version) >= Version("0.12.0"): + value = nnx.List(value) + return super().__setattr__(name, value) class EmbeddingNet(make_embedding_network(NativeNet, NativeLayer)): @@ -69,6 +82,11 @@ class NetworkCollection(NetworkCollectionDP): "fitting_network": FittingNet, } + def __setattr__(self, name: str, value: Any) -> None: + if name in {"_networks"} and Version(flax_version) >= Version("0.12.0"): + value = nnx.List([nnx.data(item) for item in value]) + return super().__setattr__(name, value) + class LayerNorm(LayerNormDP, NativeLayer): pass diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 5d4da49e08..6a3c839608 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -55,10 +55,15 @@ def deserialize_to_file(model_file: str, data: dict) -> None: def exported_whether_do_atomic_virial( do_atomic_virial: bool, has_ghost_atoms: bool - ): + ) -> "jax_export.Exported": def call_lower_with_fixed_do_atomic_virial( - coord, atype, nlist, mapping, fparam, aparam - ): + coord: jnp.ndarray, + atype: jnp.ndarray, + nlist: jnp.ndarray, + mapping: jnp.ndarray, + fparam: jnp.ndarray, + aparam: jnp.ndarray, + ) -> dict[str, jnp.ndarray]: return call_lower( coord, atype, diff --git a/deepmd/jax/utils/type_embed.py b/deepmd/jax/utils/type_embed.py index 30cd9f45a9..aff0a78a2c 100644 --- a/deepmd/jax/utils/type_embed.py +++ b/deepmd/jax/utils/type_embed.py @@ -3,12 +3,20 @@ Any, ) +from packaging.version import ( + Version, +) + from deepmd.dpmodel.utils.type_embed import TypeEmbedNet as TypeEmbedNetDP from deepmd.jax.common import ( ArrayAPIVariable, flax_module, to_jax_array, ) +from deepmd.jax.env import ( + flax_version, + nnx, +) from deepmd.jax.utils.network import ( EmbeddingNet, ) @@ -21,6 +29,8 @@ def __setattr__(self, name: str, value: Any) -> None: value = to_jax_array(value) if value is not None: value = ArrayAPIVariable(value) + elif Version(flax_version) >= Version("0.12.0"): + value = nnx.data(value) if name in {"embedding_net"}: value = EmbeddingNet.deserialize(value.serialize()) return super().__setattr__(name, value) diff --git a/deepmd/lmp.py b/deepmd/lmp.py index 15959cf243..3daf269c92 100644 --- a/deepmd/lmp.py +++ b/deepmd/lmp.py @@ -9,9 +9,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import torch # noqa: TID253 from packaging.version import ( @@ -34,7 +31,7 @@ find_libpython = None -def get_env(paths: list[Optional[str]]) -> str: +def get_env(paths: list[str | None]) -> str: """Get the environment variable from given paths.""" return ":".join(p for p in paths if p is not None) @@ -81,7 +78,6 @@ def get_library_path(module: str, filename: str) -> list[str]: pt_dir = os.path.join(torch.__path__[0], "lib") op_dir = str(SHARED_LIB_DIR) - cuda_library_paths = [] if platform.system() == "Linux": cuda_library_paths.extend( diff --git a/deepmd/lmp_check_build.py b/deepmd/lmp_check_build.py new file mode 100644 index 0000000000..dc81a31d8e --- /dev/null +++ b/deepmd/lmp_check_build.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.env import ( + GLOBAL_CONFIG, +) + +if GLOBAL_CONFIG.get("lammps_version", "") == "": + + def get_op_dir() -> str: + """Get the directory of the deepmd-kit OP library.""" + # empty + return "" +else: + from deepmd.lmp import ( + get_op_dir, + ) + +__all__ = [ + "get_op_dir", +] diff --git a/deepmd/loggers/loggers.py b/deepmd/loggers/loggers.py index 060ac1a4b0..aa5c73aa6e 100644 --- a/deepmd/loggers/loggers.py +++ b/deepmd/loggers/loggers.py @@ -44,7 +44,7 @@ class _AppFilter(logging.Filter): """Add field `app_name` to log messages.""" - def filter(self, record) -> bool: + def filter(self, record: logging.LogRecord) -> bool: record.app_name = "DEEPMD" return True @@ -56,7 +56,7 @@ def __init__(self, rank: int) -> None: super().__init__(name="MPI_rank_id") self.mpi_rank = str(rank) - def filter(self, record) -> bool: + def filter(self, record: logging.LogRecord) -> bool: record.rank = self.mpi_rank return True @@ -68,7 +68,7 @@ def __init__(self, rank: int) -> None: super().__init__(name="MPI_master_log") self.mpi_rank = rank - def filter(self, record) -> bool: + def filter(self, record: logging.LogRecord) -> bool: if self.mpi_rank == 0: return True else: @@ -135,16 +135,16 @@ def __init__( self.MPI = MPI super().__init__(filename, mode=mode, encoding=None, delay=False) - def _open(self): + def _open(self) -> "_MPIFileStream": return _MPIFileStream(self.baseFilename, self.MPI, self.mode) - def setStream(self, stream) -> NoReturn: + def setStream(self, stream: "_MPIFileStream") -> NoReturn: """Stream cannot be reasigned in MPI mode.""" raise NotImplementedError("Unable to do for MPI file handler!") def set_log_handles( - level: int, log_path: Optional["Path"] = None, mpi_log: Optional[str] = None + level: int, log_path: Optional["Path"] = None, mpi_log: str | None = None ) -> None: """Set desired level for package loggers and add file handlers. diff --git a/deepmd/loggers/training.py b/deepmd/loggers/training.py index 5de7926460..c7fe94e24d 100644 --- a/deepmd/loggers/training.py +++ b/deepmd/loggers/training.py @@ -1,14 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import datetime -from typing import ( - Optional, -) def format_training_message( batch: int, wall_time: float, - eta: Optional[int] = None, + eta: int | None = None, ) -> str: """Format a training message.""" msg = f"batch {batch:7d}: total wall time = {wall_time:.2f} s" @@ -21,7 +18,7 @@ def format_training_message_per_task( batch: int, task_name: str, rmse: dict[str, float], - learning_rate: Optional[float], + learning_rate: float | None, ) -> str: if task_name: task_name += ": " diff --git a/deepmd/main.py b/deepmd/main.py index 14c0390bdc..62118ae3c6 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -14,7 +14,7 @@ defaultdict, ) from typing import ( - Optional, + Any, ) from deepmd.backend.backend import ( @@ -63,19 +63,31 @@ class RawTextArgumentDefaultsHelpFormatter( class BackendOption(argparse.Action): """Map backend alias to unique name.""" - def __call__(self, parser, namespace, values, option_string=None): + def __call__( + self, + parser: argparse.ArgumentParser, + namespace: argparse.Namespace, + values: Any, + option_string: str | None = None, + ) -> None: setattr(namespace, self.dest, BACKEND_TABLE[values]) class DeprecateAction(argparse.Action): # See https://stackoverflow.com/a/69052677/9567349 by Ibolit under CC BY-SA 4.0 - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args: Any, **kwargs: Any) -> None: self.call_count = 0 if "help" in kwargs: kwargs["help"] = f"[DEPRECATED] {kwargs['help']}" super().__init__(*args, **kwargs) - def __call__(self, parser, namespace, values, option_string=None): + def __call__( + self, + parser: argparse.ArgumentParser, + namespace: argparse.Namespace, + values: Any, + option_string: str | None = None, + ) -> None: if self.call_count == 0: warnings.warn( f"The option `{option_string}` is deprecated. It will be ignored.", @@ -112,7 +124,7 @@ def main_parser() -> argparse.ArgumentParser: if default_backend not in BACKEND_TABLE.keys(): raise ValueError( f"Unknown backend {default_backend}. " - "Please set DP_BACKEND to either tensorflow or pytorch." + "Please set DP_BACKEND to either tensorflow, pytorch, or paddle." ) parser_backend = parser.add_mutually_exclusive_group() @@ -312,7 +324,7 @@ def main_parser() -> argparse.ArgumentParser: "--output", type=str, default="frozen_model", - help="Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth", + help="Filename (prefix) of the output model file. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth; Paddle backend: suffix is .json and .pdiparams", ) parser_frz.add_argument( "-n", @@ -371,6 +383,24 @@ def main_parser() -> argparse.ArgumentParser: type=str, help="The path to the datafile, each line of which is a path to one data system.", ) + parser_tst_subgroup.add_argument( + "--train-data", + dest="train_json", + default=None, + type=str, + help=( + "The input json file. Training data in the file will be used for testing." + ), + ) + parser_tst_subgroup.add_argument( + "--valid-data", + dest="valid_json", + default=None, + type=str, + help=( + "The input json file. Validation data in the file will be used for testing." + ), + ) parser_tst.add_argument( "-S", "--set-prefix", @@ -416,6 +446,56 @@ def main_parser() -> argparse.ArgumentParser: help="(Supported backend: PyTorch) Task head (alias: model branch) to test if in multi-task mode.", ) + # * eval_desc script *************************************************************** + parser_eval_desc = subparsers.add_parser( + "eval-desc", + parents=[parser_log], + help="evaluate descriptors using the model", + formatter_class=RawTextArgumentDefaultsHelpFormatter, + epilog=textwrap.dedent( + """\ + examples: + dp eval-desc -m graph.pb -s /path/to/system -o desc + """ + ), + ) + parser_eval_desc.add_argument( + "-m", + "--model", + default="frozen_model", + type=str, + help="Frozen model file (prefix) to import. TensorFlow backend: suffix is .pb; PyTorch backend: suffix is .pth.", + ) + parser_eval_desc_subgroup = parser_eval_desc.add_mutually_exclusive_group() + parser_eval_desc_subgroup.add_argument( + "-s", + "--system", + default=".", + type=str, + help="The system dir. Recursively detect systems in this directory", + ) + parser_eval_desc_subgroup.add_argument( + "-f", + "--datafile", + default=None, + type=str, + help="The path to the datafile, each line of which is a path to one data system.", + ) + parser_eval_desc.add_argument( + "-o", + "--output", + default="desc", + type=str, + help="Output directory for descriptor files. Descriptors will be saved as desc/(system_name).npy", + ) + parser_eval_desc.add_argument( + "--head", + "--model-branch", + default=None, + type=str, + help="(Supported backend: PyTorch) Task head (alias: model branch) to use if in multi-task mode.", + ) + # * compress model ***************************************************************** # Compress a model, which including tabulating the embedding-net. # The table is composed of fifth-order polynomial coefficients and is assembled @@ -671,12 +751,13 @@ def main_parser() -> argparse.ArgumentParser: parser_change_bias = subparsers.add_parser( "change-bias", parents=[parser_log], - help="(Supported backend: PyTorch) Change model out bias according to the input data.", + help="Change model out bias according to the input data.", formatter_class=RawTextArgumentDefaultsHelpFormatter, epilog=textwrap.dedent( """\ examples: - dp change-bias model.pt -s data -n 10 -m change + dp --pt change-bias model.pt -s data -n 10 -m change + dp --tf change-bias model.ckpt -s data -n 10 -m change """ ), ) @@ -851,13 +932,20 @@ def main_parser() -> argparse.ArgumentParser: ) parser_show.add_argument( "ATTRIBUTES", - choices=["model-branch", "type-map", "descriptor", "fitting-net", "size"], + choices=[ + "model-branch", + "type-map", + "descriptor", + "fitting-net", + "size", + "observed-type", + ], nargs="+", ) return parser -def parse_args(args: Optional[list[str]] = None) -> argparse.Namespace: +def parse_args(args: list[str] | None = None) -> argparse.Namespace: """Parse arguments and convert argument strings to objects. Parameters @@ -881,7 +969,7 @@ def parse_args(args: Optional[list[str]] = None) -> argparse.Namespace: return parsed_args -def main(args: Optional[list[str]] = None) -> None: +def main(args: list[str] | None = None) -> None: """DeePMD-kit new entry point. Parameters @@ -902,6 +990,7 @@ def main(args: Optional[list[str]] = None) -> None: if args.command in ( "test", + "eval-desc", "doc-train-input", "model-devi", "neighbor-stat", diff --git a/deepmd/pd/cxx_op.py b/deepmd/pd/cxx_op.py new file mode 100644 index 0000000000..9962bc0d24 --- /dev/null +++ b/deepmd/pd/cxx_op.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import importlib +from types import ( + ModuleType, +) + + +def load_library(module_name: str) -> tuple[bool, ModuleType]: + """Load OP library and return the module if success. + + Parameters + ---------- + module_name : str + Name of the module + + Returns + ------- + bool + Whether the library is loaded successfully + ModuleType + loaded custom operator module + """ + if importlib.util.find_spec(module_name) is not None: + module = importlib.import_module(module_name) + return True, module + + return False, None + + +ENABLE_CUSTOMIZED_OP, paddle_ops_deepmd = load_library("deepmd_op_pd") + +__all__ = [ + "ENABLE_CUSTOMIZED_OP", + "paddle_ops_deepmd", +] diff --git a/deepmd/pd/entrypoints/main.py b/deepmd/pd/entrypoints/main.py index 8d96c4e6f2..8600d73bc9 100644 --- a/deepmd/pd/entrypoints/main.py +++ b/deepmd/pd/entrypoints/main.py @@ -7,8 +7,7 @@ Path, ) from typing import ( - Optional, - Union, + Any, ) import h5py @@ -22,6 +21,7 @@ ) from deepmd.common import ( expand_sys_str, + j_loader, ) from deepmd.loggers.loggers import ( set_log_handles, @@ -79,35 +79,40 @@ def get_trainer( - config, - init_model=None, - restart_model=None, - finetune_model=None, - force_load=False, - init_frz_model=None, - shared_links=None, - finetune_links=None, -): + config: dict[str, Any], + init_model: str | None = None, + restart_model: str | None = None, + finetune_model: str | None = None, + force_load: bool = False, + init_frz_model: str | None = None, + shared_links: dict[str, Any] | None = None, + finetune_links: dict[str, Any] | None = None, +) -> training.Trainer: multi_task = "model_dict" in config.get("model", {}) # Initialize DDP world_size = dist.get_world_size() if world_size > 1: - assert paddle.version.nccl() != "0" + assert not paddle.core.is_compiled_with_nccl() or paddle.version.nccl() != "0" fleet.init(is_collective=True) def prepare_trainer_input_single( - model_params_single, data_dict_single, rank=0, seed=None - ): + model_params_single: dict[str, Any], + data_dict_single: dict[str, Any], + rank: int = 0, + seed: int | None = None, + ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data", None) validation_systems = ( validation_dataset_params["systems"] if validation_dataset_params else None ) training_systems = training_dataset_params["systems"] - training_systems = process_systems(training_systems) + trn_patterns = training_dataset_params.get("rglob_patterns", None) + training_systems = process_systems(training_systems, patterns=trn_patterns) if validation_systems is not None: - validation_systems = process_systems(validation_systems) + val_patterns = validation_dataset_params.get("rglob_patterns", None) + validation_systems = process_systems(validation_systems, val_patterns) # stat files stat_file_path_single = data_dict_single.get("stat_file", None) @@ -207,7 +212,7 @@ def get_compute_device(self) -> str: def get_ngpus(self) -> int: """Get the number of GPUs.""" - return paddle.device.cuda.device_count() + return paddle.device.device_count() def get_backend_info(self) -> dict: """Get backend information.""" @@ -222,10 +227,10 @@ def get_backend_info(self) -> dict: def train( input_file: str, - init_model: Optional[str], - restart: Optional[str], - finetune: Optional[str], - init_frz_model: Optional[str], + init_model: str | None, + restart: str | None, + finetune: str | None, + init_frz_model: str | None, model_branch: str, skip_neighbor_stat: bool = False, use_pretrain_script: bool = False, @@ -235,8 +240,7 @@ def train( log.info("Configuration path: %s", input_file) if LOCAL_RANK == 0: SummaryPrinter()() - with open(input_file) as fin: - config = json.load(fin) + config = j_loader(input_file) # ensure suffix, as in the command line help, we say "path prefix of checkpoint files" if init_model is not None and not init_model.endswith(".pd"): init_model += ".pd" @@ -341,7 +345,8 @@ def train( def freeze( model: str, output: str = "frozen_model.json", - head: Optional[str] = None, + head: str | None = None, + do_atomic_virial: bool = False, ) -> None: paddle.set_flags( { @@ -368,12 +373,13 @@ def freeze( model.forward = paddle.jit.to_static( model.forward, input_spec=[ - InputSpec([1, -1, 3], dtype="float64", name="coord"), # coord - InputSpec([1, -1], dtype="int64", name="atype"), # atype - InputSpec([1, 9], dtype="float64", name="box"), # box + InputSpec([-1, -1, 3], dtype="float64", name="coord"), # coord + InputSpec([-1, -1], dtype="int64", name="atype"), # atype + InputSpec([-1, 9], dtype="float64", name="box"), # box None, # fparam None, # aparam - True, # do_atomic_virial + # InputSpec([], dtype="bool", name="do_atomic_virial"), # do_atomic_virial + do_atomic_virial, # do_atomic_virial ], full_graph=True, ) @@ -388,17 +394,46 @@ def freeze( model.forward_lower = paddle.jit.to_static( model.forward_lower, input_spec=[ - InputSpec([1, -1, 3], dtype="float64", name="coord"), # extended_coord - InputSpec([1, -1], dtype="int32", name="atype"), # extended_atype - InputSpec([1, -1, -1], dtype="int32", name="nlist"), # nlist - InputSpec([1, -1], dtype="int64", name="mapping"), # mapping + InputSpec([-1, -1, 3], dtype="float64", name="coord"), # extended_coord + InputSpec([-1, -1], dtype="int32", name="atype"), # extended_atype + InputSpec([-1, -1, -1], dtype="int32", name="nlist"), # nlist + InputSpec([-1, -1], dtype="int64", name="mapping"), # mapping None, # fparam None, # aparam - True, # do_atomic_virial - None, # comm_dict + # InputSpec([], dtype="bool", name="do_atomic_virial"), # do_atomic_virial + do_atomic_virial, # do_atomic_virial + ( + InputSpec([-1], "int64", name="send_list"), + InputSpec([-1], "int32", name="send_proc"), + InputSpec([-1], "int32", name="recv_proc"), + InputSpec([-1], "int32", name="send_num"), + InputSpec([-1], "int32", name="recv_num"), + InputSpec([-1], "int64", name="communicator"), + # InputSpec([1], "int64", name="has_spin"), + ), # comm_dict ], full_graph=True, ) + for method_name in [ + "get_buffer_rcut", + "get_buffer_type_map", + "get_buffer_dim_fparam", + "get_buffer_dim_aparam", + "get_buffer_intensive", + "get_buffer_sel_type", + "get_buffer_numb_dos", + "get_buffer_task_dim", + ]: + if hasattr(model, method_name): + setattr( + model, + method_name, + paddle.jit.to_static( + getattr(model, method_name), + input_spec=[], + full_graph=True, + ), + ) if output.endswith(".json"): output = output[:-5] paddle.jit.save( @@ -414,12 +449,12 @@ def freeze( def change_bias( input_file: str, mode: str = "change", - bias_value: Optional[list] = None, - datafile: Optional[str] = None, + bias_value: list | None = None, + datafile: str | None = None, system: str = ".", numb_batch: int = 0, - model_branch: Optional[str] = None, - output: Optional[str] = None, + model_branch: str | None = None, + output: str | None = None, ) -> None: if input_file.endswith(".pd"): old_state_dict = paddle.load(input_file) @@ -525,7 +560,7 @@ def change_bias( log.info(f"Saved model to {output_path}") -def main(args: Optional[Union[list[str], argparse.Namespace]] = None): +def main(args: list[str] | argparse.Namespace | None = None): if not isinstance(args, argparse.Namespace): FLAGS = parse_args(args=args) else: diff --git a/deepmd/pd/infer/deep_eval.py b/deepmd/pd/infer/deep_eval.py index b0b4bab980..67c435ab3f 100644 --- a/deepmd/pd/infer/deep_eval.py +++ b/deepmd/pd/infer/deep_eval.py @@ -1,14 +1,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np import paddle +from paddle import inference as paddle_inference from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT from deepmd.dpmodel.output_def import ( @@ -16,16 +19,35 @@ OutputVariableCategory, OutputVariableDef, ) +from deepmd.infer.deep_dipole import ( + DeepDipole, +) +from deepmd.infer.deep_dos import ( + DeepDOS, +) from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper from deepmd.infer.deep_eval import ( DeepEvalBackend, ) +from deepmd.infer.deep_polar import ( + DeepGlobalPolar, + DeepPolar, +) from deepmd.infer.deep_pot import ( DeepPot, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) +from deepmd.infer.deep_wfc import ( + DeepWFC, +) from deepmd.pd.model.model import ( get_model, ) +from deepmd.pd.model.network.network import ( + TypeEmbedNetConsistent, +) from deepmd.pd.train.wrapper import ( ModelWrapper, ) @@ -42,10 +64,22 @@ to_numpy_array, to_paddle_tensor, ) +from deepmd.utils.econf_embd import ( + sort_element_type, +) +from deepmd.utils.model_branch_dict import ( + get_model_dict, +) if TYPE_CHECKING: import ase.neighborlist + from deepmd.pd.model.model.model import ( + BaseModel, + ) + +log = logging.getLogger(__name__) + class DeepEval(DeepEvalBackend): """Paddle backend implementation of DeepEval. @@ -58,7 +92,7 @@ class DeepEval(DeepEvalBackend): The output definition of the model. *args : list Positional arguments. - auto_batch_size : bool or int or AutomaticBatchSize, default: False + auto_batch_size : bool or int or AutomaticBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional @@ -73,11 +107,12 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, - head: Optional[Union[str, int]] = None, + head: str | int | None = None, + no_jit: bool = False, **kwargs: Any, - ): + ) -> None: enable_prim(True) self.output_def = output_def self.model_path = model_file @@ -110,12 +145,42 @@ def __init__( # model = paddle.jit.to_static(model) self.dp = ModelWrapper(model) self.dp.set_state_dict(state_dict) + self.rcut = self.dp.model["Default"].get_rcut() + self.type_map: list[str] = self.dp.model["Default"].get_type_map() + self.dp.eval() + self.static_model = False + elif str(self.model_path).endswith(".json"): + self.dp = paddle.jit.load(self.model_path[:-5]) + self.rcut = self.dp.get_buffer_rcut().item() + self.type_map: list[str] = "".join( + [chr(x) for x in self.dp.get_buffer_type_map().numpy()] + ).split(" ") + config = paddle_inference.Config( + self.model_path, self.model_path.replace(".json", ".pdiparams") + ) + config.enable_custom_passes( + ["add_shadow_output_after_dead_parameter_pass"], True + ) + config.enable_use_gpu(4096, 0) + config.disable_glog_info() + + self.predictor = paddle_inference.create_predictor(config) + self.coord_handle = self.predictor.get_input_handle("coord") + self.atype_handle = self.predictor.get_input_handle("atype") + self.box_handle = self.predictor.get_input_handle("box") + + self.atom_energy_handle = self.predictor.get_output_handle("fetch_name_0") + self.atom_virial_handle = self.predictor.get_output_handle("fetch_name_1") + self.energy_handle = self.predictor.get_output_handle("fetch_name_2") + self.force_handle = self.predictor.get_output_handle("fetch_name_3") + self.mask_handle = self.predictor.get_output_handle("fetch_name_4") + self.virial_handle = self.predictor.get_output_handle("fetch_name_5") + self.static_model = True else: - # self.dp = paddle.jit.load(self.model_path.split(".json")[0]) - raise ValueError(f"Unknown model file format: {self.model_path}!") - self.dp.eval() - self.rcut = self.dp.model["Default"].get_rcut() - self.type_map = self.dp.model["Default"].get_type_map() + raise NotImplementedError( + f"Only support .pd or .json format, but got {model_file}" + ) + if isinstance(auto_batch_size, bool): if auto_batch_size: self.auto_batch_size = AutoBatchSize() @@ -127,9 +192,14 @@ def __init__( self.auto_batch_size = auto_batch_size else: raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") - self._has_spin = getattr(self.dp.model["Default"], "has_spin", False) + self._has_spin = ( + getattr(self.dp.model["Default"], "has_spin", False) + if isinstance(self.dp, ModelWrapper) + else False + ) if callable(self._has_spin): - self._has_spin = self._has_spin() + self._has_spin = False + self._has_hessian = False def get_rcut(self) -> float: """Get the cutoff radius of this model.""" @@ -145,21 +215,56 @@ def get_type_map(self) -> list[str]: def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this DP.""" + if self.static_model: + return self.dp.get_buffer_dim_fparam() return self.dp.model["Default"].get_dim_fparam() def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" + if self.static_model: + return self.dp.get_buffer_dim_aparam() return self.dp.model["Default"].get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + try: + return self.dp.model["Default"].has_default_fparam() + except AttributeError: + # for compatibility with old models + return False + def get_intensive(self) -> bool: return self.dp.model["Default"].get_intensive() + def get_var_name(self) -> str: + """Get the name of the property.""" + if hasattr(self.dp.model["Default"], "get_var_name") and callable( + getattr(self.dp.model["Default"], "get_var_name") + ): + return self.dp.model["Default"].get_var_name() + else: + raise NotImplementedError + @property def model_type(self) -> type["DeepEvalWrapper"]: """The the evaluator of the model type.""" + if self.static_model: + return DeepPot model_output_type = self.dp.model["Default"].model_output_type() if "energy" in model_output_type: return DeepPot + elif "dos" in model_output_type: + return DeepDOS + elif "dipole" in model_output_type: + return DeepDipole + elif "polar" in model_output_type: + return DeepPolar + elif "global_polar" in model_output_type: + return DeepGlobalPolar + elif "wfc" in model_output_type: + return DeepWFC + elif self.get_var_name() in model_output_type: + return DeepProperty else: raise RuntimeError("Unknown model type") @@ -180,26 +285,41 @@ def get_task_dim(self) -> int: """Get the output dimension.""" return self.dp.model["Default"].get_task_dim() - def get_has_efield(self): + def get_has_efield(self) -> bool: """Check if the model has efield.""" return False - def get_ntypes_spin(self): + def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model. Only used in old implement.""" return 0 - def get_has_spin(self): + def get_has_spin(self) -> bool: """Check if the model has spin atom types.""" return self._has_spin + def get_has_hessian(self) -> bool: + """Check if the model has hessian.""" + return self._has_hessian + + def get_model_branch(self) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: + """Get the model branch information.""" + if "model_dict" in self.model_def_script: + model_alias_dict, model_branch_dict = get_model_dict( + self.model_def_script["model_dict"] + ) + return model_alias_dict, model_branch_dict + else: + # single-task model + return {"Default": "Default"}, {"Default": {"alias": [], "info": {}}} + def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -297,6 +417,7 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: OutputVariableCategory.REDU, OutputVariableCategory.DERV_R, OutputVariableCategory.DERV_C_REDU, + OutputVariableCategory.DERV_R_DERV_R, ) ] @@ -319,7 +440,7 @@ def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Calla """ if self.auto_batch_size is not None: - def eval_func(*args, **kwargs): + def eval_func(*args: Any, **kwargs: Any) -> Any: return self.auto_batch_size.execute_all( inner_func, numb_test, natoms, *args, **kwargs ) @@ -348,15 +469,133 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], ): - model = self.dp.to(DEVICE) + if not self.static_model: + model = self.dp.to(DEVICE) prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PD_FLOAT_PRECISION]] + nframes = coords.shape[0] + if len(atom_types.shape) == 1: + natoms = len(atom_types) + atom_types = np.tile(atom_types, nframes).reshape([nframes, -1]) + else: + natoms = len(atom_types[0]) + + if self.static_model: + self.coord_handle.copy_from_cpu( + coords.reshape([nframes, natoms, 3]).astype(prec) + ) + self.atype_handle.copy_from_cpu( + atom_types.astype( + NP_PRECISION_DICT[RESERVED_PRECISION_DICT[paddle.int64]] + ) + ) + if cells is not None: + self.box_handle.copy_from_cpu(cells.reshape([nframes, 3, 3])) + + if fparam is not None: + raise NotImplementedError( + "fparam_input is not supported for .json files. Please use a .pd file instead." + ) + + if aparam is not None: + raise NotImplementedError( + "aparam_input is not supported for .json files. Please use a .pd file instead." + ) + + else: + coord_input = paddle.to_tensor( + coords.reshape([nframes, natoms, 3]).astype(prec), + dtype=GLOBAL_PD_FLOAT_PRECISION, + place=DEVICE, + ) + type_input = paddle.to_tensor( + atom_types.astype( + NP_PRECISION_DICT[RESERVED_PRECISION_DICT[paddle.int64]] + ), + dtype=paddle.int64, + place=DEVICE, + ) + if cells is not None: + box_input = paddle.to_tensor( + cells.reshape([nframes, 3, 3]), + dtype=GLOBAL_PD_FLOAT_PRECISION, + place=DEVICE, + ) + else: + box_input = None + if fparam is not None: + fparam_input = to_paddle_tensor( + fparam.reshape([nframes, self.get_dim_fparam()]) + ) + else: + fparam_input = None + if aparam is not None: + aparam_input = to_paddle_tensor( + aparam.reshape([nframes, natoms, self.get_dim_aparam()]) + ) + else: + aparam_input = None + + do_atomic_virial = any( + x.category == OutputVariableCategory.DERV_C for x in request_defs + ) + if self.static_model: + self.predictor.run() + batch_output = { + "atom_energy": self.atom_energy_handle.copy_to_cpu(), + "energy": self.energy_handle.copy_to_cpu(), + "force": self.force_handle.copy_to_cpu(), + "mask": self.mask_handle.copy_to_cpu(), + "virial": self.virial_handle.copy_to_cpu(), + } + if do_atomic_virial: + batch_output["atom_virial"] = self.atom_virial_handle.copy_to_cpu() + else: + batch_output = model( + coord_input, + type_input, + box=box_input, + do_atomic_virial=do_atomic_virial, + fparam=fparam_input, + aparam=aparam_input, + ) + if isinstance(batch_output, tuple): + batch_output = batch_output[0] + + results = [] + for odef in request_defs: + pd_name = self._OUTDEF_DP2BACKEND[odef.name] + if pd_name in batch_output: + shape = self._get_output_shape(odef, nframes, natoms) + out = batch_output[pd_name].reshape(shape) + if not self.static_model: + out = out.numpy() + results.append(out) + else: + shape = self._get_output_shape(odef, nframes, natoms) + results.append( + np.full(np.abs(shape), np.nan, dtype=prec) + ) # this is kinda hacky + return tuple(results) + + def _eval_model_spin( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + spins: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + request_defs: list[OutputVariableDef], + ) -> tuple[np.ndarray, ...]: + model = self.dp.to(DEVICE) + nframes = coords.shape[0] if len(atom_types.shape) == 1: natoms = len(atom_types) @@ -365,13 +604,14 @@ def _eval_model( natoms = len(atom_types[0]) coord_input = paddle.to_tensor( - coords.reshape([nframes, natoms, 3]).astype(prec), + coords.reshape([nframes, natoms, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE, ) - type_input = paddle.to_tensor( - atom_types.astype(NP_PRECISION_DICT[RESERVED_PRECISION_DICT[paddle.int64]]), - dtype=paddle.int64, + type_input = paddle.to_tensor(atom_types, dtype=paddle.int64, place=DEVICE) + spin_input = paddle.to_tensor( + spins.reshape([nframes, natoms, 3]), + dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE, ) if cells is not None: @@ -394,12 +634,14 @@ def _eval_model( ) else: aparam_input = None + do_atomic_virial = any( - x.category == OutputVariableCategory.DERV_C for x in request_defs + x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs ) batch_output = model( coord_input, type_input, + spin=spin_input, box=box_input, do_atomic_virial=do_atomic_virial, fparam=fparam_input, @@ -413,28 +655,24 @@ def _eval_model( pd_name = self._OUTDEF_DP2BACKEND[odef.name] if pd_name in batch_output: shape = self._get_output_shape(odef, nframes, natoms) - out = batch_output[pd_name].reshape(shape).numpy() + out = batch_output[pd_name].reshape(shape).detach().cpu().numpy() results.append(out) else: shape = self._get_output_shape(odef, nframes, natoms) results.append( - np.full(np.abs(shape), np.nan, dtype=prec) + np.full( + np.abs(shape), + np.nan, + dtype=NP_PRECISION_DICT[ + RESERVED_PRECISION_DICT[GLOBAL_PD_FLOAT_PRECISION] + ], + ) ) # this is kinda hacky return tuple(results) - def _eval_model_spin( - self, - coords: np.ndarray, - cells: Optional[np.ndarray], - atom_types: np.ndarray, - spins: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], - request_defs: list[OutputVariableDef], - ): - raise NotImplementedError("_eval_model_spin is not supported yet.") - - def _get_output_shape(self, odef, nframes, natoms): + def _get_output_shape( + self, odef: OutputVariableDef, nframes: int, natoms: int + ) -> list[int]: if odef.category == OutputVariableCategory.DERV_C_REDU: # virial return [nframes, *odef.shape[:-1], 9] @@ -452,6 +690,9 @@ def _get_output_shape(self, odef, nframes, natoms): # Something wrong here? # return [nframes, *shape, natoms, 1] return [nframes, natoms, *odef.shape, 1] + elif odef.category == OutputVariableCategory.DERV_R_DERV_R: + return [nframes, 3 * natoms, 3 * natoms] + # return [nframes, *odef.shape, 3 * natoms, 3 * natoms] else: raise RuntimeError("unknown category") @@ -476,9 +717,16 @@ def eval_typeebd(self) -> np.ndarray: deepmd.pd.model.network.network.TypeEmbedNetConsistent : The type embedding network. """ - raise NotImplementedError("eval_typeebd is not supported yet.") - - def get_model_def_script(self) -> str: + out = [] + for mm in self.dp.model["Default"].sublayers(): + if mm.__class__.__name__ == TypeEmbedNetConsistent.__name__: + out.append(mm(DEVICE)) + if not out: + raise KeyError("The model has no type embedding networks.") + typeebd = paddle.concat(out, axis=1) + return to_numpy_array(typeebd) + + def get_model_def_script(self) -> dict: """Get model definition script.""" return self.model_def_script @@ -506,13 +754,39 @@ def get_model_size(self) -> dict: "total": sum_param_des + sum_param_fit, } + def get_observed_types(self) -> dict: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + dict + A dictionary containing the information of observed type in the model: + - 'type_num': the total number of observed types in this model. + - 'observed_type': a list of the observed types in this model. + """ + observed_type_list = self.dp.model["Default"].get_observed_type_list() + return { + "type_num": len(observed_type_list), + "observed_type": sort_element_type(observed_type_list), + } + + def get_model(self) -> "BaseModel": + """Get the Paddle model. + + Returns + ------- + BaseModel + The Paddle model instance. + """ + return self.dp.model["Default"] + def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate descriptors by using this DP. @@ -546,7 +820,9 @@ def eval_descriptor( descriptor Descriptors. """ - model = self.dp.model["Default"] + model = ( + self.dp.model["Default"] if isinstance(self.dp, ModelWrapper) else self.dp + ) model.set_eval_descriptor_hook(True) self.eval( coords, @@ -560,3 +836,58 @@ def eval_descriptor( descriptor = model.eval_descriptor() model.set_eval_descriptor_hook(False) return to_numpy_array(descriptor) + + def eval_fitting_last_layer( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> np.ndarray: + """Evaluate fitting before last layer by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + + Returns + ------- + fitting + Fitting output before last layer. + """ + model = self.dp.model["Default"] + model.set_eval_fitting_last_layer_hook(True) + self.eval( + coords, + cells, + atom_types, + atomic=False, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + fitting_net = model.eval_fitting_last_layer() + model.set_eval_fitting_last_layer_hook(False) + return to_numpy_array(fitting_net) diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index 09ec5ff49e..fbd806b26d 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle import paddle.nn.functional as F @@ -46,7 +43,7 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, diff --git a/deepmd/pd/model/atomic_model/base_atomic_model.py b/deepmd/pd/model/atomic_model/base_atomic_model.py index 1100813fb4..87cb18f6fc 100644 --- a/deepmd/pd/model/atomic_model/base_atomic_model.py +++ b/deepmd/pd/model/atomic_model/base_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import copy import logging -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + NoReturn, Optional, - Union, ) import numpy as np @@ -64,9 +65,9 @@ class BaseAtomicModel(paddle.nn.Layer, BaseAtomicModel_): of the atomic model. Implemented by removing the pairs from the nlist. rcond : float, optional The condition number for the regression of atomic energy. - preset_out_bias : Dict[str, list[Optional[paddle.Tensor]]], optional + preset_out_bias : dict[str, list[Optional[np.ndarray]]], optional Specifying atomic energy contribution in vacuum. Given by key:value pairs. - The value is a list specifying the bias. the elements can be None or np.array of output shape. + The value is a list specifying the bias. the elements can be None or np.ndarray of output shape. For example: [None, [2.]] means type 0 is not set, type 1 is set to [2.] The `set_davg_zero` key in the descriptor should be set. @@ -77,20 +78,30 @@ def __init__( type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], - rcond: Optional[float] = None, - preset_out_bias: Optional[dict[str, np.ndarray]] = None, - ): + rcond: float | None = None, + preset_out_bias: dict[str, np.ndarray] | None = None, + data_stat_protect: float = 1e-2, + ) -> None: paddle.nn.Layer.__init__(self) BaseAtomicModel_.__init__(self) self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) + self.ntypes = len(self.type_map) + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) self.reinit_atom_exclude(atom_exclude_types) self.reinit_pair_exclude(pair_exclude_types) self.rcond = rcond self.preset_out_bias = preset_out_bias + self.data_stat_protect = data_stat_protect - def init_out_stat(self): + def init_out_stat(self) -> None: """Initialize the output bias.""" - ntypes = self.get_ntypes() self.bias_keys: list[str] = list(self.fitting_output_def().keys()) self.max_out_size = max( [self.atomic_output_def()[kk].size for kk in self.bias_keys] @@ -104,7 +115,7 @@ def init_out_stat(self): def set_out_bias(self, out_bias: paddle.Tensor) -> None: self.out_bias = out_bias - def __setitem__(self, key, value): + def __setitem__(self, key: str, value: paddle.Tensor) -> None: if key in ["out_bias"]: self.out_bias = value elif key in ["out_std"]: @@ -112,7 +123,7 @@ def __setitem__(self, key, value): else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> paddle.Tensor: if key in ["out_bias"]: return self.out_bias elif key in ["out_std"]: @@ -124,10 +135,36 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.type_map + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + + def get_compute_stats_distinguish_types(self) -> bool: + """Get whether the fitting net computes stats which are not distinguished between different types of atoms.""" + return True + + def get_intensive(self) -> bool: + """Whether the fitting property is intensive.""" + return False + + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + def reinit_atom_exclude( self, - exclude_types: list[int] = [], - ): + exclude_types: list[int] | None = None, + ) -> None: + if exclude_types is None: + exclude_types = [] self.atom_exclude_types = exclude_types if exclude_types == []: self.atom_excl = None @@ -137,7 +174,7 @@ def reinit_atom_exclude( def reinit_pair_exclude( self, exclude_types: list[tuple[int, int]] = [], - ): + ) -> None: self.pair_exclude_types = exclude_types if exclude_types == []: self.pair_excl = None @@ -188,10 +225,10 @@ def forward_common_atomic( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ) -> dict[str, paddle.Tensor]: """Common interface for atomic inference. @@ -232,7 +269,7 @@ def forward_common_atomic( if self.pair_excl is not None: pair_mask = self.pair_excl(nlist, extended_atype) # exclude neighbors in the nlist - nlist = paddle.where(pair_mask == 1, nlist, -1) + nlist = paddle.where(pair_mask == 1, nlist, paddle.full_like(nlist, -1)) ext_atom_mask = self.make_atom_mask(extended_atype) ret_dict = self.forward_atomic( @@ -247,7 +284,6 @@ def forward_common_atomic( comm_dict=comm_dict, ) ret_dict = self.apply_out_stat(ret_dict, atype) - # nf x nloc atom_mask = ext_atom_mask[:, :nloc].astype(paddle.int32) if self.atom_excl is not None: @@ -260,10 +296,10 @@ def forward_common_atomic( out_shape2 *= ss ret_dict[kk] = ( ret_dict[kk].reshape([out_shape[0], out_shape[1], out_shape2]) - * atom_mask.unsqueeze(2).astype(ret_dict[kk].dtype) + * atom_mask[:, :, None].astype(ret_dict[kk].dtype) ).reshape(out_shape) ret_dict["mask"] = atom_mask - + # raise return ret_dict def forward( @@ -271,10 +307,10 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ) -> dict[str, paddle.Tensor]: return self.forward_common_atomic( extended_coord, @@ -287,7 +323,9 @@ def forward( ) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["BaseAtomicModel"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -332,7 +370,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "BaseAtomicModel": - data = copy.deepcopy(data) + data = data.copy() variables = data.pop("@variables", None) variables = ( {"out_bias": None, "out_std": None} if variables is None else variables @@ -352,32 +390,36 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel": def compute_or_load_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, + ) -> NoReturn: """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + Compute or load the statistics parameters of the model, + such as mean and standard deviation of descriptors or the energy bias of the fitting net. + When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update), + and saved in the `stat_file_path`(s). + When `sampled` is not provided, it will check the existence of `stat_file_path`(s) + and load the calculated statistics parameters. Parameters ---------- - merged : Union[Callable[[], list[dict]], list[dict]] - - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` - originating from the `i`-th data system. - - Callable[[], list[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. + merged + The lazy sampled function to get data frames from different data systems. + stat_file_path + The dictionary of paths to the statistics files. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ raise NotImplementedError def compute_or_load_out_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -404,7 +446,7 @@ def apply_out_stat( self, ret: dict[str, paddle.Tensor], atype: paddle.Tensor, - ): + ) -> dict[str, paddle.Tensor]: """Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. @@ -425,9 +467,9 @@ def apply_out_stat( def change_out_bias( self, - sample_merged, - stat_file_path: Optional[DPPath] = None, - bias_adjust_mode="change-by-statistic", + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + bias_adjust_mode: str = "change-by-statistic", ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -457,7 +499,6 @@ def change_out_bias( model_forward=self._get_forward_wrapper_func(), rcond=self.rcond, preset_bias=self.preset_out_bias, - atomic_output=self.atomic_output_def(), ) self._store_out_stat(delta_bias, out_std, add=True) elif bias_adjust_mode == "set-by-statistic": @@ -468,16 +509,41 @@ def change_out_bias( stat_file_path=stat_file_path, rcond=self.rcond, preset_bias=self.preset_out_bias, - atomic_output=self.atomic_output_def(), + stats_distinguish_types=self.get_compute_stats_distinguish_types(), + intensive=self.get_intensive(), ) self._store_out_stat(bias_out, std_out) else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + def compute_fitting_input_stat( + self, + sample_merged: Callable[[], list[dict]] | list[dict], + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the atomic model from packed data. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + """ + pass + def _get_forward_wrapper_func(self) -> Callable[..., paddle.Tensor]: """Get a forward wrapper of the atomic model for output bias calculation.""" - def model_forward(coord, atype, box, fparam=None, aparam=None): + def model_forward( + coord: paddle.Tensor, + atype: paddle.Tensor, + box: paddle.Tensor | None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + ) -> dict[str, paddle.Tensor]: with ( paddle.no_grad() ): # it's essential for pure paddle forward function to use auto_batchsize @@ -506,7 +572,7 @@ def model_forward(coord, atype, box, fparam=None, aparam=None): return model_forward - def _default_bias(self): + def _default_bias(self) -> paddle.Tensor: ntypes = self.get_ntypes() return paddle.zeros([self.n_out, ntypes, self.max_out_size], dtype=dtype).to( device=device @@ -544,7 +610,7 @@ def _store_out_stat( out_bias: dict[str, paddle.Tensor], out_std: dict[str, paddle.Tensor], add: bool = False, - ): + ) -> None: ntypes = self.get_ntypes() out_bias_data = paddle.clone(self.out_bias) out_std_data = paddle.clone(self.out_std) @@ -560,6 +626,12 @@ def _store_out_stat( paddle.assign(out_bias_data, self.out_bias) paddle.assign(out_std_data, self.out_std) + def get_ntypes(self): + return len(self.type_map) + + def get_buffer_ntypes(self) -> paddle.Tensor: + return self.buffer_ntypes + def _fetch_out_stat( self, keys: list[str], diff --git a/deepmd/pd/model/atomic_model/dp_atomic_model.py b/deepmd/pd/model/atomic_model/dp_atomic_model.py index 1089b93a68..c5aa8b8a56 100644 --- a/deepmd/pd/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pd/model/atomic_model/dp_atomic_model.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools import logging +from collections.abc import ( + Callable, +) from typing import ( + Any, Optional, ) @@ -47,10 +51,10 @@ class DPAtomicModel(BaseAtomicModel): def __init__( self, - descriptor, - fitting, + descriptor: BaseDescriptor, + fitting: BaseFitting, type_map: list[str], - **kwargs, + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) ntypes = len(type_map) @@ -62,17 +66,22 @@ def __init__( self.fitting_net = fitting super().init_out_stat() self.enable_eval_descriptor_hook = False + self.enable_eval_fitting_last_layer_hook = False self.eval_descriptor_list = [] + self.eval_fitting_last_layer_list = [] # register 'type_map' as buffer - def _string_to_array(s: str) -> list[int]: + def _string_to_array(s: str | list[str]) -> list[int]: return [ord(c) for c in s] - self.register_buffer( - "buffer_type_map", - paddle.to_tensor(_string_to_array(" ".join(self.type_map)), dtype="int32"), - ) - self.buffer_type_map.name = "buffer_type_map" + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor( + _string_to_array(" ".join(self.type_map)), dtype="int32" + ), + ) + self.buffer_type_map.name = "buffer_type_map" if hasattr(self.descriptor, "has_message_passing"): # register 'has_message_passing' as buffer(cast to int32 as problems may meets with vector) self.register_buffer( @@ -112,16 +121,29 @@ def _string_to_array(s: str) -> list[int]: self.buffer_aparam_nall.name = "buffer_aparam_nall" eval_descriptor_list: list[paddle.Tensor] + eval_fitting_last_layer_list: list[paddle.Tensor] def set_eval_descriptor_hook(self, enable: bool) -> None: """Set the hook for evaluating descriptor and clear the cache for descriptor list.""" self.enable_eval_descriptor_hook = enable - self.eval_descriptor_list = [] + # = [] does not work; See #4533 + self.eval_descriptor_list.clear() def eval_descriptor(self) -> paddle.Tensor: """Evaluate the descriptor.""" return paddle.concat(self.eval_descriptor_list) + def set_eval_fitting_last_layer_hook(self, enable: bool) -> None: + """Set the hook for evaluating fitting last layer output and clear the cache for fitting last layer output list.""" + self.enable_eval_fitting_last_layer_hook = enable + self.fitting_net.set_return_middle_output(enable) + # = [] does not work; See #4533 + self.eval_fitting_last_layer_list.clear() + + def eval_fitting_last_layer(self) -> paddle.Tensor: + """Evaluate the fitting last layer output.""" + return paddle.concat(self.eval_fitting_last_layer_list) + def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" return ( @@ -138,7 +160,27 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.sel - def set_case_embd(self, case_idx: int): + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + + def get_buffer_rcut(self) -> paddle.Tensor: + """Get the cut-off radius as a buffer-style Tensor.""" + return self.descriptor.get_buffer_rcut() + + def get_buffer_sel(self) -> paddle.Tensor: + """Get the neighbor selection as a buffer-style Tensor.""" + return self.descriptor.get_buffer_sel() + + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -158,7 +200,9 @@ def mixed_types(self) -> bool: return self.descriptor.mixed_types() def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["DPAtomicModel"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -199,7 +243,7 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "DPAtomicModel": + def deserialize(cls, data: dict) -> "DPAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 1) data.pop("@class", None) @@ -244,13 +288,13 @@ def enable_compression( def forward_atomic( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + extended_coord: paddle.Tensor, + extended_atype: paddle.Tensor, + nlist: paddle.Tensor, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, + comm_dict: dict[str, paddle.Tensor] | None = None, ) -> dict[str, paddle.Tensor]: """Return atomic prediction. @@ -288,7 +332,7 @@ def forward_atomic( ) assert descriptor is not None if self.enable_eval_descriptor_hook: - self.eval_descriptor_list.append(descriptor) + self.eval_descriptor_list.append(descriptor.detach()) # energy, force fit_ret = self.fitting_net( descriptor, @@ -299,6 +343,13 @@ def forward_atomic( fparam=fparam, aparam=aparam, ) + if self.enable_eval_fitting_last_layer_hook: + assert "middle_output" in fit_ret, ( + "eval_fitting_last_layer not supported for this fitting net!" + ) + self.eval_fitting_last_layer_list.append( + fit_ret.pop("middle_output").detach() + ) return fit_ret def get_out_bias(self) -> paddle.Tensor: @@ -306,8 +357,9 @@ def get_out_bias(self) -> paddle.Tensor: def compute_or_load_stat( self, - sampled_func, - stat_file_path: Optional[DPPath] = None, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, ) -> None: """ Compute or load the statistics parameters of the model, @@ -323,6 +375,9 @@ def compute_or_load_stat( The lazy sampled function to get data frames from different data systems. stat_file_path The dictionary of paths to the statistics files. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ if stat_file_path is not None and self.type_map is not None: # descriptors and fitting net with different type_map @@ -343,16 +398,55 @@ def wrapped_sampler(): return sampled self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) - self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + self.compute_fitting_input_stat(wrapped_sampler, stat_file_path) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + + def compute_fitting_input_stat( + self, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The dictionary of paths to the statistics files. + """ + self.fitting_net.compute_input_stats( + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, + ) def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.fitting_net.get_dim_fparam() + def get_buffer_dim_fparam(self) -> paddle.Tensor: + """Get the number (dimension) of frame parameters of this atomic model as a buffer-style Tensor.""" + return self.fitting_net.get_buffer_dim_fparam() + + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.fitting_net.has_default_fparam() + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.fitting_net.get_dim_aparam() + def get_buffer_dim_aparam(self) -> paddle.Tensor: + """Get the number (dimension) of atomic parameters of this atomic model as a buffer-style Tensor.""" + return self.fitting_net.get_buffer_dim_aparam() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. diff --git a/deepmd/pd/model/descriptor/descriptor.py b/deepmd/pd/model/descriptor/descriptor.py index 3050b7dca3..d4ca4bc151 100644 --- a/deepmd/pd/model/descriptor/descriptor.py +++ b/deepmd/pd/model/descriptor/descriptor.py @@ -4,11 +4,11 @@ ABC, abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, - Optional, - Union, ) import paddle @@ -101,8 +101,8 @@ def get_env_protection(self) -> float: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -177,9 +177,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Calculate DescriptorBlock.""" pass diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index 6942b096c9..65baa8daa6 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -213,7 +211,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -235,17 +233,17 @@ def __init__( concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth_type_embedding: bool = True, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented spin=None, - type: Optional[str] = None, + type: str | None = None, ) -> None: super().__init__() # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -292,10 +290,16 @@ def __init__( trainable_ln=trainable_ln, ln_eps=ln_eps, seed=child_seed(seed, 1), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd self.use_tebd_bias = use_tebd_bias self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) self.compress = False self.type_embedding = TypeEmbedNet( ntypes, @@ -305,6 +309,7 @@ def __init__( use_econf_tebd=use_econf_tebd, use_tebd_bias=use_tebd_bias, type_map=type_map, + trainable=trainable, ) self.prec = PRECISION_DICT[precision] self.tebd_dim = tebd_dim @@ -318,10 +323,18 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.se_atten.get_rcut() + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.se_atten.get_buffer_rcut() + def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.se_atten.get_rcut_smth() + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.se_atten.get_buffer_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return self.se_atten.get_nsel() @@ -338,6 +351,18 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + def get_dim_out(self) -> int: """Returns the output dimension.""" ret = self.se_atten.get_dim_out() @@ -409,8 +434,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ): """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -593,8 +618,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -662,9 +687,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index 0e3b24397f..645ad2f2fe 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -80,9 +78,9 @@ def __init__( self, ntypes: int, # args for repinit - repinit: Union[RepinitArgs, dict], + repinit: RepinitArgs | dict, # args for repformer - repformer: Union[RepformerArgs, dict], + repformer: RepformerArgs | dict, # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -90,11 +88,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: r"""The DPA-2 descriptor[1]_. @@ -184,6 +182,7 @@ def init_subclass_params(sub_data, sub_class): smooth=smooth, type_one_side=self.repinit_args.type_one_side, seed=child_seed(seed, 0), + trainable=trainable, ) self.use_three_body = self.repinit_args.use_three_body if self.use_three_body: @@ -203,6 +202,7 @@ def init_subclass_params(sub_data, sub_class): resnet_dt=self.repinit_args.resnet_dt, smooth=smooth, seed=child_seed(seed, 5), + trainable=trainable, ) else: self.repinit_three_body = None @@ -243,6 +243,7 @@ def init_subclass_params(sub_data, sub_class): g1_out_conv=self.repformer_args.g1_out_conv, g1_out_mlp=self.repformer_args.g1_out_mlp, seed=child_seed(seed, 1), + trainable=trainable, ) self.rcsl_list = [ (self.repformers.get_rcut(), self.repformers.get_nsel()), @@ -262,6 +263,11 @@ def init_subclass_params(sub_data, sub_class): self.use_econf_tebd = use_econf_tebd self.use_tebd_bias = use_tebd_bias self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) self.type_embedding = TypeEmbedNet( ntypes, self.repinit_args.tebd_dim, @@ -270,6 +276,7 @@ def init_subclass_params(sub_data, sub_class): use_econf_tebd=self.use_econf_tebd, use_tebd_bias=use_tebd_bias, type_map=type_map, + trainable=trainable, ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -295,6 +302,7 @@ def init_subclass_params(sub_data, sub_class): precision=precision, init="glorot", seed=child_seed(seed, 3), + trainable=trainable, ) self.tebd_transform = None if self.add_tebd_to_repinit_out: @@ -304,6 +312,7 @@ def init_subclass_params(sub_data, sub_class): bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) assert self.repinit.rcut > self.repformers.rcut assert self.repinit.sel[0] > self.repformers.sel[0] @@ -312,6 +321,9 @@ def init_subclass_params(sub_data, sub_class): self.rcut = self.repinit.get_rcut() self.rcut_smth = self.repinit.get_rcut_smth() self.ntypes = ntypes + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) self.sel = self.repinit.sel # set trainable for param in self.parameters(): @@ -326,6 +338,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius.""" + return self.repinit.get_buffer_rcut() + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.repinit.get_buffer_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return sum(self.sel) @@ -336,7 +356,7 @@ def get_sel(self) -> list[int]: def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" @@ -481,8 +501,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -705,8 +725,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -741,7 +761,7 @@ def forward( """ # cast the input to internal precsion - extended_coord = extended_coord.to(dtype=self.prec) + extended_coord = extended_coord.astype(dtype=self.prec) use_three_body = self.use_three_body nframes, nloc, nnei = nlist.shape @@ -762,7 +782,7 @@ def forward( type_embedding = None g1, _, _, _, _ = self.repinit( nlist_dict[ - get_multiple_nlist_key(self.repinit.get_rcut(), self.repinit.get_nsel()) + get_multiple_nlist_key(self.repinit.rcut, sum(self.repinit.sel)) ], extended_coord, extended_atype, @@ -792,14 +812,15 @@ def forward( assert self.tebd_transform is not None g1 = g1 + self.tebd_transform(g1_inp) # mapping g1 - if comm_dict is None: - assert mapping is not None + if comm_dict is None or len(comm_dict) == 0: + if paddle.in_dynamic_mode(): + assert mapping is not None mapping_ext = ( mapping.reshape([nframes, nall]) .unsqueeze(-1) .expand([-1, -1, g1.shape[-1]]) ) - g1_ext = paddle.take_along_axis(g1, mapping_ext, 1) + g1_ext = paddle.take_along_axis(g1, mapping_ext, 1, broadcast=False) g1 = g1_ext # repformer g1, g2, h2, rot_mat, sw = self.repformers( @@ -817,20 +838,20 @@ def forward( if self.concat_output_tebd: g1 = paddle.concat([g1, g1_inp], axis=-1) return ( - g1.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), - rot_mat.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), - g2.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), - h2.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), - sw.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + g1.astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + rot_mat.astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + g2.astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + h2.astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + sw.astype(dtype=env.GLOBAL_PD_FLOAT_PRECISION), ) @classmethod def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py index 99fd78c62f..205829860f 100644 --- a/deepmd/pd/model/descriptor/dpa3.py +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -91,7 +89,7 @@ class DescrptDPA3(BaseDescriptor, paddle.nn.Layer): Whether to use bias in the type embedding layer. use_loc_mapping : bool, Optional Whether to use local atom index mapping in training or non-parallel inference. - Not supported yet in Paddle. + When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation. type_map : list[str], Optional A list of strings. Give the name to each type of atoms. @@ -106,7 +104,7 @@ def __init__( self, ntypes: int, # args for repflow - repflow: Union[RepFlowArgs, dict], + repflow: RepFlowArgs | dict, # kwargs for descriptor concat_output_tebd: bool = False, activation_function: str = "silu", @@ -114,11 +112,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - use_loc_mapping: bool = False, - type_map: Optional[list[str]] = None, + use_loc_mapping: bool = True, + type_map: list[str] | None = None, ) -> None: super().__init__() @@ -160,6 +158,8 @@ def init_subclass_params(sub_data, sub_class): fix_stat_std=self.repflow_args.fix_stat_std, optim_update=self.repflow_args.optim_update, smooth_edge_update=self.repflow_args.smooth_edge_update, + edge_init_use_dist=self.repflow_args.edge_init_use_dist, + use_exp_switch=self.repflow_args.use_exp_switch, use_dynamic_sel=self.repflow_args.use_dynamic_sel, sel_reduce_factor=self.repflow_args.sel_reduce_factor, use_loc_mapping=use_loc_mapping, @@ -167,12 +167,18 @@ def init_subclass_params(sub_data, sub_class): env_protection=env_protection, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd - self.use_tebd_bias = use_tebd_bias self.use_loc_mapping = use_loc_mapping + self.use_tebd_bias = use_tebd_bias self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]), + ) self.tebd_dim = self.repflow_args.n_dim self.type_embedding = TypeEmbedNet( ntypes, @@ -182,6 +188,7 @@ def init_subclass_params(sub_data, sub_class): use_econf_tebd=self.use_econf_tebd, use_tebd_bias=use_tebd_bias, type_map=type_map, + trainable=trainable, ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -203,6 +210,9 @@ def init_subclass_params(sub_data, sub_class): self.rcut_smth = self.repflows.get_rcut_smth() self.sel = self.repflows.get_sel() self.ntypes = ntypes + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) # set trainable for param in self.parameters(): @@ -217,6 +227,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.repflows.get_buffer_rcut() + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.repflows.get_buffer_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return sum(self.sel) @@ -225,14 +243,30 @@ def get_sel(self) -> list[int]: """Returns the number of selected atoms for each type.""" return self.sel + def get_buffer_sel(self) -> paddle.Tensor: + """Returns the number of selected atoms for each type as a buffer-style Tensor.""" + return self.repflows.get_sel() + def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + def get_dim_out(self) -> int: """Returns the output dimension of this descriptor.""" ret = self.repflows.dim_out @@ -335,8 +369,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -452,8 +486,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -487,12 +521,16 @@ def forward( The smooth switch function. shape: nf x nloc x nnei """ + parallel_mode = comm_dict is not None # cast the input to internal precsion extended_coord = extended_coord.to(dtype=self.prec) nframes, nloc, nnei = nlist.shape nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 - node_ebd_ext = self.type_embedding(extended_atype) + if not parallel_mode and self.use_loc_mapping: + node_ebd_ext = self.type_embedding(extended_atype[:, :nloc]) + else: + node_ebd_ext = self.type_embedding(extended_atype) node_ebd_inp = node_ebd_ext[:, :nloc, :] # repflows node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows( @@ -517,9 +555,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/descriptor/env_mat.py b/deepmd/pd/model/descriptor/env_mat.py index 9b72da0b16..2cc23fcadf 100644 --- a/deepmd/pd/model/descriptor/env_mat.py +++ b/deepmd/pd/model/descriptor/env_mat.py @@ -3,6 +3,7 @@ import paddle from deepmd.pd.utils.preprocess import ( + compute_exp_sw, compute_smooth_weight, ) @@ -14,6 +15,7 @@ def _make_env_mat( ruct_smth: float, radial_only: bool = False, protection: float = 0.0, + use_exp_switch: bool = False, ): """Make smooth environment matrix.""" bsz, natoms, nnei = nlist.shape @@ -21,10 +23,11 @@ def _make_env_mat( nall = coord.shape[1] mask = nlist >= 0 # nlist = nlist * mask ## this impl will contribute nans in Hessian calculation. - nlist = paddle.where(mask, nlist, nall - 1) + nlist = paddle.where(mask, nlist, paddle.full_like(nlist, nall - 1)) coord_l = coord[:, :natoms].reshape([bsz, -1, 1, 3]) index = nlist.reshape([bsz, -1]).unsqueeze(-1).expand([-1, -1, 3]) - coord_r = paddle.take_along_axis(coord, axis=1, indices=index) + coord_pad = paddle.concat([coord, coord[:, -1:, :] + rcut], axis=1) + coord_r = paddle.take_along_axis(coord_pad, axis=1, indices=index, broadcast=False) coord_r = coord_r.reshape([bsz, natoms, nnei, 3]) diff = coord_r - coord_l length = paddle.linalg.norm(diff, axis=-1, keepdim=True) @@ -32,7 +35,11 @@ def _make_env_mat( length = length + (~mask.unsqueeze(-1)).astype(length.dtype) t0 = 1 / (length + protection) t1 = diff / (length + protection) ** 2 - weight = compute_smooth_weight(length, ruct_smth, rcut) + weight = ( + compute_smooth_weight(length, ruct_smth, rcut) + if not use_exp_switch + else compute_exp_sw(length, ruct_smth, rcut) + ) weight = weight * mask.unsqueeze(-1).astype(weight.dtype) if radial_only: env_mat = t0 * weight @@ -51,6 +58,7 @@ def prod_env_mat( rcut_smth: float, radial_only: bool = False, protection: float = 0.0, + use_exp_switch: bool = False, ): """Generate smooth environment matrix from atom coordinates and other context. @@ -63,6 +71,7 @@ def prod_env_mat( - rcut_smth: Smooth hyper-parameter for pair force & energy. - radial_only: Whether to return a full description or a radial-only descriptor. - protection: Protection parameter to prevent division by zero errors during calculations. + - use_exp_switch: Whether to use the exponential switch function. Returns ------- @@ -75,6 +84,7 @@ def prod_env_mat( rcut_smth, radial_only, protection=protection, + use_exp_switch=use_exp_switch, ) # shape [n_atom, dim, 4 or 1] t_avg = mean[atype] # [n_atom, dim, 4 or 1] t_std = stddev[atype] # [n_atom, dim, 4 or 1] diff --git a/deepmd/pd/model/descriptor/repflow_layer.py b/deepmd/pd/model/descriptor/repflow_layer.py index f1bdd0439d..eca1659b67 100644 --- a/deepmd/pd/model/descriptor/repflow_layer.py +++ b/deepmd/pd/model/descriptor/repflow_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle import paddle.nn as nn @@ -19,6 +15,9 @@ from deepmd.pd.model.network.mlp import ( MLPLayer, ) +from deepmd.pd.model.network.utils import ( + aggregate, +) from deepmd.pd.utils.env import ( PRECISION_DICT, ) @@ -60,7 +59,8 @@ def __init__( update_residual: float = 0.1, update_residual_init: str = "const", precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei @@ -123,6 +123,7 @@ def __init__( n_dim, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) if self.update_style == "res_residual": self.n_residual.append( @@ -132,6 +133,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) ) @@ -142,6 +144,7 @@ def __init__( n_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) if self.update_style == "res_residual": self.n_residual.append( @@ -151,6 +154,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) ) @@ -160,6 +164,7 @@ def __init__( self.n_multi_edge_message * n_dim, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) if self.update_style == "res_residual": for head_index in range(self.n_multi_edge_message): @@ -170,6 +175,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(child_seed(seed, 5), head_index), + trainable=trainable, ) ) @@ -179,6 +185,7 @@ def __init__( e_dim, precision=precision, seed=child_seed(seed, 6), + trainable=trainable, ) if self.update_style == "res_residual": self.e_residual.append( @@ -188,6 +195,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 7), + trainable=trainable, ) ) @@ -216,6 +224,7 @@ def __init__( precision=precision, bias=False, seed=child_seed(seed, 8), + trainable=trainable, ) self.a_compress_e_linear = MLPLayer( self.e_dim, @@ -223,6 +232,7 @@ def __init__( precision=precision, bias=False, seed=child_seed(seed, 9), + trainable=trainable, ) else: self.a_compress_n_linear = None @@ -234,12 +244,14 @@ def __init__( self.e_dim, precision=precision, seed=child_seed(seed, 10), + trainable=trainable, ) self.edge_angle_linear2 = MLPLayer( self.e_dim, self.e_dim, precision=precision, seed=child_seed(seed, 11), + trainable=trainable, ) if self.update_style == "res_residual": self.e_residual.append( @@ -249,6 +261,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 12), + trainable=trainable, ) ) @@ -258,6 +271,7 @@ def __init__( self.a_dim, precision=precision, seed=child_seed(seed, 13), + trainable=trainable, ) if self.update_style == "res_residual": self.a_residual.append( @@ -267,6 +281,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 14), + trainable=trainable, ) ) else: @@ -326,6 +341,61 @@ def _cal_hg( h2g2 = paddle.matmul(paddle.matrix_transpose(h2), edge_ebd) * invnnei return h2g2 + @staticmethod + def _cal_hg_dynamic( + flat_edge_ebd: paddle.Tensor, + flat_h2: paddle.Tensor, + flat_sw: paddle.Tensor, + owner: paddle.Tensor, + num_owner: int, + nb: int, + nloc: int, + scale_factor: float, + ) -> paddle.Tensor: + """ + Calculate the transposed rotation matrix. + + Parameters + ---------- + flat_edge_ebd + Flatted neighbor-wise/pair-wise invariant rep tensors, with shape n_edge x e_dim. + flat_h2 + Flatted neighbor-wise/pair-wise equivariant rep tensors, with shape n_edge x 3. + flat_sw + Flatted switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, + and remains 0 beyond rcut, with shape n_edge. + owner + The owner index of the neighbor to reduce on. + num_owner : int + The total number of the owner. + nb : int + The number of batches. + nloc : int + The number of local atoms. + scale_factor : float + The scale factor to apply after reduce. + + Returns + ------- + hg + The transposed rotation matrix, with shape nf x nloc x 3 x e_dim. + """ + n_edge, e_dim = flat_edge_ebd.shape + # n_edge x e_dim + flat_edge_ebd = flat_edge_ebd * flat_sw.unsqueeze(-1) + # n_edge x 3 x e_dim + flat_h2g2 = (flat_h2.unsqueeze(-1) * flat_edge_ebd.unsqueeze(-2)).reshape( + [-1, 3 * e_dim] + ) + # nf x nloc x 3 x e_dim + h2g2 = ( + aggregate(flat_h2g2, owner, average=False, num_owner=num_owner).reshape( + [nb, nloc, 3, e_dim] + ) + * scale_factor + ) + return h2g2 + @staticmethod def _cal_grrg(h2g2: paddle.Tensor, axis_neuron: int) -> paddle.Tensor: """ @@ -398,6 +468,63 @@ def symmetrization_op( g1_13 = self._cal_grrg(h2g2, axis_neuron) return g1_13 + def symmetrization_op_dynamic( + self, + flat_edge_ebd: paddle.Tensor, + flat_h2: paddle.Tensor, + flat_sw: paddle.Tensor, + owner: paddle.Tensor, + num_owner: int, + nb: int, + nloc: int, + scale_factor: float, + axis_neuron: int, + ) -> paddle.Tensor: + """ + Symmetrization operator to obtain atomic invariant rep. + + Parameters + ---------- + flat_edge_ebd + Flatted neighbor-wise/pair-wise invariant rep tensors, with shape n_edge x e_dim. + flat_h2 + Flatted neighbor-wise/pair-wise equivariant rep tensors, with shape n_edge x 3. + flat_sw + Flatted switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, + and remains 0 beyond rcut, with shape n_edge. + owner + The owner index of the neighbor to reduce on. + num_owner : int + The total number of the owner. + nb : int + The number of batches. + nloc : int + The number of local atoms. + scale_factor : float + The scale factor to apply after reduce. + axis_neuron + Size of the submatrix. + + Returns + ------- + grrg + Atomic invariant rep, with shape nb x nloc x (axis_neuron x e_dim) + """ + # nb x nloc x 3 x e_dim + h2g2 = self._cal_hg_dynamic( + flat_edge_ebd, + flat_h2, + flat_sw, + owner, + num_owner, + nb, + nloc, + scale_factor, + ) + # nb x nloc x (axis x e_dim) + grrg = self._cal_grrg(h2g2, axis_neuron) + return grrg + def optim_angle_update( self, angle_ebd: paddle.Tensor, @@ -419,7 +546,7 @@ def optim_angle_update( node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] # angle_dim, node_dim, edge_dim, edge_dim - sub_angle, sub_node, sub_edge_ij, sub_edge_ik = paddle.split( + sub_angle, sub_node, sub_edge_ik, sub_edge_ij = paddle.split( matrix, [angle_dim, node_dim, edge_dim, edge_dim] ) @@ -428,14 +555,66 @@ def optim_angle_update( # nf * nloc * angle_dim sub_node_update = paddle.matmul(node_ebd, sub_node) # nf * nloc * a_nnei * angle_dim - sub_edge_update_ij = paddle.matmul(edge_ebd, sub_edge_ij) sub_edge_update_ik = paddle.matmul(edge_ebd, sub_edge_ik) + sub_edge_update_ij = paddle.matmul(edge_ebd, sub_edge_ij) result_update = ( bias + sub_node_update.unsqueeze(2).unsqueeze(3) - + sub_edge_update_ij.unsqueeze(2) - + sub_edge_update_ik.unsqueeze(3) + + sub_edge_update_ik.unsqueeze(2) + + sub_edge_update_ij.unsqueeze(3) + + sub_angle_update + ) + return result_update + + def optim_angle_update_dynamic( + self, + flat_angle_ebd: paddle.Tensor, + node_ebd: paddle.Tensor, + flat_edge_ebd: paddle.Tensor, + n2a_index: paddle.Tensor, + eij2a_index: paddle.Tensor, + eik2a_index: paddle.Tensor, + feat: str = "edge", + ) -> paddle.Tensor: + if feat == "edge": + matrix, bias = self.edge_angle_linear1.matrix, self.edge_angle_linear1.bias + elif feat == "angle": + matrix, bias = self.angle_self_linear.matrix, self.angle_self_linear.bias + else: + raise NotImplementedError + nf, nloc, node_dim = node_ebd.shape + edge_dim = flat_edge_ebd.shape[-1] + angle_dim = flat_angle_ebd.shape[-1] + # angle_dim, node_dim, edge_dim, edge_dim + sub_angle, sub_node, sub_edge_ik, sub_edge_ij = paddle.split( + matrix, [angle_dim, node_dim, edge_dim, edge_dim] + ) + + # n_angle * angle_dim + sub_angle_update = paddle.matmul(flat_angle_ebd, sub_angle) + + # nf * nloc * angle_dim + sub_node_update = paddle.matmul(node_ebd, sub_node) + # n_angle * angle_dim + sub_node_update = paddle.index_select( + sub_node_update.reshape([nf * nloc, sub_node_update.shape[-1]]), + n2a_index, + 0, + ) + + # n_edge * angle_dim + sub_edge_update_ik = paddle.matmul(flat_edge_ebd, sub_edge_ik) + sub_edge_update_ij = paddle.matmul(flat_edge_ebd, sub_edge_ij) + # n_angle * angle_dim + sub_edge_update_ik = paddle.index_select(sub_edge_update_ik, eik2a_index, 0) + sub_edge_update_ij = paddle.index_select(sub_edge_update_ij, eij2a_index, 0) + + result_update = ( + bias + + sub_node_update + + sub_edge_update_ik + + sub_edge_update_ij + sub_angle_update ) return result_update @@ -475,9 +654,55 @@ def optim_edge_update( ) return result_update + def optim_edge_update_dynamic( + self, + node_ebd: paddle.Tensor, + node_ebd_ext: paddle.Tensor, + flat_edge_ebd: paddle.Tensor, + n2e_index: paddle.Tensor, + n_ext2e_index: paddle.Tensor, + feat: str = "node", + ) -> paddle.Tensor: + if feat == "node": + matrix, bias = self.node_edge_linear.matrix, self.node_edge_linear.bias + elif feat == "edge": + matrix, bias = self.edge_self_linear.matrix, self.edge_self_linear.bias + else: + raise NotImplementedError + assert bias is not None + nf, nall, node_dim = node_ebd_ext.shape + _, nloc, _ = node_ebd.shape + edge_dim = flat_edge_ebd.shape[-1] + # node_dim, node_dim, edge_dim + node, node_ext, edge = paddle.split(matrix, [node_dim, node_dim, edge_dim]) + + # nf * nloc * node/edge_dim + sub_node_update = paddle.matmul(node_ebd, node) + # n_edge * node/edge_dim + sub_node_update = paddle.index_select( + sub_node_update.reshape([nf * nloc, sub_node_update.shape[-1]]), + n2e_index, + 0, + ) + + # nf * nall * node/edge_dim + sub_node_ext_update = paddle.matmul(node_ebd_ext, node_ext) + # n_edge * node/edge_dim + sub_node_ext_update = paddle.index_select( + sub_node_ext_update.reshape([nf * nall, sub_node_update.shape[-1]]), + n_ext2e_index, + 0, + ) + + # n_edge * node/edge_dim + sub_edge_update = paddle.matmul(flat_edge_ebd, edge) + + result_update = bias + sub_node_update + sub_edge_update + sub_node_ext_update + return result_update + def forward( self, - node_ebd_ext: paddle.Tensor, # nf x nall x n_dim + node_ebd_ext: paddle.Tensor, # nf x nall x n_dim [OR] nf x nloc x n_dim when not parallel_mode edge_ebd: paddle.Tensor, # nf x nloc x nnei x e_dim h2: paddle.Tensor, # nf x nloc x nnei x 3 angle_ebd: paddle.Tensor, # nf x nloc x a_nnei x a_nnei x a_dim @@ -487,6 +712,8 @@ def forward( a_nlist: paddle.Tensor, # nf x nloc x a_nnei a_nlist_mask: paddle.Tensor, # nf x nloc x a_nnei a_sw: paddle.Tensor, # switch func, nf x nloc x a_nnei + edge_index: paddle.Tensor, # 2 x n_edge + angle_index: paddle.Tensor, # 3 x n_angle ): """ Parameters @@ -511,6 +738,18 @@ def forward( Masks of the neighbor list for angle. real nei 1 otherwise 0 a_sw : nf x nloc x a_nnei Switch function for angle. + edge_index : Optional for dynamic sel, 2 x n_edge + n2e_index : n_edge + Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i). + n_ext2e_index : n_edge + Broadcast indices from extended node(j) to edge(ij). + angle_index : Optional for dynamic sel, 3 x n_angle + n2a_index : n_angle + Broadcast indices from extended node(j) to angle(ijk). + eij2a_index : n_angle + Broadcast indices from extended edge(ij) to angle(ijk), or reduction indices from angle(ijk) to edge(ij). + eik2a_index : n_angle + Broadcast indices from extended edge(ik) to angle(ijk). Returns ------- @@ -521,15 +760,37 @@ def forward( a_updated : nf x nloc x a_nnei x a_nnei x a_dim Updated angle embedding. """ - nb, nloc, nnei, _ = edge_ebd.shape + nb, nloc, nnei = nlist.shape nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :] if paddle.in_dynamic_mode(): assert [nb, nloc] == node_ebd.shape[:2] - if paddle.in_dynamic_mode(): - assert [nb, nloc, nnei] == h2.shape[:3] + if not self.use_dynamic_sel: + if paddle.in_dynamic_mode(): + assert [nb, nloc, nnei, 3] == h2.shape + n_edge = None + else: + n_edge = h2.shape[0] del a_nlist # may be used in the future + n2e_index, n_ext2e_index = edge_index[0], edge_index[1] + n2a_index, eij2a_index, eik2a_index = ( + angle_index[0], + angle_index[1], + angle_index[2], + ) + + # nb x nloc x nnei x n_dim [OR] n_edge x n_dim + nei_node_ebd = ( + _make_nei_g1(node_ebd_ext, nlist) + if not self.use_dynamic_sel + else paddle.index_select( + node_ebd_ext.reshape([-1, self.n_dim]), + n_ext2e_index, + 0, + ) + ) + n_update_list: list[paddle.Tensor] = [node_ebd] e_update_list: list[paddle.Tensor] = [edge_ebd] a_update_list: list[paddle.Tensor] = [angle_ebd] @@ -538,8 +799,6 @@ def forward( node_self_mlp = self.act(self.node_self_mlp(node_ebd)) n_update_list.append(node_self_mlp) - nei_node_ebd = _make_nei_g1(node_ebd_ext, nlist) - # node sym (grrg + drrd) node_sym_list: list[paddle.Tensor] = [] node_sym_list.append( @@ -550,6 +809,18 @@ def forward( sw, self.axis_neuron, ) + if not self.use_dynamic_sel + else self.symmetrization_op_dynamic( + edge_ebd, + h2, + sw, + owner=n2e_index, + num_owner=nb * nloc, + nb=nb, + nloc=nloc, + scale_factor=self.dynamic_e_sel ** (-0.5), + axis_neuron=self.axis_neuron, + ) ) node_sym_list.append( self.symmetrization_op( @@ -559,20 +830,47 @@ def forward( sw, self.axis_neuron, ) + if not self.use_dynamic_sel + else self.symmetrization_op_dynamic( + nei_node_ebd, + h2, + sw, + owner=n2e_index, + num_owner=nb * nloc, + nb=nb, + nloc=nloc, + scale_factor=self.dynamic_e_sel ** (-0.5), + axis_neuron=self.axis_neuron, + ) ) node_sym = self.act(self.node_sym_linear(paddle.concat(node_sym_list, axis=-1))) n_update_list.append(node_sym) if not self.optim_update: - # nb x nloc x nnei x (n_dim * 2 + e_dim) - edge_info = paddle.concat( - [ - paddle.tile(node_ebd.unsqueeze(-2), [1, 1, self.nnei, 1]), - nei_node_ebd, - edge_ebd, - ], - axis=-1, - ) + if not self.use_dynamic_sel: + # nb x nloc x nnei x (n_dim * 2 + e_dim) + edge_info = paddle.concat( + [ + paddle.tile(node_ebd.unsqueeze(-2), [1, 1, self.nnei, 1]), + nei_node_ebd, + edge_ebd, + ], + axis=-1, + ) + else: + # n_edge x (n_dim * 2 + e_dim) + edge_info = paddle.concat( + [ + paddle.index_select( + node_ebd.reshape([-1, self.n_dim]), + n2e_index, + 0, + ), + nei_node_ebd, + edge_ebd, + ], + axis=-1, + ) else: edge_info = None @@ -592,16 +890,37 @@ def forward( nlist, "node", ) + if not self.use_dynamic_sel + else self.optim_edge_update_dynamic( + node_ebd, + node_ebd_ext, + edge_ebd, + n2e_index, + n_ext2e_index, + "node", + ) ) * sw.unsqueeze(-1) + node_edge_update = ( + (paddle.sum(node_edge_update, axis=-2) / self.nnei) + if not self.use_dynamic_sel + else ( + aggregate( + node_edge_update, + n2e_index, + average=False, + num_owner=nb * nloc, + ).reshape([nb, nloc, node_edge_update.shape[-1]]) + / self.dynamic_e_sel + ) + ) - node_edge_update = paddle.sum(node_edge_update, axis=-2) / self.nnei if self.n_multi_edge_message > 1: - # nb x nloc x nnei x h x n_dim + # nb x nloc x h x n_dim node_edge_update_mul_head = node_edge_update.reshape( [nb, nloc, self.n_multi_edge_message, self.n_dim] ) for head_index in range(self.n_multi_edge_message): - n_update_list.append(node_edge_update_mul_head[:, :, head_index, :]) + n_update_list.append(node_edge_update_mul_head[..., head_index, :]) else: n_update_list.append(node_edge_update) # update node_ebd @@ -620,6 +939,15 @@ def forward( nlist, "edge", ) + if not self.use_dynamic_sel + else self.optim_edge_update_dynamic( + node_ebd, + node_ebd_ext, + edge_ebd, + n2e_index, + n_ext2e_index, + "edge", + ) ) e_update_list.append(edge_self_update) @@ -641,48 +969,66 @@ def forward( edge_ebd_for_angle = self.a_compress_e_linear(edge_ebd) else: # use the first a_compress_dim dim for node and edge - node_ebd_for_angle = node_ebd[:, :, : self.n_a_compress_dim] - edge_ebd_for_angle = edge_ebd[:, :, :, : self.e_a_compress_dim] + node_ebd_for_angle = node_ebd[..., : self.n_a_compress_dim] + edge_ebd_for_angle = edge_ebd[..., : self.e_a_compress_dim] else: node_ebd_for_angle = node_ebd edge_ebd_for_angle = edge_ebd - # nb x nloc x a_nnei x e_dim - edge_for_angle = edge_ebd_for_angle[:, :, : self.a_sel, :] - # nb x nloc x a_nnei x e_dim - edge_for_angle = paddle.where( - a_nlist_mask.unsqueeze(-1), - edge_for_angle, - paddle.zeros_like(edge_for_angle), - ).astype(edge_for_angle.dtype) + if not self.use_dynamic_sel: + # nb x nloc x a_nnei x e_dim + edge_ebd_for_angle = edge_ebd_for_angle[..., : self.a_sel, :] + # nb x nloc x a_nnei x e_dim + edge_ebd_for_angle = edge_ebd_for_angle.masked_fill( + ~a_nlist_mask.unsqueeze(-1), 0.0 + ) if not self.optim_update: - # nb x nloc x a_nnei x a_nnei x n_dim - node_for_angle_info = paddle.tile( - node_ebd_for_angle.unsqueeze(2).unsqueeze(2), - [1, 1, self.a_sel, self.a_sel, 1], + # nb x nloc x a_nnei x a_nnei x n_dim [OR] n_angle x n_dim + node_for_angle_info = ( + paddle.tile( + node_ebd_for_angle.unsqueeze(2).unsqueeze(2), + (1, 1, self.a_sel, self.a_sel, 1), + ) + if not self.use_dynamic_sel + else paddle.index_select( + node_ebd_for_angle.reshape([-1, self.n_a_compress_dim]), + n2a_index, + 0, + ) ) - # nb x nloc x (a_nnei) x a_nnei x edge_ebd - edge_for_angle_i = paddle.tile( - edge_for_angle.unsqueeze(2), (1, 1, self.a_sel, 1, 1) + + # nb x nloc x (a_nnei) x a_nnei x e_dim [OR] n_angle x e_dim + edge_for_angle_k = ( + paddle.tile( + edge_ebd_for_angle.unsqueeze(2), (1, 1, self.a_sel, 1, 1) + ) + if not self.use_dynamic_sel + else paddle.index_select(edge_ebd_for_angle, eik2a_index, 0) ) - # nb x nloc x a_nnei x (a_nnei) x e_dim - edge_for_angle_j = paddle.tile( - edge_for_angle.unsqueeze(3), (1, 1, 1, self.a_sel, 1) + # nb x nloc x a_nnei x (a_nnei) x e_dim [OR] n_angle x e_dim + edge_for_angle_j = ( + paddle.tile( + edge_ebd_for_angle.unsqueeze(3), (1, 1, 1, self.a_sel, 1) + ) + if not self.use_dynamic_sel + else paddle.index_select(edge_ebd_for_angle, eij2a_index, 0) ) - # nb x nloc x a_nnei x a_nnei x (e_dim + e_dim) + # nb x nloc x a_nnei x a_nnei x (e_dim + e_dim) [OR] n_angle x (e_dim + e_dim) edge_for_angle_info = paddle.concat( - [edge_for_angle_i, edge_for_angle_j], axis=-1 + [edge_for_angle_k, edge_for_angle_j], axis=-1 ) angle_info_list = [angle_ebd] angle_info_list.append(node_for_angle_info) angle_info_list.append(edge_for_angle_info) # nb x nloc x a_nnei x a_nnei x (a + n_dim + e_dim*2) or (a + a/c + a/c) + # [OR] + # n_angle x (a + n_dim + e_dim*2) or (a + a/c + a/c) angle_info = paddle.concat(angle_info_list, axis=-1) else: angle_info = None # edge angle message - # nb x nloc x a_nnei x a_nnei x e_dim + # nb x nloc x a_nnei x a_nnei x e_dim [OR] n_angle x e_dim if not self.optim_update: assert angle_info is not None edge_angle_update = self.act(self.edge_angle_linear1(angle_info)) @@ -691,32 +1037,61 @@ def forward( self.optim_angle_update( angle_ebd, node_ebd_for_angle, - edge_for_angle, + edge_ebd_for_angle, + "edge", + ) + if not self.use_dynamic_sel + else self.optim_angle_update_dynamic( + angle_ebd, + node_ebd_for_angle, + edge_ebd_for_angle, + n2a_index, + eij2a_index, + eik2a_index, "edge", ) ) - # nb x nloc x a_nnei x a_nnei x e_dim - weighted_edge_angle_update = ( - a_sw[..., None, None] * a_sw[..., None, :, None] * edge_angle_update - ) - # nb x nloc x a_nnei x e_dim - reduced_edge_angle_update = paddle.sum( - weighted_edge_angle_update, axis=-2 - ) / (self.a_sel**0.5) - # nb x nloc x nnei x e_dim - padding_edge_angle_update = paddle.concat( - [ - reduced_edge_angle_update, - paddle.zeros( - [nb, nloc, self.nnei - self.a_sel, self.e_dim], - dtype=edge_ebd.dtype, - ).to(device=edge_ebd.place), - ], - axis=2, - ) + if not self.use_dynamic_sel: + # nb x nloc x a_nnei x a_nnei x e_dim + weighted_edge_angle_update = ( + a_sw.unsqueeze(-1).unsqueeze(-1) + * a_sw.unsqueeze(-2).unsqueeze(-1) + * edge_angle_update + ) + # nb x nloc x a_nnei x e_dim + reduced_edge_angle_update = paddle.sum( + weighted_edge_angle_update, axis=-2 + ) / (self.a_sel**0.5) + # nb x nloc x nnei x e_dim + padding_edge_angle_update = paddle.concat( + [ + reduced_edge_angle_update, + paddle.zeros( + [nb, nloc, self.nnei - self.a_sel, self.e_dim], + dtype=edge_ebd.dtype, + ), + ], + axis=2, + ) + else: + # n_angle x e_dim + weighted_edge_angle_update = edge_angle_update * a_sw.unsqueeze(-1) + # n_edge x e_dim + padding_edge_angle_update = aggregate( + weighted_edge_angle_update, + eij2a_index, + average=False, + num_owner=n_edge, + ) / (self.dynamic_a_sel**0.5) + if not self.smooth_edge_update: # will be deprecated in the future + # not support dynamic index, will pass anyway + if self.use_dynamic_sel: + raise NotImplementedError( + "smooth_edge_update must be True when use_dynamic_sel is True!" + ) full_mask = paddle.concat( [ a_nlist_mask, @@ -727,8 +1102,8 @@ def forward( ], axis=-1, ) - padding_edge_angle_update = paddle.where( - full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd + padding_edge_angle_update = padding_edge_angle_update.masked_fill( + ~full_mask.unsqueeze(-1), edge_ebd ) e_update_list.append( self.act(self.edge_angle_linear2(padding_edge_angle_update)) @@ -746,7 +1121,17 @@ def forward( self.optim_angle_update( angle_ebd, node_ebd_for_angle, - edge_for_angle, + edge_ebd_for_angle, + "angle", + ) + if not self.use_dynamic_sel + else self.optim_angle_update_dynamic( + angle_ebd, + node_ebd_for_angle, + edge_ebd_for_angle, + n2a_index, + eij2a_index, + eik2a_index, "angle", ) ) diff --git a/deepmd/pd/model/descriptor/repflows.py b/deepmd/pd/model/descriptor/repflows.py index 3200c26dba..21b5984ef3 100644 --- a/deepmd/pd/model/descriptor/repflows.py +++ b/deepmd/pd/model/descriptor/repflows.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -10,6 +8,10 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.pd.cxx_op import ( + ENABLE_CUSTOMIZED_OP, + paddle_ops_deepmd, +) from deepmd.pd.model.descriptor.descriptor import ( DescriptorBlock, ) @@ -19,6 +21,9 @@ from deepmd.pd.model.network.mlp import ( MLPLayer, ) +from deepmd.pd.model.network.utils import ( + get_graph_index, +) from deepmd.pd.utils import ( env, ) @@ -31,6 +36,9 @@ from deepmd.pd.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.pd.utils.spin import ( + concat_switch_virtual, +) from deepmd.pd.utils.utils import ( ActivationFn, ) @@ -45,6 +53,30 @@ RepFlowLayer, ) +if not ENABLE_CUSTOMIZED_OP: + + def border_op( + argument0, + argument1, + argument2, + argument3, + argument4, + argument5, + argument6, + argument7, + argument8, + ) -> paddle.Tensor: + raise NotImplementedError( + "The 'border_op' operator is unavailable because the custom Paddle OP library was not built when freezing the model.\n" + "To install 'border_op', run: python source/op/pd/setup.py install\n" + "For more information, please refer to the DPA3 documentation." + ) + + # Note: this hack cannot actually save a model that can be run using LAMMPS. + paddle_ops_deepmd_border_op = border_op +else: + paddle_ops_deepmd_border_op = paddle_ops_deepmd.border_op + @DescriptorBlock.register("se_repflow") class DescrptBlockRepflows(DescriptorBlock): @@ -109,12 +141,35 @@ class DescrptBlockRepflows(DescriptorBlock): smooth_edge_update : bool, optional Whether to make edge update smooth. If True, the edge update from angle message will not use self as padding. + edge_init_use_dist : bool, optional + Whether to use direct distance r to initialize the edge features instead of 1/r. + Note that when using this option, the activation function will not be used when initializing edge features. + use_exp_switch : bool, optional + Whether to use an exponential switch function instead of a polynomial one in the neighbor update. + The exponential switch function ensures neighbor contributions smoothly diminish as the interatomic distance + `r` approaches the cutoff radius `rcut`. Specifically, the function is defined as: + s(r) = \\exp(-\\exp(20 * (r - rcut_smth) / rcut_smth)) for 0 < r \\leq rcut, and s(r) = 0 for r > rcut. + Here, `rcut_smth` is an adjustable smoothing factor and `rcut_smth` should be chosen carefully + according to `rcut`, ensuring s(r) approaches zero smoothly at the cutoff. + Typical recommended values are `rcut_smth` = 5.3 for `rcut` = 6.0, and 3.5 for `rcut` = 4.0. + use_dynamic_sel : bool, optional + Whether to dynamically select neighbors within the cutoff radius. + If True, the exact number of neighbors within the cutoff radius is used + without padding to a fixed selection numbers. + When enabled, users can safely set larger values for `e_sel` or `a_sel` (e.g., 1200 or 300, respectively) + to guarantee capturing all neighbors within the cutoff radius. + Note that when using dynamic selection, the `smooth_edge_update` must be True. + sel_reduce_factor : float, optional + Reduction factor applied to neighbor-scale normalization when `use_dynamic_sel` is True. + In the dynamic selection case, neighbor-scale normalization will use `e_sel / sel_reduce_factor` + or `a_sel / sel_reduce_factor` instead of the raw `e_sel` or `a_sel` values, + accommodating larger selection numbers. + use_loc_mapping : bool, Optional + Whether to use local atom index mapping in training or non-parallel inference. + When True, local indexing and mapping are applied to neighbor lists and embeddings during descriptor computation. optim_update : bool, optional Whether to enable the optimized update method. Uses a more efficient process when enabled. Defaults to True - use_loc_mapping : bool, Optional - Whether to use local atom index mapping in training or non-parallel inference. - Not supported yet in Paddle. ntypes : int Number of element types activation_function : str, optional @@ -131,6 +186,8 @@ class DescrptBlockRepflows(DescriptorBlock): For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection. seed : int, optional Random seed for parameter initialization. + trainable : bool, default: True + Whether this block is trainable """ def __init__( @@ -162,11 +219,14 @@ def __init__( precision: str = "float64", fix_stat_std: float = 0.3, smooth_edge_update: bool = False, + edge_init_use_dist: bool = False, + use_exp_switch: bool = False, use_dynamic_sel: bool = False, sel_reduce_factor: float = 10.0, - use_loc_mapping: bool = False, + use_loc_mapping: bool = True, optim_update: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.e_rcut = float(e_rcut) @@ -176,6 +236,9 @@ def __init__( self.a_rcut_smth = float(a_rcut_smth) self.a_sel = a_sel self.ntypes = ntypes + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) self.nlayers = nlayers # for other common desciptor method sel = [e_sel] if isinstance(e_sel, int) else e_sel @@ -183,7 +246,9 @@ def __init__( self.ndescrpt = self.nnei * 4 # use full descriptor. assert len(sel) == 1 self.sel = sel + self.register_buffer("buffer_sel", paddle.to_tensor(sel)) self.rcut = e_rcut + self.register_buffer("buffer_rcut", paddle.to_tensor(self.e_rcut)) self.rcut_smth = e_rcut_smth self.sec = self.sel self.split_sel = self.sel @@ -195,13 +260,21 @@ def __init__( self.fix_stat_std = fix_stat_std self.set_stddev_constant = fix_stat_std != 0.0 self.a_compress_use_split = a_compress_use_split + self.use_loc_mapping = use_loc_mapping self.optim_update = optim_update self.smooth_edge_update = smooth_edge_update - self.use_dynamic_sel = use_dynamic_sel # not supported yet + self.edge_init_use_dist = edge_init_use_dist + self.use_exp_switch = use_exp_switch + self.use_dynamic_sel = use_dynamic_sel self.sel_reduce_factor = sel_reduce_factor - assert not self.use_dynamic_sel, "Dynamic selection is not supported yet." - self.use_loc_mapping = use_loc_mapping - assert not self.use_loc_mapping, "Local mapping is not supported yet." + if self.use_dynamic_sel and not self.smooth_edge_update: + raise NotImplementedError( + "smooth_edge_update must be True when use_dynamic_sel is True!" + ) + if self.sel_reduce_factor <= 0: + raise ValueError( + f"`sel_reduce_factor` must be > 0, got {self.sel_reduce_factor}" + ) self.n_dim = n_dim self.e_dim = e_dim @@ -223,10 +296,19 @@ def __init__( self.seed = seed self.edge_embd = MLPLayer( - 1, self.e_dim, precision=precision, seed=child_seed(seed, 0) + 1, + self.e_dim, + precision=precision, + seed=child_seed(seed, 0), + trainable=trainable, ) self.angle_embd = MLPLayer( - 1, self.a_dim, precision=precision, bias=False, seed=child_seed(seed, 1) + 1, + self.a_dim, + precision=precision, + bias=False, + seed=child_seed(seed, 1), + trainable=trainable, ) layers = [] for ii in range(nlayers): @@ -258,6 +340,7 @@ def __init__( sel_reduce_factor=self.sel_reduce_factor, smooth_edge_update=self.smooth_edge_update, seed=child_seed(child_seed(seed, 1), ii), + trainable=trainable, ) ) self.layers = paddle.nn.LayerList(layers) @@ -275,6 +358,10 @@ def get_rcut(self) -> float: """Returns the cut-off radius.""" return self.e_rcut + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.buffer_rcut + def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.e_rcut_smth @@ -287,9 +374,13 @@ def get_sel(self) -> list[int]: """Returns the number of selected atoms for each type.""" return self.sel + def get_buffer_sel(self) -> paddle.Tensor: + """Returns the number of selected atoms for each type as a buffer-style Tensor.""" + return self.buffer_sel + def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_dim_out(self) -> int: """Returns the output dimension.""" @@ -362,19 +453,20 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): - if comm_dict is None: - assert mapping is not None - assert extended_atype_embd is not None + parallel_mode = comm_dict is not None + if not parallel_mode: + if paddle.in_dynamic_mode(): + assert mapping is not None nframes, nloc, nnei = nlist.shape nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 atype = extended_atype[:, :nloc] # nb x nloc x nnei exclude_mask = self.emask(nlist, extended_atype) - nlist = paddle.where(exclude_mask != 0, nlist, -1) + nlist = paddle.where(exclude_mask != 0, nlist, paddle.full_like(nlist, -1)) # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1 dmatrix, diff, sw = prod_env_mat( extended_coord, @@ -385,36 +477,19 @@ def forward( self.e_rcut, self.e_rcut_smth, protection=self.env_protection, + use_exp_switch=self.use_exp_switch, ) nlist_mask = nlist != -1 sw = paddle.squeeze(sw, -1) # beyond the cutoff sw should be 0.0 sw = sw.masked_fill(~nlist_mask, 0.0) - # [nframes, nloc, tebd_dim] - if comm_dict is None: - if paddle.in_dynamic_mode(): - assert isinstance(extended_atype_embd, paddle.Tensor) - atype_embd = extended_atype_embd[:, :nloc, :] - if paddle.in_dynamic_mode(): - assert atype_embd.shape == [nframes, nloc, self.n_dim] - else: - atype_embd = extended_atype_embd - if paddle.in_dynamic_mode(): - assert isinstance(atype_embd, paddle.Tensor) - node_ebd = self.act(atype_embd) - n_dim = node_ebd.shape[-1] - # nb x nloc x nnei x 1, nb x nloc x nnei x 3 - edge_input, h2 = paddle.split(dmatrix, [1, 3], axis=-1) - # nb x nloc x nnei x e_dim - edge_ebd = self.act(self.edge_embd(edge_input)) - # get angle nlist (maybe smaller) a_dist_mask = (paddle.linalg.norm(diff, axis=-1) < self.a_rcut)[ :, :, : self.a_sel ] a_nlist = nlist[:, :, : self.a_sel] - a_nlist = paddle.where(a_dist_mask, a_nlist, -1) + a_nlist = paddle.where(a_dist_mask, a_nlist, paddle.full_like(a_nlist, -1)) _, a_diff, a_sw = prod_env_mat( extended_coord, a_nlist, @@ -424,13 +499,34 @@ def forward( self.a_rcut, self.a_rcut_smth, protection=self.env_protection, + use_exp_switch=self.use_exp_switch, ) a_nlist_mask = a_nlist != -1 a_sw = paddle.squeeze(a_sw, -1) # beyond the cutoff sw should be 0.0 a_sw = a_sw.masked_fill(~a_nlist_mask, 0.0) + # set all padding positions to index of 0 + # if the a neighbor is real or not is indicated by nlist_mask + nlist[nlist == -1] = 0 a_nlist[a_nlist == -1] = 0 + # get node embedding + # [nframes, nloc, tebd_dim] + assert extended_atype_embd is not None + atype_embd = extended_atype_embd[:, :nloc, :] + if paddle.in_dynamic_mode(): + assert list(atype_embd.shape) == [nframes, nloc, self.n_dim] + assert isinstance(atype_embd, paddle.Tensor) # for jit + node_ebd = self.act(atype_embd) + n_dim = node_ebd.shape[-1] + + # get edge and angle embedding input + # nb x nloc x nnei x 1, nb x nloc x nnei x 3 + edge_input, h2 = paddle.split(dmatrix, [1, 3], axis=-1) + if self.edge_init_use_dist: + # nb x nloc x nnei x 1 + edge_input = paddle.linalg.norm(diff, axis=-1, keepdim=True) + # nf x nloc x a_nnei x 3 normalized_diff_i = a_diff / ( paddle.linalg.norm(a_diff, axis=-1, keepdim=True) + 1e-6 @@ -440,19 +536,57 @@ def forward( # nf x nloc x a_nnei x a_nnei # 1 - 1e-6 for paddle.acos stability cosine_ij = paddle.matmul(normalized_diff_i, normalized_diff_j) * (1 - 1e-6) - # nf x nloc x a_nnei x a_nnei x 1 - cosine_ij = cosine_ij.unsqueeze(-1) / (paddle.pi**0.5) - # nf x nloc x a_nnei x a_nnei x a_dim - angle_ebd = self.angle_embd(cosine_ij).reshape( - [nframes, nloc, self.a_sel, self.a_sel, self.a_dim] - ) + angle_input = cosine_ij.unsqueeze(-1) / (paddle.pi**0.5) + + if not parallel_mode and self.use_loc_mapping: + if paddle.in_dynamic_mode(): + assert mapping is not None + # convert nlist from nall to nloc index + nlist = paddle.take_along_axis( + mapping, + nlist.reshape([nframes, -1]), + 1, + broadcast=False, + ).reshape(nlist.shape) + if self.use_dynamic_sel: + # get graph index + edge_index, angle_index = get_graph_index( + nlist, + nlist_mask, + a_nlist_mask, + nall, + use_loc_mapping=self.use_loc_mapping, + ) + # flat all the tensors + # n_edge x 1 + edge_input = edge_input[nlist_mask] + # n_edge x 3 + h2 = h2[nlist_mask] + # n_edge x 1 + sw = sw[nlist_mask] + # nb x nloc x a_nnei x a_nnei + a_nlist_mask = a_nlist_mask[:, :, :, None] & a_nlist_mask[:, :, None, :] + # n_angle x 1 + angle_input = angle_input[a_nlist_mask] + # n_angle x 1 + a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask] + else: + # avoid jit assertion + edge_index = paddle.zeros([2, 1], dtype=nlist.dtype) + angle_index = paddle.zeros([3, 1], dtype=nlist.dtype) + # get edge and angle embedding + # nb x nloc x nnei x e_dim [OR] n_edge x e_dim + if not self.edge_init_use_dist: + edge_ebd = self.act(self.edge_embd(edge_input)) + else: + edge_ebd = self.edge_embd(edge_input) + # nf x nloc x a_nnei x a_nnei x a_dim [OR] n_angle x a_dim + angle_ebd = self.angle_embd(angle_input) - # set all padding positions to index of 0 - # if the a neighbor is real or not is indicated by nlist_mask - nlist[nlist == -1] = 0 # nb x nall x n_dim - if comm_dict is None: - assert mapping is not None + if not parallel_mode: + if paddle.in_dynamic_mode(): + assert mapping is not None mapping = ( mapping.reshape([nframes, nall]) .unsqueeze(-1) @@ -460,14 +594,97 @@ def forward( ) for idx, ll in enumerate(self.layers): # node_ebd: nb x nloc x n_dim - # node_ebd_ext: nb x nall x n_dim - if comm_dict is None: - assert mapping is not None - node_ebd_ext = paddle.take_along_axis( - node_ebd, mapping, 1, broadcast=False + # node_ebd_ext: nb x nall x n_dim [OR] nb x nloc x n_dim when not parallel_mode + if not parallel_mode: + if paddle.in_dynamic_mode(): + assert mapping is not None + node_ebd_ext = ( + paddle.take_along_axis(node_ebd, mapping, 1, broadcast=False) + if not self.use_loc_mapping + else node_ebd ) else: - raise NotImplementedError("Not implemented") + assert len(comm_dict) >= 6 + has_spin = len(comm_dict) >= 7 + if not has_spin: + n_padding = nall - nloc + if paddle.in_dynamic_mode(): + node_ebd = paddle.nn.functional.pad( + node_ebd.squeeze(0), [0, 0, 0, n_padding], value=0.0 + ) + else: + _fill_shape = node_ebd.shape[1:] + _fill_shape[0] = n_padding + node_ebd = paddle.concat( + [ + node_ebd.squeeze(0), + paddle.zeros(_fill_shape, dtype=node_ebd.dtype), + ], + axis=0, + ) + # [nframes, nloc, tebd_dim] + real_nloc = nloc + real_nall = nall + else: + # for spin + real_nloc = nloc // 2 + real_nall = nall // 2 + real_n_padding = real_nall - real_nloc + node_ebd_real, node_ebd_virtual = paddle.split( + node_ebd, [real_nloc, real_nloc], axis=1 + ) + # mix_node_ebd: nb x real_nloc x (n_dim * 2) + mix_node_ebd = paddle.concat( + [node_ebd_real, node_ebd_virtual], axis=2 + ) + # nb x real_nall x (n_dim * 2) + if paddle.in_dynamic_mode(): + node_ebd = paddle.nn.functional.pad( + mix_node_ebd.squeeze(0), + (0, 0, 0, real_n_padding), + value=0.0, + ) + else: + _fill_shape = mix_node_ebd.shape[1:] + _fill_shape[0] = real_n_padding + node_ebd = paddle.concat( + [ + mix_node_ebd.squeeze(0), + paddle.zeros(_fill_shape, dtype=mix_node_ebd.dtype), + ], + axis=0, + ) + + assert len(comm_dict) >= 6 + ret = paddle_ops_deepmd_border_op( + comm_dict[0], + comm_dict[1], + comm_dict[2], + comm_dict[3], + comm_dict[4], + node_ebd, + comm_dict[5], + paddle.to_tensor( + real_nloc, + dtype=paddle.int32, + place=paddle.CPUPlace(), + ), # should be int of c++, placed on cpu + paddle.to_tensor( + real_nall - real_nloc, + dtype=paddle.int32, + place=paddle.CPUPlace(), + ), # should be int of c++, placed on cpu + ) + if not paddle.in_dynamic_mode(): + ret = paddle.assign(ret) + node_ebd_ext = ret.unsqueeze(0) + if has_spin: + node_ebd_real_ext, node_ebd_virtual_ext = paddle.split( + node_ebd_ext, [n_dim, n_dim], axis=2 + ) + node_ebd_ext = concat_switch_virtual( + node_ebd_real_ext, node_ebd_virtual_ext, real_nloc + ) node_ebd, edge_ebd, angle_ebd = ll.forward( node_ebd_ext, edge_ebd, @@ -479,12 +696,27 @@ def forward( a_nlist, a_nlist_mask, a_sw, + edge_index=edge_index, + angle_index=angle_index, ) # nb x nloc x 3 x e_dim - h2g2 = RepFlowLayer._cal_hg(edge_ebd, h2, nlist_mask, sw) + h2g2 = ( + RepFlowLayer._cal_hg(edge_ebd, h2, nlist_mask, sw) + if not self.use_dynamic_sel + else RepFlowLayer._cal_hg_dynamic( + edge_ebd, + h2, + sw, + owner=edge_index[0], + num_owner=nframes * nloc, + nb=nframes, + nloc=nloc, + scale_factor=(self.nnei / self.sel_reduce_factor) ** (-0.5), + ) + ) # (nb x nloc) x e_dim x 3 - rot_mat = paddle.transpose(h2g2, (0, 1, 3, 2)) + rot_mat = paddle.transpose(h2g2, [0, 1, 3, 2]) return ( node_ebd, @@ -496,8 +728,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pd/model/descriptor/repformer_layer.py b/deepmd/pd/model/descriptor/repformer_layer.py index b4d93d8301..efb1881d59 100644 --- a/deepmd/pd/model/descriptor/repformer_layer.py +++ b/deepmd/pd/model/descriptor/repformer_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle import paddle.nn as nn @@ -44,7 +40,7 @@ def get_residual( _mode: str = "norm", trainable: bool = True, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> paddle.Tensor: r""" Get residual tensor for one update vector. @@ -162,7 +158,8 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" super().__init__() @@ -175,6 +172,7 @@ def __init__( bias=False, precision=precision, seed=seed, + trainable=trainable, ) self.has_gate = has_gate self.smooth = smooth @@ -287,7 +285,8 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim @@ -298,12 +297,14 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.head_map = MLPLayer( input_dim * head_num, input_dim, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.precision = precision @@ -374,13 +375,19 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num self.head_map = MLPLayer( - head_num, 1, bias=False, precision=precision, seed=seed + head_num, + 1, + bias=False, + precision=precision, + seed=seed, + trainable=trainable, ) self.precision = precision @@ -447,7 +454,8 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim @@ -459,6 +467,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.mapkv = MLPLayer( input_dim, @@ -466,12 +475,14 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.head_map = MLPLayer( input_dim * head_num, input_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) self.smooth = smooth self.attnw_shift = attnw_shift @@ -607,11 +618,12 @@ def __init__( smooth: bool = True, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei @@ -672,6 +684,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) ) @@ -681,6 +694,7 @@ def __init__( g1_dim, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.linear2 = None self.proj_g1g2 = None @@ -697,6 +711,7 @@ def __init__( g2_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) if self.update_style == "res_residual": self.g2_residual.append( @@ -706,6 +721,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) ) if self.g1_out_mlp: @@ -714,6 +730,7 @@ def __init__( g1_dim, precision=precision, seed=child_seed(seed, 15), + trainable=trainable, ) if self.update_style == "res_residual": self.g1_residual.append( @@ -723,6 +740,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 16), + trainable=trainable, ) ) else: @@ -735,6 +753,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) else: self.proj_g1g2 = MLPLayer( @@ -743,6 +762,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) if self.update_style == "res_residual": self.g1_residual.append( @@ -752,6 +772,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 17), + trainable=trainable, ) ) if self.update_g2_has_g1g1: @@ -761,6 +782,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 5), + trainable=trainable, ) if self.update_style == "res_residual": self.g2_residual.append( @@ -770,6 +792,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 6), + trainable=trainable, ) ) if self.update_g2_has_attn or self.update_h2: @@ -781,10 +804,15 @@ def __init__( self.smooth, precision=precision, seed=child_seed(seed, 7), + trainable=trainable, ) if self.update_g2_has_attn: self.attn2_mh_apply = Atten2MultiHeadApply( - g2_dim, attn2_nhead, precision=precision, seed=child_seed(seed, 8) + g2_dim, + attn2_nhead, + precision=precision, + seed=child_seed(seed, 8), + trainable=trainable, ) self.attn2_lm = LayerNorm( g2_dim, @@ -801,12 +829,17 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 10), + trainable=trainable, ) ) if self.update_h2: self.attn2_ev_apply = Atten2EquiVarApply( - g2_dim, attn2_nhead, precision=precision, seed=child_seed(seed, 11) + g2_dim, + attn2_nhead, + precision=precision, + seed=child_seed(seed, 11), + trainable=trainable, ) if self.update_style == "res_residual": self.h2_residual.append( @@ -816,6 +849,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 12), + trainable=trainable, ) ) if self.update_g1_has_attn: @@ -826,6 +860,7 @@ def __init__( self.smooth, precision=precision, seed=child_seed(seed, 13), + trainable=trainable, ) if self.update_style == "res_residual": self.g1_residual.append( @@ -835,6 +870,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 14), + trainable=trainable, ) ) @@ -914,9 +950,7 @@ def _update_g1_conv( ).unsqueeze(-1) else: gg1 = _apply_switch(gg1, sw) - invnnei = (1.0 / float(nnei)) * paddle.ones( - (nb, nloc, 1), dtype=gg1.dtype - ).to(device=gg1.place) + invnnei = (1.0 / float(nnei)) * paddle.ones((nb, nloc, 1), dtype=gg1.dtype) if not self.g1_out_conv: # nb x nloc x ng2 g1_11 = paddle.sum(g2 * gg1, axis=2) * invnnei @@ -986,11 +1020,10 @@ def _cal_hg( if not use_sqrt_nnei: invnnei = (1.0 / float(nnei)) * paddle.ones( (nb, nloc, 1, 1), dtype=g2.dtype - ).to(device=g2.place) + ) else: invnnei = paddle.rsqrt( - float(nnei) - * paddle.ones([nb, nloc, 1, 1], dtype=g2.dtype).to(device=g2.place) + float(nnei) * paddle.ones([nb, nloc, 1, 1], dtype=g2.dtype) ) # nb x nloc x 3 x ng2 h2g2 = paddle.matmul(paddle.transpose(h2, [0, 1, 3, 2]), g2) * invnnei diff --git a/deepmd/pd/model/descriptor/repformers.py b/deepmd/pd/model/descriptor/repformers.py index 32f88dd1d3..24f92f1bee 100644 --- a/deepmd/pd/model/descriptor/repformers.py +++ b/deepmd/pd/model/descriptor/repformers.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -10,6 +8,10 @@ from deepmd.dpmodel.utils.seed import ( child_seed, ) +from deepmd.pd.cxx_op import ( + ENABLE_CUSTOMIZED_OP, + paddle_ops_deepmd, +) from deepmd.pd.model.descriptor.descriptor import ( DescriptorBlock, ) @@ -31,6 +33,9 @@ from deepmd.pd.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.pd.utils.spin import ( + concat_switch_virtual, +) from deepmd.pd.utils.utils import ( ActivationFn, ) @@ -45,6 +50,30 @@ RepformerLayer, ) +if not ENABLE_CUSTOMIZED_OP: + + def border_op( + argument0, + argument1, + argument2, + argument3, + argument4, + argument5, + argument6, + argument7, + argument8, + ) -> paddle.Tensor: + raise NotImplementedError( + "The 'border_op' operator is unavailable because the custom Paddle OP library was not built when freezing the model.\n" + "To install 'border_op', run: python source/op/pd/setup.py install\n" + "For more information, please refer to the DPA3 documentation." + ) + + # Note: this hack cannot actually save a model that can be run using LAMMPS. + paddle_ops_deepmd_border_op = border_op +else: + paddle_ops_deepmd_border_op = paddle_ops_deepmd.border_op + @DescriptorBlock.register("se_repformer") @DescriptorBlock.register("se_uni") @@ -82,11 +111,12 @@ def __init__( env_protection: float = 0.0, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, + trainable: bool = True, ) -> None: r""" The repformer descriptor block. @@ -173,11 +203,19 @@ def __init__( The epsilon value for layer normalization. seed : int, optional Random seed for parameter initialization. + trainable : bool + Whether the block is trainable """ super().__init__() self.rcut = float(rcut) + self.register_buffer("buffer_rcut", paddle.to_tensor(self.rcut)) self.rcut_smth = float(rcut_smth) + self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.rcut_smth)) self.ntypes = ntypes + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) + self.nlayers = nlayers sel = [sel] if isinstance(sel, int) else sel self.nnei = sum(sel) @@ -223,7 +261,11 @@ def __init__( self.seed = seed self.g2_embd = MLPLayer( - 1, self.g2_dim, precision=precision, seed=child_seed(seed, 0) + 1, + self.g2_dim, + precision=precision, + seed=child_seed(seed, 0), + trainable=trainable, ) layers = [] for ii in range(nlayers): @@ -261,16 +303,17 @@ def __init__( g1_out_conv=self.g1_out_conv, g1_out_mlp=self.g1_out_mlp, seed=child_seed(child_seed(seed, 1), ii), + trainable=trainable, ) ) self.layers = paddle.nn.LayerList(layers) wanted_shape = (self.ntypes, self.nnei, 4) mean = paddle.zeros(wanted_shape, dtype=env.GLOBAL_PD_FLOAT_PRECISION).to( - device=env.DEVICE + env.DEVICE ) stddev = paddle.ones(wanted_shape, dtype=env.GLOBAL_PD_FLOAT_PRECISION).to( - device=env.DEVICE + env.DEVICE ) self.register_buffer("mean", mean) self.register_buffer("stddev", stddev) @@ -284,6 +327,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.buffer_rcut + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.buffer_rcut_smth + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return sum(self.sel) @@ -294,7 +345,7 @@ def get_sel(self) -> list[int]: def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_dim_out(self) -> int: """Returns the output dimension.""" @@ -367,12 +418,12 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): - if comm_dict is None: + if (comm_dict is None or len(comm_dict) == 0) and paddle.in_dynamic_mode(): assert mapping is not None assert extended_atype_embd is not None nframes, nloc, nnei = nlist.shape @@ -398,7 +449,7 @@ def forward( sw = sw.masked_fill(~nlist_mask, 0.0) # [nframes, nloc, tebd_dim] - if comm_dict is None: + if comm_dict is None or len(comm_dict) == 0: if paddle.in_dynamic_mode(): assert isinstance(extended_atype_embd, paddle.Tensor) # for jit atype_embd = extended_atype_embd[:, :nloc, :] @@ -406,8 +457,8 @@ def forward( assert list(atype_embd.shape) == [nframes, nloc, self.g1_dim] else: atype_embd = extended_atype_embd - if paddle.in_dynamic_mode(): - assert isinstance(atype_embd, paddle.Tensor) # for jit + if paddle.in_dynamic_mode(): + assert isinstance(atype_embd, paddle.Tensor) # for jit g1 = self.act(atype_embd) ng1 = g1.shape[-1] # nb x nloc x nnei x 1, nb x nloc x nnei x 3 @@ -424,8 +475,9 @@ def forward( # if the a neighbor is real or not is indicated by nlist_mask nlist[nlist == -1] = 0 # nb x nall x ng1 - if comm_dict is None: - assert mapping is not None + if comm_dict is None or len(comm_dict) == 0: + if paddle.in_dynamic_mode(): + assert mapping is not None mapping = ( mapping.reshape([nframes, nall]) .unsqueeze(-1) @@ -434,14 +486,91 @@ def forward( for idx, ll in enumerate(self.layers): # g1: nb x nloc x ng1 # g1_ext: nb x nall x ng1 - if comm_dict is None: - assert mapping is not None + if comm_dict is None or len(comm_dict) == 0: + if paddle.in_dynamic_mode(): + assert mapping is not None g1_ext = paddle.take_along_axis( g1, axis=1, indices=mapping, broadcast=False ) else: - raise NotImplementedError("Not implemented yet") - + has_spin = len(comm_dict) >= 7 + if not has_spin: + n_padding = nall - nloc + if paddle.in_dynamic_mode(): + g1 = paddle.nn.functional.pad( + g1.squeeze(0), + (0, 0, 0, n_padding), + value=0.0, + pad_from_left_axis=False, + ) + else: + _fill_shape = g1.shape[1:] + _fill_shape[0] = n_padding + g1 = paddle.concat( + [g1.squeeze(0), paddle.zeros(_fill_shape, dtype=g1.dtype)], + axis=0, + ) + real_nloc = nloc + real_nall = nall + else: + # for spin + real_nloc = nloc // 2 + real_nall = nall // 2 + real_n_padding = real_nall - real_nloc + g1_real, g1_virtual = paddle.split( + g1, [real_nloc, real_nloc], axis=1 + ) + # mix_g1: nb x real_nloc x (ng1 * 2) + mix_g1 = paddle.concat([g1_real, g1_virtual], axis=2) + # nb x real_nall x (ng1 * 2) + if paddle.in_dynamic_mode(): + g1 = paddle.nn.functional.pad( + mix_g1.squeeze(0), + (0, 0, 0, real_n_padding), + value=0.0, + pad_from_left_axis=False, + ) + else: + _fill_shape = mix_g1.shape[1:] + _fill_shape[0] = real_n_padding + g1 = paddle.concat( + [ + mix_g1.squeeze(0), + paddle.zeros(_fill_shape, dtype=mix_g1.dtype), + ], + axis=0, + ) + + assert len(comm_dict) >= 6 + ret = paddle_ops_deepmd_border_op( + comm_dict[0], + comm_dict[1], + comm_dict[2], + comm_dict[3], + comm_dict[4], + g1, + comm_dict[5], + paddle.to_tensor( + [real_nloc], + dtype=paddle.int32, + place=paddle.CPUPlace(), + ), # should be int of c++, placed on cpu + paddle.to_tensor( + [real_nall - real_nloc], + dtype=paddle.int32, + place=paddle.CPUPlace(), + ), # should be int of c++, placed on cpu + ) + if not paddle.in_dynamic_mode(): + ret = paddle.assign(ret) + g1_ext = ret.unsqueeze(0) + if has_spin: + g1_real_ext, g1_virtual_ext = paddle.split( + g1_ext, [ng1, ng1], axis=2 + ) + g1_ext = concat_switch_virtual( + g1_real_ext, g1_virtual_ext, real_nloc + ) g1, g2, h2 = ll.forward( g1_ext, g2, @@ -468,8 +597,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -503,11 +632,11 @@ def compute_input_stats( mean, stddev = env_mat_stat() if not self.set_davg_zero: paddle.assign( - paddle.to_tensor(mean, dtype=self.mean.dtype).to(device=env.DEVICE), + paddle.to_tensor(mean, dtype=self.mean.dtype).to(env.DEVICE), self.mean, ) # pylint: disable=no-explicit-dtype paddle.assign( - paddle.to_tensor(stddev, dtype=self.stddev.dtype).to(device=env.DEVICE), + paddle.to_tensor(stddev, dtype=self.stddev.dtype).to(env.DEVICE), self.stddev, ) # pylint: disable=no-explicit-dtype diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index 7b70a742ce..17ef5d67c4 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( ClassVar, - Optional, - Union, ) import numpy as np @@ -84,9 +84,9 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - ntypes: Optional[int] = None, # to be compat with input - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + ntypes: int | None = None, # to be compat with input + type_map: list[str] | None = None, # not implemented spin=None, ) -> None: @@ -95,6 +95,11 @@ def __init__( raise NotImplementedError("old implementation of spin is not supported.") super().__init__() self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) self.compress = False self.prec = PRECISION_DICT[precision] self.sea = DescrptBlockSeA( @@ -122,6 +127,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.sea.get_rcut_smth() + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.sea.get_buffer_rcut() + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.sea.get_buffer_rcut_smth() + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return self.sea.get_nsel() @@ -138,6 +151,18 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + def get_dim_out(self) -> int: """Returns the output dimension.""" return self.sea.get_dim_out() @@ -201,8 +226,8 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ): """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -259,8 +284,8 @@ def forward( coord_ext: paddle.Tensor, atype_ext: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -376,9 +401,9 @@ def t_cvt(xx): def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -424,7 +449,7 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, **kwargs, ) -> None: """Construct an embedding net of type `se_a`. @@ -438,7 +463,9 @@ def __init__( """ super().__init__() self.rcut = float(rcut) + self.register_buffer("buffer_rcut", paddle.to_tensor(self.rcut)) self.rcut_smth = float(rcut_smth) + self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.rcut_smth)) self.neuron = neuron self.filter_neuron = self.neuron self.axis_neuron = axis_neuron @@ -449,6 +476,9 @@ def __init__( self.resnet_dt = resnet_dt self.env_protection = env_protection self.ntypes = len(sel) + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) self.type_one_side = type_one_side self.seed = seed # order matters, placed after the assignment of self.ntypes @@ -481,6 +511,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, ii), + trainable=trainable, ) self.filter_layers = filter_layers self.stats = None @@ -512,6 +543,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.buffer_rcut + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.buffer_rcut_smth + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return sum(self.sel) @@ -522,7 +561,7 @@ def get_sel(self) -> list[int]: def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_dim_out(self) -> int: """Returns the output dimension.""" @@ -584,8 +623,8 @@ def __getitem__(self, key): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -645,7 +684,7 @@ def reinit_exclude( def enable_compression( self, table_data: dict[str, paddle.Tensor], - table_config: list[Union[int, float]], + table_config: list[int | float], lower: dict[str, int], upper: dict[str, int], ) -> None: @@ -683,9 +722,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Calculate decoded embedding for each atom. diff --git a/deepmd/pd/model/descriptor/se_atten.py b/deepmd/pd/model/descriptor/se_atten.py index 6bec47b12e..2c93c35ef8 100644 --- a/deepmd/pd/model/descriptor/se_atten.py +++ b/deepmd/pd/model/descriptor/se_atten.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -56,7 +54,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -78,9 +76,10 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, - type: Optional[str] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, + type: str | None = None, + trainable: bool = True, ) -> None: r"""Construct an embedding net of type `se_atten`. @@ -146,11 +145,15 @@ def __init__( If not None, the scaling of attention weights is `temperature` itself. seed : int, Optional Random seed for parameter initialization. + trainable : bool, default: True + Whether this block is trainable """ super().__init__() del type self.rcut = float(rcut) + self.register_buffer("buffer_rcut", paddle.to_tensor(self.rcut)) self.rcut_smth = float(rcut_smth) + self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.rcut_smth)) self.neuron = neuron self.filter_neuron = self.neuron self.axis_neuron = axis_neuron @@ -182,6 +185,10 @@ def __init__( sel = [sel] self.ntypes = ntypes + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) + self.sel = sel self.sec = self.sel self.split_sel = self.sel @@ -205,6 +212,7 @@ def __init__( smooth=self.smooth, precision=self.precision, seed=child_seed(self.seed, 0), + trainable=trainable, ) wanted_shape = (self.ntypes, self.nnei, 4) @@ -229,6 +237,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 1), + trainable=trainable, ) self.filter_layers = filter_layers if self.tebd_input_mode in ["strip"]: @@ -242,6 +251,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 2), + trainable=trainable, ) self.filter_layers_strip = filter_layers_strip self.stats = None @@ -272,6 +282,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.buffer_rcut + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.buffer_rcut_smth + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return sum(self.sel) @@ -282,7 +300,7 @@ def get_sel(self) -> list[int]: def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_dim_in(self) -> int: """Returns the input dimension.""" @@ -349,8 +367,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -439,9 +457,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Compute the descriptor. @@ -519,7 +537,10 @@ def forward( index = nlist.reshape([nb, nloc * nnei]).unsqueeze(-1).expand([-1, -1, nt]) # nb x (nloc x nnei) x nt atype_tebd_nlist = paddle.take_along_axis( - atype_tebd_ext, axis=1, indices=index + atype_tebd_ext, + axis=1, + indices=index, + broadcast=False, ) # j # nb x nloc x nnei x nt atype_tebd_nlist = atype_tebd_nlist.reshape([nb, nloc, nnei, nt]) @@ -557,14 +578,16 @@ def forward( nlist_index = nlist.reshape([nb, nloc * nnei]) # nf x (nl x nnei) nei_type = paddle.take_along_axis( - extended_atype, indices=nlist_index, axis=1 + extended_atype, indices=nlist_index, axis=1, broadcast=False ) # (nf x nl x nnei) x ng nei_type_index = nei_type.reshape([-1, 1]).expand([-1, ng]).to(paddle.int64) if self.type_one_side: tt_full = self.filter_layers_strip.networks[0](type_embedding) # (nf x nl x nnei) x ng - gg_t = paddle.take_along_axis(tt_full, indices=nei_type_index, axis=0) + gg_t = paddle.take_along_axis( + tt_full, indices=nei_type_index, axis=0, broadcast=False + ) else: idx_i = paddle.tile( atype.reshape([-1, 1]) * ntypes_with_padding, [1, nnei] @@ -588,7 +611,9 @@ def forward( ).reshape([-1, nt * 2]) tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) # (nf x nl x nnei) x ng - gg_t = paddle.take_along_axis(tt_full, axis=0, indices=idx) + gg_t = paddle.take_along_axis( + tt_full, axis=0, indices=idx, broadcast=False + ) # (nf x nl) x nnei x ng gg_t = gg_t.reshape([nfnl, nnei, ng]) if self.smooth: @@ -649,12 +674,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a neighbor-wise attention net.""" super().__init__() @@ -690,6 +716,7 @@ def __init__( smooth=smooth, precision=precision, seed=child_seed(seed, i), + trainable=trainable, ) ) self.attention_layers = nn.LayerList(attention_layers) @@ -698,8 +725,8 @@ def forward( self, input_G, nei_mask, - input_r: Optional[paddle.Tensor] = None, - sw: Optional[paddle.Tensor] = None, + input_r: paddle.Tensor | None = None, + sw: paddle.Tensor | None = None, ): """Compute the multi-layer gated self-attention. @@ -791,12 +818,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, smooth: bool = True, trainable_ln: bool = True, ln_eps: float = 1e-5, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a neighbor-wise attention layer.""" super().__init__() @@ -824,6 +852,7 @@ def __init__( smooth=smooth, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.attn_layer_norm = LayerNorm( self.embed_dim, @@ -837,8 +866,8 @@ def forward( self, x, nei_mask, - input_r: Optional[paddle.Tensor] = None, - sw: Optional[paddle.Tensor] = None, + input_r: paddle.Tensor | None = None, + sw: paddle.Tensor | None = None, ): residual = x x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) @@ -899,11 +928,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a multi-head neighbor-wise attention net.""" super().__init__() @@ -936,6 +966,7 @@ def __init__( stddev=1.0, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.out_proj = MLPLayer( hidden_dim, @@ -946,14 +977,15 @@ def __init__( stddev=1.0, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) def forward( self, query, nei_mask, - input_r: Optional[paddle.Tensor] = None, - sw: Optional[paddle.Tensor] = None, + input_r: paddle.Tensor | None = None, + sw: paddle.Tensor | None = None, attnw_shift: float = 20.0, ): """Compute the multi-head gated self-attention. diff --git a/deepmd/pd/model/descriptor/se_atten_v2.py b/deepmd/pd/model/descriptor/se_atten_v2.py index 6c90846f9a..8855fb3037 100644 --- a/deepmd/pd/model/descriptor/se_atten_v2.py +++ b/deepmd/pd/model/descriptor/se_atten_v2.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle @@ -40,7 +36,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -61,16 +57,16 @@ def __init__( concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented spin=None, - type: Optional[str] = None, + type: str | None = None, ) -> None: r"""Construct smooth version of embedding net of type `se_atten_v2`. diff --git a/deepmd/pd/model/descriptor/se_t_tebd.py b/deepmd/pd/model/descriptor/se_t_tebd.py index 2898283f0c..b0f409f2bb 100644 --- a/deepmd/pd/model/descriptor/se_t_tebd.py +++ b/deepmd/pd/model/descriptor/se_t_tebd.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, ) import paddle @@ -124,7 +122,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, @@ -136,8 +134,8 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias=False, @@ -160,10 +158,16 @@ def __init__( env_protection=env_protection, smooth=smooth, seed=child_seed(seed, 1), + trainable=trainable, ) self.prec = PRECISION_DICT[precision] self.use_econf_tebd = use_econf_tebd self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(type_map)]), + ) self.smooth = smooth self.type_embedding = TypeEmbedNet( ntypes, @@ -173,6 +177,7 @@ def __init__( use_econf_tebd=use_econf_tebd, type_map=type_map, use_tebd_bias=use_tebd_bias, + trainable=trainable, ) self.tebd_dim = tebd_dim self.tebd_input_mode = tebd_input_mode @@ -206,6 +211,18 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + def get_dim_out(self) -> int: """Returns the output dimension.""" ret = self.se_ttebd.get_dim_out() @@ -277,8 +294,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ): """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -412,8 +429,8 @@ def forward( extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, nlist: paddle.Tensor, - mapping: Optional[paddle.Tensor] = None, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + mapping: paddle.Tensor | None = None, + comm_dict: list[paddle.Tensor] | None = None, ): """Compute the descriptor. @@ -481,9 +498,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -516,7 +533,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, @@ -528,11 +545,14 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.rcut = float(rcut) + self.register_buffer("buffer_rcut", paddle.to_tensor(self.rcut)) self.rcut_smth = float(rcut_smth) + self.register_buffer("buffer_rcut_smth", paddle.to_tensor(self.rcut_smth)) self.neuron = neuron self.filter_neuron = self.neuron self.tebd_dim = tebd_dim @@ -550,6 +570,10 @@ def __init__( sel = [sel] self.ntypes = ntypes + self.register_buffer( + "buffer_ntypes", paddle.to_tensor(self.ntypes, dtype="int64") + ) + self.sel = sel self.sec = self.sel self.split_sel = self.sel @@ -585,6 +609,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 1), + trainable=trainable, ) self.filter_layers = filter_layers if self.tebd_input_mode in ["strip"]: @@ -598,6 +623,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 2), + trainable=trainable, ) self.filter_layers_strip = filter_layers_strip self.stats = None @@ -610,6 +636,14 @@ def get_rcut_smth(self) -> float: """Returns the radius where the neighbor information starts to smoothly decay to 0.""" return self.rcut_smth + def get_buffer_rcut(self) -> paddle.Tensor: + """Returns the cut-off radius as a buffer-style Tensor.""" + return self.buffer_rcut + + def get_buffer_rcut_smth(self) -> paddle.Tensor: + """Returns the radius where the neighbor information starts to smoothly decay to 0 as a buffer-style Tensor.""" + return self.buffer_rcut_smth + def get_nsel(self) -> int: """Returns the number of selected atoms in the cut-off radius.""" return sum(self.sel) @@ -620,7 +654,7 @@ def get_sel(self) -> list[int]: def get_ntypes(self) -> int: """Returns the number of element types.""" - return self.ntypes + return self.ntypes if paddle.in_dynamic_mode() else self.buffer_ntypes def get_dim_in(self) -> int: """Returns the input dimension.""" @@ -683,8 +717,8 @@ def dim_emb(self): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -746,9 +780,9 @@ def forward( nlist: paddle.Tensor, extended_coord: paddle.Tensor, extended_atype: paddle.Tensor, - extended_atype_embd: Optional[paddle.Tensor] = None, - mapping: Optional[paddle.Tensor] = None, - type_embedding: Optional[paddle.Tensor] = None, + extended_atype_embd: paddle.Tensor | None = None, + mapping: paddle.Tensor | None = None, + type_embedding: paddle.Tensor | None = None, ): """Compute the descriptor. @@ -840,7 +874,7 @@ def forward( # nb x (nloc x nnei) x nt # atype_tebd_nlist = paddle.take_along_axis(atype_tebd_ext, axis=1, index=index) atype_tebd_nlist = paddle.take_along_axis( - atype_tebd_ext, axis=1, indices=index + atype_tebd_ext, axis=1, indices=index, broadcast=False ) # nb x nloc x nnei x nt atype_tebd_nlist = atype_tebd_nlist.reshape([nb, nloc, nnei, nt]) @@ -864,7 +898,7 @@ def forward( nlist_index = nlist.reshape([nb, nloc * nnei]) # nf x (nl x nnei) nei_type = paddle.take_along_axis( - extended_atype, indices=nlist_index, axis=1 + extended_atype, indices=nlist_index, axis=1, broadcast=False ) # nfnl x nnei nei_type = nei_type.reshape([nfnl, nnei]) @@ -897,7 +931,7 @@ def forward( ).reshape([-1, nt * 2]) tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) # (nfnl x nt_i x nt_j) x ng - gg_t = paddle.take_along_axis(tt_full, indices=idx, axis=0) + gg_t = paddle.take_along_axis(tt_full, indices=idx, axis=0, broadcast=False) # (nfnl x nt_i x nt_j) x ng gg_t = gg_t.reshape([nfnl, nnei, nnei, ng]) if self.smooth: diff --git a/deepmd/pd/model/model/dp_model.py b/deepmd/pd/model/model/dp_model.py index e014be5b68..fe107263c4 100644 --- a/deepmd/pd/model/model/dp_model.py +++ b/deepmd/pd/model/model/dp_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -20,9 +17,9 @@ class DPModelCommon: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/model/ener_model.py b/deepmd/pd/model/model/ener_model.py index a5b1b9d4b3..8111ae434d 100644 --- a/deepmd/pd/model/model/ener_model.py +++ b/deepmd/pd/model/model/ener_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -34,6 +31,18 @@ def __init__( DPModelCommon.__init__(self) DPEnergyModel_.__init__(self, *args, **kwargs) + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return super().get_buffer_type_map() + def translated_output_def(self): out_def_data = self.model_output_def().get_data() output_def = { @@ -56,9 +65,9 @@ def forward( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: model_ret = self.forward_common( @@ -81,6 +90,10 @@ def forward( model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze( -3 ) + else: + model_predict["atom_virial"] = paddle.zeros( + [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64 + ) else: model_predict["force"] = model_ret["dforce"] if "mask" in model_ret: @@ -95,11 +108,11 @@ def forward_lower( extended_coord, extended_atype, nlist, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + comm_dict: list[paddle.Tensor] | None = None, ): model_ret = self.forward_common_lower( extended_coord, @@ -124,6 +137,10 @@ def forward_lower( model_predict["extended_virial"] = model_ret[ "energy_derv_c" ].squeeze(-3) + else: + model_predict["extended_virial"] = paddle.zeros( + [model_predict["energy"].shape[0], 1, 9], dtype=paddle.float64 + ) else: assert model_ret["dforce"] is not None model_predict["dforce"] = model_ret["dforce"] diff --git a/deepmd/pd/model/model/frozen.py b/deepmd/pd/model/model/frozen.py index e8128c6bd1..f585cb1e24 100644 --- a/deepmd/pd/model/model/frozen.py +++ b/deepmd/pd/model/model/frozen.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json -from typing import ( - Optional, -) import paddle @@ -101,9 +98,9 @@ def forward( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: return self.model.forward( @@ -123,7 +120,7 @@ def get_model_def_script(self) -> str: # be a problem return self.model.get_model_def_script() - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.model.get_min_nbor_dist() @@ -154,9 +151,9 @@ def get_nsel(self) -> int: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pd/model/model/make_model.py b/deepmd/pd/model/model/make_model.py index acb237b5ac..03aaf621f8 100644 --- a/deepmd/pd/model/model/make_model.py +++ b/deepmd/pd/model/model/make_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -67,7 +64,7 @@ def __init__( self, *args, # underscore to prevent conflict with normal inputs - atomic_model_: Optional[T_AtomicModel] = None, + atomic_model_: T_AtomicModel | None = None, **kwargs, ) -> None: super().__init__(*args, **kwargs) @@ -129,9 +126,9 @@ def forward_common( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: """Return model prediction. @@ -228,17 +225,19 @@ def change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, ) + if bias_adjust_mode == "set-by-statistic": + self.atomic_model.compute_fitting_input_stat(merged) def forward_common_lower( self, extended_coord, extended_atype, nlist, - mapping: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + mapping: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, paddle.Tensor]] = None, + comm_dict: list[paddle.Tensor] | None = None, extra_nlist_sort: bool = False, ): """Return model prediction. Lower interface that takes @@ -304,14 +303,14 @@ def forward_common_lower( def input_type_cast( self, coord: paddle.Tensor, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ) -> tuple[ paddle.Tensor, - Optional[paddle.Tensor], - Optional[paddle.Tensor], - Optional[paddle.Tensor], + paddle.Tensor | None, + paddle.Tensor | None, + paddle.Tensor | None, str, ]: """Cast the input data to global float type.""" @@ -326,7 +325,7 @@ def input_type_cast( # " does not match" # f" that of the coordinate {input_prec}" # ) - _lst: list[Optional[paddle.Tensor]] = [ + _lst: list[paddle.Tensor | None] = [ vv.astype(coord.dtype) if vv is not None else None for vv in [box, fparam, aparam] ] @@ -364,7 +363,7 @@ def output_type_cast( continue if check_operation_applied(odef[kk], OutputVariableOperation.REDU): model_ret[kk] = ( - model_ret[kk].to(self.global_pd_ener_float_precision) + model_ret[kk].astype(self.global_pd_ener_float_precision) if model_ret[kk] is not None else None ) @@ -445,7 +444,7 @@ def _format_nlist( * paddle.ones( [n_nf, n_nloc, nnei - n_nnei], dtype=nlist.dtype, - ).to(nlist.place), + ), ], axis=-1, ) @@ -458,17 +457,21 @@ def _format_nlist( coord0 = extended_coord[:, :n_nloc, :] # nf x (nloc x nnei) x 3 index = nlist.reshape([n_nf, n_nloc * n_nnei, 1]).expand([-1, -1, 3]) - coord1 = paddle.take_along_axis(extended_coord, axis=1, indices=index) + coord1 = paddle.take_along_axis( + extended_coord, axis=1, indices=index, broadcast=False + ) # nf x nloc x nnei x 3 coord1 = coord1.reshape([n_nf, n_nloc, n_nnei, 3]) # nf x nloc x nnei rr = paddle.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) - rr = paddle.where(m_real_nei, rr, float("inf")) + rr = paddle.where(m_real_nei, rr, paddle.full_like(rr, float("inf"))) rr, nlist_mapping = ( paddle.sort(rr, axis=-1), paddle.argsort(rr, axis=-1), ) - nlist = paddle.take_along_axis(nlist, axis=2, indices=nlist_mapping) + nlist = paddle.take_along_axis( + nlist, axis=2, indices=nlist_mapping, broadcast=False + ) nlist = paddle.where(rr > rcut, paddle.full_like(nlist, -1), nlist) nlist = nlist[..., :nnei] else: # not extra_nlist_sort and n_nnei <= nnei: @@ -478,7 +481,7 @@ def _format_nlist( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -487,7 +490,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. @@ -525,6 +528,14 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.atomic_model.get_dim_aparam() + def get_buffer_dim_fparam(self) -> paddle.Tensor: + """Get the number (dimension) of frame parameters of this atomic model as a buffer-style Tensor.""" + return self.atomic_model.get_buffer_dim_fparam() + + def get_buffer_dim_aparam(self) -> paddle.Tensor: + """Get the number (dimension) of atomic parameters of this atomic model as a buffer-style Tensor.""" + return self.atomic_model.get_buffer_dim_aparam() + def get_sel_type(self) -> list[int]: """Get the selected atom types of this model. @@ -549,6 +560,22 @@ def get_type_map(self) -> list[str]: """Get the type map.""" return self.atomic_model.get_type_map() + def get_buffer_rcut(self) -> paddle.Tensor: + """Get the cut-off radius as a buffer-style Tensor.""" + return self.atomic_model.get_buffer_rcut() + + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.atomic_model.get_buffer_type_map() + def get_nsel(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" return self.atomic_model.get_nsel() @@ -564,7 +591,7 @@ def atomic_output_def(self) -> FittingOutputDef: def compute_or_load_stat( self, sampled_func, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ): """Compute or load the statistics.""" return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path) @@ -597,9 +624,9 @@ def forward( self, coord, atype, - box: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, paddle.Tensor]: # directly call the forward_common method when no specific transform rule diff --git a/deepmd/pd/model/model/model.py b/deepmd/pd/model/model/model.py index 06a2c6910f..5027590a9e 100644 --- a/deepmd/pd/model/model/model.py +++ b/deepmd/pd/model/model/model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import paddle @@ -23,7 +20,7 @@ def __init__(self, *args, **kwargs): def compute_or_load_stat( self, sampled_func, - stat_file_path: Optional[DPPath] = None, + stat_file_path: DPPath | None = None, ): """ Compute or load the statistics parameters of the model, @@ -46,10 +43,14 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" return self.min_nbor_dist def get_ntypes(self): """Returns the number of element types.""" return len(self.get_type_map()) + + def get_buffer_ntypes(self) -> paddle.Tensor: + """Returns the number of element types as a buffer-style Tensor.""" + return self.get_buffer_ntypes() diff --git a/deepmd/pd/model/model/transform_output.py b/deepmd/pd/model/model/transform_output.py index 469bfd3168..47004265c7 100644 --- a/deepmd/pd/model/model/transform_output.py +++ b/deepmd/pd/model/model/transform_output.py @@ -223,9 +223,7 @@ def communicate_extended_output( mapping = mapping.reshape(mldims + [1] * len(derv_r_ext_dims)).expand( [-1] * len(mldims) + derv_r_ext_dims ) - force = paddle.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype).to( - device=vv.place - ) + force = paddle.zeros(vldims + derv_r_ext_dims, dtype=vv.dtype) # nf x nloc x nvar x 3 new_ret[kk_derv_r] = decomp.scatter_reduce( force, @@ -242,9 +240,7 @@ def communicate_extended_output( mapping, [1] * (len(mldims) + len(vdef.shape)) + [3], ) - virial = paddle.zeros(vldims + derv_c_ext_dims, dtype=vv.dtype).to( - device=vv.place - ) + virial = paddle.zeros(vldims + derv_c_ext_dims, dtype=vv.dtype) # nf x nloc x nvar x 9 new_ret[kk_derv_c] = decomp.scatter_reduce( virial, @@ -254,9 +250,9 @@ def communicate_extended_output( reduce="sum", ) new_ret[kk_derv_c + "_redu"] = paddle.sum( - new_ret[kk_derv_c].to(redu_prec), axis=1 + new_ret[kk_derv_c].astype(redu_prec), axis=1 ) - if not do_atomic_virial: + if not do_atomic_virial and paddle.in_dynamic_mode(): # pop atomic virial, because it is not correctly calculated. new_ret.pop(kk_derv_c) return new_ret diff --git a/deepmd/pd/model/network/layernorm.py b/deepmd/pd/model/network/layernorm.py index 7cfffe1b93..d3d00beacd 100644 --- a/deepmd/pd/model/network/layernorm.py +++ b/deepmd/pd/model/network/layernorm.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np import paddle @@ -45,7 +41,7 @@ def __init__( stddev: float = 1.0, precision: str = DEFAULT_PRECISION, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ): super().__init__() self.eps = eps diff --git a/deepmd/pd/model/network/mlp.py b/deepmd/pd/model/network/mlp.py index 41286fbbae..ee408b8719 100644 --- a/deepmd/pd/model/network/mlp.py +++ b/deepmd/pd/model/network/mlp.py @@ -85,8 +85,10 @@ def __init__( precision: str = DEFAULT_PRECISION, init: str = "default", seed: int | list[int] | None = None, + trainable: bool = True, ): super().__init__() + self.trainable = trainable # only use_timestep when skip connection is established. self.use_timestep = use_timestep and ( num_out == num_in or num_out == num_in * 2 @@ -251,6 +253,7 @@ def serialize(self) -> dict: activation_function=self.activate_name, resnet=self.resnet, precision=self.precision, + trainable=self.trainable, ) nl.w, nl.b, nl.idt = ( to_numpy_array(self.matrix), @@ -277,6 +280,7 @@ def deserialize(cls, data: dict) -> MLPLayer: activation_function=nl["activation_function"], resnet=nl["resnet"], precision=nl["precision"], + trainable=nl["trainable"], ) prec = PRECISION_DICT[obj.precision] diff --git a/deepmd/pd/model/network/network.py b/deepmd/pd/model/network/network.py index 9cdb7b3adc..81e2dad710 100644 --- a/deepmd/pd/model/network/network.py +++ b/deepmd/pd/model/network/network.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle import paddle.nn as nn @@ -41,10 +37,11 @@ def __init__( bavg=0.0, stddev=1.0, precision="default", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd=False, use_tebd_bias: bool = False, type_map=None, + trainable: bool = True, ) -> None: """Construct a type embedding net.""" super().__init__() @@ -65,6 +62,7 @@ def __init__( type_map=type_map, precision=precision, seed=seed, + trainable=trainable, ) # init.normal_(self.embedding.weight[:-1], mean=bavg, std=stddev) @@ -160,11 +158,11 @@ def __init__( activation_function: str = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: """Construct a type embedding net.""" super().__init__() @@ -195,6 +193,7 @@ def __init__( self.precision, self.seed, bias=self.use_tebd_bias, + trainable=trainable, ) for param in self.parameters(): param.stop_gradient = not trainable @@ -208,9 +207,7 @@ def forward(self, device: str): Type embedding network. """ if not self.use_econf_tebd: - embed = self.embedding_net( - paddle.eye(self.ntypes, dtype=self.prec).to(device=device) - ) + embed = self.embedding_net(paddle.eye(self.ntypes, dtype=self.prec)) else: assert self.econf_tebd is not None embed = self.embedding_net(self.econf_tebd.to(device)) @@ -218,9 +215,7 @@ def forward(self, device: str): embed = paddle.concat( [ embed, - paddle.zeros([1, embed.shape[1]], dtype=self.prec).to( - device=device - ), + paddle.zeros([1, embed.shape[1]], dtype=self.prec), ] ) return embed diff --git a/deepmd/pd/model/network/utils.py b/deepmd/pd/model/network/utils.py new file mode 100644 index 0000000000..cfd6753638 --- /dev/null +++ b/deepmd/pd/model/network/utils.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import paddle + + +def aggregate( + data: paddle.Tensor, + owners: paddle.Tensor, + average: bool = True, + num_owner: int | None = None, +) -> paddle.Tensor: + """ + Aggregate rows in data by specifying the owners. + + Parameters + ---------- + data : data tensor to aggregate [n_row, feature_dim] + owners : specify the owner of each row [n_row, 1] + average : if True, average the rows, if False, sum the rows. + Default = True + num_owner : the number of owners, this is needed if the + max idx of owner is not presented in owners tensor + Default = None + + Returns + ------- + output: [num_owner, feature_dim] + """ + if num_owner is None or average: + # requires bincount + bin_count = paddle.bincount(owners) + bin_count = bin_count.where(bin_count != 0, paddle.ones_like(bin_count)) + + if (num_owner is not None) and (bin_count.shape[0] != num_owner): + difference = num_owner - bin_count.shape[0] + bin_count = paddle.concat( + [bin_count, paddle.ones([difference], dtype=bin_count.dtype)] + ) + else: + bin_count = None + + # make sure this operation is done on the same device of data and owners + output = paddle.zeros([num_owner, data.shape[1]]) + output = output.index_add_(owners, 0, data.astype(output.dtype)) + if average: + assert bin_count is not None + output = (output.T / bin_count).T + return output + + +def get_graph_index( + nlist: paddle.Tensor, + nlist_mask: paddle.Tensor, + a_nlist_mask: paddle.Tensor, + nall: int, + use_loc_mapping: bool = True, +): + """ + Get the index mapping for edge graph and angle graph, ready in `aggregate` or `index_select`. + + Parameters + ---------- + nlist : nf x nloc x nnei + Neighbor list. (padded neis are set to 0) + nlist_mask : nf x nloc x nnei + Masks of the neighbor list. real nei 1 otherwise 0 + a_nlist_mask : nf x nloc x a_nnei + Masks of the neighbor list for angle. real nei 1 otherwise 0 + nall + The number of extended atoms. + + Returns + ------- + edge_index : 2 x n_edge + n2e_index : n_edge + Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i). + n_ext2e_index : n_edge + Broadcast indices from extended node(j) to edge(ij). + angle_index : 3 x n_angle + n2a_index : n_angle + Broadcast indices from extended node(j) to angle(ijk). + eij2a_index : n_angle + Broadcast indices from extended edge(ij) to angle(ijk), or reduction indices from angle(ijk) to edge(ij). + eik2a_index : n_angle + Broadcast indices from extended edge(ik) to angle(ijk). + """ + nf, nloc, nnei = nlist.shape + _, _, a_nnei = a_nlist_mask.shape + # nf x nloc x nnei x nnei + # nlist_mask_3d = nlist_mask[:, :, :, None] & nlist_mask[:, :, None, :] + a_nlist_mask_3d = a_nlist_mask[:, :, :, None] & a_nlist_mask[:, :, None, :] + n_edge = nlist_mask.sum().item() + # n_angle = a_nlist_mask_3d.sum().item() + + # following: get n2e_index, n_ext2e_index, n2a_index, eij2a_index, eik2a_index + + # 1. atom graph + # node(i) to edge(ij) index_select; edge(ij) to node aggregate + nlist_loc_index = paddle.arange(0, nf * nloc, dtype=nlist.dtype).to(nlist.place) + # nf x nloc x nnei + n2e_index = nlist_loc_index.reshape([nf, nloc, 1]).expand([-1, -1, nnei]) + # n_edge + n2e_index = n2e_index[nlist_mask] # graph node index, atom_graph[:, 0] + + # node_ext(j) to edge(ij) index_select + frame_shift = paddle.arange(0, nf, dtype=nlist.dtype) * ( + nall if not use_loc_mapping else nloc + ) + shifted_nlist = nlist + frame_shift[:, None, None] + # n_edge + n_ext2e_index = shifted_nlist[nlist_mask] # graph neighbor index, atom_graph[:, 1] + + # 2. edge graph + # node(i) to angle(ijk) index_select + n2a_index = nlist_loc_index.reshape([nf, nloc, 1, 1]).expand( + [-1, -1, a_nnei, a_nnei] + ) + # n_angle + n2a_index = n2a_index[a_nlist_mask_3d] + + # edge(ij) to angle(ijk) index_select; angle(ijk) to edge(ij) aggregate + edge_id = paddle.arange(0, n_edge, dtype=nlist.dtype) + # nf x nloc x nnei + edge_index = paddle.zeros([nf, nloc, nnei], dtype=nlist.dtype) + edge_index[nlist_mask] = edge_id + # only cut a_nnei neighbors, to avoid nnei x nnei + edge_index = edge_index[:, :, :a_nnei] + edge_index_ij = edge_index.unsqueeze(-1).expand([-1, -1, -1, a_nnei]) + # n_angle + eij2a_index = edge_index_ij[a_nlist_mask_3d] + + # edge(ik) to angle(ijk) index_select + edge_index_ik = edge_index.unsqueeze(-2).expand([-1, -1, a_nnei, -1]) + # n_angle + eik2a_index = edge_index_ik[a_nlist_mask_3d] + + edge_index_result = paddle.stack([n2e_index, n_ext2e_index], axis=0) + angle_index_result = paddle.stack([n2a_index, eij2a_index, eik2a_index], axis=0) + + return edge_index_result, angle_index_result diff --git a/deepmd/pd/model/task/ener.py b/deepmd/pd/model/task/ener.py index 789ef75066..3d2025bcc2 100644 --- a/deepmd/pd/model/task/ener.py +++ b/deepmd/pd/model/task/ener.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging -from typing import ( - Optional, - Union, -) import paddle @@ -38,7 +34,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[paddle.Tensor] = None, + bias_atom_e: paddle.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -46,8 +42,8 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, **kwargs, ): super().__init__( @@ -72,7 +68,7 @@ def __init__( @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) diff --git a/deepmd/pd/model/task/fitting.py b/deepmd/pd/model/task/fitting.py index a478c12f97..d92ad6e947 100644 --- a/deepmd/pd/model/task/fitting.py +++ b/deepmd/pd/model/task/fitting.py @@ -3,9 +3,8 @@ from abc import ( abstractmethod, ) -from typing import ( - Optional, - Union, +from collections.abc import ( + Callable, ) import numpy as np @@ -39,6 +38,9 @@ get_index_between_two_maps, map_atom_exclude_types, ) +from deepmd.utils.path import ( + DPPath, +) dtype = env.GLOBAL_PD_FLOAT_PRECISION device = env.DEVICE @@ -71,6 +73,96 @@ def share_params(self, base_class, shared_level, resume=False) -> None: else: raise NotImplementedError + def compute_input_stats( + self, + merged: Callable[[], list[dict]] | list[dict], + protection: float = 1e-2, + stat_file_path: DPPath | None = None, + ) -> None: + """ + Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + + Parameters + ---------- + merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + protection : float + Divided-by-zero protection + stat_file_path : Optional[DPPath] + The path to the stat file. + """ + if self.numb_fparam == 0 and self.numb_aparam == 0: + # skip data statistics + return + if callable(merged): + sampled = merged() + else: + sampled = merged + # stat fparam + if self.numb_fparam > 0: + cat_data = paddle.concat([frame["fparam"] for frame in sampled], axis=0) + cat_data = paddle.reshape(cat_data, [-1, self.numb_fparam]) + fparam_avg = paddle.mean(cat_data, axis=0) + fparam_std = paddle.std(cat_data, axis=0, unbiased=False) + fparam_std = paddle.where( + fparam_std < protection, + paddle.to_tensor(protection, dtype=fparam_std.dtype), + fparam_std, + ) + fparam_inv_std = 1.0 / fparam_std + paddle.assign( + paddle.to_tensor( + fparam_avg, place=env.DEVICE, dtype=self.fparam_avg.dtype + ), + self.fparam_avg, + ) + paddle.assign( + paddle.to_tensor( + fparam_inv_std, place=env.DEVICE, dtype=self.fparam_inv_std.dtype + ), + self.fparam_inv_std, + ) + # stat aparam + if self.numb_aparam > 0: + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in [frame["aparam"] for frame in sampled]: + ss = paddle.reshape(ss_, [-1, self.numb_aparam]) + sys_sumv.append(paddle.sum(ss, axis=0)) + sys_sumv2.append(paddle.sum(ss * ss, axis=0)) + sys_sumn.append(ss.shape[0]) + sumv = paddle.sum(paddle.stack(sys_sumv), axis=0) + sumv2 = paddle.sum(paddle.stack(sys_sumv2), axis=0) + sumn = sum(sys_sumn) + aparam_avg = sumv / sumn + aparam_std = paddle.sqrt(sumv2 / sumn - (sumv / sumn) ** 2) + aparam_std = paddle.where( + aparam_std < protection, + paddle.to_tensor( + protection, dtype=aparam_std.dtype, place=aparam_std.device + ), + aparam_std, + ) + aparam_inv_std = 1.0 / aparam_std + paddle.assign( + paddle.to_tensor( + aparam_avg, place=env.DEVICE, dtype=self.aparam_avg.dtype + ), + self.aparam_avg, + ) + paddle.assign( + paddle.to_tensor( + aparam_inv_std, place=env.DEVICE, dtype=self.aparam_inv_std.dtype + ), + self.aparam_inv_std, + ) + class GeneralFitting(Fitting): """Construct a general fitting net. @@ -95,6 +187,10 @@ class GeneralFitting(Fitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. + This parameter is not supported in PaddlePaddle. dim_case_embd : int Dimension of case specific embedding. activation_function : str @@ -130,7 +226,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[paddle.Tensor] = None, + bias_atom_e: paddle.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -138,13 +234,14 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - trainable: Union[bool, list[bool]] = True, - remove_vaccum_contribution: Optional[list[bool]] = None, - type_map: Optional[list[str]] = None, + trainable: bool | list[bool] = True, + remove_vaccum_contribution: list[bool] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, + default_fparam: list[float] | None = None, **kwargs, ) -> None: super().__init__() @@ -155,14 +252,26 @@ def __init__( self.mixed_types = mixed_types self.resnet_dt = resnet_dt self.numb_fparam = numb_fparam + self.register_buffer( + "buffer_numb_fparam", paddle.to_tensor([numb_fparam], dtype=paddle.int64) + ) self.numb_aparam = numb_aparam + self.register_buffer( + "buffer_numb_aparam", paddle.to_tensor([numb_aparam], dtype=paddle.int64) + ) self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] self.rcond = rcond self.seed = seed self.type_map = type_map + if type_map is not None: + self.register_buffer( + "buffer_type_map", + paddle.to_tensor([ord(c) for c in " ".join(self.type_map)]), + ) self.use_aparam_as_mask = use_aparam_as_mask # order matters, should be place after the assignment of ntypes self.reinit_exclude(exclude_types) @@ -246,6 +355,8 @@ def __init__( for param in self.parameters(): param.stop_gradient = not self.trainable + self.eval_return_middle_output = False + def reinit_exclude( self, exclude_types: list[int] = [], @@ -282,7 +393,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 3, + "@version": 4, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -291,6 +402,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function, "precision": self.precision, "mixed_types": self.mixed_types, @@ -338,6 +450,14 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam + def get_buffer_dim_fparam(self) -> paddle.Tensor: + """Get the number (dimension) of frame parameters of this atomic model as a buffer-style Tensor.""" + return self.buffer_numb_fparam + + def get_buffer_dim_aparam(self) -> paddle.Tensor: + """Get the number (dimension) of atomic parameters of this atomic model as a buffer-style Tensor.""" + return self.buffer_numb_aparam + # make jit happy exclude_types: list[int] @@ -359,6 +479,18 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map + def get_buffer_type_map(self) -> paddle.Tensor: + """ + Return the type map as a buffer-style Tensor for JIT saving. + + The original type map (e.g., ['Ni', 'O']) is first joined into a single space-separated string + (e.g., "Ni O"). Each character in this string is then converted to its ASCII code using `ord()`, + and the resulting integer sequence is stored as a 1D paddle.Tensor of dtype int. + + This format allows the type map to be serialized as a raw byte buffer during JIT model saving. + """ + return self.buffer_type_map + def set_case_embd(self, case_idx: int): """ Set the case embedding of this fitting net by the given case_idx, @@ -368,6 +500,9 @@ def set_case_embd(self, case_idx: int): case_idx ] + def set_return_middle_output(self, return_middle_output: bool = True) -> None: + self.eval_return_middle_output = return_middle_output + def __setitem__(self, key, value) -> None: if key in ["bias_atom_e"]: value = value.reshape([self.ntypes, self._net_out_dim()]) @@ -420,16 +555,16 @@ def _forward_common( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): # cast the input to internal precsion - xx = descriptor.to(self.prec) - fparam = fparam.to(self.prec) if fparam is not None else None - aparam = aparam.to(self.prec) if aparam is not None else None + xx = descriptor.astype(self.prec) + fparam = fparam.astype(self.prec) if fparam is not None else None + aparam = aparam.astype(self.prec) if aparam is not None else None if self.remove_vaccum_contribution is not None: # TODO: compute the input for vaccm when remove_vaccum_contribution is set @@ -514,15 +649,37 @@ def _forward_common( outs = paddle.zeros( (nf, nloc, net_dim_out), dtype=env.GLOBAL_PD_FLOAT_PRECISION, - ).to(device=descriptor.place) + ) + results = {} + if self.mixed_types: - atom_property = self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] + atom_property = self.filter_layers.networks[0](xx) + if self.eval_return_middle_output: + results["middle_output"] = self.filter_layers.networks[ + 0 + ].call_until_last(xx) if xx_zeros is not None: atom_property -= self.filter_layers.networks[0](xx_zeros) outs = ( - outs + atom_property + self.bias_atom_e[atype].to(self.prec) + outs + atom_property + self.bias_atom_e[atype].astype(self.prec) ) # Shape is [nframes, natoms[0], net_dim_out] else: + if self.eval_return_middle_output: + outs_middle = paddle.zeros( + (nf, nloc, self.neuron[-1]), + dtype=self.prec, + ).to(device=descriptor.place) # jit assertion + for type_i, ll in enumerate(self.filter_layers.networks): + mask = (atype == type_i).unsqueeze(-1) + mask = paddle.tile(mask, (1, 1, net_dim_out)) + middle_output_type = ll.call_until_last(xx) + middle_output_type = paddle.where( + paddle.tile(mask, (1, 1, self.neuron[-1])), + middle_output_type, + paddle.zeros_like(middle_output_type), + ) + outs_middle = outs_middle + middle_output_type + results["middle_output"] = outs_middle for type_i, ll in enumerate(self.filter_layers.networks): mask = (atype == type_i).unsqueeze(-1) mask.stop_gradient = True @@ -537,12 +694,15 @@ def _forward_common( ): atom_property -= ll(xx_zeros) atom_property = atom_property + self.bias_atom_e[type_i] - atom_property = paddle.where(mask, atom_property, 0.0) + atom_property = paddle.where( + mask, atom_property, paddle.full_like(atom_property, 0.0) + ) outs = ( outs + atom_property ) # Shape is [nframes, natoms[0], net_dim_out] # nf x nloc - mask = self.emask(atype).to("bool") + mask = self.emask(atype).astype("bool") # nf x nloc x nod - outs = paddle.where(mask[:, :, None], outs, 0.0) - return {self.var_name: outs.astype(env.GLOBAL_PD_FLOAT_PRECISION)} + outs = paddle.where(mask[:, :, None], outs, paddle.zeros_like(outs)) + results.update({self.var_name: outs}) + return results diff --git a/deepmd/pd/model/task/invar_fitting.py b/deepmd/pd/model/task/invar_fitting.py index b92c862dc8..6dd5b29cb9 100644 --- a/deepmd/pd/model/task/invar_fitting.py +++ b/deepmd/pd/model/task/invar_fitting.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging -from typing import ( - Optional, - Union, -) import paddle @@ -91,7 +87,7 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[paddle.Tensor] = None, + bias_atom_e: paddle.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -99,11 +95,11 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - atom_ener: Optional[list[Optional[paddle.Tensor]]] = None, - type_map: Optional[list[str]] = None, + atom_ener: list[paddle.Tensor | None] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, **kwargs, ): @@ -147,7 +143,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) return super().deserialize(data) def output_def(self) -> FittingOutputDef: @@ -167,11 +163,11 @@ def forward( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): """Based on embedding net output, alculate total energy. diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index d72c270667..4e5fea081f 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import contextlib import functools import logging import time @@ -18,6 +19,7 @@ from paddle.distributed import ( fleet, ) +from paddle.distributed.fleet.utils import hybrid_parallel_util as hpu from paddle.framework import ( core, ) @@ -52,6 +54,7 @@ ) from deepmd.pd.utils.env import ( CINN, + CINN_ALLOW_DYNAMIC_SHAPE, DEFAULT_PRECISION, DEVICE, JIT, @@ -130,6 +133,9 @@ def __init__( # Iteration config self.num_steps = training_params["numb_steps"] + self.acc_freq: int = training_params.get( + "acc_freq", 1 + ) # gradient accumulation steps self.disp_file = training_params.get("disp_file", "lcurve.out") self.disp_freq = training_params.get("disp_freq", 1000) self.save_ckpt = training_params.get("save_ckpt", "model.ckpt") @@ -607,40 +613,65 @@ def warm_up_linear(step, warmup_steps): ) backend = "CINN" if CINN else None - self.wrapper.forward = jit.to_static( - backend=backend, - input_spec=[ - static.InputSpec([1, -1, 3], "float64", name="coord"), # coord - static.InputSpec([1, -1], "int32", name="atype"), # atype - None, # spin - static.InputSpec([1, 9], "float64", name="box"), # box - static.InputSpec([], "float64", name="cur_lr"), # cur_lr - { - "find_box": np.float32(1.0), - "find_coord": np.float32(1.0), - "find_numb_copy": np.float32(0.0), - "numb_copy": static.InputSpec( - [1, 1], "int64", name="numb_copy" - ), - "find_energy": np.float32(1.0), - "energy": static.InputSpec([1, 1], "float64", name="energy"), - "find_force": np.float32(1.0), - "force": static.InputSpec([1, -1, 3], "float64", name="force"), - "natoms": static.InputSpec([1, -1], "int32", name="natoms"), - }, # label, - # None, # task_key - # False, # inference_only - # False, # do_atomic_virial - # None, # fparam - # None, # aparam - ], - full_graph=True, - )(self.wrapper.forward) + if CINN_ALLOW_DYNAMIC_SHAPE: + # Build spec only for keys present in sample data + # NOTE: This is a trick to decide the right input_spec for wrapper.forward + _, label_dict, _ = self.get_data(is_train=True) + # Define specification templates + spec_templates = { + "find_box": np.float32(1.0), + "find_coord": np.float32(1.0), + "find_numb_copy": np.float32(0.0), + "numb_copy": static.InputSpec([1, 1], "int64", name="numb_copy"), + "find_energy": np.float32(1.0), + "energy": static.InputSpec([1, 1], "float64", name="energy"), + "find_force": np.float32(1.0), + "force": static.InputSpec([1, -1, 3], "float64", name="force"), + "find_virial": np.float32(0.0), + "virial": static.InputSpec([1, 9], "float64", name="virial"), + "natoms": static.InputSpec([1, -1], "int32", name="natoms"), + } + label_dict_spec = { + k: spec_templates[k] + for k in label_dict.keys() + if k in spec_templates + } + self.wrapper.forward = jit.to_static( + backend=backend, + input_spec=[ + static.InputSpec([1, -1, 3], "float64", name="coord"), # coord + static.InputSpec([1, -1], "int32", name="atype"), # atype + None, # spin + static.InputSpec([1, 9], "float64", name="box"), # box + static.InputSpec([], "float64", name="cur_lr"), # cur_lr + label_dict_spec, # label, + # None, # task_key + # False, # inference_only + # False, # do_atomic_virial + # None, # fparam + # None, # aparam + ], + full_graph=True, + )(self.wrapper.forward) + else: + self.wrapper.forward = jit.to_static(full_graph=True, backend=backend)( + self.wrapper.forward + ) log.info( - "Enable CINN during training, there may be some additional " - "compilation time in the first traning step." + "[CINN] Enable CINN during training, there may be some additional " + "compilation time in the first training step." ) + if not CINN_ALLOW_DYNAMIC_SHAPE: + log.info( + "[CINN] Dynamic shape is disabled (CINN_ALLOW_DYNAMIC_SHAPE=0). " + "Make sure the input batch shapes are fixed during training. " + "This is recommended for optimal performance, e.g., as in examples/water." + ) + log.info( + "[CINN] If batch data from your dataset(s) has varying input shapes, consider setting " + "CINN_ALLOW_DYNAMIC_SHAPE=1 to enable dynamic shape support." + ) if dist.is_available() and dist.is_initialized(): # DDP will guarantee the model parameters are identical across all processes @@ -716,7 +747,6 @@ def step(_step_id, task_key="Default") -> None: _lr = self.lr_exp cur_lr = _lr.value(_step_id) pref_lr = cur_lr - self.optimizer.clear_grad(set_to_zero=False) with nvprof_context(enable_profiling, "Fetching data"): input_dict, label_dict, log_dict = self.get_data( @@ -732,28 +762,47 @@ def step(_step_id, task_key="Default") -> None: pref_lr = _lr.start_lr else: pref_lr = cur_lr - with nvprof_context(enable_profiling, "Forward pass"): - model_pred, loss, more_loss = self.wrapper( - **input_dict, - cur_lr=paddle.full([], pref_lr, DEFAULT_PRECISION), - label=label_dict, - task_key=task_key, - ) - with nvprof_context(enable_profiling, "Backward pass"): - loss.backward() + # disable synchronization in forward-backward manually + # as derivatives exist in model forward + no_sync_context = ( + self.wrapper.no_sync + if self.world_size > 1 + else contextlib.nullcontext + ) + with no_sync_context(): + with nvprof_context(enable_profiling, "Forward pass"): + model_pred, loss, more_loss = self.wrapper( + **input_dict, + cur_lr=paddle.full([], pref_lr, DEFAULT_PRECISION), + label=label_dict, + task_key=task_key, + ) + + with nvprof_context(enable_profiling, "Backward pass"): + loss.backward() - if self.gradient_max_norm > 0.0: - with nvprof_context(enable_profiling, "Gradient clip"): - paddle.nn.utils.clip_grad_norm_( - self.wrapper.parameters(), - self.gradient_max_norm, - error_if_nonfinite=True, + # gradient accumulation + if (_step_id + 1) % self.acc_freq == 0: + # fuse + allreduce manually before optimization if use DDP + no_sync + # details in https://github.com/PaddlePaddle/Paddle/issues/48898#issuecomment-1343838622 + if self.world_size > 1: + hpu.fused_allreduce_gradients( + list(self.wrapper.parameters()), None ) - with nvprof_context(enable_profiling, "Adam update"): - self.optimizer.step() - self.scheduler.step() + if self.gradient_max_norm > 0.0: + with nvprof_context(enable_profiling, "Gradient clip"): + paddle.nn.utils.clip_grad_norm_( + self.wrapper.parameters(), + self.gradient_max_norm, + error_if_nonfinite=True, + ) + + with nvprof_context(enable_profiling, "Adam update"): + self.optimizer.step() + self.optimizer.clear_grad(set_to_zero=False) + self.scheduler.step() else: raise ValueError(f"Not supported optimizer type '{self.opt_type}'") @@ -876,7 +925,9 @@ def log_loss_valid(_task_key="Default"): self.t0 = current_time if self.rank == 0 and self.timing_in_training: eta = int( - (self.num_steps - display_step_id) / self.disp_freq * train_time + (self.num_steps - display_step_id) + / min(self.disp_freq, display_step_id - self.start_step) + * train_time ) log.info( format_training_message( diff --git a/deepmd/pd/train/wrapper.py b/deepmd/pd/train/wrapper.py index bd28b17c88..6c8db691c4 100644 --- a/deepmd/pd/train/wrapper.py +++ b/deepmd/pd/train/wrapper.py @@ -7,13 +7,10 @@ from collections import ( OrderedDict, ) -from typing import ( - Union, -) import paddle -_StateDict = Union[dict[str, paddle.Tensor], OrderedDict[str, paddle.Tensor]] +_StateDict = dict[str, paddle.Tensor] | OrderedDict[str, paddle.Tensor] log = logging.getLogger(__name__) diff --git a/deepmd/pd/utils/auto_batch_size.py b/deepmd/pd/utils/auto_batch_size.py index 0431fb80ae..f45746ed95 100644 --- a/deepmd/pd/utils/auto_batch_size.py +++ b/deepmd/pd/utils/auto_batch_size.py @@ -36,7 +36,7 @@ def is_gpu_available(self) -> bool: bool True if GPU is available """ - return paddle.device.cuda.device_count() > 0 + return paddle.device.device_count() > 0 def is_oom_error(self, e: Exception) -> bool: """Check if the exception is an OOM error. @@ -51,6 +51,6 @@ def is_oom_error(self, e: Exception) -> bool: # (the meaningless error message should be considered as a bug in cusolver) if isinstance(e, MemoryError) and ("ResourceExhaustedError" in e.args[0]): # Release all unoccupied cached memory - paddle.device.cuda.empty_cache() + paddle.device.empty_cache() return True return False diff --git a/deepmd/pd/utils/dataset.py b/deepmd/pd/utils/dataset.py index 1f0533d8fc..fa9106044c 100644 --- a/deepmd/pd/utils/dataset.py +++ b/deepmd/pd/utils/dataset.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) - from paddle.io import ( Dataset, ) @@ -16,7 +12,7 @@ class DeepmdDataSetForLoader(Dataset): - def __init__(self, system: str, type_map: Optional[list[str]] = None): + def __init__(self, system: str, type_map: list[str] | None = None): """Construct DeePMD-style dataset containing frames cross different systems. Args: diff --git a/deepmd/pd/utils/decomp.py b/deepmd/pd/utils/decomp.py index 3b7bddbcd1..ab9a57dbfd 100644 --- a/deepmd/pd/utils/decomp.py +++ b/deepmd/pd/utils/decomp.py @@ -112,10 +112,7 @@ def masked_add__decomp( """ assert mask.dtype == paddle.bool, f"mask must be bool type, but got {mask.dtype}" # indices is bool mask - mask_coord = paddle.concat( - paddle.nonzero(mask, as_tuple=True), - axis=1, - ) # [nz, dim] + mask_coord = paddle.nonzero(mask, as_tuple=False) # [nz, dim] if not paddle.is_tensor(v): v = paddle.full([mask_coord.shape[0]], v, dtype=x.dtype) t = paddle.scatter_nd_add( diff --git a/deepmd/pd/utils/env.py b/deepmd/pd/utils/env.py index cf5b1f835c..8715b59f54 100644 --- a/deepmd/pd/utils/env.py +++ b/deepmd/pd/utils/env.py @@ -27,12 +27,12 @@ ncpus = os.cpu_count() NUM_WORKERS = int(os.environ.get("NUM_WORKERS", min(0, ncpus))) # Make sure DDP uses correct device if applicable -LOCAL_RANK = paddle.distributed.get_rank() +LOCAL_RANK = int(os.environ.get("PADDLE_LOCAL_RANK", 0)) -if os.environ.get("DEVICE") == "cpu" or paddle.device.cuda.device_count() <= 0: +if os.environ.get("DEVICE") == "cpu" or paddle.device.device_count() <= 0: DEVICE = "cpu" else: - DEVICE = f"gpu:{LOCAL_RANK}" + DEVICE = paddle.device.get_device() paddle.device.set_device(DEVICE) @@ -69,8 +69,17 @@ def to_bool(flag: int | bool | str) -> bool: "installation or recompiling with CINN enabled." ) +# NOTE: Allow the CINN compiler to optimize inputs with dynamic shapes, +# may lead to a slight performance decrease compared to static shapes. + +# If you can confirm that the shape of the input tensors will not change, +# you can set it to False to further enhance performance. +# Otherwise, please use the default value(True) to improve runtime compatibility. +CINN_ALLOW_DYNAMIC_SHAPE = to_bool(os.environ.get("CINN_ALLOW_DYNAMIC_SHAPE", True)) + CACHE_PER_SYS = 5 # keep at most so many sets per sys in memory ENERGY_BIAS_TRAINABLE = True +CUSTOM_OP_USE_JIT = to_bool(os.environ.get("CUSTOM_OP_USE_JIT", False)) PRECISION_DICT = { "float16": paddle.float16, @@ -198,6 +207,8 @@ def enable_prim(enable: bool = True): __all__ = [ "CACHE_PER_SYS", "CINN", + "CINN_ALLOW_DYNAMIC_SHAPE", + "CUSTOM_OP_USE_JIT", "DEFAULT_PRECISION", "DEVICE", "ENERGY_BIAS_TRAINABLE", diff --git a/deepmd/pd/utils/env_mat_stat.py b/deepmd/pd/utils/env_mat_stat.py index 4ea4850ba7..0e41243924 100644 --- a/deepmd/pd/utils/env_mat_stat.py +++ b/deepmd/pd/utils/env_mat_stat.py @@ -4,7 +4,6 @@ ) from typing import ( TYPE_CHECKING, - Union, ) import numpy as np @@ -77,7 +76,7 @@ def __init__(self, descriptor: "DescriptorBlock"): ) # se_r=1, se_a=4 def iter( - self, data: list[dict[str, Union[paddle.Tensor, list[tuple[int, int]]]]] + self, data: list[dict[str, paddle.Tensor | list[tuple[int, int]]]] ) -> Iterator[dict[str, StatItem]]: """Get the iterator of the environment matrix. diff --git a/deepmd/pd/utils/exclude_mask.py b/deepmd/pd/utils/exclude_mask.py index 29c9cc3501..cde8730c9a 100644 --- a/deepmd/pd/utils/exclude_mask.py +++ b/deepmd/pd/utils/exclude_mask.py @@ -58,7 +58,7 @@ def forward( """ nf, natom = atype.shape - return self.type_mask[atype].reshape([nf, natom]).to(atype.place) + return self.type_mask[atype].reshape([nf, natom]) class PairExcludeMask(paddle.nn.Layer): @@ -126,31 +126,25 @@ def forward( """ if self.no_exclusion: # safely return 1 if nothing is excluded. - return paddle.ones_like(nlist, dtype=paddle.int32).to(device=nlist.place) + return paddle.ones_like(nlist, dtype=paddle.int32) nf, nloc, nnei = nlist.shape nall = atype_ext.shape[1] # add virtual atom of type ntypes. nf x nall+1 ae = paddle.concat( [ atype_ext, - self.ntypes - * paddle.ones([nf, 1], dtype=atype_ext.dtype).to(atype_ext.place), + self.ntypes * paddle.ones([nf, 1], dtype=atype_ext.dtype), ], axis=-1, ) type_i = atype_ext[:, :nloc].reshape([nf, nloc]) * (self.ntypes + 1) # nf x nloc x nnei index = paddle.where(nlist == -1, nall, nlist).reshape([nf, nloc * nnei]) - type_j = paddle.take_along_axis(ae, axis=1, indices=index).reshape( - [nf, nloc, nnei] - ) + type_j = paddle.take_along_axis( + ae, axis=1, indices=index, broadcast=False + ).reshape([nf, nloc, nnei]) type_ij = type_i[:, :, None] + type_j # nf x (nloc x nnei) type_ij = type_ij.reshape([nf, nloc * nnei]) - mask = ( - self.type_mask[type_ij] - .reshape([nf, nloc, nnei]) - .to(atype_ext.place) - .astype("bool") - ) + mask = self.type_mask[type_ij].reshape([nf, nloc, nnei]).astype("bool") return mask diff --git a/deepmd/pd/utils/neighbor_stat.py b/deepmd/pd/utils/neighbor_stat.py index 9bc0079475..f569999bfc 100644 --- a/deepmd/pd/utils/neighbor_stat.py +++ b/deepmd/pd/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np import paddle @@ -52,7 +49,7 @@ def forward( self, coord: paddle.Tensor, atype: paddle.Tensor, - cell: Optional[paddle.Tensor], + cell: paddle.Tensor | None, ) -> tuple[paddle.Tensor, paddle.Tensor]: """Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix. @@ -173,7 +170,7 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, ): """Execute the operation. diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py index 707cbd125b..cd041a345c 100644 --- a/deepmd/pd/utils/nlist.py +++ b/deepmd/pd/utils/nlist.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import paddle @@ -22,7 +18,7 @@ def extend_input_and_build_neighbor_list( rcut: float, sel: list[int], mixed_types: bool = False, - box: Optional[paddle.Tensor] = None, + box: paddle.Tensor | None = None, ): nframes, nloc = atype.shape[:2] if box is not None: @@ -54,7 +50,7 @@ def build_neighbor_list( atype: paddle.Tensor, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> paddle.Tensor: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -94,36 +90,36 @@ def build_neighbor_list( """ batch_size = coord.shape[0] - coord = coord.reshape([batch_size, -1]) nall = coord.shape[1] // 3 # fill virtual atoms with large coords so they are not neighbors of any # real atom. - # NOTE: control flow with double backward is not supported well yet by paddle.jit if not paddle.in_dynamic_mode() or decomp.numel(coord) > 0: xmax = paddle.max(coord) + 2.0 * rcut else: xmax = paddle.zeros([], dtype=coord.dtype).to(device=coord.place) + 2.0 * rcut + coord_xyz = coord.reshape([batch_size, nall, 3]) # nf x nall is_vir = atype < 0 - coord1 = paddle.where( - is_vir[:, :, None], xmax, coord.reshape([batch_size, nall, 3]) - ).reshape([batch_size, nall * 3]) + # batch_size x nall x 3 + vcoord_xyz = paddle.where(is_vir.unsqueeze(2), xmax, coord_xyz) if isinstance(sel, int): sel = [sel] - # nloc x 3 - coord0 = coord1[:, : nloc * 3] - # nloc x nall x 3 - diff = coord1.reshape([batch_size, -1, 3]).unsqueeze(1) - coord0.reshape( - [batch_size, -1, 3] - ).unsqueeze(2) + + # Get the coordinates for the local atoms (first nloc atoms) + # batch_size x nloc x 3 + vcoord_local_xyz = vcoord_xyz[:, :nloc, :] + + # Calculate displacement vectors. + diff = vcoord_xyz.unsqueeze(1) - vcoord_local_xyz.unsqueeze(2) if paddle.in_dynamic_mode(): assert list(diff.shape) == [batch_size, nloc, nall, 3] # nloc x nall rr = paddle.linalg.norm(diff, axis=-1) # if central atom has two zero distances, sorting sometimes can not exclude itself - rr = rr - paddle.eye(nloc, nall, dtype=rr.dtype).to(device=rr.place).unsqueeze(0) + rr = rr - paddle.eye(nloc, nall, dtype=rr.dtype).unsqueeze(0) rr, nlist = paddle.sort(rr, axis=-1), paddle.argsort(rr, axis=-1) + # nloc x (nall-1) rr = rr[:, :, 1:] nlist = nlist[:, :, 1:] @@ -155,20 +151,13 @@ def _trim_mask_distinguish_nlist( rr = paddle.concat( [ rr, - paddle.ones([batch_size, nloc, nsel - nnei]).to( - device=rr.place, dtype=rr.dtype - ) + paddle.ones([batch_size, nloc, nsel - nnei]).astype(dtype=rr.dtype) + rcut, ], axis=-1, ) nlist = paddle.concat( - [ - nlist, - paddle.ones([batch_size, nloc, nsel - nnei], dtype=nlist.dtype).to( - device=rr.place - ), - ], + [nlist, paddle.ones([batch_size, nloc, nsel - nnei], dtype=nlist.dtype)], axis=-1, ) if paddle.in_dynamic_mode(): @@ -188,7 +177,7 @@ def build_directional_neighbor_list( coord_neig: paddle.Tensor, atype_neig: paddle.Tensor, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> paddle.Tensor: """Build directional neighbor list. @@ -318,7 +307,11 @@ def nlist_distinguish_types( paddle.argsort(pick_mask, axis=-1, descending=True, stable=True), ) # nloc x s(nsel) - inlist = paddle.take_along_axis(nlist, axis=2, indices=imap) + inlist = paddle.take_along_axis( + nlist, + axis=2, + indices=imap, + ) inlist = inlist.masked_fill(~(pick_mask.to(paddle.bool)), -1) # nloc x nsel[ii] ret_nlist.append(paddle.split(inlist, [ss, snsel - ss], axis=-1)[0]) @@ -377,7 +370,7 @@ def build_multiple_neighbor_list( pad = -paddle.ones( [nb, nloc, nsels[-1] - nsel], dtype=nlist.dtype, - ).to(device=nlist.place) + ) # nb x nloc x nsel nlist = paddle.concat([nlist, pad], axis=-1) if paddle.is_tensor(nsel): @@ -399,9 +392,11 @@ def build_multiple_neighbor_list( .expand([-1, -1, 3]) ) # nb x nloc x nsel x 3 - coord2 = paddle.take_along_axis(coord1, axis=1, indices=index).reshape( - [nb, nloc, nsel, 3] - ) + coord2 = paddle.take_along_axis( + coord1, + axis=1, + indices=index, + ).reshape([nb, nloc, nsel, 3]) # nb x nloc x nsel x 3 diff = coord2 - coord0[:, :, None, :] # nb x nloc x nsel @@ -418,9 +413,9 @@ def build_multiple_neighbor_list( def extend_coord_with_ghosts( coord: paddle.Tensor, atype: paddle.Tensor, - cell: Optional[paddle.Tensor], + cell: paddle.Tensor | None, rcut: float, - cell_cpu: Optional[paddle.Tensor] = None, + cell_cpu: paddle.Tensor | None = None, ) -> tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor]: """Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors @@ -452,7 +447,7 @@ def extend_coord_with_ghosts( device = coord.place nf, nloc = atype.shape[:2] # int64 for index - aidx = paddle.tile(paddle.arange(nloc).to(device=device).unsqueeze(0), [nf, 1]) # pylint: disable=no-explicit-dtype + aidx = paddle.tile(paddle.arange(nloc).unsqueeze(0), [nf, 1]) # pylint: disable=no-explicit-dtype if cell is None: nall = nloc extend_coord = coord.clone() @@ -496,14 +491,12 @@ def extend_coord_with_ghosts( # .cpu() ) # pylint: disable=no-explicit-dtype eye_3 = ( - paddle.eye(3, dtype=env.GLOBAL_PD_FLOAT_PRECISION).to( - dtype=env.GLOBAL_PD_FLOAT_PRECISION - ) + paddle.eye(3) # .cpu() - ) - xyz = xi.reshape([-1, 1, 1, 1]) * eye_3[0] - xyz = xyz + yi.reshape([1, -1, 1, 1]) * eye_3[1] - xyz = xyz + zi.reshape([1, 1, -1, 1]) * eye_3[2] + ).to(dtype=env.GLOBAL_PD_FLOAT_PRECISION) + xyz = xi.reshape([-1, 1, 1, 1]).astype(eye_3.dtype) * eye_3[0] + xyz = xyz + yi.reshape([1, -1, 1, 1]).astype(eye_3.dtype) * eye_3[1] + xyz = xyz + zi.reshape([1, 1, -1, 1]).astype(eye_3.dtype) * eye_3[2] xyz = xyz.reshape([-1, 3]) # xyz = xyz.to(device=device) # ns x 3 @@ -519,7 +512,7 @@ def extend_coord_with_ghosts( # nf x ns x nloc extend_aidx = paddle.tile(aidx.unsqueeze(-2), [1, ns, 1]) return ( - extend_coord.reshape([nf, nall * 3]).to(device), - extend_atype.reshape([nf, nall]).to(device), - extend_aidx.reshape([nf, nall]).to(device), + extend_coord.reshape([nf, nall * 3]), + extend_atype.reshape([nf, nall]), + extend_aidx.reshape([nf, nall]), ) diff --git a/deepmd/pd/utils/preprocess.py b/deepmd/pd/utils/preprocess.py index 3e047c1b8b..3be42b522e 100644 --- a/deepmd/pd/utils/preprocess.py +++ b/deepmd/pd/utils/preprocess.py @@ -10,9 +10,20 @@ def compute_smooth_weight(distance, rmin: float, rmax: float): """Compute smooth weight for descriptor elements.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") - min_mask = distance <= rmin - max_mask = distance >= rmax - mid_mask = paddle.logical_not(paddle.logical_or(min_mask, max_mask)) + distance = paddle.clip(distance, min=rmin, max=rmax) uu = (distance - rmin) / (rmax - rmin) - vv = uu * uu * uu * (-6 * uu * uu + 15 * uu - 10) + 1 - return vv * mid_mask.astype(vv.dtype) + min_mask.astype(vv.dtype) + uu2 = uu * uu + vv = uu2 * uu * (-6 * uu2 + 15 * uu - 10) + 1 + return vv + + +def compute_exp_sw(distance, rmin: float, rmax: float): + """Compute the exponential switch function for neighbor update.""" + if rmin >= rmax: + raise ValueError("rmin should be less than rmax.") + distance = paddle.clip(distance, min=0.0, max=rmax) + C = 20 + a = C / rmin + b = rmin + exp_sw = paddle.exp(-paddle.exp(a * (distance - b))) + return exp_sw diff --git a/deepmd/pd/utils/serialization.py b/deepmd/pd/utils/serialization.py index f4206ce993..bd70deb75c 100644 --- a/deepmd/pd/utils/serialization.py +++ b/deepmd/pd/utils/serialization.py @@ -69,9 +69,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None: model.forward, full_graph=True, input_spec=[ - InputSpec([1, -1, 3], dtype="float64", name="coord"), - InputSpec([1, -1], dtype="int64", name="atype"), - InputSpec([1, 9], dtype="float64", name="box"), + InputSpec([-1, -1, 3], dtype="float64", name="coord"), + InputSpec([-1, -1], dtype="int64", name="atype"), + InputSpec([-1, 9], dtype="float64", name="box"), None, None, True, @@ -88,9 +88,9 @@ def deserialize_to_file(model_file: str, data: dict) -> None: model.forward_lower, full_graph=True, input_spec=[ - InputSpec([1, -1, 3], dtype="float64", name="coord"), - InputSpec([1, -1], dtype="int32", name="atype"), - InputSpec([1, -1, -1], dtype="int32", name="nlist"), + InputSpec([-1, -1, 3], dtype="float64", name="coord"), + InputSpec([-1, -1], dtype="int32", name="atype"), + InputSpec([-1, -1, -1], dtype="int32", name="nlist"), None, None, None, @@ -101,4 +101,5 @@ def deserialize_to_file(model_file: str, data: dict) -> None: paddle.jit.save( model, model_file.split(".json")[0], + skip_prune_program=True, ) diff --git a/deepmd/pd/utils/spin.py b/deepmd/pd/utils/spin.py index 934fb3762a..27bc355877 100644 --- a/deepmd/pd/utils/spin.py +++ b/deepmd/pd/utils/spin.py @@ -21,7 +21,6 @@ def concat_switch_virtual( extended_tensor_updated = paddle.zeros( out_shape, dtype=extended_tensor.dtype, - device=extended_tensor.place, ) extended_tensor_updated[:, :nloc] = extended_tensor[:, :nloc] extended_tensor_updated[:, nloc : nloc + nloc] = extended_tensor_virtual[:, :nloc] diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py index e0abb1b289..3642d309b8 100644 --- a/deepmd/pd/utils/stat.py +++ b/deepmd/pd/utils/stat.py @@ -3,18 +3,16 @@ from collections import ( defaultdict, ) -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import numpy as np import paddle -from deepmd.dpmodel.output_def import ( - FittingOutputDef, -) from deepmd.pd.utils import ( AtomExcludeMask, ) @@ -27,6 +25,7 @@ to_paddle_tensor, ) from deepmd.utils.out_stat import ( + compute_stats_do_not_distinguish_types, compute_stats_from_atomic, compute_stats_from_redu, ) @@ -37,7 +36,9 @@ log = logging.getLogger(__name__) -def make_stat_input(datasets, dataloaders, nbatches): +def make_stat_input( + datasets: list[Any], dataloaders: list[Any], nbatches: int +) -> dict[str, Any]: """Pack data for statistics. Args: @@ -61,6 +62,14 @@ def make_stat_input(datasets, dataloaders, nbatches): except StopIteration: iterator = iter(dataloaders[i]) stat_data = next(iterator) + if ( + "find_fparam" in stat_data + and "fparam" in stat_data + and stat_data["find_fparam"] == 0.0 + ): + # for model using default fparam + stat_data.pop("fparam") + stat_data.pop("find_fparam") for dd in stat_data: if stat_data[dd] is None: sys_stat[dd] = None @@ -88,7 +97,7 @@ def make_stat_input(datasets, dataloaders, nbatches): def _restore_from_file( stat_file_path: DPPath, keys: list[str] = ["energy"], -) -> Optional[dict]: +) -> dict | None: if stat_file_path is None: return None, None stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] @@ -117,7 +126,7 @@ def _save_to_file( stat_file_path: DPPath, bias_out: dict, std_out: dict, -): +) -> None: assert stat_file_path is not None stat_file_path.mkdir(exist_ok=True, parents=True) for kk, vv in bias_out.items(): @@ -129,26 +138,31 @@ def _save_to_file( def _post_process_stat( - out_bias, - out_std, -): + out_bias: paddle.Tensor, + out_std: paddle.Tensor, +) -> tuple[paddle.Tensor, paddle.Tensor]: """Post process the statistics. For global statistics, we do not have the std for each type of atoms, thus fake the output std by ones for all the types. + If the shape of out_std is already the same as out_bias, + we do not need to do anything. """ new_std = {} for kk, vv in out_bias.items(): - new_std[kk] = np.ones_like(vv) + if vv.shape == out_std[kk].shape: + new_std[kk] = out_std[kk] + else: + new_std[kk] = np.ones_like(vv) return out_bias, new_std def _compute_model_predict( - sampled: Union[Callable[[], list[dict]], list[dict]], + sampled: Callable[[], list[dict]] | list[dict], keys: list[str], model_forward: Callable[..., paddle.Tensor], -): +) -> dict[str, list[paddle.Tensor]]: auto_batch_size = AutoBatchSize() model_predict = {kk: [] for kk in keys} for system in sampled: @@ -185,8 +199,8 @@ def model_forward_auto_batch_size(*args, **kwargs): def _make_preset_out_bias( ntypes: int, - ibias: list[Optional[np.ndarray]], -) -> Optional[np.ndarray]: + ibias: list[np.ndarray | None], +) -> np.ndarray | None: """Make preset out bias. output: @@ -209,9 +223,9 @@ def _make_preset_out_bias( def _fill_stat_with_global( - atomic_stat: Union[np.ndarray, None], + atomic_stat: np.ndarray | None, global_stat: np.ndarray, -): +) -> np.ndarray | None: """This function is used to fill atomic stat with global stat. Parameters @@ -235,15 +249,16 @@ def _fill_stat_with_global( def compute_output_stats( - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], ntypes: int, - keys: Union[str, list[str]] = ["energy"], - stat_file_path: Optional[DPPath] = None, - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_forward: Optional[Callable[..., paddle.Tensor]] = None, - atomic_output: Optional[FittingOutputDef] = None, -): + keys: str | list[str] = ["energy"], + stat_file_path: DPPath | None = None, + rcond: float | None = None, + preset_bias: dict[str, list[np.ndarray | None]] | None = None, + model_forward: Callable[..., paddle.Tensor] | None = None, + stats_distinguish_types: bool = True, + intensive: bool = False, +) -> dict[str, Any]: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -272,8 +287,10 @@ def compute_output_stats( If not None, the model will be utilized to generate the original energy prediction, which will be subtracted from the energy label of the data. The difference will then be used to calculate the delta complement energy bias for each type. - atomic_output : FittingOutputDef, optional - The output of atomic model. + stats_distinguish_types : bool, optional + Whether to distinguish different element types in the statistics. + intensive : bool, optional + Whether the fitting target is intensive. """ # try to restore the bias from stat file bias_atom_e, std_atom_e = _restore_from_file(stat_file_path, keys) @@ -362,7 +379,8 @@ def compute_output_stats( rcond, preset_bias, model_pred_g, - atomic_output, + stats_distinguish_types, + intensive, ) bias_atom_a, std_atom_a = compute_output_stats_atomic( sampled, @@ -402,11 +420,12 @@ def compute_output_stats_global( sampled: list[dict], ntypes: int, keys: list[str], - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[paddle.Tensor]]]] = None, - model_pred: Optional[dict[str, np.ndarray]] = None, - atomic_output: Optional[FittingOutputDef] = None, -): + rcond: float | None = None, + preset_bias: dict[str, list[paddle.Tensor | None]] | None = None, + model_pred: dict[str, np.ndarray] | None = None, + stats_distinguish_types: bool = True, + intensive: bool = False, +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: """This function only handle stat computation from reduced global labels.""" # return directly if model predict is empty for global if model_pred == {}: @@ -469,26 +488,31 @@ def compute_output_stats_global( # subtract the model bias and output the delta bias stats_input = { - kk: merged_output[kk] - model_pred[kk] for kk in keys if kk in merged_output + kk: merged_output[kk] - model_pred[kk].reshape(merged_output[kk].shape) + for kk in keys + if kk in merged_output } bias_atom_e = {} std_atom_e = {} for kk in keys: if kk in stats_input: - if atomic_output is not None and atomic_output.get_data()[kk].intensive: - task_dim = stats_input[kk].shape[1] - assert merged_natoms[kk].shape == (nf[kk], ntypes) - stats_input[kk] = ( - merged_natoms[kk].sum(axis=1).reshape([-1, 1]) * stats_input[kk] + if not stats_distinguish_types: + bias_atom_e[kk], std_atom_e[kk] = ( + compute_stats_do_not_distinguish_types( + stats_input[kk], + merged_natoms[kk], + assigned_bias=assigned_atom_ener[kk], + intensive=intensive, + ) + ) + else: + bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( + stats_input[kk], + merged_natoms[kk], + assigned_bias=assigned_atom_ener[kk], + rcond=rcond, ) - assert stats_input[kk].shape == (nf[kk], task_dim) - bias_atom_e[kk], std_atom_e[kk] = compute_stats_from_redu( - stats_input[kk], - merged_natoms[kk], - assigned_bias=assigned_atom_ener[kk], - rcond=rcond, - ) else: # this key does not have global labels, skip it. continue @@ -509,16 +533,16 @@ def compute_output_stats_global( } atom_numbs = {kk: merged_natoms[kk].sum(-1) for kk in bias_atom_e.keys()} - def rmse(x): + def rmse(x: np.ndarray) -> float: return np.sqrt(np.mean(np.square(x))) for kk in bias_atom_e.keys(): rmse_ae = rmse( ( - unbias_e[kk].reshape([nf[kk], -1]).astype(merged_output[kk].dtype) + unbias_e[kk].reshape([nf[kk], -1]) - merged_output[kk].reshape([nf[kk], -1]) ) - / atom_numbs[kk][:, None].astype(merged_output[kk].dtype) + / atom_numbs[kk][:, None] ) log.info( f"RMSE of {kk} per atom after linear regression is: {rmse_ae} in the unit of {kk}." @@ -530,8 +554,8 @@ def compute_output_stats_atomic( sampled: list[dict], ntypes: int, keys: list[str], - model_pred: Optional[dict[str, np.ndarray]] = None, -): + model_pred: dict[str, np.ndarray] | None = None, +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { kk: [ @@ -549,7 +573,17 @@ def compute_output_stats_atomic( ] for kk in keys } - # shape: (nframes, nloc, ndim) + # reshape outputs [nframes, nloc * ndim] --> reshape to [nframes * nloc, 1, ndim] for concatenation + # reshape natoms [nframes, nloc] --> reshape to [nframes * nolc, 1] for concatenation + natoms = {k: [sys_v.reshape([-1, 1]) for sys_v in v] for k, v in natoms.items()} + outputs = { + k: [ + sys.reshape([natoms[k][sys_idx].shape[0], 1, -1]) + for sys_idx, sys in enumerate(v) + ] + for k, v in outputs.items() + } + merged_output = { kk: to_numpy_array(paddle.concat(outputs[kk])) for kk in keys diff --git a/deepmd/pd/utils/utils.py b/deepmd/pd/utils/utils.py index a756491a8d..7224547805 100644 --- a/deepmd/pd/utils/utils.py +++ b/deepmd/pd/utils/utils.py @@ -3,6 +3,7 @@ annotations, ) +import warnings from contextlib import ( contextmanager, ) @@ -20,6 +21,9 @@ ) from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT +from deepmd.pd.utils import ( + env, +) from .env import ( DEVICE, @@ -32,15 +36,126 @@ ) +def silut_forward( + x: paddle.Tensor, threshold: float, slope: float, const_val: float +) -> paddle.Tensor: + sig = F.sigmoid(x) + silu = x * sig + tanh = paddle.tanh(slope * (x - threshold)) + const_val + return paddle.where(x >= threshold, tanh, silu) + + +def silut_backward( + x: paddle.Tensor, grad_output: paddle.Tensor, threshold: float, slope: float +) -> paddle.Tensor: + sig = F.sigmoid(x) + grad_silu = sig * (1 + x * (1 - sig)) + + tanh = paddle.tanh(slope * (x - threshold)) + grad_tanh = slope * (1 - tanh * tanh) + + grad = paddle.where(x >= threshold, grad_tanh, grad_silu) + return grad * grad_output + + +def silut_double_backward( + x: paddle.Tensor, + grad_grad_output: paddle.Tensor, + grad_output: paddle.Tensor, + threshold: float, + slope: float, +) -> tuple[paddle.Tensor, paddle.Tensor]: + # SiLU branch + sig = F.sigmoid(x) + + sig_prime = sig * (1 - sig) + grad_silu = sig + x * sig_prime + grad_grad_silu = sig_prime * (2 + x * (1 - 2 * sig)) + + # Tanh branch + tanh = paddle.tanh(slope * (x - threshold)) + tanh_square = tanh * tanh # .square is slow for jit.script! + grad_tanh = slope * (1 - tanh_square) + grad_grad_tanh = -2 * slope * tanh * grad_tanh + + grad = paddle.where(x >= threshold, grad_tanh, grad_silu) + grad_grad = paddle.where(x >= threshold, grad_grad_tanh, grad_grad_silu) + return grad_output * grad_grad * grad_grad_output, grad * grad_grad_output + + +class SiLUTScript(paddle.nn.Layer): + def __init__(self, threshold: float = 3.0): + super().__init__() + self.threshold = threshold + + # Precompute parameters for the tanh replacement + sigmoid_threshold = 1 / (1 + np.exp(-threshold)) + self.slope = float( + sigmoid_threshold + threshold * sigmoid_threshold * (1 - sigmoid_threshold) + ) + self.const_val = float(threshold * sigmoid_threshold) + self.get_script_code() + + def get_script_code(self): + silut_forward_script = paddle.jit.to_static(silut_forward, full_graph=True) + silut_backward_script = paddle.jit.to_static(silut_backward, full_graph=True) + silut_double_backward_script = paddle.jit.to_static( + silut_double_backward, full_graph=True + ) + + class SiLUTFunction(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, x, threshold, slope, const_val): + ctx.save_for_backward(x) + ctx.threshold = threshold + ctx.slope = slope + ctx.const_val = const_val + return silut_forward_script(x, threshold, slope, const_val) + + @staticmethod + def backward(ctx, grad_output): + (x,) = ctx.saved_tensor() + threshold = ctx.threshold + slope = ctx.slope + + grad_input = SiLUTGradFunction.apply(x, grad_output, threshold, slope) + return grad_input + + class SiLUTGradFunction(paddle.autograd.PyLayer): + @staticmethod + def forward(ctx, x, grad_output, threshold, slope): + ctx.threshold = threshold + ctx.slope = slope + grad_input = silut_backward_script(x, grad_output, threshold, slope) + ctx.save_for_backward(x, grad_output) + return grad_input + + @staticmethod + def backward(ctx, grad_grad_output): + (x, grad_output) = ctx.saved_tensor() + threshold = ctx.threshold + slope = ctx.slope + + grad_input, grad_mul_grad_grad_output = silut_double_backward_script( + x, grad_grad_output, grad_output, threshold, slope + ) + return grad_input, grad_mul_grad_grad_output + + self.SiLUTFunction = SiLUTFunction + + def forward(self, x): + return self.SiLUTFunction.apply(x, self.threshold, self.slope, self.const_val) + + class SiLUT(paddle.nn.Layer): def __init__(self, threshold=3.0): super().__init__() def sigmoid(x): - return paddle.nn.functional.sigmoid(x) + return F.sigmoid(x) def silu(x): - return paddle.nn.functional.silu(x) + return F.silu(x) def silu_grad(x): sig = sigmoid(x) @@ -76,7 +191,11 @@ def __init__(self, activation: str | None): threshold = ( float(self.activation.split(":")[-1]) if ":" in self.activation else 3.0 ) - self.silut = SiLUT(threshold=threshold) + if env.CUSTOM_OP_USE_JIT: + # for efficient training but can not be jit + self.silut = SiLUTScript(threshold=threshold) + else: + self.silut = SiLUT(threshold=threshold) else: self.silut = None @@ -227,8 +346,21 @@ def get_generator( generator = paddle.framework.core.default_cuda_generator( int(DEVICE.split("gpu:")[1]) ) + elif DEVICE == "xpu": + generator = paddle.framework.core.default_xpu_generator(0) + elif DEVICE.startswith("xpu:"): + generator = paddle.framework.core.default_xpu_generator( + int(DEVICE.split("xpu:")[1]) + ) else: - raise ValueError("DEVICE should be cpu or gpu or gpu:x") + # return none for compability in different devices + warnings.warn( + f"DEVICE is {DEVICE}, which is not supported. Returning None.", + category=UserWarning, + stacklevel=2, + ) + return None + # raise ValueError("DEVICE should be cpu or gpu or gpu:x or xpu or xpu:x") generator.manual_seed(seed) return generator else: diff --git a/deepmd/pt/cxx_op.py b/deepmd/pt/cxx_op.py index f7922a5c52..1106a8887f 100644 --- a/deepmd/pt/cxx_op.py +++ b/deepmd/pt/cxx_op.py @@ -1,5 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import platform +from ctypes import ( + CDLL, + RTLD_GLOBAL, +) +from importlib import ( + metadata, +) import torch from packaging.version import ( @@ -87,6 +94,29 @@ def load_library(module_name: str) -> bool: return False +def load_mpi_library() -> None: + """Load MPI library. + + When building with cibuildwheel, the link to the MPI library is lost + after the wheel is repaired. + """ + if platform.system() == "Linux": + libname = "libmpi.so.*" + elif platform.system() == "Darwin": + libname = "libmpi.*.dylib" + else: + raise RuntimeError("Unsupported platform") + MPI_LIB = next(p for p in metadata.files("mpich") if p.match(libname)).locate() + # use CDLL to load the library + CDLL(MPI_LIB, mode=RTLD_GLOBAL) + + +if GLOBAL_CONFIG.get("cibuildwheel", "0") == "1" and platform.system() in ( + "Linux", + "Darwin", +): + load_mpi_library() + ENABLE_CUSTOMIZED_OP = load_library("deepmd_op_pt") __all__ = [ diff --git a/deepmd/pt/entrypoints/compress.py b/deepmd/pt/entrypoints/compress.py index 3a4efaeea5..83f2999971 100644 --- a/deepmd/pt/entrypoints/compress.py +++ b/deepmd/pt/entrypoints/compress.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json import logging -from typing import ( - Optional, -) import torch @@ -35,7 +32,7 @@ def enable_compression( stride: float = 0.01, extrapolate: int = 5, check_frequency: int = -1, - training_script: Optional[str] = None, + training_script: str | None = None, ) -> None: saved_model = torch.jit.load(input_file, map_location="cpu") model_def_script = json.loads(saved_model.model_def_script) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 0e248583ec..d349d519ba 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -8,8 +8,7 @@ Path, ) from typing import ( - Optional, - Union, + Any, ) import h5py @@ -25,6 +24,7 @@ ) from deepmd.common import ( expand_sys_str, + j_loader, ) from deepmd.env import ( GLOBAL_CONFIG, @@ -94,20 +94,23 @@ def get_trainer( - config, - init_model=None, - restart_model=None, - finetune_model=None, - force_load=False, - init_frz_model=None, - shared_links=None, - finetune_links=None, -): + config: dict[str, Any], + init_model: str | None = None, + restart_model: str | None = None, + finetune_model: str | None = None, + force_load: bool = False, + init_frz_model: str | None = None, + shared_links: dict[str, Any] | None = None, + finetune_links: dict[str, Any] | None = None, +) -> training.Trainer: multi_task = "model_dict" in config.get("model", {}) def prepare_trainer_input_single( - model_params_single, data_dict_single, rank=0, seed=None - ): + model_params_single: dict[str, Any], + data_dict_single: dict[str, Any], + rank: int = 0, + seed: int | None = None, + ) -> tuple[DpLoaderSet, DpLoaderSet | None, DPPath | None]: training_dataset_params = data_dict_single["training_data"] validation_dataset_params = data_dict_single.get("validation_data", None) validation_systems = ( @@ -240,10 +243,10 @@ def get_backend_info(self) -> dict: def train( input_file: str, - init_model: Optional[str], - restart: Optional[str], - finetune: Optional[str], - init_frz_model: Optional[str], + init_model: str | None, + restart: str | None, + finetune: str | None, + init_frz_model: str | None, model_branch: str, skip_neighbor_stat: bool = False, use_pretrain_script: bool = False, @@ -254,8 +257,7 @@ def train( env.CUSTOM_OP_USE_JIT = True if LOCAL_RANK == 0: SummaryPrinter()() - with open(input_file) as fin: - config = json.load(fin) + config = j_loader(input_file) # ensure suffix, as in the command line help, we say "path prefix of checkpoint files" if init_model is not None and not init_model.endswith(".pt"): init_model += ".pt" @@ -368,7 +370,7 @@ def train( def freeze( model: str, output: str = "frozen_model.pth", - head: Optional[str] = None, + head: str | None = None, ) -> None: model = inference.Tester(model, head=head).model model.eval() @@ -385,12 +387,12 @@ def freeze( def change_bias( input_file: str, mode: str = "change", - bias_value: Optional[list] = None, - datafile: Optional[str] = None, + bias_value: list | None = None, + datafile: str | None = None, system: str = ".", numb_batch: int = 0, - model_branch: Optional[str] = None, - output: Optional[str] = None, + model_branch: str | None = None, + output: str | None = None, ) -> None: if input_file.endswith(".pt"): old_state_dict = torch.load( @@ -514,7 +516,7 @@ def change_bias( @record -def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: +def main(args: list[str] | argparse.Namespace | None = None) -> None: if not isinstance(args, argparse.Namespace): FLAGS = parse_args(args=args) else: diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index c4d5d028ce..2726b61152 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +import logging +from collections.abc import ( + Callable, +) from typing import ( TYPE_CHECKING, Any, - Callable, Optional, - Union, ) import numpy as np @@ -64,10 +66,22 @@ to_numpy_array, to_torch_tensor, ) +from deepmd.utils.econf_embd import ( + sort_element_type, +) +from deepmd.utils.model_branch_dict import ( + get_model_dict, +) if TYPE_CHECKING: import ase.neighborlist + from deepmd.pt.model.model.model import ( + BaseModel, + ) + +log = logging.getLogger(__name__) + class DeepEval(DeepEvalBackend): """PyTorch backend implementation of DeepEval. @@ -80,7 +94,7 @@ class DeepEval(DeepEvalBackend): The output definition of the model. *args : list Positional arguments. - auto_batch_size : bool or int or AutomaticBatchSize, default: False + auto_batch_size : bool or int or AutomaticBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional @@ -95,9 +109,10 @@ def __init__( model_file: str, output_def: ModelOutputDef, *args: Any, - auto_batch_size: Union[bool, int, AutoBatchSize] = True, + auto_batch_size: bool | int | AutoBatchSize = True, neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, - head: Optional[Union[str, int]] = None, + head: str | int | None = None, + no_jit: bool = False, **kwargs: Any, ) -> None: self.output_def = output_def @@ -112,15 +127,36 @@ def __init__( self.model_def_script = self.input_param self.multi_task = "model_dict" in self.input_param if self.multi_task: + model_alias_dict, model_branch_dict = get_model_dict( + self.input_param["model_dict"] + ) model_keys = list(self.input_param["model_dict"].keys()) + if head is None and "Default" in model_alias_dict: + head = "Default" + log.info( + f"Using default head {model_alias_dict[head]} for multitask model." + ) if isinstance(head, int): head = model_keys[0] assert head is not None, ( - f"Head must be set for multitask model! Available heads are: {model_keys}" + f"Head must be set for multitask model! Available heads are: {model_keys}, " + f"use `dp --pt show your_model.pt model-branch` to show detail information." ) - assert head in model_keys, ( - f"No head named {head} in model! Available heads are: {model_keys}" + if head not in model_alias_dict: + # preprocess with potentially case-insensitive input + head_lower = head.lower() + for mk in model_alias_dict: + if mk.lower() == head_lower: + # mapped the first matched head + head = mk + break + # replace with alias + assert head in model_alias_dict, ( + f"No head or alias named {head} in model! Available heads are: {model_keys}," + f"use `dp --pt show your_model.pt model-branch` to show detail information." ) + head = model_alias_dict[head] + self.input_param = self.input_param["model_dict"][head] state_dict_head = {"_extra_state": state_dict["_extra_state"]} for item in state_dict: @@ -130,7 +166,7 @@ def __init__( ] = state_dict[item].clone() state_dict = state_dict_head model = get_model(self.input_param).to(DEVICE) - if not self.input_param.get("hessian_mode"): + if not self.input_param.get("hessian_mode") and not no_jit: model = torch.jit.script(model) self.dp = ModelWrapper(model) self.dp.load_state_dict(state_dict) @@ -183,6 +219,14 @@ def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this DP.""" return self.dp.model["Default"].get_dim_aparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + try: + return self.dp.model["Default"].has_default_fparam() + except AttributeError: + # for compatibility with old models + return False + def get_intensive(self) -> bool: return self.dp.model["Default"].get_intensive() @@ -205,7 +249,7 @@ def model_type(self) -> type["DeepEvalWrapper"]: return DeepDOS elif "dipole" in model_output_type: return DeepDipole - elif "polar" in model_output_type: + elif "polar" in model_output_type or "polarizability" in model_output_type: return DeepPolar elif "global_polar" in model_output_type: return DeepGlobalPolar @@ -241,22 +285,33 @@ def get_ntypes_spin(self) -> int: """Get the number of spin atom types of this model. Only used in old implement.""" return 0 - def get_has_spin(self): + def get_has_spin(self) -> bool: """Check if the model has spin atom types.""" return self._has_spin - def get_has_hessian(self): + def get_has_hessian(self) -> bool: """Check if the model has hessian.""" return self._has_hessian + def get_model_branch(self) -> tuple[dict[str, str], dict[str, dict[str, Any]]]: + """Get the model branch information.""" + if "model_dict" in self.model_def_script: + model_alias_dict, model_branch_dict = get_model_dict( + self.model_def_script["model_dict"] + ) + return model_alias_dict, model_branch_dict + else: + # single-task model + return {"Default": "Default"}, {"Default": {"alias": [], "info": {}}} + def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -343,9 +398,9 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: The requested output definitions. """ if atomic: - return list(self.output_def.var_defs.values()) + output_defs = list(self.output_def.var_defs.values()) else: - return [ + output_defs = [ x for x in self.output_def.var_defs.values() if x.category @@ -357,6 +412,13 @@ def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: OutputVariableCategory.DERV_R_DERV_R, ) ] + if not self.get_has_hessian(): + output_defs = [ + x + for x in output_defs + if x.category != OutputVariableCategory.DERV_R_DERV_R + ] + return output_defs def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: """Wrapper method with auto batch size. @@ -377,7 +439,7 @@ def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Calla """ if self.auto_batch_size is not None: - def eval_func(*args, **kwargs): + def eval_func(*args: Any, **kwargs: Any) -> Any: return self.auto_batch_size.execute_all( inner_func, numb_test, natoms, *args, **kwargs ) @@ -406,12 +468,12 @@ def _get_natoms_and_nframes( def _eval_model( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], - ): + ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) prec = NP_PRECISION_DICT[RESERVED_PRECISION_DICT[GLOBAL_PT_FLOAT_PRECISION]] @@ -483,13 +545,13 @@ def _eval_model( def _eval_model_spin( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, spins: np.ndarray, - fparam: Optional[np.ndarray], - aparam: Optional[np.ndarray], + fparam: np.ndarray | None, + aparam: np.ndarray | None, request_defs: list[OutputVariableDef], - ): + ) -> tuple[np.ndarray, ...]: model = self.dp.to(DEVICE) nframes = coords.shape[0] @@ -566,7 +628,9 @@ def _eval_model_spin( ) # this is kinda hacky return tuple(results) - def _get_output_shape(self, odef, nframes, natoms): + def _get_output_shape( + self, odef: OutputVariableDef, nframes: int, natoms: int + ) -> list[int]: if odef.category == OutputVariableCategory.DERV_C_REDU: # virial return [nframes, *odef.shape[:-1], 9] @@ -620,7 +684,7 @@ def eval_typeebd(self) -> np.ndarray: typeebd = torch.cat(out, dim=1) return to_numpy_array(typeebd) - def get_model_def_script(self) -> str: + def get_model_def_script(self) -> dict: """Get model definition script.""" return self.model_def_script @@ -648,13 +712,39 @@ def get_model_size(self) -> dict: "total": sum_param_des + sum_param_fit, } + def get_observed_types(self) -> dict: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + dict + A dictionary containing the information of observed type in the model: + - 'type_num': the total number of observed types in this model. + - 'observed_type': a list of the observed types in this model. + """ + observed_type_list = self.dp.model["Default"].get_observed_type_list() + return { + "type_num": len(observed_type_list), + "observed_type": sort_element_type(observed_type_list), + } + + def get_model(self) -> "BaseModel": + """Get the PyTorch model. + + Returns + ------- + BaseModel + The PyTorch model instance. + """ + return self.dp.model["Default"] + def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, **kwargs: Any, ) -> np.ndarray: """Evaluate descriptors by using this DP. @@ -702,3 +792,58 @@ def eval_descriptor( descriptor = model.eval_descriptor() model.set_eval_descriptor_hook(False) return to_numpy_array(descriptor) + + def eval_fitting_last_layer( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> np.ndarray: + """Evaluate fitting before last layer by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + + Returns + ------- + fitting + Fitting output before last layer. + """ + model = self.dp.model["Default"] + model.set_eval_fitting_last_layer_hook(True) + self.eval( + coords, + cells, + atom_types, + atomic=False, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + fitting_net = model.eval_fitting_last_layer() + model.set_eval_fitting_last_layer_hook(False) + return to_numpy_array(fitting_net) diff --git a/deepmd/pt/infer/inference.py b/deepmd/pt/infer/inference.py index dd0e7eaccb..4c49abeef8 100644 --- a/deepmd/pt/infer/inference.py +++ b/deepmd/pt/infer/inference.py @@ -25,8 +25,8 @@ class Tester: def __init__( self, - model_ckpt, - head=None, + model_ckpt: str | torch.nn.Module, + head: str | None = None, ) -> None: """Construct a DeePMD tester. diff --git a/deepmd/pt/loss/denoise.py b/deepmd/pt/loss/denoise.py index 574210adb6..c8eeff6185 100644 --- a/deepmd/pt/loss/denoise.py +++ b/deepmd/pt/loss/denoise.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + import torch import torch.nn.functional as F @@ -13,15 +17,15 @@ class DenoiseLoss(TaskLoss): def __init__( self, - ntypes, - masked_token_loss=1.0, - masked_coord_loss=1.0, - norm_loss=0.01, - use_l1=True, - beta=1.00, - mask_loss_coord=True, - mask_loss_token=True, - **kwargs, + ntypes: int, + masked_token_loss: float = 1.0, + masked_coord_loss: float = 1.0, + norm_loss: float = 0.01, + use_l1: bool = True, + beta: float = 1.00, + mask_loss_coord: bool = True, + mask_loss_token: bool = True, + **kwargs: Any, ) -> None: """Construct a layer to compute loss on coord, and type reconstruction.""" super().__init__() @@ -38,7 +42,14 @@ def __init__( self.mask_loss_coord = mask_loss_coord self.mask_loss_token = mask_loss_token - def forward(self, model_pred, label, natoms, learning_rate, mae=False): + def forward( + self, + model_pred: dict[str, torch.Tensor], + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float, + mae: bool = False, + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Return loss on coord and type denoise. Returns diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index 493cc85694..bc77f34437 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import torch @@ -26,8 +29,8 @@ def __init__( limit_pref_ados: float = 0.0, start_pref_acdf: float = 0.0, limit_pref_acdf: float = 0.0, - inference=False, - **kwargs, + inference: bool = False, + **kwargs: Any, ) -> None: r"""Construct a loss for local and global tensors. @@ -85,7 +88,15 @@ def __init__( ) ) - def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False): + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float = 0.0, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: """Return loss on local and global tensors. Parameters diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 10e2bf9971..1778806343 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -23,7 +23,9 @@ ) -def custom_huber_loss(predictions, targets, delta=1.0): +def custom_huber_loss( + predictions: torch.Tensor, targets: torch.Tensor, delta: float = 1.0 +) -> torch.Tensor: error = targets - predictions abs_error = torch.abs(error) quadratic_loss = 0.5 * torch.pow(error, 2) @@ -35,27 +37,27 @@ def custom_huber_loss(predictions, targets, delta=1.0): class EnergyStdLoss(TaskLoss): def __init__( self, - starter_learning_rate=1.0, - start_pref_e=0.0, - limit_pref_e=0.0, - start_pref_f=0.0, - limit_pref_f=0.0, - start_pref_v=0.0, - limit_pref_v=0.0, + starter_learning_rate: float = 1.0, + start_pref_e: float = 0.0, + limit_pref_e: float = 0.0, + start_pref_f: float = 0.0, + limit_pref_f: float = 0.0, + start_pref_v: float = 0.0, + limit_pref_v: float = 0.0, start_pref_ae: float = 0.0, limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, numb_generalized_coord: int = 0, use_l1_all: bool = False, - inference=False, - use_huber=False, - huber_delta=0.01, - **kwargs, + inference: bool = False, + use_huber: bool = False, + huber_delta: float = 0.01, + **kwargs: Any, ) -> None: r"""Construct a layer to compute loss on energy, force and virial. @@ -149,7 +151,15 @@ def __init__( "Huber loss is not implemented for force with atom_pref, generalized force and relative force. " ) - def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: """Return loss on energy and force. Parameters @@ -528,10 +538,10 @@ def deserialize(cls, data: dict) -> "TaskLoss": class EnergyHessianStdLoss(EnergyStdLoss): def __init__( self, - start_pref_h=0.0, - limit_pref_h=0.0, - **kwargs, - ): + start_pref_h: float = 0.0, + limit_pref_h: float = 0.0, + **kwargs: Any, + ) -> None: r"""Enable the layer to compute loss on hessian. Parameters @@ -549,7 +559,15 @@ def __init__( self.start_pref_h = start_pref_h self.limit_pref_h = limit_pref_h - def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: model_pred, loss, more_loss = super().forward( input_dict, model, label, natoms, learning_rate, mae=mae ) diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 6a926f4051..9b87d4234f 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import torch import torch.nn.functional as F @@ -20,21 +23,21 @@ class EnergySpinLoss(TaskLoss): def __init__( self, - starter_learning_rate=1.0, - start_pref_e=0.0, - limit_pref_e=0.0, - start_pref_fr=0.0, - limit_pref_fr=0.0, - start_pref_fm=0.0, - limit_pref_fm=0.0, - start_pref_v=0.0, - limit_pref_v=0.0, + starter_learning_rate: float = 1.0, + start_pref_e: float = 0.0, + limit_pref_e: float = 0.0, + start_pref_fr: float = 0.0, + limit_pref_fr: float = 0.0, + start_pref_fm: float = 0.0, + limit_pref_fm: float = 0.0, + start_pref_v: float = 0.0, + limit_pref_v: float = 0.0, start_pref_ae: float = 0.0, limit_pref_ae: float = 0.0, enable_atom_ener_coeff: bool = False, use_l1_all: bool = False, - inference=False, - **kwargs, + inference: bool = False, + **kwargs: Any, ) -> None: r"""Construct a layer to compute loss on energy, real force, magnetic force and virial. @@ -93,7 +96,15 @@ def __init__( self.use_l1_all = use_l1_all self.inference = inference - def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: """Return energy loss with magnetic labels. Parameters diff --git a/deepmd/pt/loss/loss.py b/deepmd/pt/loss/loss.py index d1777a29b3..51eb7a9319 100644 --- a/deepmd/pt/loss/loss.py +++ b/deepmd/pt/loss/loss.py @@ -4,6 +4,7 @@ abstractmethod, ) from typing import ( + Any, NoReturn, ) @@ -18,11 +19,18 @@ class TaskLoss(torch.nn.Module, ABC, make_plugin_registry("loss")): - def __init__(self, **kwargs) -> None: + def __init__(self, **kwargs: Any) -> None: """Construct loss.""" super().__init__() - def forward(self, input_dict, model, label, natoms, learning_rate) -> NoReturn: + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float | torch.Tensor, + ) -> NoReturn: """Return loss .""" raise NotImplementedError diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 9d42c81b45..189bcb2a4a 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Union, + Any, ) import torch @@ -23,15 +23,15 @@ class PropertyLoss(TaskLoss): def __init__( self, - task_dim, + task_dim: int, var_name: str, loss_func: str = "smooth_mae", - metric: list = ["mae"], + metric: list[str] = ["mae"], beta: float = 1.00, - out_bias: Union[list, None] = None, - out_std: Union[list, None] = None, + out_bias: list | None = None, + out_std: list | None = None, intensive: bool = False, - **kwargs, + **kwargs: Any, ) -> None: r"""Construct a layer to compute loss on property. @@ -42,7 +42,7 @@ def __init__( var_name : str The atomic property to fit, 'energy', 'dipole', and 'polar'. loss_func : str - The loss function, such as "smooth_mae", "mae", "rmse". + The loss function, such as "smooth_mae", "mae", "rmse", "mape". metric : list The metric such as mae, rmse which will be printed. beta : float @@ -66,7 +66,15 @@ def __init__( self.intensive = intensive self.var_name = var_name - def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False): + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float = 0.0, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: """Return loss on properties . Parameters @@ -151,6 +159,12 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False reduction="mean", ) ) + elif self.loss_func == "mape": + loss += torch.mean( + torch.abs( + (label[var_name] - model_pred[var_name]) / (label[var_name] + 1e-3) + ) + ) else: raise RuntimeError(f"Unknown loss function : {self.loss_func}") @@ -182,6 +196,12 @@ def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False reduction="mean", ) ).detach() + if "mape" in self.metric: + more_loss["mape"] = torch.mean( + torch.abs( + (label[var_name] - model_pred[var_name]) / (label[var_name] + 1e-3) + ) + ).detach() return model_pred, loss, more_loss diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 0acc3989be..625a9b30bc 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import torch @@ -21,9 +24,9 @@ def __init__( label_name: str, pref_atomic: float = 0.0, pref: float = 0.0, - inference=False, + inference: bool = False, enable_atomic_weight: bool = False, - **kwargs, + **kwargs: Any, ) -> None: r"""Construct a loss for local and global tensors. @@ -64,7 +67,15 @@ def __init__( "Can not assian zero weight both to `pref` and `pref_atomic`" ) - def forward(self, input_dict, model, label, natoms, learning_rate=0.0, mae=False): + def forward( + self, + input_dict: dict[str, torch.Tensor], + model: torch.nn.Module, + label: dict[str, torch.Tensor], + natoms: int, + learning_rate: float = 0.0, + mae: bool = False, + ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]: """Return loss on local and global tensors. Parameters diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 56af5f4f43..0ccf539757 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -1,11 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( NoReturn, Optional, - Union, ) import numpy as np @@ -77,8 +78,8 @@ def __init__( type_map: list[str], atom_exclude_types: list[int] = [], pair_exclude_types: list[tuple[int, int]] = [], - rcond: Optional[float] = None, - preset_out_bias: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_out_bias: dict[str, np.ndarray] | None = None, data_stat_protect: float = 1e-2, ) -> None: torch.nn.Module.__init__(self) @@ -106,7 +107,7 @@ def init_out_stat(self) -> None: def set_out_bias(self, out_bias: torch.Tensor) -> None: self.out_bias = out_bias - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: torch.Tensor) -> None: if key in ["out_bias"]: self.out_bias = value elif key in ["out_std"]: @@ -114,7 +115,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> torch.Tensor: if key in ["out_bias"]: return self.out_bias elif key in ["out_std"]: @@ -135,6 +136,10 @@ def get_intensive(self) -> bool: """Whether the fitting property is intensive.""" return False + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return False + def reinit_atom_exclude( self, exclude_types: list[int] = [], @@ -199,10 +204,10 @@ def forward_common_atomic( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: """Common interface for atomic inference. @@ -280,10 +285,10 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: return self.forward_common_atomic( extended_coord, @@ -296,7 +301,9 @@ def forward( ) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["BaseAtomicModel"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -361,31 +368,35 @@ def deserialize(cls, data: dict) -> "BaseAtomicModel": def compute_or_load_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, ) -> NoReturn: """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + Compute or load the statistics parameters of the model, + such as mean and standard deviation of descriptors or the energy bias of the fitting net. + When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update), + and saved in the `stat_file_path`(s). + When `sampled` is not provided, it will check the existence of `stat_file_path`(s) + and load the calculated statistics parameters. Parameters ---------- - merged : Union[Callable[[], list[dict]], list[dict]] - - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` - originating from the `i`-th data system. - - Callable[[], list[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. + merged + The lazy sampled function to get data frames from different data systems. + stat_file_path + The dictionary of paths to the statistics files. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ raise NotImplementedError def compute_or_load_out_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -413,7 +424,7 @@ def apply_out_stat( self, ret: dict[str, torch.Tensor], atype: torch.Tensor, - ): + ) -> dict[str, torch.Tensor]: """Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. @@ -434,9 +445,9 @@ def apply_out_stat( def change_out_bias( self, - sample_merged, - stat_file_path: Optional[DPPath] = None, - bias_adjust_mode="change-by-statistic", + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + bias_adjust_mode: str = "change-by-statistic", ) -> None: """Change the output bias according to the input data and the pretrained model. @@ -483,10 +494,34 @@ def change_out_bias( else: raise RuntimeError("Unknown bias_adjust_mode mode: " + bias_adjust_mode) + def compute_fitting_input_stat( + self, + sample_merged: Callable[[], list[dict]] | list[dict], + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the atomic model from packed data. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + """ + pass + def _get_forward_wrapper_func(self) -> Callable[..., torch.Tensor]: """Get a forward wrapper of the atomic model for output bias calculation.""" - def model_forward(coord, atype, box, fparam=None, aparam=None): + def model_forward( + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: with ( torch.no_grad() ): # it's essential for pure torch forward function to use auto_batchsize @@ -515,13 +550,13 @@ def model_forward(coord, atype, box, fparam=None, aparam=None): return model_forward - def _default_bias(self): + def _default_bias(self) -> torch.Tensor: ntypes = self.get_ntypes() return torch.zeros( [self.n_out, ntypes, self.max_out_size], dtype=dtype, device=device ) - def _default_std(self): + def _default_std(self) -> torch.Tensor: ntypes = self.get_ntypes() return torch.ones( [self.n_out, ntypes, self.max_out_size], dtype=dtype, device=device diff --git a/deepmd/pt/model/atomic_model/dipole_atomic_model.py b/deepmd/pt/model/atomic_model/dipole_atomic_model.py index 3796aa2e83..c9badefcad 100644 --- a/deepmd/pt/model/atomic_model/dipole_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dipole_atomic_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import torch @@ -12,7 +15,9 @@ class DPDipoleAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: Any, **kwargs: Any + ) -> None: if not isinstance(fitting, DipoleFittingNet): raise TypeError( "fitting must be an instance of DipoleFittingNet for DPDipoleAtomicModel" @@ -23,6 +28,6 @@ def apply_out_stat( self, ret: dict[str, torch.Tensor], atype: torch.Tensor, - ): + ) -> dict[str, torch.Tensor]: # dipole not applying bias return ret diff --git a/deepmd/pt/model/atomic_model/dos_atomic_model.py b/deepmd/pt/model/atomic_model/dos_atomic_model.py index 2af1a4e052..7bc0108fc5 100644 --- a/deepmd/pt/model/atomic_model/dos_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dos_atomic_model.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + from deepmd.pt.model.task.dos import ( DOSFittingNet, ) @@ -9,7 +13,9 @@ class DPDOSAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: Any, **kwargs: Any + ) -> None: if not isinstance(fitting, DOSFittingNet): raise TypeError( "fitting must be an instance of DOSFittingNet for DPDOSAtomicModel" diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 5a5655b72c..af2e8954df 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -1,7 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools import logging +from collections.abc import ( + Callable, +) from typing import ( + Any, Optional, ) @@ -47,10 +51,10 @@ class DPAtomicModel(BaseAtomicModel): def __init__( self, - descriptor, - fitting, + descriptor: BaseDescriptor, + fitting: BaseFitting, type_map: list[str], - **kwargs, + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) ntypes = len(type_map) @@ -60,11 +64,16 @@ def __init__( self.rcut = self.descriptor.get_rcut() self.sel = self.descriptor.get_sel() self.fitting_net = fitting + if hasattr(self.fitting_net, "reinit_exclude"): + self.fitting_net.reinit_exclude(self.atom_exclude_types) super().init_out_stat() self.enable_eval_descriptor_hook = False + self.enable_eval_fitting_last_layer_hook = False self.eval_descriptor_list = [] + self.eval_fitting_last_layer_list = [] eval_descriptor_list: list[torch.Tensor] + eval_fitting_last_layer_list: list[torch.Tensor] def set_eval_descriptor_hook(self, enable: bool) -> None: """Set the hook for evaluating descriptor and clear the cache for descriptor list.""" @@ -76,6 +85,17 @@ def eval_descriptor(self) -> torch.Tensor: """Evaluate the descriptor.""" return torch.concat(self.eval_descriptor_list) + def set_eval_fitting_last_layer_hook(self, enable: bool) -> None: + """Set the hook for evaluating fitting last layer output and clear the cache for fitting last layer output list.""" + self.enable_eval_fitting_last_layer_hook = enable + self.fitting_net.set_return_middle_output(enable) + # = [] does not work; See #4533 + self.eval_fitting_last_layer_list.clear() + + def eval_fitting_last_layer(self) -> torch.Tensor: + """Evaluate the fitting last layer output.""" + return torch.concat(self.eval_fitting_last_layer_list) + @torch.jit.export def fitting_output_def(self) -> FittingOutputDef: """Get the output def of the fitting net.""" @@ -94,7 +114,7 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.sel - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -114,7 +134,9 @@ def mixed_types(self) -> bool: return self.descriptor.mixed_types() def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["DPAtomicModel"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -131,6 +153,9 @@ def change_type_map( else None, ) self.fitting_net.change_type_map(type_map=type_map) + # Reinitialize fitting to get correct sel_type + if hasattr(self.fitting_net, "reinit_exclude"): + self.fitting_net.reinit_exclude(self.atom_exclude_types) def has_message_passing(self) -> bool: """Returns whether the atomic model has message passing.""" @@ -155,7 +180,7 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "DPAtomicModel": + def deserialize(cls, data: dict) -> "DPAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 1) data.pop("@class", None) @@ -200,13 +225,13 @@ def enable_compression( def forward_atomic( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: """Return atomic prediction. @@ -255,6 +280,13 @@ def forward_atomic( fparam=fparam, aparam=aparam, ) + if self.enable_eval_fitting_last_layer_hook: + assert "middle_output" in fit_ret, ( + "eval_fitting_last_layer not supported for this fitting net!" + ) + self.eval_fitting_last_layer_list.append( + fit_ret.pop("middle_output").detach() + ) return fit_ret def get_out_bias(self) -> torch.Tensor: @@ -262,8 +294,9 @@ def get_out_bias(self) -> torch.Tensor: def compute_or_load_stat( self, - sampled_func, - stat_file_path: Optional[DPPath] = None, + sampled_func: Callable[[], list[dict]], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, ) -> None: """ Compute or load the statistics parameters of the model, @@ -279,6 +312,9 @@ def compute_or_load_stat( The lazy sampled function to get data frames from different data systems. stat_file_path The dictionary of paths to the statistics files. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ if stat_file_path is not None and self.type_map is not None: # descriptors and fitting net with different type_map @@ -286,7 +322,7 @@ def compute_or_load_stat( stat_file_path /= " ".join(self.type_map) @functools.lru_cache - def wrapped_sampler(): + def wrapped_sampler() -> list[dict]: sampled = sampled_func() if self.pair_excl is not None: pair_exclude_types = self.pair_excl.get_exclude_types() @@ -296,18 +332,58 @@ def wrapped_sampler(): atom_exclude_types = self.atom_excl.get_exclude_types() for sample in sampled: sample["atom_exclude_types"] = list(atom_exclude_types) + if ( + "find_fparam" not in sampled[0] + and "fparam" not in sampled[0] + and self.has_default_fparam() + ): + default_fparam = self.get_default_fparam() + for sample in sampled: + nframe = sample["atype"].shape[0] + sample["fparam"] = default_fparam.repeat(nframe, 1) return sampled self.descriptor.compute_input_stats(wrapped_sampler, stat_file_path) + self.compute_fitting_input_stat(wrapped_sampler, stat_file_path) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) + + def compute_fitting_input_stat( + self, + sample_merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + ) -> None: + """Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. + + Parameters + ---------- + sample_merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The dictionary of paths to the statistics files. + """ self.fitting_net.compute_input_stats( - wrapped_sampler, protection=self.data_stat_protect + sample_merged, + protection=self.data_stat_protect, + stat_file_path=stat_file_path, ) - self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.fitting_net.get_dim_fparam() + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.fitting_net.has_default_fparam() + + def get_default_fparam(self) -> torch.Tensor | None: + return self.fitting_net.get_default_fparam() + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.fitting_net.get_dim_aparam() diff --git a/deepmd/pt/model/atomic_model/energy_atomic_model.py b/deepmd/pt/model/atomic_model/energy_atomic_model.py index 6d894b4aab..9f513fc53d 100644 --- a/deepmd/pt/model/atomic_model/energy_atomic_model.py +++ b/deepmd/pt/model/atomic_model/energy_atomic_model.py @@ -1,4 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + from deepmd.pt.model.task.ener import ( EnergyFittingNet, EnergyFittingNetDirect, @@ -11,7 +15,9 @@ class DPEnergyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: Any, **kwargs: Any + ) -> None: if not ( isinstance(fitting, EnergyFittingNet) or isinstance(fitting, EnergyFittingNetDirect) diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 3d894dc3a0..96b3baf6ec 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +import functools +from collections.abc import ( Callable, +) +from typing import ( + Any, Optional, - Union, ) import torch @@ -55,8 +58,8 @@ def __init__( self, models: list[BaseAtomicModel], type_map: list[str], - weights: Optional[Union[str, list[float]]] = "mean", - **kwargs, + weights: str | list[float] | None = "mean", + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() @@ -135,7 +138,9 @@ def get_type_map(self) -> list[str]: return self.type_map def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["LinearEnergyAtomicModel"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -158,7 +163,7 @@ def get_model_rcuts(self) -> list[float]: def get_sel(self) -> list[int]: return [max([model.get_nsel() for model in self.models])] - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -228,10 +233,10 @@ def forward_atomic( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: """Return atomic prediction. @@ -290,6 +295,7 @@ def forward_atomic( mapping, fparam, aparam, + comm_dict=comm_dict, )["energy"] ) weights = self._compute_weight(extended_coord, extended_atype, nlists_) @@ -306,7 +312,7 @@ def apply_out_stat( self, ret: dict[str, torch.Tensor], atype: torch.Tensor, - ): + ) -> dict[str, torch.Tensor]: """Apply the stat to each atomic output. The developer may override the method to define how the bias is applied to the atomic output of the model. @@ -319,6 +325,10 @@ def apply_out_stat( The atom types. nf x nloc """ + out_bias, out_std = self._fetch_out_stat(self.bias_keys) + for kk in self.bias_keys: + # nf x nloc x odims, out_bias: ntypes x odims + ret[kk] = ret[kk] + out_bias[kk][atype] return ret @staticmethod @@ -464,34 +474,11 @@ def is_aparam_nall(self) -> bool: """ return False - def compute_or_load_out_stat( - self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, - ) -> None: - """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. - - Parameters - ---------- - merged : Union[Callable[[], list[dict]], list[dict]] - - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` - originating from the `i`-th data system. - - Callable[[], list[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. - - """ - for md in self.models: - md.compute_or_load_out_stat(merged, stat_file_path) - def compute_or_load_stat( self, - sampled_func, - stat_file_path: Optional[DPPath] = None, + sampled_func: Callable[[], list[dict[str, Any]]], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, ) -> None: """ Compute or load the statistics parameters of the model, @@ -507,9 +494,34 @@ def compute_or_load_stat( The lazy sampled function to get data frames from different data systems. stat_file_path The dictionary of paths to the statistics files. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ for md in self.models: - md.compute_or_load_stat(sampled_func, stat_file_path) + md.compute_or_load_stat( + sampled_func, stat_file_path, compute_or_load_out_stat=False + ) + + if stat_file_path is not None and self.type_map is not None: + # descriptors and fitting net with different type_map + # should not share the same parameters + stat_file_path /= " ".join(self.type_map) + + @functools.lru_cache + def wrapped_sampler() -> list[dict[str, Any]]: + sampled = sampled_func() + if self.pair_excl is not None: + pair_exclude_types = self.pair_excl.get_exclude_types() + for sample in sampled: + sample["pair_exclude_types"] = list(pair_exclude_types) + if self.atom_excl is not None: + atom_exclude_types = self.atom_excl.get_exclude_types() + for sample in sampled: + sample["atom_exclude_types"] = list(atom_exclude_types) + return sampled + + self.compute_or_load_out_stat(wrapped_sampler, stat_file_path) class DPZBLLinearEnergyAtomicModel(LinearEnergyAtomicModel): @@ -540,8 +552,8 @@ def __init__( sw_rmin: float, sw_rmax: float, type_map: list[str], - smin_alpha: Optional[float] = 0.1, - **kwargs, + smin_alpha: float | None = 0.1, + **kwargs: Any, ) -> None: models = [dp_model, zbl_model] kwargs["models"] = models @@ -569,7 +581,7 @@ def serialize(self) -> dict: ) return dd - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -578,7 +590,7 @@ def set_case_embd(self, case_idx: int): self.models[0].set_case_embd(case_idx) @classmethod - def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": + def deserialize(cls, data: dict[str, Any]) -> "DPZBLLinearEnergyAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 1) models = [ diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 62b47afb32..6933fdc19a 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, Optional, - Union, ) import torch @@ -66,9 +68,9 @@ def __init__( self, tab_file: str, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], type_map: list[str], - **kwargs, + **kwargs: Any, ) -> None: super().__init__(type_map, **kwargs) super().init_out_stat() @@ -141,7 +143,7 @@ def get_type_map(self) -> list[str]: def get_sel(self) -> list[int]: return [self.sel] - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -175,7 +177,9 @@ def need_sorted_nlist_for_lower(self) -> bool: return False def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["PairTabAtomicModel"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -202,7 +206,7 @@ def serialize(self) -> dict: return dd @classmethod - def deserialize(cls, data) -> "PairTabAtomicModel": + def deserialize(cls, data: dict[str, Any]) -> "PairTabAtomicModel": data = data.copy() check_version_compatibility(data.pop("@version", 1), 2, 1) tab = PairTab.deserialize(data.pop("tab")) @@ -224,37 +228,42 @@ def deserialize(cls, data) -> "PairTabAtomicModel": def compute_or_load_stat( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + sampled_func: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, + compute_or_load_out_stat: bool = True, ) -> None: """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + Compute or load the statistics parameters of the model, + such as mean and standard deviation of descriptors or the energy bias of the fitting net. + When `sampled` is provided, all the statistics parameters will be calculated (or re-calculated for update), + and saved in the `stat_file_path`(s). + When `sampled` is not provided, it will check the existence of `stat_file_path`(s) + and load the calculated statistics parameters. Parameters ---------- - merged : Union[Callable[[], list[dict]], list[dict]] - - list[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` - originating from the `i`-th data system. - - Callable[[], list[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. + sampled_func + The lazy sampled function to get data frames from different data systems. + stat_file_path + The dictionary of paths to the statistics files. + compute_or_load_out_stat : bool + Whether to compute the output statistics. + If False, it will only compute the input statistics (e.g. mean and standard deviation of descriptors). """ - self.compute_or_load_out_stat(merged, stat_file_path) + if compute_or_load_out_stat: + self.compute_or_load_out_stat(sampled_func, stat_file_path) def forward_atomic( self, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, ) -> dict[str, torch.Tensor]: nframes, nloc, nnei = nlist.shape extended_coord = extended_coord.view(nframes, -1, 3) diff --git a/deepmd/pt/model/atomic_model/polar_atomic_model.py b/deepmd/pt/model/atomic_model/polar_atomic_model.py index 6bd063591f..4484d1945b 100644 --- a/deepmd/pt/model/atomic_model/polar_atomic_model.py +++ b/deepmd/pt/model/atomic_model/polar_atomic_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import torch @@ -12,7 +15,9 @@ class DPPolarAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: Any, **kwargs: Any + ) -> None: if not isinstance(fitting, PolarFittingNet): raise TypeError( "fitting must be an instance of PolarFittingNet for DPPolarAtomicModel" @@ -23,7 +28,7 @@ def apply_out_stat( self, ret: dict[str, torch.Tensor], atype: torch.Tensor, - ): + ) -> dict[str, torch.Tensor]: """Apply the stat to each atomic output. Parameters diff --git a/deepmd/pt/model/atomic_model/property_atomic_model.py b/deepmd/pt/model/atomic_model/property_atomic_model.py index 3622c9f476..baf9c5b7fc 100644 --- a/deepmd/pt/model/atomic_model/property_atomic_model.py +++ b/deepmd/pt/model/atomic_model/property_atomic_model.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) import torch @@ -12,7 +15,9 @@ class DPPropertyAtomicModel(DPAtomicModel): - def __init__(self, descriptor, fitting, type_map, **kwargs): + def __init__( + self, descriptor: Any, fitting: Any, type_map: Any, **kwargs: Any + ) -> None: if not isinstance(fitting, PropertyFittingNet): raise TypeError( "fitting must be an instance of PropertyFittingNet for DPPropertyAtomicModel" @@ -31,7 +36,7 @@ def apply_out_stat( self, ret: dict[str, torch.Tensor], atype: torch.Tensor, - ): + ) -> dict[str, torch.Tensor]: """Apply the stat to each atomic output. In property fitting, each output will be multiplied by label std and then plus the label average value. diff --git a/deepmd/pt/model/descriptor/descriptor.py b/deepmd/pt/model/descriptor/descriptor.py index 3b374751c7..2a6e857a6f 100644 --- a/deepmd/pt/model/descriptor/descriptor.py +++ b/deepmd/pt/model/descriptor/descriptor.py @@ -4,11 +4,12 @@ ABC, abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, NoReturn, - Optional, - Union, ) import torch @@ -43,7 +44,7 @@ class DescriptorBlock(torch.nn.Module, ABC, make_plugin_registry("DescriptorBloc local_cluster = False - def __new__(cls, *args, **kwargs): + def __new__(cls, *args: Any, **kwargs: Any) -> "DescriptorBlock": if cls is DescriptorBlock: try: descrpt_type = kwargs["type"] @@ -101,8 +102,8 @@ def get_env_protection(self) -> float: def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -126,7 +127,9 @@ def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" raise NotImplementedError - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: "DescriptorBlock", shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -175,10 +178,16 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Calculate DescriptorBlock.""" pass @@ -192,14 +201,18 @@ def need_sorted_nlist_for_lower(self) -> bool: def make_default_type_embedding( - ntypes, -): + ntypes: int, +) -> tuple[TypeEmbedNet, dict[str, Any]]: aux = {} aux["tebd_dim"] = 8 return TypeEmbedNet(ntypes, aux["tebd_dim"]), aux -def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None: +def extend_descrpt_stat( + des: DescriptorBlock, + type_map: list[str], + des_with_stat: DescriptorBlock | None = None, +) -> None: r""" Extend the statistics of a descriptor block with types from newly provided `type_map`. diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 9c1e144f48..359bf2f084 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +import warnings +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -219,7 +221,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -236,22 +238,22 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, scaling_factor: int = 1.0, - normalize=True, - temperature=None, + normalize: bool = True, + temperature: float | None = None, concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, smooth_type_embedding: bool = True, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented - spin=None, - type: Optional[str] = None, + spin: Any | None = None, + type: str | None = None, ) -> None: super().__init__() # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -298,11 +300,13 @@ def __init__( trainable_ln=trainable_ln, ln_eps=ln_eps, seed=child_seed(seed, 1), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd self.use_tebd_bias = use_tebd_bias self.type_map = type_map - self.compress = False + self.tebd_compress = False + self.geo_compress = False self.type_embedding = TypeEmbedNet( ntypes, tebd_dim, @@ -311,6 +315,7 @@ def __init__( use_econf_tebd=use_econf_tebd, use_tebd_bias=use_tebd_bias, type_map=type_map, + trainable=trainable, ) self.prec = PRECISION_DICT[precision] self.tebd_dim = tebd_dim @@ -378,7 +383,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_atten.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -402,18 +409,18 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -446,7 +453,7 @@ def get_stat_mean_and_stddev(self) -> tuple[torch.Tensor, torch.Tensor]: return self.se_atten.mean, self.se_atten.stddev def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -546,7 +553,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": data["use_tebd_bias"] = True obj = cls(**data) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.se_atten.prec, device=env.DEVICE) obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize( @@ -587,12 +594,17 @@ def enable_compression( check_frequency The overflow check frequency """ - # do some checks before the mocel compression process - if self.compress: + # do some checks before the model compression process + if self.tebd_compress or self.geo_compress: raise ValueError("Compression is already enabled.") + + if self.tebd_input_mode != "strip": + raise RuntimeError("Type embedding compression only works in strip mode") + assert not self.se_atten.resnet_dt, ( "Model compression error: descriptor resnet_dt must be false!" ) + for tt in self.se_atten.exclude_types: if (tt[0] not in range(self.se_atten.ntypes)) or ( tt[1] not in range(self.se_atten.ntypes) @@ -604,6 +616,7 @@ def enable_compression( + str(self.se_atten.ntypes) + "!" ) + if ( self.se_atten.ntypes * self.se_atten.ntypes - len(self.se_atten.exclude_types) @@ -613,43 +626,53 @@ def enable_compression( "Empty embedding-nets are not supported in model compression!" ) - if self.se_atten.attn_layer != 0: - raise RuntimeError("Cannot compress model when attention layer is not 0.") - - if self.tebd_input_mode != "strip": - raise RuntimeError("Cannot compress model when tebd_input_mode == 'concat'") - - data = self.serialize() - self.table = DPTabulate( - self, - data["neuron"], - data["type_one_side"], - data["exclude_types"], - ActivationFn(data["activation_function"]), - ) - self.table_config = [ - table_extrapolate, - table_stride_1, - table_stride_2, - check_frequency, - ] - self.lower, self.upper = self.table.build( - min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2 - ) + # Enable type embedding compression + self.se_atten.type_embedding_compression(self.type_embedding) + self.tebd_compress = True + + if self.se_atten.attn_layer == 0: + data = self.serialize() + self.table = DPTabulate( + self, + data["neuron"], + data["type_one_side"], + data["exclude_types"], + ActivationFn(data["activation_function"]), + ) + self.table_config = [ + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ] + self.lower, self.upper = self.table.build( + min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2 + ) - self.se_atten.enable_compression( - self.table.data, self.table_config, self.lower, self.upper - ) - self.compress = True + self.se_atten.enable_compression( + self.table.data, self.table_config, self.lower, self.upper + ) + self.geo_compress = True + else: + warnings.warn( + "Attention layer is not 0, only type embedding is compressed. Geometric part is not compressed.", + UserWarning, + ) def forward( self, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -706,19 +729,21 @@ def forward( return ( g1.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) + if rot_mat is not None + else None, g2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if g2 is not None else None, - h2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + h2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if h2 is not None else None, + sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if sw is not None else None, ) @classmethod def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 969fdca5fc..30d1987b97 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -1,8 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +import warnings +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -84,9 +86,9 @@ def __init__( self, ntypes: int, # args for repinit - repinit: Union[RepinitArgs, dict], + repinit: RepinitArgs | dict, # args for repformer - repformer: Union[RepformerArgs, dict], + repformer: RepformerArgs | dict, # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -94,11 +96,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, add_tebd_to_repinit_out: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: r"""The DPA-2 descriptor[1]_. @@ -155,7 +157,7 @@ def __init__( """ super().__init__() - def init_subclass_params(sub_data, sub_class): + def init_subclass_params(sub_data: Any, sub_class: Any) -> Any: if isinstance(sub_data, dict): return sub_class(**sub_data) elif isinstance(sub_data, sub_class): @@ -188,6 +190,7 @@ def init_subclass_params(sub_data, sub_class): smooth=smooth, type_one_side=self.repinit_args.type_one_side, seed=child_seed(seed, 0), + trainable=trainable, ) self.use_three_body = self.repinit_args.use_three_body if self.use_three_body: @@ -207,6 +210,7 @@ def init_subclass_params(sub_data, sub_class): resnet_dt=self.repinit_args.resnet_dt, smooth=smooth, seed=child_seed(seed, 5), + trainable=trainable, ) else: self.repinit_three_body = None @@ -247,6 +251,7 @@ def init_subclass_params(sub_data, sub_class): g1_out_conv=self.repformer_args.g1_out_conv, g1_out_mlp=self.repformer_args.g1_out_mlp, seed=child_seed(seed, 1), + trainable=trainable, ) self.rcsl_list = [ (self.repformers.get_rcut(), self.repformers.get_nsel()), @@ -274,6 +279,7 @@ def init_subclass_params(sub_data, sub_class): use_econf_tebd=self.use_econf_tebd, use_tebd_bias=use_tebd_bias, type_map=type_map, + trainable=trainable, ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -299,6 +305,7 @@ def init_subclass_params(sub_data, sub_class): precision=precision, init="glorot", seed=child_seed(seed, 3), + trainable=trainable, ) self.tebd_transform = None if self.add_tebd_to_repinit_out: @@ -308,6 +315,7 @@ def init_subclass_params(sub_data, sub_class): bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) assert self.repinit.rcut > self.repformers.rcut assert self.repinit.sel[0] > self.repformers.sel[0] @@ -384,7 +392,9 @@ def get_env_protection(self) -> float: # the env_protection of repinit is the same as that of the repformer return self.repinit.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -416,7 +426,7 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -471,18 +481,18 @@ def change_type_map( repinit_three_body["dstd"] = repinit_three_body["dstd"][remap_index] @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -650,7 +660,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": if obj.repinit.dim_out != obj.repformers.dim_in: obj.g1_shape_tranform = MLPLayer.deserialize(g1_shape_tranform) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.repinit.prec, device=env.DEVICE) # deserialize repinit @@ -703,9 +713,15 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -814,19 +830,21 @@ def forward( g1 = torch.cat([g1, g1_inp], dim=-1) return ( g1.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - g2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - h2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) + if rot_mat is not None + else None, + g2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if g2 is not None else None, + h2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if h2 is not None else None, + sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if sw is not None else None, ) @classmethod def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -921,36 +939,41 @@ def enable_compression( "Repinit empty embedding-nets are not supported in model compression!" ) - if self.repinit.attn_layer != 0: + if self.repinit.tebd_input_mode != "strip": raise RuntimeError( - "Cannot compress model when repinit attention layer is not 0." + "Cannot compress model when repinit tebd_input_mode != 'strip'" ) - if self.repinit.tebd_input_mode != "strip": - raise RuntimeError( - "Cannot compress model when repinit tebd_input_mode == 'concat'" + if self.repinit.attn_layer == 0: + # repinit doesn't have a serialize method + data = self.serialize() + self.table = DPTabulate( + self, + data["repinit_args"]["neuron"], + data["repinit_args"]["type_one_side"], + data["exclude_types"], + ActivationFn(data["repinit_args"]["activation_function"]), + ) + self.table_config = [ + table_extrapolate, + table_stride_1, + table_stride_2, + check_frequency, + ] + self.lower, self.upper = self.table.build( + min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2 ) - # repinit doesn't have a serialize method - data = self.serialize() - self.table = DPTabulate( - self, - data["repinit_args"]["neuron"], - data["repinit_args"]["type_one_side"], - data["exclude_types"], - ActivationFn(data["repinit_args"]["activation_function"]), - ) - self.table_config = [ - table_extrapolate, - table_stride_1, - table_stride_2, - check_frequency, - ] - self.lower, self.upper = self.table.build( - min_nbor_dist, table_extrapolate, table_stride_1, table_stride_2 - ) + self.repinit.enable_compression( + self.table.data, self.table_config, self.lower, self.upper + ) + else: + warnings.warn( + "Attention layer is not 0, only type embedding is compressed. Geometric part is not compressed.", + UserWarning, + ) + + # Enable type embedding compression for repinit (se_atten) + self.repinit.type_embedding_compression(self.type_embedding) - self.repinit.enable_compression( - self.table.data, self.table_config, self.lower, self.upper - ) self.compress = True diff --git a/deepmd/pt/model/descriptor/dpa3.py b/deepmd/pt/model/descriptor/dpa3.py index dd2da9a3c8..136527123e 100644 --- a/deepmd/pt/model/descriptor/dpa3.py +++ b/deepmd/pt/model/descriptor/dpa3.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -106,7 +107,7 @@ def __init__( self, ntypes: int, # args for repflow - repflow: Union[RepFlowArgs, dict], + repflow: RepFlowArgs | dict, # kwargs for descriptor concat_output_tebd: bool = False, activation_function: str = "silu", @@ -114,15 +115,15 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, use_loc_mapping: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__() - def init_subclass_params(sub_data, sub_class): + def init_subclass_params(sub_data: Any, sub_class: Any) -> Any: if isinstance(sub_data, dict): return sub_class(**sub_data) elif isinstance(sub_data, sub_class): @@ -169,6 +170,7 @@ def init_subclass_params(sub_data, sub_class): env_protection=env_protection, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.use_econf_tebd = use_econf_tebd @@ -184,6 +186,7 @@ def init_subclass_params(sub_data, sub_class): use_econf_tebd=self.use_econf_tebd, use_tebd_bias=use_tebd_bias, type_map=type_map, + trainable=trainable, ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -270,7 +273,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.repflows.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -294,7 +299,7 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -323,18 +328,18 @@ def change_type_map( repflow["dstd"] = repflow["dstd"][remap_index] @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -425,7 +430,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA3": type_embedding ) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.repflows.prec, device=env.DEVICE) # deserialize repflow @@ -448,9 +453,15 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -507,19 +518,23 @@ def forward( node_ebd = torch.cat([node_ebd, node_ebd_inp], dim=-1) return ( node_ebd.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - edge_ebd.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - h2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) + if rot_mat is not None + else None, + edge_ebd.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) + if edge_ebd is not None + else None, + h2.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if h2 is not None else None, + sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if sw is not None else None, ) @classmethod def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/env_mat.py b/deepmd/pt/model/descriptor/env_mat.py index c57ae209fd..0ffdbb7dbb 100644 --- a/deepmd/pt/model/descriptor/env_mat.py +++ b/deepmd/pt/model/descriptor/env_mat.py @@ -9,14 +9,14 @@ def _make_env_mat( - nlist, - coord, + nlist: torch.Tensor, + coord: torch.Tensor, rcut: float, ruct_smth: float, radial_only: bool = False, protection: float = 0.0, use_exp_switch: bool = False, -): +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Make smooth environment matrix.""" bsz, natoms, nnei = nlist.shape coord = coord.view(bsz, -1, 3) @@ -49,17 +49,17 @@ def _make_env_mat( def prod_env_mat( - extended_coord, - nlist, - atype, - mean, - stddev, + extended_coord: torch.Tensor, + nlist: torch.Tensor, + atype: torch.Tensor, + mean: torch.Tensor, + stddev: torch.Tensor, rcut: float, rcut_smth: float, radial_only: bool = False, protection: float = 0.0, use_exp_switch: bool = False, -): +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Generate smooth environment matrix from atom coordinates and other context. Args: diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index e13b014037..0f001bc4c8 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -3,7 +3,6 @@ from typing import ( Any, Optional, - Union, ) import numpy as np @@ -44,8 +43,8 @@ class DescrptHybrid(BaseDescriptor, torch.nn.Module): def __init__( self, - list: list[Union[BaseDescriptor, dict[str, Any]]], - **kwargs, + list: list[BaseDescriptor | dict[str, Any]], + **kwargs: Any, ) -> None: super().__init__() # warning: list is conflict with built-in list @@ -140,7 +139,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension.""" return sum([descrpt.get_dim_emb() for descrpt in self.descrpt_list]) - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -164,7 +163,9 @@ def get_env_protection(self) -> float: ) return all_protection[0] - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: "DescrptHybrid", shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -182,7 +183,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["DescrptHybrid"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -196,7 +199,7 @@ def change_type_map( ) def compute_input_stats( - self, merged: list[dict], path: Optional[DPPath] = None + self, merged: list[dict], path: DPPath | None = None ) -> None: """Update mean and stddev for descriptor elements.""" for descrpt in self.descrpt_list: @@ -204,8 +207,8 @@ def compute_input_stats( def set_stat_mean_and_stddev( self, - mean: list[Union[torch.Tensor, list[torch.Tensor]]], - stddev: list[Union[torch.Tensor, list[torch.Tensor]]], + mean: list[torch.Tensor | list[torch.Tensor]], + stddev: list[torch.Tensor | list[torch.Tensor]], ) -> None: """Update mean and stddev for descriptor.""" for ii, descrpt in enumerate(self.descrpt_list): @@ -214,8 +217,8 @@ def set_stat_mean_and_stddev( def get_stat_mean_and_stddev( self, ) -> tuple[ - list[Union[torch.Tensor, list[torch.Tensor]]], - list[Union[torch.Tensor, list[torch.Tensor]]], + list[torch.Tensor | list[torch.Tensor]], + list[torch.Tensor | list[torch.Tensor]], ]: """Get mean and stddev for descriptor.""" mean_list = [] @@ -263,9 +266,15 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -299,9 +308,9 @@ def forward( """ out_descriptor = [] out_gr = [] - out_g2: Optional[torch.Tensor] = None - out_h2: Optional[torch.Tensor] = None - out_sw: Optional[torch.Tensor] = None + out_g2: torch.Tensor | None = None + out_h2: torch.Tensor | None = None + out_sw: torch.Tensor | None = None if self.sel_no_mixed_types is not None: nl_distinguish_types = nlist_distinguish_types( nlist, @@ -334,9 +343,9 @@ def forward( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 37d4f07bb4..338f48b060 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import torch import torch.nn as nn @@ -63,7 +59,8 @@ def __init__( update_residual: float = 0.1, update_residual_init: str = "const", precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei @@ -126,6 +123,7 @@ def __init__( n_dim, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) if self.update_style == "res_residual": self.n_residual.append( @@ -135,6 +133,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) ) @@ -145,6 +144,7 @@ def __init__( n_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) if self.update_style == "res_residual": self.n_residual.append( @@ -154,6 +154,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) ) @@ -163,6 +164,7 @@ def __init__( self.n_multi_edge_message * n_dim, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) if self.update_style == "res_residual": for head_index in range(self.n_multi_edge_message): @@ -173,6 +175,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(child_seed(seed, 5), head_index), + trainable=trainable, ) ) @@ -182,6 +185,7 @@ def __init__( e_dim, precision=precision, seed=child_seed(seed, 6), + trainable=trainable, ) if self.update_style == "res_residual": self.e_residual.append( @@ -191,6 +195,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 7), + trainable=trainable, ) ) @@ -219,6 +224,7 @@ def __init__( precision=precision, bias=False, seed=child_seed(seed, 8), + trainable=trainable, ) self.a_compress_e_linear = MLPLayer( self.e_dim, @@ -226,6 +232,7 @@ def __init__( precision=precision, bias=False, seed=child_seed(seed, 9), + trainable=trainable, ) else: self.a_compress_n_linear = None @@ -237,12 +244,14 @@ def __init__( self.e_dim, precision=precision, seed=child_seed(seed, 10), + trainable=trainable, ) self.edge_angle_linear2 = MLPLayer( self.e_dim, self.e_dim, precision=precision, seed=child_seed(seed, 11), + trainable=trainable, ) if self.update_style == "res_residual": self.e_residual.append( @@ -252,6 +261,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 12), + trainable=trainable, ) ) @@ -261,6 +271,7 @@ def __init__( self.a_dim, precision=precision, seed=child_seed(seed, 13), + trainable=trainable, ) if self.update_style == "res_residual": self.a_residual.append( @@ -270,6 +281,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 14), + trainable=trainable, ) ) else: @@ -370,7 +382,7 @@ def _cal_hg_dynamic( # n_edge x e_dim flat_edge_ebd = flat_edge_ebd * flat_sw.unsqueeze(-1) # n_edge x 3 x e_dim - flat_h2g2 = (flat_h2[..., None] * flat_edge_ebd[:, None, :]).reshape( + flat_h2g2 = (flat_h2.unsqueeze(-1) * flat_edge_ebd.unsqueeze(-2)).reshape( -1, 3 * e_dim ) # nf x nloc x 3 x e_dim @@ -694,9 +706,9 @@ def forward( a_nlist: torch.Tensor, # nf x nloc x a_nnei a_nlist_mask: torch.Tensor, # nf x nloc x a_nnei a_sw: torch.Tensor, # switch func, nf x nloc x a_nnei - edge_index: torch.Tensor, # n_edge x 2 - angle_index: torch.Tensor, # n_angle x 3 - ): + edge_index: torch.Tensor, # 2 x n_edge + angle_index: torch.Tensor, # 3 x n_angle + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Parameters ---------- @@ -720,12 +732,12 @@ def forward( Masks of the neighbor list for angle. real nei 1 otherwise 0 a_sw : nf x nloc x a_nnei Switch function for angle. - edge_index : Optional for dynamic sel, n_edge x 2 + edge_index : Optional for dynamic sel, 2 x n_edge n2e_index : n_edge Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i). n_ext2e_index : n_edge Broadcast indices from extended node(j) to edge(ij). - angle_index : Optional for dynamic sel, n_angle x 3 + angle_index : Optional for dynamic sel, 3 x n_angle n2a_index : n_angle Broadcast indices from extended node(j) to angle(ijk). eij2a_index : n_angle @@ -745,19 +757,21 @@ def forward( nb, nloc, nnei = nlist.shape nall = node_ebd_ext.shape[1] node_ebd = node_ebd_ext[:, :nloc, :] - n_edge = int(nlist_mask.sum().item()) assert (nb, nloc) == node_ebd.shape[:2] if not self.use_dynamic_sel: assert (nb, nloc, nnei, 3) == h2.shape + n_edge = None else: - assert (n_edge, 3) == h2.shape + # n_edge = int(nlist_mask.sum().item()) + # assert (n_edge, 3) == h2.shape + n_edge = h2.shape[0] del a_nlist # may be used in the future - n2e_index, n_ext2e_index = edge_index[:, 0], edge_index[:, 1] + n2e_index, n_ext2e_index = edge_index[0], edge_index[1] n2a_index, eij2a_index, eik2a_index = ( - angle_index[:, 0], - angle_index[:, 1], - angle_index[:, 2], + angle_index[0], + angle_index[1], + angle_index[2], ) # nb x nloc x nnei x n_dim [OR] n_edge x n_dim @@ -1026,7 +1040,9 @@ def forward( if not self.use_dynamic_sel: # nb x nloc x a_nnei x a_nnei x e_dim weighted_edge_angle_update = ( - a_sw[..., None, None] * a_sw[..., None, :, None] * edge_angle_update + a_sw.unsqueeze(-1).unsqueeze(-1) + * a_sw.unsqueeze(-2).unsqueeze(-1) + * edge_angle_update ) # nb x nloc x a_nnei x e_dim reduced_edge_angle_update = torch.sum( diff --git a/deepmd/pt/model/descriptor/repflows.py b/deepmd/pt/model/descriptor/repflows.py index 5408c49482..433897860f 100644 --- a/deepmd/pt/model/descriptor/repflows.py +++ b/deepmd/pt/model/descriptor/repflows.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -54,15 +55,15 @@ if not hasattr(torch.ops.deepmd, "border_op"): def border_op( - argument0, - argument1, - argument2, - argument3, - argument4, - argument5, - argument6, - argument7, - argument8, + argument0: Any, + argument1: Any, + argument2: Any, + argument3: Any, + argument4: Any, + argument5: Any, + argument6: Any, + argument7: Any, + argument8: Any, ) -> torch.Tensor: raise NotImplementedError( "border_op is not available since customized PyTorch OP library is not built when freezing the model. " @@ -181,15 +182,17 @@ class DescrptBlockRepflows(DescriptorBlock): For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection. seed : int, optional Random seed for parameter initialization. + trainable : bool, default: True + Whether this block is trainable """ def __init__( self, - e_rcut, - e_rcut_smth, + e_rcut: float, + e_rcut_smth: float, e_sel: int, - a_rcut, - a_rcut_smth, + a_rcut: float, + a_rcut_smth: float, a_sel: int, ntypes: int, nlayers: int = 6, @@ -218,7 +221,8 @@ def __init__( sel_reduce_factor: float = 10.0, use_loc_mapping: bool = True, optim_update: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.e_rcut = float(e_rcut) @@ -283,10 +287,19 @@ def __init__( self.seed = seed self.edge_embd = MLPLayer( - 1, self.e_dim, precision=precision, seed=child_seed(seed, 0) + 1, + self.e_dim, + precision=precision, + seed=child_seed(seed, 0), + trainable=trainable, ) self.angle_embd = MLPLayer( - 1, self.a_dim, precision=precision, bias=False, seed=child_seed(seed, 1) + 1, + self.a_dim, + precision=precision, + bias=False, + seed=child_seed(seed, 1), + trainable=trainable, ) layers = [] for ii in range(nlayers): @@ -318,6 +331,7 @@ def __init__( sel_reduce_factor=self.sel_reduce_factor, smooth_edge_update=self.smooth_edge_update, seed=child_seed(child_seed(seed, 1), ii), + trainable=trainable, ) ) self.layers = torch.nn.ModuleList(layers) @@ -363,7 +377,7 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension e_dim.""" return self.e_dim - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -371,7 +385,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -396,17 +410,17 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.n_dim @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.n_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension e_dim.""" return self.get_dim_emb() @@ -422,10 +436,16 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: parallel_mode = comm_dict is not None if not parallel_mode: assert mapping is not None @@ -537,9 +557,8 @@ def forward( a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask] else: # avoid jit assertion - edge_index = angle_index = torch.zeros( - [1, 3], device=nlist.device, dtype=nlist.dtype - ) + edge_index = torch.zeros([2, 1], device=nlist.device, dtype=nlist.dtype) + angle_index = torch.zeros([3, 1], device=nlist.device, dtype=nlist.dtype) # get edge and angle embedding # nb x nloc x nnei x e_dim [OR] n_edge x e_dim if not self.edge_init_use_dist: @@ -646,7 +665,7 @@ def forward( edge_ebd, h2, sw, - owner=edge_index[:, 0], + owner=edge_index[0], num_owner=nframes * nloc, nb=nframes, nloc=nloc, @@ -660,8 +679,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pt/model/descriptor/repformer_layer.py b/deepmd/pt/model/descriptor/repformer_layer.py index 1e2cba66d6..fd72dbc0a3 100644 --- a/deepmd/pt/model/descriptor/repformer_layer.py +++ b/deepmd/pt/model/descriptor/repformer_layer.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import torch import torch.nn as nn @@ -43,7 +39,7 @@ def get_residual( _mode: str = "norm", trainable: bool = True, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> torch.Tensor: r""" Get residual tensor for one update vector. @@ -159,7 +155,8 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" super().__init__() @@ -172,6 +169,7 @@ def __init__( bias=False, precision=precision, seed=seed, + trainable=trainable, ) self.has_gate = has_gate self.smooth = smooth @@ -284,7 +282,8 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim @@ -295,12 +294,14 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.head_map = MLPLayer( input_dim * head_num, input_dim, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.precision = precision @@ -369,13 +370,19 @@ def __init__( input_dim: int, head_num: int, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim self.head_num = head_num self.head_map = MLPLayer( - head_num, 1, bias=False, precision=precision, seed=seed + head_num, + 1, + bias=False, + precision=precision, + seed=seed, + trainable=trainable, ) self.precision = precision @@ -442,7 +449,8 @@ def __init__( smooth: bool = True, attnw_shift: float = 20.0, precision: str = "float64", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.input_dim = input_dim @@ -454,6 +462,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.mapkv = MLPLayer( input_dim, @@ -461,12 +470,14 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.head_map = MLPLayer( input_dim * head_num, input_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) self.smooth = smooth self.attnw_shift = attnw_shift @@ -570,12 +581,12 @@ def deserialize(cls, data: dict) -> "LocalAtten": class RepformerLayer(torch.nn.Module): def __init__( self, - rcut, - rcut_smth, + rcut: float, + rcut_smth: float, sel: int, ntypes: int, - g1_dim=128, - g2_dim=16, + g1_dim: int = 128, + g2_dim: int = 16, axis_neuron: int = 4, update_chnnl_2: bool = True, update_g1_has_conv: bool = True, @@ -597,11 +608,12 @@ def __init__( smooth: bool = True, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.epsilon = 1e-4 # protection of 1./nnei @@ -662,6 +674,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) ) @@ -671,6 +684,7 @@ def __init__( g1_dim, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) self.linear2 = None self.proj_g1g2 = None @@ -687,6 +701,7 @@ def __init__( g2_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, ) if self.update_style == "res_residual": self.g2_residual.append( @@ -696,6 +711,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 3), + trainable=trainable, ) ) if self.g1_out_mlp: @@ -704,6 +720,7 @@ def __init__( g1_dim, precision=precision, seed=child_seed(seed, 15), + trainable=trainable, ) if self.update_style == "res_residual": self.g1_residual.append( @@ -713,6 +730,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 16), + trainable=trainable, ) ) else: @@ -725,6 +743,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) else: self.proj_g1g2 = MLPLayer( @@ -733,6 +752,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 4), + trainable=trainable, ) if self.update_style == "res_residual": self.g1_residual.append( @@ -742,6 +762,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 17), + trainable=trainable, ) ) if self.update_g2_has_g1g1: @@ -751,6 +772,7 @@ def __init__( bias=False, precision=precision, seed=child_seed(seed, 5), + trainable=trainable, ) if self.update_style == "res_residual": self.g2_residual.append( @@ -760,6 +782,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 6), + trainable=trainable, ) ) if self.update_g2_has_attn or self.update_h2: @@ -771,10 +794,15 @@ def __init__( self.smooth, precision=precision, seed=child_seed(seed, 7), + trainable=trainable, ) if self.update_g2_has_attn: self.attn2_mh_apply = Atten2MultiHeadApply( - g2_dim, attn2_nhead, precision=precision, seed=child_seed(seed, 8) + g2_dim, + attn2_nhead, + precision=precision, + seed=child_seed(seed, 8), + trainable=trainable, ) self.attn2_lm = LayerNorm( g2_dim, @@ -791,12 +819,17 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 10), + trainable=trainable, ) ) if self.update_h2: self.attn2_ev_apply = Atten2EquiVarApply( - g2_dim, attn2_nhead, precision=precision, seed=child_seed(seed, 11) + g2_dim, + attn2_nhead, + precision=precision, + seed=child_seed(seed, 11), + trainable=trainable, ) if self.update_style == "res_residual": self.h2_residual.append( @@ -806,6 +839,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 12), + trainable=trainable, ) ) if self.update_g1_has_attn: @@ -816,6 +850,7 @@ def __init__( self.smooth, precision=precision, seed=child_seed(seed, 13), + trainable=trainable, ) if self.update_style == "res_residual": self.g1_residual.append( @@ -825,6 +860,7 @@ def __init__( self.update_residual_init, precision=precision, seed=child_seed(seed, 14), + trainable=trainable, ) ) @@ -1101,7 +1137,7 @@ def forward( nlist: torch.Tensor, # nf x nloc x nnei nlist_mask: torch.Tensor, # nf x nloc x nnei sw: torch.Tensor, # switch func, nf x nloc x nnei - ): + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Parameters ---------- diff --git a/deepmd/pt/model/descriptor/repformers.py b/deepmd/pt/model/descriptor/repformers.py index 82773d1a78..69f2cc4eaa 100644 --- a/deepmd/pt/model/descriptor/repformers.py +++ b/deepmd/pt/model/descriptor/repformers.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -51,15 +52,15 @@ if not hasattr(torch.ops.deepmd, "border_op"): def border_op( - argument0, - argument1, - argument2, - argument3, - argument4, - argument5, - argument6, - argument7, - argument8, + argument0: Any, + argument1: Any, + argument2: Any, + argument3: Any, + argument4: Any, + argument5: Any, + argument6: Any, + argument7: Any, + argument8: Any, ) -> torch.Tensor: raise NotImplementedError( "border_op is not available since customized PyTorch OP library is not built when freezing the model. " @@ -75,13 +76,13 @@ def border_op( class DescrptBlockRepformers(DescriptorBlock): def __init__( self, - rcut, - rcut_smth, + rcut: float, + rcut_smth: float, sel: int, ntypes: int, nlayers: int = 3, - g1_dim=128, - g2_dim=16, + g1_dim: int = 128, + g2_dim: int = 16, axis_neuron: int = 4, direct_dist: bool = False, update_g1_has_conv: bool = True, @@ -106,11 +107,12 @@ def __init__( env_protection: float = 0.0, precision: str = "float64", trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, use_sqrt_nnei: bool = True, g1_out_conv: bool = True, g1_out_mlp: bool = True, + trainable: bool = True, ) -> None: r""" The repformer descriptor block. @@ -197,6 +199,8 @@ def __init__( The epsilon value for layer normalization. seed : int, optional Random seed for parameter initialization. + trainable : bool + Whether the block is trainable """ super().__init__() self.rcut = float(rcut) @@ -247,7 +251,11 @@ def __init__( self.seed = seed self.g2_embd = MLPLayer( - 1, self.g2_dim, precision=precision, seed=child_seed(seed, 0) + 1, + self.g2_dim, + precision=precision, + seed=child_seed(seed, 0), + trainable=trainable, ) layers = [] for ii in range(nlayers): @@ -285,6 +293,7 @@ def __init__( g1_out_conv=self.g1_out_conv, g1_out_mlp=self.g1_out_mlp, seed=child_seed(child_seed(seed, 1), ii), + trainable=trainable, ) ) self.layers = torch.nn.ModuleList(layers) @@ -328,7 +337,7 @@ def get_dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.g2_dim - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -336,7 +345,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -361,17 +370,17 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.g1_dim @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.g1_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the embedding dimension g2.""" return self.get_dim_emb() @@ -387,11 +396,17 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: if comm_dict is None: assert mapping is not None assert extended_atype_embd is not None @@ -531,8 +546,8 @@ def forward( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index fc3e14bd25..59c165ddb0 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, ClassVar, - Optional, - Union, ) import numpy as np @@ -93,11 +94,11 @@ def tabulate_fusion_se_a( class DescrptSeA(BaseDescriptor, torch.nn.Module): def __init__( self, - rcut, - rcut_smth, - sel, - neuron=[25, 50, 100], - axis_neuron=16, + rcut: float, + rcut_smth: float, + sel: list[int] | int, + neuron: list[int] = [25, 50, 100], + axis_neuron: int = 16, set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "float64", @@ -106,11 +107,11 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - ntypes: Optional[int] = None, # to be compat with input - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + ntypes: int | None = None, # to be compat with input + type_map: list[str] | None = None, # not implemented - spin=None, + spin: Any | None = None, ) -> None: del ntypes if spin is not None: @@ -168,7 +169,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension.""" return self.sea.get_dim_emb() - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -186,7 +187,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.sea.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -205,12 +208,12 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.sea.dim_out def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -223,9 +226,9 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -303,9 +306,15 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -345,10 +354,12 @@ def forward( ) return ( g1.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), - rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + rot_mat.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) + if rot_mat is not None + else None, None, None, - sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if sw is not None else None, ) def set_stat_mean_and_stddev( @@ -408,7 +419,7 @@ def deserialize(cls, data: dict) -> "DescrptSeA": env_mat = data.pop("env_mat") obj = cls(**data) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.sea.prec, device=env.DEVICE) obj.sea["davg"] = t_cvt(variables["davg"]) @@ -420,9 +431,9 @@ def t_cvt(xx): def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -455,11 +466,11 @@ class DescrptBlockSeA(DescriptorBlock): def __init__( self, - rcut, - rcut_smth, - sel, - neuron=[25, 50, 100], - axis_neuron=16, + rcut: float, + rcut_smth: float, + sel: int | list[int], + neuron: list[int] = [25, 50, 100], + axis_neuron: int = 16, set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "float64", @@ -468,8 +479,8 @@ def __init__( env_protection: float = 0.0, type_one_side: bool = True, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - **kwargs, + seed: int | list[int] | None = None, + **kwargs: Any, ) -> None: """Construct an embedding net of type `se_a`. @@ -525,6 +536,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, ii), + trainable=trainable, ) self.filter_layers = filter_layers self.stats = None @@ -601,7 +613,7 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] * self.axis_neuron @@ -610,7 +622,7 @@ def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return 0 - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: torch.Tensor) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -618,7 +630,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> torch.Tensor: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -628,8 +640,8 @@ def __getitem__(self, key): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -687,7 +699,7 @@ def reinit_exclude( def enable_compression( self, table_data: dict[str, torch.Tensor], - table_config: list[Union[int, float]], + table_config: list[int | float], lower: dict[str, int], upper: dict[str, int], ) -> None: @@ -725,10 +737,16 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Calculate decoded embedding for each atom. Args: diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 1ce6ad4583..32c8830b1e 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -26,6 +27,9 @@ MLPLayer, NetworkCollection, ) +from deepmd.pt.model.network.network import ( + TypeEmbedNet, +) from deepmd.pt.utils import ( env, ) @@ -75,7 +79,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -86,20 +90,21 @@ def __init__( attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, - activation_function="tanh", + activation_function: str = "tanh", precision: str = "float64", resnet_dt: bool = False, - scaling_factor=1.0, - normalize=True, - temperature=None, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float | None = None, smooth: bool = True, type_one_side: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, - seed: Optional[Union[int, list[int]]] = None, - type: Optional[str] = None, + ln_eps: float | None = 1e-5, + seed: int | list[int] | None = None, + type: str | None = None, + trainable: bool = True, ) -> None: r"""Construct an embedding net of type `se_atten`. @@ -165,6 +170,8 @@ def __init__( If not None, the scaling of attention weights is `temperature` itself. seed : int, Optional Random seed for parameter initialization. + trainable : bool, default: True + Whether this block is trainable """ super().__init__() del type @@ -224,6 +231,7 @@ def __init__( smooth=self.smooth, precision=self.precision, seed=child_seed(self.seed, 0), + trainable=trainable, ) wanted_shape = (self.ntypes, self.nnei, 4) @@ -248,6 +256,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 1), + trainable=trainable, ) self.filter_layers = filter_layers if self.tebd_input_mode in ["strip"]: @@ -261,19 +270,25 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 2), + trainable=trainable, ) self.filter_layers_strip = filter_layers_strip self.stats = None - # add for compression - self.compress = False + self.tebd_compress = False + self.geo_compress = False self.is_sorted = False + # For geometric compression self.compress_info = nn.ParameterList( [nn.Parameter(torch.zeros(0, dtype=self.prec, device="cpu"))] ) self.compress_data = nn.ParameterList( [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))] ) + # For type embedding compression + self.register_buffer( + "type_embd_data", torch.zeros(0, dtype=self.prec, device=env.DEVICE) + ) def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -311,7 +326,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -319,7 +334,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -344,24 +359,24 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] * self.axis_neuron @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.tebd_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -419,10 +434,10 @@ def reinit_exclude( def enable_compression( self, - table_data, - table_config, - lower, - upper, + table_data: dict, + table_config: dict, + lower: dict, + upper: dict, ) -> None: net = "filter_net" self.compress_info[0] = torch.as_tensor( @@ -438,17 +453,75 @@ def enable_compression( device="cpu", ) self.compress_data[0] = table_data[net].to(device=env.DEVICE, dtype=self.prec) - self.compress = True + self.geo_compress = True + + def type_embedding_compression(self, type_embedding_net: TypeEmbedNet) -> None: + """Enable type embedding compression for strip mode. + + Precomputes embedding network outputs for all type combinations: + - One-side: (ntypes+1) combinations (neighbor types only) + - Two-side: (ntypes+1)² combinations (neighbor x center type pairs) + + Parameters + ---------- + type_embedding_net : TypeEmbedNet + The type embedding network that provides get_full_embedding() method + """ + if self.tebd_input_mode != "strip": + raise RuntimeError("Type embedding compression only works in strip mode") + if self.filter_layers_strip is None: + raise RuntimeError( + "filter_layers_strip must be initialized for type embedding compression" + ) + + with torch.no_grad(): + # Get full type embedding: (ntypes+1) x tebd_dim + full_embd = type_embedding_net.get_full_embedding(env.DEVICE) + nt, t_dim = full_embd.shape + + if self.type_one_side: + # One-side: only neighbor types, much simpler! + # Precompute for all (ntypes+1) neighbor types + embd_tensor = self.filter_layers_strip.networks[0](full_embd).detach() + if hasattr(self, "type_embd_data"): + del self.type_embd_data + self.register_buffer("type_embd_data", embd_tensor) + else: + # Two-side: all (ntypes+1)² type pair combinations + # Create [neighbor, center] combinations + # for a fixed row i, all columns j have different neighbor types + embd_nei = full_embd.view(1, nt, t_dim).expand(nt, nt, t_dim) + # for a fixed row i, all columns j share the same center type i + embd_center = full_embd.view(nt, 1, t_dim).expand(nt, nt, t_dim) + two_side_embd = torch.cat([embd_nei, embd_center], dim=-1).reshape( + -1, t_dim * 2 + ) + # Precompute for all type pairs + # Index formula: idx = center_type * nt + neighbor_type + embd_tensor = self.filter_layers_strip.networks[0]( + two_side_embd + ).detach() + if hasattr(self, "type_embd_data"): + del self.type_embd_data + self.register_buffer("type_embd_data", embd_tensor) + + self.tebd_compress = True def forward( self, nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -559,47 +632,49 @@ def forward( nlist_index = nlist.reshape(nb, nloc * nnei) # nf x (nl x nnei) nei_type = torch.gather(extended_atype, dim=1, index=nlist_index) - # (nf x nl x nnei) x ng - nei_type_index = nei_type.view(-1, 1).expand(-1, ng).type(torch.long) if self.type_one_side: - tt_full = self.filter_layers_strip.networks[0](type_embedding) - # (nf x nl x nnei) x ng - gg_t = torch.gather(tt_full, dim=0, index=nei_type_index) + if self.tebd_compress: + tt_full = self.type_embd_data + else: + # (ntypes+1, tebd_dim) -> (ntypes+1, ng) + tt_full = self.filter_layers_strip.networks[0](type_embedding) + # (nf*nl*nnei,) -> (nf*nl*nnei, ng) + gg_t = tt_full[nei_type.view(-1).type(torch.long)] else: idx_i = torch.tile( atype.reshape(-1, 1) * ntypes_with_padding, [1, nnei] ).view(-1) idx_j = nei_type.view(-1) + # (nf x nl x nnei) + idx = (idx_i + idx_j).to(torch.long) + if self.tebd_compress: + # ((ntypes+1)^2, ng) + tt_full = self.type_embd_data + else: + # ((ntypes+1)^2) * (ntypes+1)^2 * nt + type_embedding_nei = torch.tile( + type_embedding.view(1, ntypes_with_padding, nt), + [ntypes_with_padding, 1, 1], + ) + # (ntypes+1)^2 * ((ntypes+1)^2) * nt + type_embedding_center = torch.tile( + type_embedding.view(ntypes_with_padding, 1, nt), + [1, ntypes_with_padding, 1], + ) + # ((ntypes+1)^2 * (ntypes+1)^2) * (nt+nt) + two_side_type_embedding = torch.cat( + [type_embedding_nei, type_embedding_center], -1 + ).reshape(-1, nt * 2) + tt_full = self.filter_layers_strip.networks[0]( + two_side_type_embedding + ) # (nf x nl x nnei) x ng - idx = ( - (idx_i + idx_j) - .view(-1, 1) - .expand(-1, ng) - .type(torch.long) - .to(torch.long) - ) - # (ntypes) * ntypes * nt - type_embedding_nei = torch.tile( - type_embedding.view(1, ntypes_with_padding, nt), - [ntypes_with_padding, 1, 1], - ) - # ntypes * (ntypes) * nt - type_embedding_center = torch.tile( - type_embedding.view(ntypes_with_padding, 1, nt), - [1, ntypes_with_padding, 1], - ) - # (ntypes * ntypes) * (nt+nt) - two_side_type_embedding = torch.cat( - [type_embedding_nei, type_embedding_center], -1 - ).reshape(-1, nt * 2) - tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) - # (nf x nl x nnei) x ng - gg_t = torch.gather(tt_full, dim=0, index=idx) + gg_t = tt_full[idx] # (nf x nl) x nnei x ng gg_t = gg_t.reshape(nfnl, nnei, ng) if self.smooth: gg_t = gg_t * sw.reshape(-1, self.nnei, 1) - if self.compress: + if self.geo_compress: ss = ss.reshape(-1, 1) gg_t = gg_t.reshape(-1, gg_t.size(-1)) xyz_scatter = torch.ops.deepmd.tabulate_fusion_se_atten( @@ -647,7 +722,7 @@ def forward( return ( result.view(nframes, nloc, self.filter_neuron[-1] * self.axis_neuron), gg.view(nframes, nloc, self.nnei, self.filter_neuron[-1]) - if not self.compress + if not self.geo_compress else None, dmatrix.view(nframes, nloc, self.nnei, 4)[..., 1:], rot_mat.view(nframes, nloc, self.filter_neuron[-1], 3), @@ -674,12 +749,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, ln_eps: float = 1e-5, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a neighbor-wise attention net.""" super().__init__() @@ -715,17 +791,18 @@ def __init__( smooth=smooth, precision=precision, seed=child_seed(seed, i), + trainable=trainable, ) ) self.attention_layers = nn.ModuleList(attention_layers) def forward( self, - input_G, - nei_mask, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, - ): + input_G: torch.Tensor, + nei_mask: torch.Tensor, + input_r: torch.Tensor | None = None, + sw: torch.Tensor | None = None, + ) -> torch.Tensor: """Compute the multi-layer gated self-attention. Parameters @@ -745,13 +822,13 @@ def forward( out = layer(out, nei_mask, input_r=input_r, sw=sw) return out - def __getitem__(self, key): + def __getitem__(self, key: int) -> Any: if isinstance(key, int): return self.attention_layers[key] else: raise TypeError(key) - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: int, value: Any) -> None: if not isinstance(key, int): raise TypeError(key) if isinstance(value, self.network_type): @@ -817,12 +894,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, smooth: bool = True, trainable_ln: bool = True, ln_eps: float = 1e-5, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a neighbor-wise attention layer.""" super().__init__() @@ -850,6 +928,7 @@ def __init__( smooth=smooth, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.attn_layer_norm = LayerNorm( self.embed_dim, @@ -861,11 +940,11 @@ def __init__( def forward( self, - x, - nei_mask, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, - ): + x: torch.Tensor, + nei_mask: torch.Tensor, + input_r: torch.Tensor | None = None, + sw: torch.Tensor | None = None, + ) -> torch.Tensor: residual = x x, _ = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) x = residual + x @@ -925,11 +1004,12 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: """Construct a multi-head neighbor-wise attention net.""" super().__init__() @@ -962,6 +1042,7 @@ def __init__( stddev=1.0, precision=precision, seed=child_seed(seed, 0), + trainable=trainable, ) self.out_proj = MLPLayer( hidden_dim, @@ -972,16 +1053,17 @@ def __init__( stddev=1.0, precision=precision, seed=child_seed(seed, 1), + trainable=trainable, ) def forward( self, - query, - nei_mask, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, + query: torch.Tensor, + nei_mask: torch.Tensor, + input_r: torch.Tensor | None = None, + sw: torch.Tensor | None = None, attnw_shift: float = 20.0, - ): + ) -> tuple[torch.Tensor, torch.Tensor]: """Compute the multi-head gated self-attention. Parameters diff --git a/deepmd/pt/model/descriptor/se_atten_v2.py b/deepmd/pt/model/descriptor/se_atten_v2.py index 533d7887e0..5fa031b8bf 100644 --- a/deepmd/pt/model/descriptor/se_atten_v2.py +++ b/deepmd/pt/model/descriptor/se_atten_v2.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, ) import torch @@ -40,7 +39,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], axis_neuron: int = 16, @@ -56,21 +55,21 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, scaling_factor: int = 1.0, - normalize=True, - temperature=None, + normalize: bool = True, + temperature: float | None = None, concat_output_tebd: bool = True, trainable: bool = True, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-5, + ln_eps: float | None = 1e-5, type_one_side: bool = False, - stripped_type_embedding: Optional[bool] = None, - seed: Optional[Union[int, list[int]]] = None, + stripped_type_embedding: bool | None = None, + seed: int | list[int] | None = None, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, # not implemented - spin=None, - type: Optional[str] = None, + spin: Any | None = None, + type: str | None = None, ) -> None: r"""Construct smooth version of embedding net of type `se_atten_v2`. @@ -257,7 +256,7 @@ def deserialize(cls, data: dict) -> "DescrptSeAttenV2": data["use_tebd_bias"] = True obj = cls(**data) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.se_atten.prec, device=env.DEVICE) obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize( diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index a91757460c..1d5a8fc1a8 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import numpy as np @@ -81,10 +82,10 @@ def tabulate_fusion_se_r( class DescrptSeR(BaseDescriptor, torch.nn.Module): def __init__( self, - rcut, - rcut_smth, - sel, - neuron=[25, 50, 100], + rcut: float, + rcut_smth: float, + sel: list[int] | int, + neuron: list[int] = [25, 50, 100], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "float64", @@ -92,9 +93,9 @@ def __init__( exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - **kwargs, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + **kwargs: Any, ) -> None: super().__init__() self.rcut = float(rcut) @@ -142,6 +143,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, ii), + trainable=trainable, ) self.filter_layers = filter_layers self.stats = None @@ -225,7 +227,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.env_protection - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -267,7 +271,7 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -280,8 +284,8 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -329,7 +333,7 @@ def get_stats(self) -> dict[str, StatItem]: ) return self.stats - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -337,7 +341,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -421,9 +425,15 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -574,7 +584,7 @@ def deserialize(cls, data: dict) -> "DescrptSeR": env_mat = data.pop("env_mat") obj = cls(**data) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.prec, device=env.DEVICE) obj["davg"] = t_cvt(variables["davg"]) @@ -586,9 +596,9 @@ def t_cvt(xx): def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index 6e075a04e4..9771d8fe6f 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, ClassVar, - Optional, - Union, ) import numpy as np @@ -142,11 +143,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - ntypes: Optional[int] = None, # to be compat with input + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + ntypes: int | None = None, # to be compat with input # not implemented - spin=None, + spin: dict | None = None, ) -> None: del ntypes if spin is not None: @@ -202,7 +203,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension.""" return self.seat.get_dim_emb() - def mixed_types(self): + def mixed_types(self) -> bool: """Returns if the descriptor requires a neighbor list that distinguish different atomic types or not. """ @@ -220,7 +221,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.seat.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -239,12 +242,12 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.seat.dim_out def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -257,9 +260,9 @@ def change_type_map( def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -338,9 +341,15 @@ def forward( coord_ext: torch.Tensor, atype_ext: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -384,7 +393,7 @@ def forward( None, None, None, - sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if sw is not None else None, ) def set_stat_mean_and_stddev( @@ -439,7 +448,7 @@ def deserialize(cls, data: dict) -> "DescrptSeT": env_mat = data.pop("env_mat") obj = cls(**data) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.seat.prec, device=env.DEVICE) obj.seat["davg"] = t_cvt(variables["davg"]) @@ -451,9 +460,9 @@ def t_cvt(xx): def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -497,7 +506,7 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: r"""Construct an embedding net of type `se_e3`. @@ -575,6 +584,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, ii), + trainable=trainable, ) self.filter_layers = filter_layers self.stats = None @@ -647,7 +657,7 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] @@ -656,7 +666,7 @@ def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return 0 - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -664,7 +674,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -674,8 +684,8 @@ def __getitem__(self, key): def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -732,10 +742,10 @@ def reinit_exclude( def enable_compression( self, - table_data, - table_config, - lower, - upper, + table_data: dict, + table_config: dict, + lower: dict, + upper: dict, ) -> None: for embedding_idx, ll in enumerate(self.filter_layers.networks): ti = embedding_idx % self.ntypes @@ -764,10 +774,16 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters diff --git a/deepmd/pt/model/descriptor/se_t_tebd.py b/deepmd/pt/model/descriptor/se_t_tebd.py index 7e27805bd5..d639cc94bc 100644 --- a/deepmd/pt/model/descriptor/se_t_tebd.py +++ b/deepmd/pt/model/descriptor/se_t_tebd.py @@ -1,11 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch +import torch.nn as nn from deepmd.dpmodel.utils import EnvMat as DPEnvMat from deepmd.dpmodel.utils.seed import ( @@ -38,9 +40,15 @@ from deepmd.pt.utils.exclude_mask import ( PairExcludeMask, ) +from deepmd.pt.utils.tabulate import ( + DPTabulate, +) from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.pt.utils.utils import ( + ActivationFn, +) from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -124,7 +132,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [2, 4, 8], tebd_dim: int = 8, @@ -136,11 +144,11 @@ def __init__( exclude_types: list[tuple[int, int]] = [], precision: str = "float64", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, concat_output_tebd: bool = True, use_econf_tebd: bool = False, - use_tebd_bias=False, + use_tebd_bias: bool = False, smooth: bool = True, ) -> None: super().__init__() @@ -160,6 +168,7 @@ def __init__( env_protection=env_protection, smooth=smooth, seed=child_seed(seed, 1), + trainable=trainable, ) self.prec = PRECISION_DICT[precision] self.use_econf_tebd = use_econf_tebd @@ -170,6 +179,7 @@ def __init__( tebd_dim, precision=precision, seed=child_seed(seed, 2), + trainable=trainable, use_econf_tebd=use_econf_tebd, type_map=type_map, use_tebd_bias=use_tebd_bias, @@ -178,6 +188,7 @@ def __init__( self.tebd_input_mode = tebd_input_mode self.concat_output_tebd = concat_output_tebd self.trainable = trainable + self.compress = False # set trainable for param in self.parameters(): param.requires_grad = trainable @@ -240,7 +251,9 @@ def get_env_protection(self) -> float: """Returns the protection of building environment matrix.""" return self.se_ttebd.get_env_protection() - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -264,18 +277,18 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError @property - def dim_out(self): + def dim_out(self) -> int: return self.get_dim_out() @property - def dim_emb(self): + def dim_emb(self) -> int: return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, - ): + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, + ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -308,7 +321,7 @@ def get_stat_mean_and_stddev(self) -> tuple[torch.Tensor, torch.Tensor]: return self.se_ttebd.mean, self.se_ttebd.stddev def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -388,7 +401,7 @@ def deserialize(cls, data: dict) -> "DescrptSeTTebd": embeddings_strip = None obj = cls(**data) - def t_cvt(xx): + def t_cvt(xx: Any) -> torch.Tensor: return torch.tensor(xx, dtype=obj.se_ttebd.prec, device=env.DEVICE) obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize( @@ -408,9 +421,15 @@ def forward( extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, - mapping: Optional[torch.Tensor] = None, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + mapping: torch.Tensor | None = None, + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -470,16 +489,16 @@ def forward( None, None, None, - sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION), + sw.to(dtype=env.GLOBAL_PT_FLOAT_PRECISION) if sw is not None else None, ) @classmethod def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -505,6 +524,86 @@ def update_sel( local_jdata_cpy["sel"] = sel[0] return local_jdata_cpy, min_nbor_dist + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + # do some checks before the model compression process + if self.compress: + raise ValueError("Compression is already enabled.") + assert not self.se_ttebd.resnet_dt, ( + "Model compression error: descriptor resnet_dt must be false!" + ) + if self.tebd_input_mode != "strip": + raise RuntimeError("Cannot compress model when tebd_input_mode != 'strip'") + for tt in self.se_ttebd.exclude_types: + if (tt[0] not in range(self.se_ttebd.ntypes)) or ( + tt[1] not in range(self.se_ttebd.ntypes) + ): + raise RuntimeError( + "exclude types" + + str(tt) + + " must within the number of atomic types " + + str(self.se_ttebd.ntypes) + + "!" + ) + if ( + self.se_ttebd.ntypes * self.se_ttebd.ntypes + - len(self.se_ttebd.exclude_types) + == 0 + ): + raise RuntimeError( + "Empty embedding-nets are not supported in model compression!" + ) + + data = self.serialize() + self.table = DPTabulate( + self, + data["neuron"], + exclude_types=data["exclude_types"], + activation_fn=ActivationFn(data["activation_function"]), + ) + # Scale the stride values for SE_T descriptor + stride_1_scaled = table_stride_1 * 10 + stride_2_scaled = table_stride_2 * 10 + self.table_config = [ + table_extrapolate, + stride_1_scaled, + stride_2_scaled, + check_frequency, + ] + self.lower, self.upper = self.table.build( + min_nbor_dist, table_extrapolate, stride_1_scaled, stride_2_scaled + ) + + self.se_ttebd.enable_compression( + self.type_embedding, + self.table.data, + self.table_config, + self.lower, + self.upper, + ) + self.compress = True + @DescriptorBlock.register("se_ttebd") class DescrptBlockSeTTebd(DescriptorBlock): @@ -512,19 +611,20 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list = [25, 50, 100], tebd_dim: int = 8, tebd_input_mode: str = "concat", set_davg_zero: bool = True, - activation_function="tanh", + activation_function: str = "tanh", precision: str = "float64", resnet_dt: bool = False, exclude_types: list[tuple[int, int]] = [], env_protection: float = 0.0, smooth: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() self.rcut = float(rcut) @@ -577,6 +677,7 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 1), + trainable=trainable, ) self.filter_layers = filter_layers if self.tebd_input_mode in ["strip"]: @@ -590,9 +691,23 @@ def __init__( precision=self.precision, resnet_dt=self.resnet_dt, seed=child_seed(self.seed, 2), + trainable=trainable, ) self.filter_layers_strip = filter_layers_strip self.stats = None + # compression related variables + self.compress = False + # For geometric compression + self.compress_info = nn.ParameterList( + [nn.Parameter(torch.zeros(0, dtype=self.prec, device="cpu"))] + ) + self.compress_data = nn.ParameterList( + [nn.Parameter(torch.zeros(0, dtype=self.prec, device=env.DEVICE))] + ) + # For type embedding compression + self.register_buffer( + "type_embd_data", torch.zeros(0, dtype=self.prec, device=env.DEVICE) + ) def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -626,7 +741,7 @@ def get_dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.filter_neuron[-1] - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ("avg", "data_avg", "davg"): self.mean = value elif key in ("std", "data_std", "dstd"): @@ -634,7 +749,7 @@ def __setitem__(self, key, value) -> None: else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ("avg", "data_avg", "davg"): return self.mean elif key in ("std", "data_std", "dstd"): @@ -659,24 +774,24 @@ def get_env_protection(self) -> float: return self.env_protection @property - def dim_out(self): + def dim_out(self) -> int: """Returns the output dimension of this descriptor.""" return self.filter_neuron[-1] @property - def dim_in(self): + def dim_in(self) -> int: """Returns the atomic input dimension of this descriptor.""" return self.tebd_dim @property - def dim_emb(self): + def dim_emb(self) -> int: """Returns the output dimension of embedding.""" return self.get_dim_emb() def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -736,10 +851,16 @@ def forward( nlist: torch.Tensor, extended_coord: torch.Tensor, extended_atype: torch.Tensor, - extended_atype_embd: Optional[torch.Tensor] = None, - mapping: Optional[torch.Tensor] = None, - type_embedding: Optional[torch.Tensor] = None, - ): + extended_atype_embd: torch.Tensor | None = None, + mapping: torch.Tensor | None = None, + type_embedding: torch.Tensor | None = None, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: """Compute the descriptor. Parameters @@ -791,6 +912,7 @@ def forward( self.rcut_smth, protection=self.env_protection, ) + # dmatrix: [1/r, dx/r^2, dy/r^2, dz/r^2], sw: distance weighting # nb x nloc x nnei exclude_mask = self.emask(nlist, extended_atype) nlist = torch.where(exclude_mask != 0, nlist, -1) @@ -811,11 +933,13 @@ def forward( rr = dmatrix rr = rr * exclude_mask[:, :, None] - # nfnl x nt_i x 3 + # nfnl x nt_i x 3: direction vectors + # nt_i = nnei + # nt_j = nnei rr_i = rr[:, :, 1:] # nfnl x nt_j x 3 rr_j = rr[:, :, 1:] - # nfnl x nt_i x nt_j + # nfnl x nt_i x nt_j: three-body angular correlations (cos theta_ij) env_ij = torch.einsum("ijm,ikm->ijk", rr_i, rr_j) # nfnl x nt_i x nt_j x 1 ss = env_ij.unsqueeze(-1) @@ -837,8 +961,24 @@ def forward( # nfnl x nt_i x nt_j x ng gg = self.filter_layers.networks[0](ss) elif self.tebd_input_mode in ["strip"]: - # nfnl x nt_i x nt_j x ng - gg_s = self.filter_layers.networks[0](ss) + if self.compress: + # Tabulated geometric embedding from angular features + # using SE_T_TEBD specific function + ebd_env_ij = env_ij.view(-1, 1) + gg_s = torch.ops.deepmd.tabulate_fusion_se_t_tebd( + self.compress_data[0].contiguous(), + self.compress_info[0].cpu().contiguous(), + ebd_env_ij.contiguous(), # em_x: (nfnl * nt_i * nt_j, 1) + env_ij.contiguous(), # em: (nfnl, nt_i, nt_j) + self.filter_neuron[-1], + )[0] + # SE_T_TEBD tabulation preserves the full neighbor structure + # nfnl x nt_i x nt_j x ng + gg_s = gg_s.view(nfnl, nnei, nnei, self.filter_neuron[-1]) + else: + # nfnl x nt_i x nt_j x ng + gg_s = self.filter_layers.networks[0](ss) + assert self.filter_layers_strip is not None assert type_embedding is not None ng = self.filter_neuron[-1] @@ -854,44 +994,40 @@ def forward( nei_type_j = nei_type.unsqueeze(1).expand([-1, nnei, -1]) idx_i = nei_type_i * ntypes_with_padding idx_j = nei_type_j - # (nf x nl x nt_i x nt_j) x ng - idx = ( - (idx_i + idx_j) - .view(-1, 1) - .expand(-1, ng) - .type(torch.long) - .to(torch.long) - ) - # ntypes * (ntypes) * nt - type_embedding_i = torch.tile( - type_embedding.view(ntypes_with_padding, 1, nt), - [1, ntypes_with_padding, 1], - ) - # (ntypes) * ntypes * nt - type_embedding_j = torch.tile( - type_embedding.view(1, ntypes_with_padding, nt), - [ntypes_with_padding, 1, 1], - ) - # (ntypes * ntypes) * (nt+nt) - two_side_type_embedding = torch.cat( - [type_embedding_i, type_embedding_j], -1 - ).reshape(-1, nt * 2) - tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) + idx = (idx_i + idx_j).reshape(-1).to(torch.long) + if self.compress: + tt_full = self.type_embd_data + else: + type_embedding_i = torch.tile( + type_embedding.view(ntypes_with_padding, 1, nt), + [1, ntypes_with_padding, 1], + ) + type_embedding_j = torch.tile( + type_embedding.view(1, ntypes_with_padding, nt), + [ntypes_with_padding, 1, 1], + ) + two_side_type_embedding = torch.cat( + [type_embedding_i, type_embedding_j], -1 + ).reshape(-1, nt * 2) + tt_full = self.filter_layers_strip.networks[0](two_side_type_embedding) # (nfnl x nt_i x nt_j) x ng - gg_t = torch.gather(tt_full, dim=0, index=idx) + gg_t = tt_full[idx] # (nfnl x nt_i x nt_j) x ng gg_t = gg_t.reshape(nfnl, nnei, nnei, ng) if self.smooth: + # Apply distance weighting to type features gg_t = ( gg_t * sw.reshape(nfnl, self.nnei, 1, 1) * sw.reshape(nfnl, 1, self.nnei, 1) ) + # Combine geometric and type embeddings: gg_s * (1 + gg_t) # nfnl x nt_i x nt_j x ng gg = gg_s * gg_t + gg_s else: raise NotImplementedError + # Contract angular correlations with learned features # nfnl x ng res_ij = torch.einsum("ijk,ijkm->im", env_ij, gg) res_ij = res_ij * (1.0 / float(self.nnei) / float(self.nnei)) @@ -905,6 +1041,72 @@ def forward( sw, ) + def enable_compression( + self, + type_embedding_net: TypeEmbedNet, + table_data: dict, + table_config: dict, + lower: dict, + upper: dict, + ) -> None: + """Enable compression for the SE_T_TEBD descriptor block. + + Parameters + ---------- + type_embedding_net : TypeEmbedNet + The type embedding network + table_data : dict + The tabulated data from DPTabulate + table_config : dict + Configuration for table compression + lower : dict + Lower bounds for compression + upper : dict + Upper bounds for compression + """ + if self.tebd_input_mode != "strip": + raise RuntimeError("Type embedding compression only works in strip mode") + if self.filter_layers_strip is None: + raise RuntimeError( + "filter_layers_strip must exist for type embedding compression" + ) + + # Compress the main geometric embedding network (self.filter_layers) + net_key = "filter_net" + self.compress_info[0] = torch.as_tensor( + [ + lower[net_key], + upper[net_key], + upper[net_key] * table_config[0], + table_config[1], + table_config[2], + table_config[3], + ], + dtype=self.prec, + device="cpu", + ) + self.compress_data[0] = table_data[net_key].to( + device=env.DEVICE, dtype=self.prec + ) + + # Compress the type embedding network (self.filter_layers_strip) + with torch.no_grad(): + full_embd = type_embedding_net.get_full_embedding(env.DEVICE) + nt, t_dim = full_embd.shape + type_embedding_i = full_embd.view(nt, 1, t_dim).expand(nt, nt, t_dim) + type_embedding_j = full_embd.view(1, nt, t_dim).expand(nt, nt, t_dim) + two_side_type_embedding = torch.cat( + [type_embedding_i, type_embedding_j], dim=-1 + ).reshape(-1, t_dim * 2) + embd_tensor = self.filter_layers_strip.networks[0]( + two_side_type_embedding + ).detach() + if hasattr(self, "type_embd_data"): + del self.type_embd_data + self.register_buffer("type_embd_data", embd_tensor) + + self.compress = True + def has_message_passing(self) -> bool: """Returns whether the descriptor block has message passing.""" return False diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 8d451f087f..b8f7b171d4 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -14,7 +14,7 @@ import copy import json from typing import ( - Optional, + Any, ) import numpy as np @@ -75,7 +75,7 @@ ) -def _get_standard_model_components(model_params, ntypes): +def _get_standard_model_components(model_params: dict, ntypes: int) -> tuple: if "type_embedding" in model_params: raise ValueError( "In the PyTorch backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." @@ -102,7 +102,7 @@ def _get_standard_model_components(model_params, ntypes): return descriptor, fitting, fitting_net["type"] -def get_spin_model(model_params): +def get_spin_model(model_params: dict) -> SpinModel: model_params = copy.deepcopy(model_params) if not model_params["spin"]["use_spin"] or isinstance( model_params["spin"]["use_spin"][0], int @@ -138,7 +138,7 @@ def get_spin_model(model_params): return SpinEnergyModel(backbone_model=backbone_model, spin=spin) -def get_linear_model(model_params): +def get_linear_model(model_params: dict) -> LinearEnergyModel: model_params = copy.deepcopy(model_params) weights = model_params.get("weights", "mean") list_of_models = [] @@ -178,7 +178,7 @@ def get_linear_model(model_params): ) -def get_zbl_model(model_params): +def get_zbl_model(model_params: dict) -> DPZBLModel: model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) descriptor, fitting, _ = _get_standard_model_components(model_params, ntypes) @@ -209,7 +209,7 @@ def get_zbl_model(model_params): return model -def _can_be_converted_to_float(value) -> Optional[bool]: +def _can_be_converted_to_float(value: Any) -> bool | None: try: float(value) return True @@ -218,7 +218,9 @@ def _can_be_converted_to_float(value) -> Optional[bool]: return False -def _convert_preset_out_bias_to_array(preset_out_bias, type_map): +def _convert_preset_out_bias_to_array( + preset_out_bias: dict | None, type_map: list[str] +) -> dict | None: if preset_out_bias is not None: for kk in preset_out_bias: if len(preset_out_bias[kk]) != len(type_map): @@ -241,7 +243,7 @@ def _convert_preset_out_bias_to_array(preset_out_bias, type_map): return preset_out_bias -def get_standard_model(model_params): +def get_standard_model(model_params: dict) -> BaseModel: model_params_old = model_params model_params = copy.deepcopy(model_params) ntypes = len(model_params["type_map"]) @@ -284,7 +286,7 @@ def get_standard_model(model_params): return model -def get_model(model_params): +def get_model(model_params: dict) -> Any: model_type = model_params.get("type", "standard") if model_type == "standard": if "spin" in model_params: diff --git a/deepmd/pt/model/model/dipole_model.py b/deepmd/pt/model/model/dipole_model.py index a24820b74a..c6813ce079 100644 --- a/deepmd/pt/model/model/dipole_model.py +++ b/deepmd/pt/model/model/dipole_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -28,13 +28,13 @@ class DipoleModel(DPModelCommon, DPDipoleModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPDipoleModel_.__init__(self, *args, **kwargs) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, Any]: out_def_data = self.model_output_def().get_data() output_def = { "dipole": out_def_data["dipole"], @@ -54,11 +54,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -91,14 +91,15 @@ def forward( @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, @@ -107,6 +108,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: diff --git a/deepmd/pt/model/model/dos_model.py b/deepmd/pt/model/model/dos_model.py index 41d85a559e..75f89c141a 100644 --- a/deepmd/pt/model/model/dos_model.py +++ b/deepmd/pt/model/model/dos_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -28,13 +28,13 @@ class DOSModel(DPModelCommon, DPDOSModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPDOSModel_.__init__(self, *args, **kwargs) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, Any]: out_def_data = self.model_output_def().get_data() output_def = { "atom_dos": out_def_data["dos"], @@ -46,11 +46,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -81,14 +81,15 @@ def get_numb_dos(self) -> int: @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, @@ -97,6 +98,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: diff --git a/deepmd/pt/model/model/dp_linear_model.py b/deepmd/pt/model/model/dp_linear_model.py index fe4487a495..b43f849258 100644 --- a/deepmd/pt/model/model/dp_linear_model.py +++ b/deepmd/pt/model/model/dp_linear_model.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch +from deepmd.dpmodel.output_def import ( + OutputVariableDef, +) from deepmd.pt.model.atomic_model import ( LinearEnergyAtomicModel, ) @@ -31,12 +34,12 @@ class LinearEnergyModel(DPLinearModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, OutputVariableDef]: out_def_data = self.model_output_def().get_data() output_def = { "atom_energy": out_def_data["energy"], @@ -56,11 +59,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -90,14 +93,15 @@ def forward( @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, @@ -106,6 +110,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) @@ -129,9 +134,9 @@ def forward_lower( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index e71c5e08de..a8b5b55584 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import torch @@ -20,9 +17,9 @@ class DPModelCommon: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -47,11 +44,12 @@ def update_sel( ) return local_jdata_cpy, min_nbor_dist - def get_fitting_net(self): + # sadly, use -> BaseFitting here will not make torchscript happy + def get_fitting_net(self): # noqa: ANN201 """Get the fitting network.""" return self.atomic_model.fitting_net - def get_descriptor(self): + def get_descriptor(self): # noqa: ANN201 """Get the descriptor.""" return self.atomic_model.descriptor @@ -64,3 +62,13 @@ def set_eval_descriptor_hook(self, enable: bool) -> None: def eval_descriptor(self) -> torch.Tensor: """Evaluate the descriptor.""" return self.atomic_model.eval_descriptor() + + @torch.jit.export + def set_eval_fitting_last_layer_hook(self, enable: bool) -> None: + """Set the hook for evaluating fitting_last_layer and clear the cache for fitting_last_layer list.""" + self.atomic_model.set_eval_fitting_last_layer_hook(enable) + + @torch.jit.export + def eval_fitting_last_layer(self) -> torch.Tensor: + """Evaluate the fitting_last_layer.""" + return self.atomic_model.eval_fitting_last_layer() diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index 3c019b9376..d533cbe125 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -31,12 +31,12 @@ class DPZBLModel(DPZBLModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, Any]: out_def_data = self.model_output_def().get_data() output_def = { "atom_energy": out_def_data["energy"], @@ -56,11 +56,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -90,14 +90,15 @@ def forward( @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, @@ -106,6 +107,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) @@ -129,9 +131,9 @@ def forward_lower( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/ener_model.py b/deepmd/pt/model/model/ener_model.py index 8064d3eac7..8f8a3cbad7 100644 --- a/deepmd/pt/model/model/ener_model.py +++ b/deepmd/pt/model/model/ener_model.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -31,20 +31,46 @@ class EnergyModel(DPModelCommon, DPEnergyModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPEnergyModel_.__init__(self, *args, **kwargs) self._hessian_enabled = False - def enable_hessian(self): + def enable_hessian(self) -> None: self.__class__ = make_hessian_model(type(self)) self.hess_fitting_def = super(type(self), self).atomic_output_def() self.requires_hessian("energy") self._hessian_enabled = True - def translated_output_def(self): + @torch.jit.export + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + observed_type_list: a list of the observed types in this model. + """ + type_map = self.get_type_map() + out_bias = self.atomic_model.get_out_bias()[0] + + assert out_bias is not None, "No out_bias found in the model." + assert out_bias.dim() == 2, "The supported out_bias should be a 2D tensor." + assert out_bias.size(0) == len(type_map), ( + "The out_bias shape does not match the type_map length." + ) + bias_mask = ( + torch.gt(torch.abs(out_bias), 1e-6).any(dim=-1).detach().cpu() + ) # 1e-6 for stability + + observed_type_list: list[str] = [] + for i in range(len(type_map)): + if bias_mask[i]: + observed_type_list.append(type_map[i]) + return observed_type_list + + def translated_output_def(self) -> dict[str, Any]: out_def_data = self.model_output_def().get_data() output_def = { "atom_energy": out_def_data["energy"], @@ -66,11 +92,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -107,15 +133,15 @@ def forward( @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index 27284ec276..402e78e95c 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -2,8 +2,8 @@ import json import tempfile from typing import ( + Any, NoReturn, - Optional, ) import torch @@ -32,7 +32,7 @@ class FrozenModel(BaseModel): The path to the frozen model """ - def __init__(self, model_file: str, **kwargs) -> None: + def __init__(self, model_file: str, **kwargs: Any) -> None: super().__init__(**kwargs) self.model_file = model_file if model_file.endswith(".pth"): @@ -116,11 +116,11 @@ def need_sorted_nlist_for_lower(self) -> bool: @torch.jit.export def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: return self.model.forward( @@ -142,7 +142,7 @@ def get_model_def_script(self) -> str: return self.model.get_model_def_script() @torch.jit.export - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.model.get_min_nbor_dist() @@ -175,9 +175,9 @@ def get_nsel(self) -> int: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/pt/model/model/make_hessian_model.py b/deepmd/pt/model/model/make_hessian_model.py index 000b9abea4..1b1bc3feba 100644 --- a/deepmd/pt/model/model/make_hessian_model.py +++ b/deepmd/pt/model/model/make_hessian_model.py @@ -2,8 +2,7 @@ import copy import math from typing import ( - Optional, - Union, + Any, ) import torch @@ -11,9 +10,12 @@ from deepmd.dpmodel import ( get_hessian_name, ) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, +) -def make_hessian_model(T_Model): +def make_hessian_model(T_Model: type) -> type: """Make a model that can compute Hessian. LIMITATION: this model is not jitable due to the restrictions of torch jit script. @@ -34,8 +36,8 @@ def make_hessian_model(T_Model): class CM(T_Model): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: super().__init__( *args, @@ -45,7 +47,7 @@ def __init__( def requires_hessian( self, - keys: Union[str, list[str]], + keys: str | list[str], ) -> None: """Set which output variable(s) requires hessian.""" if isinstance(keys, str): @@ -54,17 +56,17 @@ def requires_hessian( if kk in keys: self.hess_fitting_def[kk].r_hessian = True - def atomic_output_def(self): + def atomic_output_def(self) -> FittingOutputDef: """Get the fitting output def.""" return self.hess_fitting_def def forward_common( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -117,9 +119,9 @@ def _cal_hessian_all( self, coord: torch.Tensor, atype: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: nf, nloc = atype.shape coord = coord.view([nf, (nloc * 3)]) @@ -159,12 +161,12 @@ def _cal_hessian_all( def _cal_hessian_one_component( self, - ci, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + ci: int, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> torch.Tensor: # coord, # (nloc x 3) # atype, # nloc @@ -185,9 +187,9 @@ def __init__( obj: CM, ci: int, atype: torch.Tensor, - box: Optional[torch.Tensor], - fparam: Optional[torch.Tensor], - aparam: Optional[torch.Tensor], + box: torch.Tensor | None, + fparam: torch.Tensor | None, + aparam: torch.Tensor | None, ) -> None: self.atype, self.box, self.fparam, self.aparam = atype, box, fparam, aparam self.ci = ci @@ -195,8 +197,8 @@ def __init__( def __call__( self, - xx, - ): + xx: torch.Tensor, + ) -> torch.Tensor: ci = self.ci atype, box, fparam, aparam = self.atype, self.box, self.fparam, self.aparam res = super(CM, self.obj).forward_common( diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index c32abaa095..c958a62bf6 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( - Optional, + Any, ) import torch @@ -39,7 +42,7 @@ ) -def make_model(T_AtomicModel: type[BaseAtomicModel]): +def make_model(T_AtomicModel: type[BaseAtomicModel]) -> type: """Make a model as a derived class of an atomic model. The model provide two interfaces. @@ -65,10 +68,10 @@ def make_model(T_AtomicModel: type[BaseAtomicModel]): class CM(BaseModel): def __init__( self, - *args, + *args: Any, # underscore to prevent conflict with normal inputs - atomic_model_: Optional[T_AtomicModel] = None, - **kwargs, + atomic_model_: T_AtomicModel | None = None, + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) if atomic_model_ is not None: @@ -80,7 +83,7 @@ def __init__( self.global_pt_float_precision = GLOBAL_PT_FLOAT_PRECISION self.global_pt_ener_float_precision = GLOBAL_PT_ENER_FLOAT_PRECISION - def model_output_def(self): + def model_output_def(self) -> ModelOutputDef: """Get the output def for the model.""" return ModelOutputDef(self.atomic_output_def()) @@ -129,11 +132,11 @@ def enable_compression( # cannot use the name forward. torch script does not work def forward_common( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: """Return model prediction. @@ -206,8 +209,8 @@ def set_out_bias(self, out_bias: torch.Tensor) -> None: def change_out_bias( self, - merged, - bias_adjust_mode="change-by-statistic", + merged: Any, + bias_adjust_mode: str = "change-by-statistic", ) -> None: """Change the output bias of atomic model according to the input data and the pretrained model. @@ -230,19 +233,21 @@ def change_out_bias( merged, bias_adjust_mode=bias_adjust_mode, ) + if bias_adjust_mode == "set-by-statistic": + self.atomic_model.compute_fitting_input_stat(merged) def forward_common_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, - ): + ) -> dict[str, torch.Tensor]: """Return model prediction. Lower interface that takes extended atomic coordinates and types, nlist, and mapping as input, and returns the predictions on the extended region. @@ -299,6 +304,7 @@ def forward_common_lower( cc_ext, do_atomic_virial=do_atomic_virial, create_graph=self.training, + mask=atomic_ret["mask"] if "mask" in atomic_ret else None, ) model_predict = self.output_type_cast(model_predict, input_prec) return model_predict @@ -306,14 +312,14 @@ def forward_common_lower( def input_type_cast( self, coord: torch.Tensor, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, - Optional[torch.Tensor], - Optional[torch.Tensor], - Optional[torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, str, ]: """Cast the input data to global float type.""" @@ -328,7 +334,7 @@ def input_type_cast( # " does not match" # f" that of the coordinate {input_prec}" # ) - _lst: list[Optional[torch.Tensor]] = [ + _lst: list[torch.Tensor | None] = [ vv.to(coord.dtype) if vv is not None else None for vv in [box, fparam, aparam] ] @@ -382,7 +388,7 @@ def format_nlist( extended_atype: torch.Tensor, nlist: torch.Tensor, extra_nlist_sort: bool = False, - ): + ) -> torch.Tensor: """Format the neighbor list. 1. If the number of neighbors in the `nlist` is equal to sum(self.sel), @@ -433,7 +439,7 @@ def _format_nlist( nlist: torch.Tensor, nnei: int, extra_nlist_sort: bool = False, - ): + ) -> torch.Tensor: n_nf, n_nloc, n_nnei = nlist.shape # nf x nall x 3 extended_coord = extended_coord.view([n_nf, -1, 3]) @@ -478,7 +484,7 @@ def _format_nlist( def do_grad_r( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is r_differentiable. if var_name is None, returns if any of the variable is r_differentiable. @@ -487,7 +493,7 @@ def do_grad_r( def do_grad_c( self, - var_name: Optional[str] = None, + var_name: str | None = None, ) -> bool: """Tell if the output variable `var_name` is c_differentiable. if var_name is None, returns if any of the variable is c_differentiable. @@ -495,7 +501,7 @@ def do_grad_c( return self.atomic_model.do_grad_c(var_name) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -511,10 +517,10 @@ def serialize(self) -> dict: return self.atomic_model.serialize() @classmethod - def deserialize(cls, data) -> "CM": + def deserialize(cls, data: Any) -> "CM": return cls(atomic_model_=T_AtomicModel.deserialize(data)) - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: self.atomic_model.set_case_embd(case_idx) @torch.jit.export @@ -522,6 +528,14 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.atomic_model.get_dim_fparam() + @torch.jit.export + def has_default_fparam(self) -> bool: + """Check if the model has default frame parameters.""" + return self.atomic_model.has_default_fparam() + + def get_default_fparam(self) -> torch.Tensor | None: + return self.atomic_model.get_default_fparam() + @torch.jit.export def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" @@ -571,9 +585,9 @@ def atomic_output_def(self) -> FittingOutputDef: def compute_or_load_stat( self, - sampled_func, - stat_file_path: Optional[DPPath] = None, - ): + sampled_func: Callable[[], Any], + stat_file_path: DPPath | None = None, + ) -> None: """Compute or load the statistics.""" return self.atomic_model.compute_or_load_stat(sampled_func, stat_file_path) @@ -604,11 +618,11 @@ def need_sorted_nlist_for_lower(self) -> bool: def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: # directly call the forward_common method when no specific transform rule diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index 44600ab804..d9169fdc54 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + Any, NoReturn, - Optional, ) import torch @@ -18,7 +18,7 @@ class BaseModel(torch.nn.Module, make_base_model()): - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args: Any, **kwargs: Any) -> None: """Construct a basic model for different tasks.""" torch.nn.Module.__init__(self) self.model_def_script = "" @@ -28,8 +28,8 @@ def __init__(self, *args, **kwargs) -> None: def compute_or_load_stat( self, - sampled_func, - stat_file_path: Optional[DPPath] = None, + sampled_func: Any, + stat_file_path: DPPath | None = None, ) -> NoReturn: """ Compute or load the statistics parameters of the model, @@ -48,19 +48,29 @@ def compute_or_load_stat( """ raise NotImplementedError + @torch.jit.export + def get_observed_type_list(self) -> list[str]: + """Get observed types (elements) of the model during data statistics. + + Returns + ------- + observed_type_list: a list of the observed types in this model. + """ + raise NotImplementedError + @torch.jit.export def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script @torch.jit.export - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum distance between two atoms.""" if self.min_nbor_dist.item() == -1.0: return None return self.min_nbor_dist.item() @torch.jit.export - def get_ntypes(self): + def get_ntypes(self) -> int: """Returns the number of element types.""" return len(self.get_type_map()) diff --git a/deepmd/pt/model/model/polar_model.py b/deepmd/pt/model/model/polar_model.py index cb72532366..7210117823 100644 --- a/deepmd/pt/model/model/polar_model.py +++ b/deepmd/pt/model/model/polar_model.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch +from deepmd.dpmodel.output_def import ( + OutputVariableDef, +) from deepmd.pt.model.atomic_model import ( DPPolarAtomicModel, ) @@ -28,13 +31,13 @@ class PolarModel(DPModelCommon, DPPolarModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPPolarModel_.__init__(self, *args, **kwargs) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, OutputVariableDef]: out_def_data = self.model_output_def().get_data() output_def = { "polar": out_def_data["polarizability"], @@ -46,11 +49,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -75,14 +78,15 @@ def forward( @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, @@ -91,6 +95,7 @@ def forward_lower( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, + comm_dict=comm_dict, extra_nlist_sort=self.need_sorted_nlist_for_lower(), ) if self.get_fitting_net() is not None: diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 7c50c75ff1..c24ca7fa64 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch +from deepmd.dpmodel.output_def import ( + OutputVariableDef, +) from deepmd.pt.model.atomic_model import ( DPPropertyAtomicModel, ) @@ -28,13 +31,13 @@ class PropertyModel(DPModelCommon, DPPropertyModel_): def __init__( self, - *args, - **kwargs, + *args: Any, + **kwargs: Any, ) -> None: DPModelCommon.__init__(self) DPPropertyModel_.__init__(self, *args, **kwargs) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, OutputVariableDef]: out_def_data = self.model_output_def().get_data() output_def = { f"atom_{self.get_var_name()}": out_def_data[self.get_var_name()], @@ -46,11 +49,11 @@ def translated_output_def(self): def forward( self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -86,15 +89,15 @@ def get_var_name(self) -> str: @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index ac94668039..0eb7014416 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import functools +from collections.abc import ( + Callable, +) from copy import ( deepcopy, ) from typing import ( - Optional, + Any, ) import torch @@ -38,7 +41,7 @@ class SpinModel(torch.nn.Module): def __init__( self, - backbone_model, + backbone_model: DPAtomicModel, spin: Spin, ) -> None: super().__init__() @@ -48,7 +51,9 @@ def __init__( self.virtual_scale_mask = to_torch_tensor(self.spin.get_virtual_scale_mask()) self.spin_mask = to_torch_tensor(self.spin.get_spin_mask()) - def process_spin_input(self, coord, atype, spin): + def process_spin_input( + self, coord: torch.Tensor, atype: torch.Tensor, spin: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: """Generate virtual coordinates and types, concat into the input.""" nframes, nloc = atype.shape coord = coord.reshape(nframes, nloc, 3) @@ -62,12 +67,12 @@ def process_spin_input(self, coord, atype, spin): def process_spin_input_lower( self, - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping: Optional[torch.Tensor] = None, - ): + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + extended_spin: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: """ Add `extended_spin` into `extended_coord` to generate virtual atoms, and extend `nlist` and `mapping`. Note that the final `extended_coord_updated` with shape [nframes, nall + nall, 3] has the following order: @@ -103,8 +108,12 @@ def process_spin_input_lower( ) def process_spin_output( - self, atype, out_tensor, add_mag: bool = True, virtual_scale: bool = True - ): + self, + atype: torch.Tensor, + out_tensor: torch.Tensor, + add_mag: bool = True, + virtual_scale: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Split the output both real and virtual atoms, and scale the latter. add_mag: whether to add magnetic tensor onto the real tensor. @@ -132,12 +141,12 @@ def process_spin_output( def process_spin_output_lower( self, - extended_atype, - extended_out_tensor, + extended_atype: torch.Tensor, + extended_out_tensor: torch.Tensor, nloc: int, add_mag: bool = True, virtual_scale: bool = True, - ): + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Split the extended output of both real and virtual atoms with switch, and scale the latter. add_mag: whether to add magnetic tensor onto the real tensor. @@ -177,7 +186,7 @@ def process_spin_output_lower( return extended_out_real, extended_out_mag, atomic_mask > 0.0 @staticmethod - def extend_nlist(extended_atype, nlist): + def extend_nlist(extended_atype: torch.Tensor, nlist: torch.Tensor) -> torch.Tensor: nframes, nloc, nnei = nlist.shape nall = extended_atype.shape[1] nlist_mask = nlist != -1 @@ -207,7 +216,7 @@ def extend_nlist(extended_atype, nlist): return extended_nlist @staticmethod - def expand_aparam(aparam, nloc: int): + def expand_aparam(aparam: torch.Tensor, nloc: int) -> torch.Tensor: """Expand the atom parameters for virtual atoms if necessary.""" nframes, natom, numb_aparam = aparam.shape if natom == nloc: # good @@ -239,22 +248,22 @@ def get_type_map(self) -> list[str]: return tmap[:ntypes] @torch.jit.export - def get_ntypes(self): + def get_ntypes(self) -> int: """Returns the number of element types.""" return len(self.get_type_map()) @torch.jit.export - def get_rcut(self): + def get_rcut(self) -> float: """Get the cut-off radius.""" return self.backbone_model.get_rcut() @torch.jit.export - def get_dim_fparam(self): + def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.backbone_model.get_dim_fparam() @torch.jit.export - def get_dim_aparam(self): + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.backbone_model.get_dim_aparam() @@ -285,7 +294,7 @@ def get_model_def_script(self) -> str: return self.backbone_model.get_model_def_script() @torch.jit.export - def get_min_nbor_dist(self) -> Optional[float]: + def get_min_nbor_dist(self) -> float | None: """Get the minimum neighbor distance.""" return self.backbone_model.get_min_nbor_dist() @@ -320,7 +329,7 @@ def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the model needs sorted nlist when using `forward_lower`.""" return self.backbone_model.need_sorted_nlist_for_lower() - def model_output_def(self): + def model_output_def(self) -> ModelOutputDef: """Get the output def for the model.""" model_output_type = self.backbone_model.model_output_type() if "mask" in model_output_type: @@ -330,7 +339,7 @@ def model_output_def(self): backbone_model_atomic_output_def[var_name].magnetic = True return ModelOutputDef(backbone_model_atomic_output_def) - def __getattr__(self, name): + def __getattr__(self, name: str) -> Any: """Get attribute from the wrapped model.""" if ( name == "backbone_model" @@ -343,8 +352,8 @@ def __getattr__(self, name): def compute_or_load_stat( self, - sampled_func, - stat_file_path: Optional[DPPath] = None, + sampled_func: Callable[[], list[dict[str, Any]]], + stat_file_path: DPPath | None = None, ) -> None: """ Compute or load the statistics parameters of the model, @@ -363,7 +372,7 @@ def compute_or_load_stat( """ @functools.lru_cache - def spin_sampled_func(): + def spin_sampled_func() -> list[dict[str, Any]]: sampled = sampled_func() spin_sampled = [] for sys in sampled: @@ -389,12 +398,12 @@ def spin_sampled_func(): def forward_common( self, - coord, - atype, - spin, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + spin: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: nframes, nloc = atype.shape @@ -437,17 +446,17 @@ def forward_common( def forward_common_lower( self, - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + extended_spin: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, + comm_dict: dict[str, torch.Tensor] | None = None, extra_nlist_sort: bool = False, - ): + ) -> dict[str, torch.Tensor]: nframes, nloc = nlist.shape[:2] ( extended_coord_updated, @@ -506,7 +515,7 @@ def serialize(self) -> dict: } @classmethod - def deserialize(cls, data) -> "SpinModel": + def deserialize(cls, data: dict[str, Any]) -> "SpinModel": backbone_model_obj = make_model(DPAtomicModel).deserialize( data["backbone_model"] ) @@ -524,12 +533,12 @@ class SpinEnergyModel(SpinModel): def __init__( self, - backbone_model, + backbone_model: DPAtomicModel, spin: Spin, ) -> None: super().__init__(backbone_model, spin) - def translated_output_def(self): + def translated_output_def(self) -> dict[str, Any]: out_def_data = self.model_output_def().get_data() output_def = { "atom_energy": out_def_data["energy"], @@ -545,12 +554,12 @@ def translated_output_def(self): def forward( self, - coord, - atype, - spin, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + coord: torch.Tensor, + atype: torch.Tensor, + spin: torch.Tensor, + box: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, ) -> dict[str, torch.Tensor]: model_ret = self.forward_common( @@ -575,16 +584,16 @@ def forward( @torch.jit.export def forward_lower( self, - extended_coord, - extended_atype, - extended_spin, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + extended_coord: torch.Tensor, + extended_atype: torch.Tensor, + extended_spin: torch.Tensor, + nlist: torch.Tensor, + mapping: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, - comm_dict: Optional[dict[str, torch.Tensor]] = None, - ): + comm_dict: dict[str, torch.Tensor] | None = None, + ) -> dict[str, torch.Tensor]: model_ret = self.forward_common_lower( extended_coord, extended_atype, diff --git a/deepmd/pt/model/model/transform_output.py b/deepmd/pt/model/model/transform_output.py index b8f1e024e0..839a028fcf 100644 --- a/deepmd/pt/model/model/transform_output.py +++ b/deepmd/pt/model/model/transform_output.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import torch @@ -20,7 +17,7 @@ def atomic_virial_corr( extended_coord: torch.Tensor, atom_energy: torch.Tensor, -): +) -> torch.Tensor: nall = extended_coord.shape[1] nloc = atom_energy.shape[1] coord, _ = torch.split(extended_coord, [nloc, nall - nloc], dim=1) @@ -29,7 +26,7 @@ def atomic_virial_corr( ce = coord * atom_energy sumce0, sumce1, sumce2 = torch.split(torch.sum(ce, dim=1), [1, 1, 1], dim=-1) faked_grad = torch.ones_like(sumce0) - lst = torch.jit.annotate(list[Optional[torch.Tensor]], [faked_grad]) + lst = torch.jit.annotate(list[torch.Tensor | None], [faked_grad]) extended_virial_corr0 = torch.autograd.grad( [sumce0], [extended_coord], @@ -72,9 +69,9 @@ def task_deriv_one( do_virial: bool = True, do_atomic_virial: bool = False, create_graph: bool = True, -): +) -> tuple[torch.Tensor, torch.Tensor | None]: faked_grad = torch.ones_like(energy) - lst = torch.jit.annotate(list[Optional[torch.Tensor]], [faked_grad]) + lst = torch.jit.annotate(list[torch.Tensor | None], [faked_grad]) extended_force = torch.autograd.grad( [energy], [extended_coord], @@ -102,7 +99,7 @@ def task_deriv_one( def get_leading_dims( vv: torch.Tensor, vdef: OutputVariableDef, -): +) -> list[int]: """Get the dimensions of nf x nloc.""" vshape = vv.shape return list(vshape[: (len(vshape) - len(vdef.shape))]) @@ -116,7 +113,7 @@ def take_deriv( do_virial: bool = False, do_atomic_virial: bool = False, create_graph: bool = True, -): +) -> tuple[torch.Tensor, torch.Tensor | None]: size = 1 for ii in vdef.shape: size *= ii @@ -158,6 +155,7 @@ def fit_output_to_model_output( coord_ext: torch.Tensor, do_atomic_virial: bool = False, create_graph: bool = True, + mask: torch.Tensor | None = None, ) -> dict[str, torch.Tensor]: """Transform the output of the fitting network to the model output. @@ -172,7 +170,12 @@ def fit_output_to_model_output( if vdef.reducible: kk_redu = get_reduce_name(kk) if vdef.intensive: - model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis) + if mask is not None: + model_ret[kk_redu] = torch.sum( + vv.to(redu_prec), dim=atom_axis + ) / torch.sum(mask, dim=-1, keepdim=True) + else: + model_ret[kk_redu] = torch.mean(vv.to(redu_prec), dim=atom_axis) else: model_ret[kk_redu] = torch.sum(vv.to(redu_prec), dim=atom_axis) if vdef.r_differentiable: diff --git a/deepmd/pt/model/network/init.py b/deepmd/pt/model/network/init.py index 53e2c70892..136690b749 100644 --- a/deepmd/pt/model/network/init.py +++ b/deepmd/pt/model/network/init.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import math import warnings -from typing import Optional as _Optional import torch from torch import ( @@ -18,19 +17,36 @@ # functions that use `with torch.no_grad()`. The JIT doesn't support context # managers, so these need to be implemented as builtins. Using these wrappers # lets us keep those builtins small and reusable. -def _no_grad_uniform_(tensor, a, b, generator=None): +def _no_grad_uniform_( + tensor: torch.Tensor, + a: float, + b: float, + generator: torch.Generator | None = None, +) -> torch.Tensor: with torch.no_grad(): return tensor.uniform_(a, b, generator=generator) -def _no_grad_normal_(tensor, mean, std, generator=None): +def _no_grad_normal_( + tensor: torch.Tensor, + mean: float, + std: float, + generator: torch.Generator | None = None, +) -> torch.Tensor: with torch.no_grad(): return tensor.normal_(mean, std, generator=generator) -def _no_grad_trunc_normal_(tensor, mean, std, a, b, generator=None): +def _no_grad_trunc_normal_( + tensor: torch.Tensor, + mean: float, + std: float, + a: float, + b: float, + generator: torch.Generator | None = None, +) -> torch.Tensor: # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf - def norm_cdf(x): + def norm_cdf(x: float) -> float: # Computes standard normal cumulative distribution function return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 @@ -65,17 +81,17 @@ def norm_cdf(x): return tensor -def _no_grad_zero_(tensor): +def _no_grad_zero_(tensor: torch.Tensor) -> torch.Tensor: with torch.no_grad(): return tensor.zero_() -def _no_grad_fill_(tensor, val): +def _no_grad_fill_(tensor: torch.Tensor, val: float) -> torch.Tensor: with torch.no_grad(): return tensor.fill_(val) -def calculate_gain(nonlinearity, param=None): +def calculate_gain(nonlinearity: str, param: float | None = None) -> float: r"""Return the recommended gain value for the given nonlinearity function. The values are as follows: @@ -146,7 +162,7 @@ def calculate_gain(nonlinearity, param=None): raise ValueError(f"Unsupported nonlinearity {nonlinearity}") -def _calculate_fan_in_and_fan_out(tensor): +def _calculate_fan_in_and_fan_out(tensor: torch.Tensor) -> tuple[int, int]: dimensions = tensor.dim() if dimensions < 2: raise ValueError( @@ -167,7 +183,7 @@ def _calculate_fan_in_and_fan_out(tensor): return fan_in, fan_out -def _calculate_correct_fan(tensor, mode): +def _calculate_correct_fan(tensor: torch.Tensor, mode: str) -> int: mode = mode.lower() valid_modes = ["fan_in", "fan_out"] if mode not in valid_modes: @@ -228,7 +244,7 @@ def normal_( tensor: Tensor, mean: float = 0.0, std: float = 1.0, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input Tensor with values drawn from the normal distribution. @@ -258,7 +274,7 @@ def trunc_normal_( std: float = 1.0, a: float = -2.0, b: float = 2.0, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input Tensor with values drawn from a truncated normal distribution. @@ -289,8 +305,8 @@ def kaiming_uniform_( a: float = 0, mode: str = "fan_in", nonlinearity: str = "leaky_relu", - generator: _Optional[torch.Generator] = None, -): + generator: torch.Generator | None = None, +) -> Tensor: r"""Fill the input `Tensor` with values using a Kaiming uniform distribution. The method is described in `Delving deep into rectifiers: Surpassing @@ -347,8 +363,8 @@ def kaiming_normal_( a: float = 0, mode: str = "fan_in", nonlinearity: str = "leaky_relu", - generator: _Optional[torch.Generator] = None, -): + generator: torch.Generator | None = None, +) -> Tensor: r"""Fill the input `Tensor` with values using a Kaiming normal distribution. The method is described in `Delving deep into rectifiers: Surpassing @@ -389,7 +405,7 @@ def kaiming_normal_( def xavier_uniform_( - tensor: Tensor, gain: float = 1.0, generator: _Optional[torch.Generator] = None + tensor: Tensor, gain: float = 1.0, generator: torch.Generator | None = None ) -> Tensor: r"""Fill the input `Tensor` with values using a Xavier uniform distribution. @@ -423,7 +439,7 @@ def xavier_uniform_( def xavier_normal_( tensor: Tensor, gain: float = 1.0, - generator: _Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> Tensor: r"""Fill the input `Tensor` with values using a Xavier normal distribution. diff --git a/deepmd/pt/model/network/layernorm.py b/deepmd/pt/model/network/layernorm.py index 89bd16d569..c7ffa3e205 100644 --- a/deepmd/pt/model/network/layernorm.py +++ b/deepmd/pt/model/network/layernorm.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np import torch @@ -30,21 +26,21 @@ device = env.DEVICE -def empty_t(shape, precision): +def empty_t(shape: tuple[int, ...], precision: torch.dtype) -> torch.Tensor: return torch.empty(shape, dtype=precision, device=device) class LayerNorm(nn.Module): def __init__( self, - num_in, + num_in: int, eps: float = 1e-5, uni_init: bool = True, bavg: float = 0.0, stddev: float = 1.0, precision: str = DEFAULT_PRECISION, trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, ) -> None: super().__init__() self.eps = eps @@ -141,7 +137,7 @@ def deserialize(cls, data: dict) -> "LayerNorm": ) prec = PRECISION_DICT[obj.precision] - def check_load_param(ss): + def check_load_param(ss: str) -> nn.Parameter | None: return ( nn.Parameter(data=to_torch_tensor(nl[ss])) if nl[ss] is not None diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index 22675d6163..02f7611429 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + Any, ClassVar, - Optional, - Union, ) import numpy as np @@ -43,7 +42,7 @@ ) -def empty_t(shape, precision): +def empty_t(shape: tuple[int, ...], precision: torch.dtype) -> torch.Tensor: return torch.empty(shape, dtype=precision, device=device) @@ -72,19 +71,21 @@ def deserialize(cls, data: dict) -> "Identity": class MLPLayer(nn.Module): def __init__( self, - num_in, - num_out, + num_in: int, + num_out: int, bias: bool = True, use_timestep: bool = False, - activation_function: Optional[str] = None, + activation_function: str | None = None, resnet: bool = False, bavg: float = 0.0, stddev: float = 1.0, precision: str = DEFAULT_PRECISION, init: str = "default", - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, + trainable: bool = True, ) -> None: super().__init__() + self.trainable = trainable # only use_timestep when skip connection is established. self.use_timestep = use_timestep and ( num_out == num_in or num_out == num_in * 2 @@ -130,7 +131,7 @@ def __init__( def check_type_consistency(self) -> None: precision = self.precision - def check_var(var) -> None: + def check_var(var: torch.Tensor | None) -> None: if var is not None: # assertion "float64" == "double" would fail assert PRECISION_DICT[var.dtype.name] is PRECISION_DICT[precision] @@ -149,7 +150,7 @@ def _default_normal_init( self, bavg: float = 0.0, stddev: float = 1.0, - generator: Optional[torch.Generator] = None, + generator: torch.Generator | None = None, ) -> None: normal_( self.matrix.data, @@ -162,7 +163,7 @@ def _default_normal_init( normal_(self.idt.data, mean=0.1, std=0.001, generator=generator) def _trunc_normal_init( - self, scale=1.0, generator: Optional[torch.Generator] = None + self, scale: float = 1.0, generator: torch.Generator | None = None ) -> None: # Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.) TRUNCATED_NORMAL_STDDEV_FACTOR = 0.87962566103423978 @@ -171,17 +172,17 @@ def _trunc_normal_init( std = (scale**0.5) / TRUNCATED_NORMAL_STDDEV_FACTOR trunc_normal_(self.matrix, mean=0.0, std=std, generator=generator) - def _glorot_uniform_init(self, generator: Optional[torch.Generator] = None) -> None: + def _glorot_uniform_init(self, generator: torch.Generator | None = None) -> None: xavier_uniform_(self.matrix, gain=1, generator=generator) - def _zero_init(self, use_bias=True) -> None: + def _zero_init(self, use_bias: bool = True) -> None: with torch.no_grad(): self.matrix.fill_(0.0) if use_bias and self.bias is not None: with torch.no_grad(): self.bias.fill_(1.0) - def _normal_init(self, generator: Optional[torch.Generator] = None) -> None: + def _normal_init(self, generator: torch.Generator | None = None) -> None: kaiming_normal_(self.matrix, nonlinearity="linear", generator=generator) def forward( @@ -233,6 +234,7 @@ def serialize(self) -> dict: activation_function=self.activate_name, resnet=self.resnet, precision=self.precision, + trainable=self.trainable, ) nl.w, nl.b, nl.idt = ( to_numpy_array(self.matrix), @@ -259,10 +261,11 @@ def deserialize(cls, data: dict) -> "MLPLayer": activation_function=nl["activation_function"], resnet=nl["resnet"], precision=nl["precision"], + trainable=nl["trainable"], ) prec = PRECISION_DICT[obj.precision] - def check_load_param(ss): + def check_load_param(ss: str) -> nn.Parameter | None: return ( nn.Parameter(data=to_torch_tensor(nl[ss])) if nl[ss] is not None @@ -279,7 +282,7 @@ def check_load_param(ss): class MLP(MLP_): - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.layers = torch.nn.ModuleList(self.layers) @@ -300,7 +303,7 @@ class NetworkCollection(DPNetworkCollection, nn.Module): "fitting_network": FittingNet, } - def __init__(self, *args, **kwargs) -> None: + def __init__(self, *args: Any, **kwargs: Any) -> None: # init both two base classes DPNetworkCollection.__init__(self, *args, **kwargs) nn.Module.__init__(self) diff --git a/deepmd/pt/model/network/network.py b/deepmd/pt/model/network/network.py index ab01a90774..a9662acfec 100644 --- a/deepmd/pt/model/network/network.py +++ b/deepmd/pt/model/network/network.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( + Any, Final, - Optional, - Union, ) import numpy as np @@ -32,7 +31,7 @@ ) -def Tensor(*shape): +def Tensor(*shape: int) -> torch.Tensor: return torch.empty(shape, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) @@ -41,12 +40,12 @@ class SimpleLinear(nn.Module): def __init__( self, - num_in, - num_out, - bavg=0.0, - stddev=1.0, - use_timestep=False, - activate=None, + num_in: int, + num_out: int, + bavg: float = 0.0, + stddev: float = 1.0, + use_timestep: bool = False, + activate: str | None = None, bias: bool = True, ) -> None: """Construct a linear layer. @@ -74,7 +73,7 @@ def __init__( self.idt = nn.Parameter(data=Tensor(1, num_out)) nn.init.normal_(self.idt.data, mean=0.1, std=0.001) - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: """Return X*W+b.""" xw = torch.matmul(inputs, self.matrix) hidden = xw + self.bias if self.bias is not None else xw @@ -121,7 +120,7 @@ def __init__( else: raise ValueError("Invalid init method.") - def _trunc_normal_init(self, scale=1.0) -> None: + def _trunc_normal_init(self, scale: float = 1.0) -> None: # Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.) TRUNCATED_NORMAL_STDDEV_FACTOR = 0.87962566103423978 _, fan_in = self.weight.shape @@ -132,7 +131,7 @@ def _trunc_normal_init(self, scale=1.0) -> None: def _glorot_uniform_init(self) -> None: nn.init.xavier_uniform_(self.weight, gain=1) - def _zero_init(self, use_bias=True) -> None: + def _zero_init(self, use_bias: bool = True) -> None: with torch.no_grad(): self.weight.fill_(0.0) if use_bias: @@ -144,13 +143,19 @@ def _normal_init(self) -> None: class NonLinearHead(nn.Module): - def __init__(self, input_dim, out_dim, activation_fn, hidden=None) -> None: + def __init__( + self, + input_dim: int, + out_dim: int, + activation_fn: str, + hidden: int | None = None, + ) -> None: super().__init__() hidden = input_dim if not hidden else hidden self.linear1 = SimpleLinear(input_dim, hidden, activate=activation_fn) self.linear2 = SimpleLinear(hidden, out_dim) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.linear1(x) x = self.linear2(x) return x @@ -159,7 +164,13 @@ def forward(self, x): class MaskLMHead(nn.Module): """Head for masked language modeling.""" - def __init__(self, embed_dim, output_dim, activation_fn, weight=None) -> None: + def __init__( + self, + embed_dim: int, + output_dim: int, + activation_fn: str, + weight: torch.Tensor | None = None, + ) -> None: super().__init__() self.dense = SimpleLinear(embed_dim, embed_dim) self.activation_fn = ActivationFn(activation_fn) @@ -174,7 +185,12 @@ def __init__(self, embed_dim, output_dim, activation_fn, weight=None) -> None: torch.zeros(output_dim, dtype=env.GLOBAL_PT_FLOAT_PRECISION) # pylint: disable=no-explicit-dtype,no-explicit-device ) - def forward(self, features, masked_tokens: Optional[torch.Tensor] = None, **kwargs): + def forward( + self, + features: torch.Tensor, + masked_tokens: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: # Only project the masked tokens while training, # saves both memory and computation if masked_tokens is not None: @@ -190,7 +206,13 @@ def forward(self, features, masked_tokens: Optional[torch.Tensor] = None, **kwar class ResidualDeep(nn.Module): def __init__( - self, type_id, embedding_width, neuron, bias_atom_e, out_dim=1, resnet_dt=False + self, + type_id: int, + embedding_width: int, + neuron: list[int], + bias_atom_e: float, + out_dim: int = 1, + resnet_dt: bool = False, ) -> None: """Construct a filter on the given element as neighbor. @@ -221,7 +243,7 @@ def __init__( bias_atom_e = 0 self.final_layer = SimpleLinear(self.neuron[-1], self.out_dim, bias_atom_e) - def forward(self, inputs): + def forward(self, inputs: torch.Tensor) -> torch.Tensor: """Calculate decoded embedding for each atom. Args: @@ -244,15 +266,16 @@ def forward(self, inputs): class TypeEmbedNet(nn.Module): def __init__( self, - type_nums, - embed_dim, - bavg=0.0, - stddev=1.0, - precision="default", - seed: Optional[Union[int, list[int]]] = None, - use_econf_tebd=False, + type_nums: int, + embed_dim: int, + bavg: float = 0.0, + stddev: float = 1.0, + precision: str = "default", + seed: int | list[int] | None = None, + use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map=None, + type_map: list[str] | None = None, + trainable: bool = True, ) -> None: """Construct a type embedding net.""" super().__init__() @@ -273,10 +296,11 @@ def __init__( type_map=type_map, precision=precision, seed=seed, + trainable=trainable, ) # nn.init.normal_(self.embedding.weight[:-1], mean=bavg, std=stddev) - def forward(self, atype): + def forward(self, atype: torch.Tensor) -> torch.Tensor: """ Args: atype: Type of each input, [nframes, nloc] or [nframes, nloc, nnei]. @@ -288,7 +312,7 @@ def forward(self, atype): """ return torch.embedding(self.embedding(atype.device), atype) - def get_full_embedding(self, device: torch.device): + def get_full_embedding(self, device: torch.device) -> torch.Tensor: """ Get the type embeddings of all types. @@ -305,7 +329,9 @@ def get_full_embedding(self, device: torch.device): """ return self.embedding(device) - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, base_class: Any, shared_level: int, resume: bool = False + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -322,7 +348,7 @@ def share_params(self, base_class, shared_level, resume=False) -> None: raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -368,11 +394,11 @@ def __init__( activation_function: str = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[Union[int, list[int]]] = None, + seed: int | list[int] | None = None, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: """Construct a type embedding net.""" super().__init__() @@ -407,7 +433,7 @@ def __init__( for param in self.parameters(): param.requires_grad = trainable - def forward(self, device: torch.device): + def forward(self, device: torch.device) -> torch.Tensor: """Caulate type embedding network. Returns @@ -429,7 +455,7 @@ def forward(self, device: torch.device): return embed def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -491,7 +517,7 @@ def change_type_map( self.ntypes = len(type_map) @classmethod - def deserialize(cls, data: dict): + def deserialize(cls, data: dict) -> "TypeEmbedNetConsistent": """Deserialize the model. Parameters diff --git a/deepmd/pt/model/network/utils.py b/deepmd/pt/model/network/utils.py index 2047efec2b..0613ca3e4d 100644 --- a/deepmd/pt/model/network/utils.py +++ b/deepmd/pt/model/network/utils.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import torch @@ -11,7 +8,7 @@ def aggregate( data: torch.Tensor, owners: torch.Tensor, average: bool = True, - num_owner: Optional[int] = None, + num_owner: int | None = None, ) -> torch.Tensor: """ Aggregate rows in data by specifying the owners. @@ -57,7 +54,7 @@ def get_graph_index( a_nlist_mask: torch.Tensor, nall: int, use_loc_mapping: bool = True, -): +) -> tuple[torch.Tensor, torch.Tensor]: """ Get the index mapping for edge graph and angle graph, ready in `aggregate` or `index_select`. @@ -74,12 +71,12 @@ def get_graph_index( Returns ------- - edge_index : n_edge x 2 + edge_index : 2 x n_edge n2e_index : n_edge Broadcast indices from node(i) to edge(ij), or reduction indices from edge(ij) to node(i). n_ext2e_index : n_edge Broadcast indices from extended node(j) to edge(ij). - angle_index : n_angle x 3 + angle_index : 3 x n_angle n2a_index : n_angle Broadcast indices from extended node(j) to angle(ijk). eij2a_index : n_angle @@ -135,9 +132,7 @@ def get_graph_index( # n_angle eik2a_index = edge_index_ik[a_nlist_mask_3d] - return torch.cat( - [n2e_index.unsqueeze(-1), n_ext2e_index.unsqueeze(-1)], dim=-1 - ), torch.cat( - [n2a_index.unsqueeze(-1), eij2a_index.unsqueeze(-1), eik2a_index.unsqueeze(-1)], - dim=-1, - ) + edge_index_result = torch.stack([n2e_index, n_ext2e_index], dim=0) + angle_index_result = torch.stack([n2a_index, eij2a_index, eik2a_index], dim=0) + + return edge_index_result, angle_index_result diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index fc9e8943e9..bb155322e9 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -26,12 +26,12 @@ class DenoiseNet(Fitting): def __init__( self, - feature_dim, - ntypes, - attn_head=8, - prefactor=[0.5, 0.5], - activation_function="gelu", - **kwargs, + feature_dim: int, + ntypes: int, + attn_head: int = 8, + prefactor: list[float] = [0.5, 0.5], + activation_function: str = "gelu", + **kwargs: Any, ) -> None: """Construct a denoise net. @@ -71,7 +71,7 @@ def __init__( self.pair2coord_proj.append(_pair2coord_proj) self.pair2coord_proj = torch.nn.ModuleList(self.pair2coord_proj) - def output_def(self): + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ OutputVariableDef( @@ -93,13 +93,13 @@ def output_def(self): def forward( self, - pair_weights, - diff, - nlist_mask, - features, - sw, - masked_tokens: Optional[torch.Tensor] = None, - ): + pair_weights: torch.Tensor, + diff: torch.Tensor, + nlist_mask: torch.Tensor, + features: torch.Tensor, + sw: torch.Tensor, + masked_tokens: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: """Calculate the updated coord. Args: - coord: Input noisy coord with shape [nframes, nloc, 3]. diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 65b64220ae..6e9b0340ff 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import torch @@ -72,6 +73,9 @@ class DipoleFittingNet(GeneralFitting): Only reducible variable are differentiable. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -87,13 +91,14 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], r_differentiable: bool = True, c_differentiable: bool = True, - type_map: Optional[list[str]] = None, - **kwargs, + type_map: list[str] | None = None, + default_fparam: list | None = None, + **kwargs: Any, ) -> None: self.embedding_width = embedding_width self.r_differentiable = r_differentiable @@ -114,10 +119,11 @@ def __init__( seed=seed, exclude_types=exclude_types, type_map=type_map, + default_fparam=default_fparam, **kwargs, ) - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" return self.embedding_width @@ -132,7 +138,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name", None) return super().deserialize(data) @@ -151,8 +157,8 @@ def output_def(self) -> FittingOutputDef: def compute_output_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], - stat_file_path: Optional[DPPath] = None, + merged: Callable[[], list[dict]] | list[dict], + stat_file_path: DPPath | None = None, ) -> None: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -176,12 +182,12 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: nframes, nloc, _ = descriptor.shape assert gr is not None, "Must provide the rotation matrix for dipole fitting." # cast the input to internal precsion diff --git a/deepmd/pt/model/task/dos.py b/deepmd/pt/model/task/dos.py index 568ef81c92..b11bfd7c7f 100644 --- a/deepmd/pt/model/task/dos.py +++ b/deepmd/pt/model/task/dos.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, - Union, -) import torch @@ -48,15 +44,16 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, - bias_dos: Optional[torch.Tensor] = None, - trainable: Union[bool, list[bool]] = True, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + bias_dos: torch.Tensor | None = None, + trainable: bool | list[bool] = True, + seed: int | list[int] | None = None, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, exclude_types: list[int] = [], mixed_types: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, ) -> None: if bias_dos is not None: self.bias_dos = bias_dos @@ -83,6 +80,7 @@ def __init__( exclude_types=exclude_types, trainable=trainable, type_map=type_map, + default_fparam=default_fparam, ) def output_def(self) -> FittingOutputDef: @@ -101,7 +99,7 @@ def output_def(self) -> FittingOutputDef: @classmethod def deserialize(cls, data: dict) -> "DOSFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("@class", None) data.pop("var_name", None) data.pop("tot_ener_zero", None) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 07351b33f6..39dabf99f8 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, - Union, + Any, ) import numpy as np @@ -46,7 +45,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -54,9 +53,10 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - seed: Optional[Union[int, list[int]]] = None, - type_map: Optional[list[str]] = None, - **kwargs, + seed: int | list[int] | None = None, + type_map: list[str] | None = None, + default_fparam: list | None = None, + **kwargs: Any, ) -> None: super().__init__( "energy", @@ -74,13 +74,14 @@ def __init__( mixed_types=mixed_types, seed=seed, type_map=type_map, + default_fparam=default_fparam, **kwargs, ) @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data.pop("var_name") data.pop("dim_out") return super().deserialize(data) @@ -102,15 +103,15 @@ def serialize(self) -> dict: class EnergyFittingNetDirect(Fitting): def __init__( self, - ntypes, - dim_descrpt, - neuron, - bias_atom_e=None, - out_dim=1, - resnet_dt=True, - use_tebd=True, - return_energy=False, - **kwargs, + ntypes: int, + dim_descrpt: int, + neuron: list[int], + bias_atom_e: torch.Tensor | None = None, + out_dim: int = 1, + resnet_dt: bool = True, + use_tebd: bool = True, + return_energy: bool = False, + **kwargs: Any, ) -> None: """Construct a fitting net for energy. @@ -160,7 +161,7 @@ def __init__( filter_layers.append(one) self.filter_layers = torch.nn.ModuleList(filter_layers) - def output_def(self): + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ OutputVariableDef( @@ -187,7 +188,7 @@ def deserialize(self) -> "EnergyFittingNetDirect": raise NotImplementedError def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: raise NotImplementedError @@ -198,11 +199,11 @@ def forward( self, inputs: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ) -> tuple[torch.Tensor, None]: """Based on embedding net output, alculate total energy. diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 0865b61f52..f46bc0fb7d 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -3,10 +3,12 @@ from abc import ( abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, +) +from typing import ( + Any, Optional, - Union, ) import numpy as np @@ -36,10 +38,16 @@ to_numpy_array, to_torch_tensor, ) +from deepmd.utils.env_mat_stat import ( + StatItem, +) from deepmd.utils.finetune import ( get_index_between_two_maps, map_atom_exclude_types, ) +from deepmd.utils.path import ( + DPPath, +) dtype = env.GLOBAL_PT_FLOAT_PRECISION device = env.DEVICE @@ -50,12 +58,19 @@ class Fitting(torch.nn.Module, BaseFitting): # plugin moved to BaseFitting - def __new__(cls, *args, **kwargs): + def __new__(cls, *args: Any, **kwargs: Any) -> "Fitting": if cls is Fitting: return BaseFitting.__new__(BaseFitting, *args, **kwargs) return super().__new__(cls) - def share_params(self, base_class, shared_level, resume=False) -> None: + def share_params( + self, + base_class: "Fitting", + shared_level: int, + model_prob: float = 1.0, + protection: float = 1e-2, + resume: bool = False, + ) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -66,16 +81,164 @@ def share_params(self, base_class, shared_level, resume=False) -> None: ) if shared_level == 0: # only not share the bias_atom_e and the case_embd + # link fparam buffers + if self.numb_fparam > 0: + if not resume: + base_fparam = base_class.stats["fparam"] + assert len(base_fparam) == self.numb_fparam + for ii in range(self.numb_fparam): + base_fparam[ii] += self.get_stats()["fparam"][ii] * model_prob + fparam_avg = np.array([ii.compute_avg() for ii in base_fparam]) + fparam_std = np.array( + [ii.compute_std(protection=protection) for ii in base_fparam] + ) + fparam_inv_std = 1.0 / fparam_std + base_class.fparam_avg.copy_( + torch.tensor( + fparam_avg, + device=env.DEVICE, + dtype=base_class.fparam_avg.dtype, + ) + ) + base_class.fparam_inv_std.copy_( + torch.tensor( + fparam_inv_std, + device=env.DEVICE, + dtype=base_class.fparam_inv_std.dtype, + ) + ) + self.fparam_avg = base_class.fparam_avg + self.fparam_inv_std = base_class.fparam_inv_std + + # link aparam buffers + if self.numb_aparam > 0: + if not resume: + base_aparam = base_class.stats["aparam"] + assert len(base_aparam) == self.numb_aparam + for ii in range(self.numb_aparam): + base_aparam[ii] += self.get_stats()["aparam"][ii] * model_prob + aparam_avg = np.array([ii.compute_avg() for ii in base_aparam]) + aparam_std = np.array( + [ii.compute_std(protection=protection) for ii in base_aparam] + ) + aparam_inv_std = 1.0 / aparam_std + base_class.aparam_avg.copy_( + torch.tensor( + aparam_avg, + device=env.DEVICE, + dtype=base_class.aparam_avg.dtype, + ) + ) + base_class.aparam_inv_std.copy_( + torch.tensor( + aparam_inv_std, + device=env.DEVICE, + dtype=base_class.aparam_inv_std.dtype, + ) + ) + self.aparam_avg = base_class.aparam_avg + self.aparam_inv_std = base_class.aparam_inv_std # the following will successfully link all the params except buffers, which need manually link. for item in self._modules: self._modules[item] = base_class._modules[item] else: raise NotImplementedError + def save_to_file_fparam( + self, + stat_file_path: DPPath, + ) -> None: + """Save the statistics of fparam. + + Parameters + ---------- + stat_file_path : DPPath + The path to save the statistics of fparam. + """ + assert stat_file_path is not None + stat_file_path.mkdir(exist_ok=True, parents=True) + if len(self.stats) == 0: + raise ValueError("The statistics hasn't been computed.") + fp = stat_file_path / "fparam" + _fparam_stat = [] + for ii in range(self.numb_fparam): + _tmp_stat = self.stats["fparam"][ii] + _fparam_stat.append( + [_tmp_stat.number, _tmp_stat.sum, _tmp_stat.squared_sum] + ) + _fparam_stat = np.array(_fparam_stat) + fp.save_numpy(_fparam_stat) + log.info(f"Save fparam stats to {fp}.") + + def save_to_file_aparam( + self, + stat_file_path: DPPath, + ) -> None: + """Save the statistics of aparam. + + Parameters + ---------- + stat_file_path : DPPath + The path to save the statistics of aparam. + """ + assert stat_file_path is not None + stat_file_path.mkdir(exist_ok=True, parents=True) + if len(self.stats) == 0: + raise ValueError("The statistics hasn't been computed.") + fp = stat_file_path / "aparam" + _aparam_stat = [] + for ii in range(self.numb_aparam): + _tmp_stat = self.stats["aparam"][ii] + _aparam_stat.append( + [_tmp_stat.number, _tmp_stat.sum, _tmp_stat.squared_sum] + ) + _aparam_stat = np.array(_aparam_stat) + fp.save_numpy(_aparam_stat) + log.info(f"Save aparam stats to {fp}.") + + def restore_fparam_from_file(self, stat_file_path: DPPath) -> None: + """Load the statistics of fparam. + + Parameters + ---------- + stat_file_path : DPPath + The path to load the statistics of fparam. + """ + fp = stat_file_path / "fparam" + arr = fp.load_numpy() + assert arr.shape == (self.numb_fparam, 3) + _fparam_stat = [] + for ii in range(self.numb_fparam): + _fparam_stat.append( + StatItem(number=arr[ii][0], sum=arr[ii][1], squared_sum=arr[ii][2]) + ) + self.stats["fparam"] = _fparam_stat + log.info(f"Load fparam stats from {fp}.") + + def restore_aparam_from_file(self, stat_file_path: DPPath) -> None: + """Load the statistics of aparam. + + Parameters + ---------- + stat_file_path : DPPath + The path to load the statistics of aparam. + """ + fp = stat_file_path / "aparam" + arr = fp.load_numpy() + assert arr.shape == (self.numb_aparam, 3) + _aparam_stat = [] + for ii in range(self.numb_aparam): + _aparam_stat.append( + StatItem(number=arr[ii][0], sum=arr[ii][1], squared_sum=arr[ii][2]) + ) + self.stats["aparam"] = _aparam_stat + log.info(f"Load aparam stats from {fp}.") + def compute_input_stats( self, - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], protection: float = 1e-2, + stat_file_path: DPPath | None = None, ) -> None: """ Compute the input statistics (e.g. mean and stddev) for the fittings from packed data. @@ -91,67 +254,101 @@ def compute_input_stats( the lazy function helps by only sampling once. protection : float Divided-by-zero protection + stat_file_path : Optional[DPPath] + The path to the stat file. """ if self.numb_fparam == 0 and self.numb_aparam == 0: # skip data statistics + self.stats = None return - if callable(merged): - sampled = merged() - else: - sampled = merged + + self.stats = {} + # stat fparam if self.numb_fparam > 0: - cat_data = torch.cat([frame["fparam"] for frame in sampled], dim=0) - cat_data = torch.reshape(cat_data, [-1, self.numb_fparam]) - fparam_avg = torch.mean(cat_data, dim=0) - fparam_std = torch.std(cat_data, dim=0, unbiased=False) - fparam_std = torch.where( - fparam_std < protection, - torch.tensor( - protection, dtype=fparam_std.dtype, device=fparam_std.device - ), - fparam_std, - ) - fparam_inv_std = 1.0 / fparam_std - self.fparam_avg.copy_( - torch.tensor(fparam_avg, device=env.DEVICE, dtype=self.fparam_avg.dtype) - ) - self.fparam_inv_std.copy_( - torch.tensor( - fparam_inv_std, device=env.DEVICE, dtype=self.fparam_inv_std.dtype + if ( + stat_file_path is not None + and stat_file_path.is_dir() + and (stat_file_path / "fparam").is_file() + ): + self.restore_fparam_from_file(stat_file_path) + else: + sampled = merged() if callable(merged) else merged + self.stats["fparam"] = [] + cat_data = to_numpy_array( + torch.cat([frame["fparam"] for frame in sampled], dim=0) ) + cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) + sumv = np.sum(cat_data, axis=0) + sumv2 = np.sum(cat_data * cat_data, axis=0) + sumn = cat_data.shape[0] + for ii in range(self.numb_fparam): + self.stats["fparam"].append( + StatItem( + number=sumn, + sum=sumv[ii], + squared_sum=sumv2[ii], + ) + ) + if stat_file_path is not None: + self.save_to_file_fparam(stat_file_path) + + fparam_avg = np.array([ii.compute_avg() for ii in self.stats["fparam"]]) + fparam_std = np.array( + [ii.compute_std(protection=protection) for ii in self.stats["fparam"]] ) + fparam_inv_std = 1.0 / fparam_std + log.info(f"fparam_avg is {fparam_avg}, fparam_inv_std is {fparam_inv_std}") + self.fparam_avg.copy_(to_torch_tensor(fparam_avg)) + self.fparam_inv_std.copy_(to_torch_tensor(fparam_inv_std)) + # stat aparam if self.numb_aparam > 0: - sys_sumv = [] - sys_sumv2 = [] - sys_sumn = [] - for ss_ in [frame["aparam"] for frame in sampled]: - ss = torch.reshape(ss_, [-1, self.numb_aparam]) - sys_sumv.append(torch.sum(ss, dim=0)) - sys_sumv2.append(torch.sum(ss * ss, dim=0)) - sys_sumn.append(ss.shape[0]) - sumv = torch.sum(torch.stack(sys_sumv), dim=0) - sumv2 = torch.sum(torch.stack(sys_sumv2), dim=0) - sumn = sum(sys_sumn) - aparam_avg = sumv / sumn - aparam_std = torch.sqrt(sumv2 / sumn - (sumv / sumn) ** 2) - aparam_std = torch.where( - aparam_std < protection, - torch.tensor( - protection, dtype=aparam_std.dtype, device=aparam_std.device - ), - aparam_std, + if ( + stat_file_path is not None + and stat_file_path.is_dir() + and (stat_file_path / "aparam").is_file() + ): + self.restore_aparam_from_file(stat_file_path) + else: + sampled = merged() if callable(merged) else merged + self.stats["aparam"] = [] + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in [frame["aparam"] for frame in sampled]: + ss = np.reshape(to_numpy_array(ss_), [-1, self.numb_aparam]) + sys_sumv.append(np.sum(ss, axis=0)) + sys_sumv2.append(np.sum(ss * ss, axis=0)) + sys_sumn.append(ss.shape[0]) + sumv = np.sum(np.stack(sys_sumv), axis=0) + sumv2 = np.sum(np.stack(sys_sumv2), axis=0) + sumn = sum(sys_sumn) + for ii in range(self.numb_aparam): + self.stats["aparam"].append( + StatItem( + number=sumn, + sum=sumv[ii], + squared_sum=sumv2[ii], + ) + ) + if stat_file_path is not None: + self.save_to_file_aparam(stat_file_path) + + aparam_avg = np.array([ii.compute_avg() for ii in self.stats["aparam"]]) + aparam_std = np.array( + [ii.compute_std(protection=protection) for ii in self.stats["aparam"]] ) aparam_inv_std = 1.0 / aparam_std - self.aparam_avg.copy_( - torch.tensor(aparam_avg, device=env.DEVICE, dtype=self.aparam_avg.dtype) - ) - self.aparam_inv_std.copy_( - torch.tensor( - aparam_inv_std, device=env.DEVICE, dtype=self.aparam_inv_std.dtype - ) - ) + log.info(f"aparam_avg is {aparam_avg}, aparam_inv_std is {aparam_inv_std}") + self.aparam_avg.copy_(to_torch_tensor(aparam_avg)) + self.aparam_inv_std.copy_(to_torch_tensor(aparam_inv_std)) + + def get_stats(self) -> dict[str, list[StatItem]]: + """Get the statistics of the fitting_net.""" + if self.stats is None: + raise RuntimeError("The statistics of fitting net has not been computed.") + return self.stats class GeneralFitting(Fitting): @@ -204,6 +401,9 @@ class GeneralFitting(Fitting): A list of strings. Give the name to each type of atoms. use_aparam_as_mask: bool If True, the aparam will not be used in fitting net for embedding. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -212,7 +412,7 @@ def __init__( ntypes: int, dim_descrpt: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -220,14 +420,15 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - trainable: Union[bool, list[bool]] = True, - remove_vaccum_contribution: Optional[list[bool]] = None, - type_map: Optional[list[str]] = None, + trainable: bool | list[bool] = True, + remove_vaccum_contribution: list[bool] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, - **kwargs, + default_fparam: list[float] | None = None, + **kwargs: Any, ) -> None: super().__init__() self.var_name = var_name @@ -238,6 +439,7 @@ def __init__( self.resnet_dt = resnet_dt self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam + self.default_fparam = default_fparam self.dim_case_embd = dim_case_embd self.activation_function = activation_function self.precision = precision @@ -299,6 +501,20 @@ def __init__( else: self.case_embd = None + if self.default_fparam is not None: + if self.numb_fparam > 0: + assert len(self.default_fparam) == self.numb_fparam, ( + "default_fparam length mismatch!" + ) + self.register_buffer( + "default_fparam_tensor", + torch.tensor( + np.array(self.default_fparam), dtype=self.prec, device=device + ), + ) + else: + self.default_fparam_tensor = None + in_dim = ( self.dim_descrpt + self.numb_fparam @@ -320,6 +536,7 @@ def __init__( self.precision, bias_out=True, seed=child_seed(self.seed, ii), + trainable=trainable, ) for ii in range(self.ntypes if not self.mixed_types else 1) ], @@ -328,6 +545,8 @@ def __init__( for param in self.parameters(): param.requires_grad = self.trainable + self.eval_return_middle_output = False + def reinit_exclude( self, exclude_types: list[int] = [], @@ -336,7 +555,9 @@ def reinit_exclude( self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, + type_map: list[str], + model_with_new_type_stat: Optional["GeneralFitting"] = None, ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -363,7 +584,7 @@ def serialize(self) -> dict: """Serialize the fitting to dict.""" return { "@class": "Fitting", - "@version": 3, + "@version": 4, "var_name": self.var_name, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -372,6 +593,7 @@ def serialize(self) -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function, "precision": self.precision, "mixed_types": self.mixed_types, @@ -415,6 +637,13 @@ def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" return self.numb_fparam + def has_default_fparam(self) -> bool: + """Check if the fitting has default frame parameters.""" + return self.default_fparam is not None + + def get_default_fparam(self) -> torch.Tensor | None: + return self.default_fparam_tensor + def get_dim_aparam(self) -> int: """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam @@ -440,7 +669,7 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def set_case_embd(self, case_idx: int): + def set_case_embd(self, case_idx: int) -> None: """ Set the case embedding of this fitting net by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. @@ -449,7 +678,10 @@ def set_case_embd(self, case_idx: int): case_idx ] - def __setitem__(self, key, value) -> None: + def set_return_middle_output(self, return_middle_output: bool = True) -> None: + self.eval_return_middle_output = return_middle_output + + def __setitem__(self, key: str, value: torch.Tensor) -> None: if key in ["bias_atom_e"]: value = value.view([self.ntypes, self._net_out_dim()]) self.bias_atom_e = value @@ -465,10 +697,12 @@ def __setitem__(self, key, value) -> None: self.case_embd = value elif key in ["scale"]: self.scale = value + elif key in ["default_fparam_tensor"]: + self.default_fparam_tensor = value else: raise KeyError(key) - def __getitem__(self, key): + def __getitem__(self, key: str) -> torch.Tensor: if key in ["bias_atom_e"]: return self.bias_atom_e elif key in ["fparam_avg"]: @@ -483,11 +717,13 @@ def __getitem__(self, key): return self.case_embd elif key in ["scale"]: return self.scale + elif key in ["default_fparam_tensor"]: + return self.default_fparam_tensor else: raise KeyError(key) @abstractmethod - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" pass @@ -501,14 +737,21 @@ def _forward_common( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: # cast the input to internal precsion xx = descriptor.to(self.prec) + nf, nloc, nd = xx.shape + + if self.numb_fparam > 0 and fparam is None: + # use default fparam + assert self.default_fparam_tensor is not None + fparam = torch.tile(self.default_fparam_tensor.unsqueeze(0), [nf, 1]) + fparam = fparam.to(self.prec) if fparam is not None else None aparam = aparam.to(self.prec) if aparam is not None else None @@ -521,7 +764,6 @@ def _forward_common( xx_zeros = torch.zeros_like(xx) else: xx_zeros = None - nf, nloc, nd = xx.shape net_dim_out = self._net_out_dim() if nd != self.dim_descrpt: @@ -597,14 +839,37 @@ def _forward_common( dtype=self.prec, device=descriptor.device, ) # jit assertion + results = {} + if self.mixed_types: atom_property = self.filter_layers.networks[0](xx) + if self.eval_return_middle_output: + results["middle_output"] = self.filter_layers.networks[ + 0 + ].call_until_last(xx) if xx_zeros is not None: atom_property -= self.filter_layers.networks[0](xx_zeros) outs = ( outs + atom_property + self.bias_atom_e[atype].to(self.prec) ) # Shape is [nframes, natoms[0], net_dim_out] else: + if self.eval_return_middle_output: + outs_middle = torch.zeros( + (nf, nloc, self.neuron[-1]), + dtype=self.prec, + device=descriptor.device, + ) # jit assertion + for type_i, ll in enumerate(self.filter_layers.networks): + mask = (atype == type_i).unsqueeze(-1) + mask = torch.tile(mask, (1, 1, net_dim_out)) + middle_output_type = ll.call_until_last(xx) + middle_output_type = torch.where( + torch.tile(mask, (1, 1, self.neuron[-1])), + middle_output_type, + 0.0, + ) + outs_middle = outs_middle + middle_output_type + results["middle_output"] = outs_middle for type_i, ll in enumerate(self.filter_layers.networks): mask = (atype == type_i).unsqueeze(-1) mask = torch.tile(mask, (1, 1, net_dim_out)) @@ -626,4 +891,10 @@ def _forward_common( mask = self.emask(atype).to(torch.bool) # nf x nloc x nod outs = torch.where(mask[:, :, None], outs, 0.0) - return {self.var_name: outs} + results.update({self.var_name: outs}) + return results + + @torch.jit.export + def get_task_dim(self) -> int: + """Get the output dimension of the fitting net.""" + return self._net_out_dim() diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index b1599eac60..a8953fcd2b 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, - Union, + Any, ) import torch @@ -80,6 +79,9 @@ class InvarFitting(GeneralFitting): A list of strings. Give the name to each type of atoms. use_aparam_as_mask: bool If True, the aparam will not be used in fitting net for embedding. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -89,7 +91,7 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: list[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_e: torch.Tensor | None = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -97,13 +99,14 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], - atom_ener: Optional[list[Optional[torch.Tensor]]] = None, - type_map: Optional[list[str]] = None, + atom_ener: list[torch.Tensor | None] | None = None, + type_map: list[str] | None = None, use_aparam_as_mask: bool = False, - **kwargs, + default_fparam: list[float] | None = None, + **kwargs: Any, ) -> None: self.dim_out = dim_out self.atom_ener = atom_ener @@ -128,10 +131,11 @@ def __init__( else [x is not None for x in atom_ener], type_map=type_map, use_aparam_as_mask=use_aparam_as_mask, + default_fparam=default_fparam, **kwargs, ) - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" return self.dim_out @@ -145,7 +149,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) return super().deserialize(data) def output_def(self) -> FittingOutputDef: @@ -165,12 +169,12 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: """Based on embedding net output, alculate total energy. Args: @@ -181,10 +185,17 @@ def forward( ------- - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. """ - out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[ - self.var_name - ] - return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} + out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam) + result = {self.var_name: out[self.var_name].to(env.GLOBAL_PT_FLOAT_PRECISION)} + if "middle_output" in out: + result.update( + { + "middle_output": out["middle_output"].to( + env.GLOBAL_PT_FLOAT_PRECISION + ) + } + ) + return result # make jit happy with torch 2.0.0 exclude_types: list[int] diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index a326802918..c3a7ed52a1 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, - Union, + Any, ) import torch @@ -75,7 +74,9 @@ class PolarFittingNet(GeneralFitting): Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. - + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. """ def __init__( @@ -91,14 +92,15 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[Union[int, list[int]]] = None, + rcond: float | None = None, + seed: int | list[int] | None = None, exclude_types: list[int] = [], fit_diag: bool = True, - scale: Optional[Union[list[float], float]] = None, + scale: list[float] | float | None = None, shift_diag: bool = True, - type_map: Optional[list[str]] = None, - **kwargs, + type_map: list[str] | None = None, + default_fparam: list | None = None, + **kwargs: Any, ) -> None: self.embedding_width = embedding_width self.fit_diag = fit_diag @@ -139,10 +141,11 @@ def __init__( seed=seed, exclude_types=exclude_types, type_map=type_map, + default_fparam=default_fparam, **kwargs, ) - def _net_out_dim(self): + def _net_out_dim(self) -> int: """Set the FittingNet output dim.""" return ( self.embedding_width @@ -150,20 +153,20 @@ def _net_out_dim(self): else self.embedding_width * self.embedding_width ) - def __setitem__(self, key, value) -> None: + def __setitem__(self, key: str, value: Any) -> None: if key in ["constant_matrix"]: self.constant_matrix = value else: super().__setitem__(key, value) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: if key in ["constant_matrix"]: return self.constant_matrix else: return super().__getitem__(key) def change_type_map( - self, type_map: list[str], model_with_new_type_stat=None + self, type_map: list[str], model_with_new_type_stat: Any | None = None ) -> None: """Change the type related params to new ones, according to `type_map` and the original one in the model. If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. @@ -195,7 +198,7 @@ def change_type_map( def serialize(self) -> dict: data = super().serialize() data["type"] = "polar" - data["@version"] = 4 + data["@version"] = 5 data["embedding_width"] = self.embedding_width data["fit_diag"] = self.fit_diag data["shift_diag"] = self.shift_diag @@ -206,7 +209,7 @@ def serialize(self) -> dict: @classmethod def deserialize(cls, data: dict) -> "GeneralFitting": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("var_name", None) return super().deserialize(data) @@ -227,12 +230,12 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> dict[str, torch.Tensor]: nframes, nloc, _ = descriptor.shape assert gr is not None, ( "Must provide the rotation matrix for polarizability fitting." diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 5ef0cd0233..4d0b791cfa 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, - Union, + Any, ) import torch @@ -80,7 +79,7 @@ def __init__( property_name: str, task_dim: int = 1, neuron: list[int] = [128, 128, 128], - bias_atom_p: Optional[torch.Tensor] = None, + bias_atom_p: torch.Tensor | None = None, intensive: bool = False, resnet_dt: bool = True, numb_fparam: int = 0, @@ -89,9 +88,10 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - trainable: Union[bool, list[bool]] = True, - seed: Optional[int] = None, - **kwargs, + trainable: bool | list[bool] = True, + seed: int | None = None, + default_fparam: list | None = None, + **kwargs: Any, ) -> None: self.task_dim = task_dim self.intensive = intensive @@ -111,6 +111,7 @@ def __init__( mixed_types=mixed_types, trainable=trainable, seed=seed, + default_fparam=default_fparam, **kwargs, ) @@ -135,7 +136,7 @@ def get_intensive(self) -> bool: @classmethod def deserialize(cls, data: dict) -> "PropertyFittingNet": data = data.copy() - check_version_compatibility(data.pop("@version", 1), 4, 1) + check_version_compatibility(data.pop("@version", 1), 5, 1) data.pop("dim_out") data["property_name"] = data.pop("var_name") obj = super().deserialize(data) @@ -150,7 +151,7 @@ def serialize(self) -> dict: "task_dim": self.task_dim, "intensive": self.intensive, } - dd["@version"] = 4 + dd["@version"] = 5 return dd diff --git a/deepmd/pt/model/task/type_predict.py b/deepmd/pt/model/task/type_predict.py index e4a980c3ea..a5b2ff8884 100644 --- a/deepmd/pt/model/task/type_predict.py +++ b/deepmd/pt/model/task/type_predict.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, + Any, ) import torch @@ -15,7 +15,11 @@ class TypePredictNet(Fitting): def __init__( - self, feature_dim, ntypes, activation_function="gelu", **kwargs + self, + feature_dim: int, + ntypes: int, + activation_function: str = "gelu", + **kwargs: Any, ) -> None: """Construct a type predict net. @@ -34,7 +38,9 @@ def __init__( weight=None, ) - def forward(self, features, masked_tokens: Optional[torch.Tensor] = None): + def forward( + self, features: torch.Tensor, masked_tokens: torch.Tensor | None = None + ) -> torch.Tensor: """Calculate the predicted logits. Args: - features: Input features with shape [nframes, nloc, feature_dim]. diff --git a/deepmd/pt/optimizer/LKF.py b/deepmd/pt/optimizer/LKF.py index c342960e5b..c79e875f3e 100644 --- a/deepmd/pt/optimizer/LKF.py +++ b/deepmd/pt/optimizer/LKF.py @@ -1,6 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import math +from typing import ( + Any, +) import torch import torch.distributed as dist @@ -9,7 +12,7 @@ ) -def distribute_indices(total_length, num_workers): +def distribute_indices(total_length: int, num_workers: int) -> list[tuple[int, int]]: indices_per_worker = total_length // num_workers remainder = total_length % num_workers @@ -27,10 +30,10 @@ def distribute_indices(total_length, num_workers): class LKFOptimizer(Optimizer): def __init__( self, - params, - kalman_lambda=0.98, - kalman_nue=0.9987, - block_size=5120, + params: Any, + kalman_lambda: float = 0.98, + kalman_nue: float = 0.9987, + block_size: int = 5120, ) -> None: defaults = {"lr": 0.1, "kalman_nue": kalman_nue, "block_size": block_size} @@ -158,13 +161,13 @@ def __init_P(self) -> None: self._state.setdefault("weights_num", len(P)) self._state.setdefault("params_packed_index", params_packed_index) - def __get_blocksize(self): + def __get_blocksize(self) -> int: return self.param_groups[0]["block_size"] - def __get_nue(self): + def __get_nue(self) -> float: return self.param_groups[0]["kalman_nue"] - def __split_weights(self, weight): + def __split_weights(self, weight: torch.Tensor) -> list[torch.Tensor]: block_size = self.__get_blocksize() param_num = weight.nelement() res = [] @@ -179,7 +182,9 @@ def __split_weights(self, weight): res.append(weight[i * block_size :]) return res - def __update(self, H, error, weights) -> None: + def __update( + self, H: torch.Tensor, error: torch.Tensor, weights: torch.Tensor + ) -> None: P = self._state.get("P") kalman_lambda = self._state.get("kalman_lambda") weights_num = self._state.get("weights_num") @@ -253,10 +258,10 @@ def __update(self, H, error, weights) -> None: i += 1 param.data = tmp_weight.reshape(param.data.T.shape).T.contiguous() - def set_grad_prefactor(self, grad_prefactor) -> None: + def set_grad_prefactor(self, grad_prefactor: float) -> None: self.grad_prefactor = grad_prefactor - def step(self, error) -> None: + def step(self, error: torch.Tensor) -> None: params_packed_index = self._state.get("params_packed_index") weights = [] @@ -313,7 +318,7 @@ def step(self, error) -> None: self.__update(H, error, weights) - def get_device_id(self, index): + def get_device_id(self, index: int) -> int | None: for i, (start, end) in enumerate(self.dindex): if start <= index < end: return i diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 7a6ff0ebde..24440e19de 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -3,6 +3,8 @@ import logging import time from collections.abc import ( + Callable, + Generator, Iterable, ) from copy import ( @@ -50,6 +52,7 @@ dp_random, ) from deepmd.pt.utils.dataloader import ( + DpLoaderSet, get_sampler_from_params, ) from deepmd.pt.utils.env import ( @@ -92,16 +95,16 @@ class Trainer: def __init__( self, config: dict[str, Any], - training_data, - stat_file_path=None, - validation_data=None, - init_model=None, - restart_model=None, - finetune_model=None, - force_load=False, - shared_links=None, - finetune_links=None, - init_frz_model=None, + training_data: DpLoaderSet, + stat_file_path: str | None = None, + validation_data: DpLoaderSet | None = None, + init_model: str | None = None, + restart_model: str | None = None, + finetune_model: str | None = None, + force_load: bool = False, + shared_links: dict[str, str] | None = None, + finetune_links: dict[str, str] | None = None, + init_frz_model: str | None = None, ) -> None: """Construct a DeePMD trainer. @@ -140,6 +143,7 @@ def __init__( self.num_steps = training_params["numb_steps"] self.disp_file = training_params.get("disp_file", "lcurve.out") self.disp_freq = training_params.get("disp_freq", 1000) + self.disp_avg = training_params.get("disp_avg", False) self.save_ckpt = training_params.get("save_ckpt", "model.ckpt") self.save_freq = training_params.get("save_freq", 1000) self.max_ckpt_keep = training_params.get("max_ckpt_keep", 5) @@ -150,7 +154,7 @@ def __init__( ) self.lcurve_should_print_header = True - def get_opt_param(params): + def get_opt_param(params: dict[str, Any]) -> tuple[str, dict[str, Any]]: opt_type = params.get("opt_type", "Adam") opt_param = { "kf_blocksize": params.get("kf_blocksize", 5120), @@ -162,7 +166,7 @@ def get_opt_param(params): } return opt_type, opt_param - def cycle_iterator(iterable: Iterable): + def cycle_iterator(iterable: Iterable) -> Generator[Any, None, None]: """ Produces an infinite iterator by repeatedly cycling through the given iterable. @@ -178,8 +182,20 @@ def cycle_iterator(iterable: Iterable): it = iter(iterable) yield from it - def get_data_loader(_training_data, _validation_data, _training_params): - def get_dataloader_and_iter(_data, _params): + def get_data_loader( + _training_data: DpLoaderSet, + _validation_data: DpLoaderSet | None, + _training_params: dict[str, Any], + ) -> tuple[ + DataLoader, + Generator[Any, None, None], + DataLoader | None, + Generator[Any, None, None] | None, + int, + ]: + def get_dataloader_and_iter( + _data: DpLoaderSet, _params: dict[str, Any] + ) -> tuple[DataLoader, Generator[Any, None, None]]: _sampler = get_sampler_from_params(_data, _params) if _sampler is None: log.warning( @@ -194,7 +210,7 @@ def get_dataloader_and_iter(_data, _params): else 0, # setting to 0 diverges the behavior of its iterator; should be >=1 drop_last=False, collate_fn=lambda batch: batch, # prevent extra conversion - pin_memory=True, + pin_memory=(DEVICE != "cpu"), # pin memory only if not on CPU ) _data_iter = cycle_iterator(_dataloader) return _dataloader, _data_iter @@ -226,21 +242,21 @@ def get_dataloader_and_iter(_data, _params): ) def single_model_stat( - _model, - _data_stat_nbatch, - _training_data, - _validation_data, - _stat_file_path, - _data_requirement, - finetune_has_new_type=False, - ): + _model: Any, + _data_stat_nbatch: int, + _training_data: DpLoaderSet, + _validation_data: DpLoaderSet | None, + _stat_file_path: str | None, + _data_requirement: list[DataRequirementItem], + finetune_has_new_type: bool = False, + ) -> Callable[[], Any]: _data_requirement += get_additional_data_requirement(_model) _training_data.add_data_requirement(_data_requirement) if _validation_data is not None: _validation_data.add_data_requirement(_data_requirement) @functools.lru_cache - def get_sample(): + def get_sample() -> Any: sampled = make_stat_input( _training_data.systems, _training_data.dataloaders, @@ -257,7 +273,7 @@ def get_sample(): _stat_file_path.root.close() return get_sample - def get_lr(lr_params): + def get_lr(lr_params: dict[str, Any]) -> LearningRateExp: assert lr_params.get("type", "exp") == "exp", ( "Only learning rate `exp` is supported!" ) @@ -495,11 +511,11 @@ def get_lr(lr_params): state_dict = pretrained_model_wrapper.state_dict() def collect_single_finetune_params( - _model_key, - _finetune_rule_single, - _new_state_dict, - _origin_state_dict, - _random_state_dict, + _model_key: str, + _finetune_rule_single: Any, + _new_state_dict: dict[str, Any], + _origin_state_dict: dict[str, Any], + _random_state_dict: dict[str, Any], ) -> None: _new_fitting = _finetune_rule_single.get_random_fitting() _model_key_from = _finetune_rule_single.get_model_branch() @@ -509,15 +525,31 @@ def collect_single_finetune_params( if i != "_extra_state" and f".{_model_key}." in i ] for item_key in target_keys: - if _new_fitting and (".descriptor." not in item_key): + new_key = item_key.replace( + f".{_model_key}.", f".{_model_key_from}." + ) + use_random_initialization = _new_fitting and ( + ".descriptor." not in item_key + ) + if ( + not use_random_initialization + and new_key not in _origin_state_dict + ): + # for ZBL models finetuning from standard models + if ".models.0." in new_key: + new_key = new_key.replace(".models.0.", ".") + elif ".models.1." in new_key: + use_random_initialization = True + else: + raise KeyError( + f"Key {new_key} not found in pretrained model." + ) + if use_random_initialization: # print(f'Keep {item_key} in old model!') _new_state_dict[item_key] = ( _random_state_dict[item_key].clone().detach() ) else: - new_key = item_key.replace( - f".{_model_key}.", f".{_model_key_from}." - ) # print(f'Replace {item_key} with {new_key} in pretrained_model!') _new_state_dict[item_key] = ( _origin_state_dict[new_key].clone().detach() @@ -544,10 +576,10 @@ def collect_single_finetune_params( if finetune_model is not None: def single_model_finetune( - _model, - _finetune_rule_single, - _sample_func, - ): + _model: Any, + _finetune_rule_single: Any, + _sample_func: Callable, + ) -> Any: _model = model_change_out_bias( _model, _sample_func, @@ -581,13 +613,45 @@ def single_model_finetune( if init_frz_model is not None: frz_model = torch.jit.load(init_frz_model, map_location=DEVICE) - self.model.load_state_dict(frz_model.state_dict()) + state = frz_model.state_dict() + missing, unexpected = self.model.load_state_dict(state, strict=False) + if missing or unexpected: + log.warning( + "Checkpoint loaded non-strictly. " + f"Missing keys: {missing}, Unexpected keys: {unexpected}" + ) + + # Get model prob for multi-task + if self.multi_task: + self.model_prob = np.array([0.0 for key in self.model_keys]) + if training_params.get("model_prob", None) is not None: + model_prob = training_params["model_prob"] + for ii, model_key in enumerate(self.model_keys): + if model_key in model_prob: + self.model_prob[ii] += float(model_prob[model_key]) + else: + for ii, model_key in enumerate(self.model_keys): + self.model_prob[ii] += float(len(self.training_data[model_key])) + sum_prob = np.sum(self.model_prob) + assert sum_prob > 0.0, "Sum of model prob must be larger than 0!" + self.model_prob = self.model_prob / sum_prob # Multi-task share params if shared_links is not None: + _data_stat_protect = np.array( + [ + model_params["model_dict"][ii].get("data_stat_protect", 1e-2) + for ii in model_params["model_dict"] + ] + ) + assert np.allclose(_data_stat_protect, _data_stat_protect[0]), ( + "Model key 'data_stat_protect' must be the same in each branch when multitask!" + ) self.wrapper.share_params( shared_links, resume=(resuming and not self.finetune_update_stat) or self.rank != 0, + model_key_prob_map=dict(zip(self.model_keys, self.model_prob)), + data_stat_protect=_data_stat_protect[0], ) if dist.is_available() and dist.is_initialized(): @@ -602,7 +666,7 @@ def single_model_finetune( # TODO add lr warmups for multitask # author: iProzd - def warm_up_linear(step, warmup_steps): + def warm_up_linear(step: int, warmup_steps: int) -> float: if step < warmup_steps: return step / warmup_steps else: @@ -637,21 +701,6 @@ def warm_up_linear(step, warmup_steps): else: raise ValueError(f"Not supported optimizer type '{self.opt_type}'") - # Get model prob for multi-task - if self.multi_task: - self.model_prob = np.array([0.0 for key in self.model_keys]) - if training_params.get("model_prob", None) is not None: - model_prob = training_params["model_prob"] - for ii, model_key in enumerate(self.model_keys): - if model_key in model_prob: - self.model_prob[ii] += float(model_prob[model_key]) - else: - for ii, model_key in enumerate(self.model_keys): - self.model_prob[ii] += float(len(self.training_data[model_key])) - sum_prob = np.sum(self.model_prob) - assert sum_prob > 0.0, "Sum of model prob must be larger than 0!" - self.model_prob = self.model_prob / sum_prob - # Tensorboard self.enable_tensorboard = training_params.get("tensorboard", False) self.tensorboard_log_dir = training_params.get("tensorboard_log_dir", "log") @@ -695,7 +744,7 @@ def run(self) -> None: ) prof.start() - def step(_step_id, task_key="Default") -> None: + def step(_step_id: int, task_key: str = "Default") -> None: if self.multi_task: model_index = dp_random.choice( np.arange(self.num_model, dtype=np.int_), @@ -769,7 +818,7 @@ def step(_step_id, task_key="Default") -> None: else self.wrapper ) - def fake_model(): + def fake_model() -> dict: return model_pred _, loss, more_loss = module.loss[task_key]( @@ -808,6 +857,33 @@ def fake_model(): else: raise ValueError(f"Not supported optimizer type '{self.opt_type}'") + if self.disp_avg: + # Accumulate loss for averaging over display interval + self.step_count_in_interval += 1 + if not self.multi_task: + # Accumulate loss for single task + if not self.train_loss_accu: + # Initialize accumulator with current loss structure + for item in more_loss: + if "l2_" not in item: + self.train_loss_accu[item] = 0.0 + for item in more_loss: + if "l2_" not in item: + self.train_loss_accu[item] += more_loss[item] + else: + # Accumulate loss for multi-task + if task_key not in self.train_loss_accu: + self.train_loss_accu[task_key] = {} + if task_key not in self.step_count_per_task: + self.step_count_per_task[task_key] = 0 + self.step_count_per_task[task_key] += 1 + + for item in more_loss: + if "l2_" not in item: + if item not in self.train_loss_accu[task_key]: + self.train_loss_accu[task_key][item] = 0.0 + self.train_loss_accu[task_key][item] += more_loss[item] + # Log and persist display_step_id = _step_id + 1 if self.display_in_training and ( @@ -815,18 +891,47 @@ def fake_model(): ): self.wrapper.eval() # Will set to train mode before fininshing validation - def log_loss_train(_loss, _more_loss, _task_key="Default"): - results = {} - rmse_val = { - item: _more_loss[item] - for item in _more_loss - if "l2_" not in item - } - for item in sorted(rmse_val.keys()): - results[item] = rmse_val[item] - return results + if self.disp_avg: + + def log_loss_train( + _loss: Any, _more_loss: Any, _task_key: str = "Default" + ) -> dict: + results = {} + if not self.multi_task: + # Use accumulated average loss for single task + for item in self.train_loss_accu: + results[item] = ( + self.train_loss_accu[item] + / self.step_count_in_interval + ) + else: + # Use accumulated average loss for multi-task + if ( + _task_key in self.train_loss_accu + and _task_key in self.step_count_per_task + ): + for item in self.train_loss_accu[_task_key]: + results[item] = ( + self.train_loss_accu[_task_key][item] + / self.step_count_per_task[_task_key] + ) + return results + else: - def log_loss_valid(_task_key="Default"): + def log_loss_train( + _loss: Any, _more_loss: Any, _task_key: str = "Default" + ) -> dict: + results = {} + rmse_val = { + item: _more_loss[item] + for item in _more_loss + if "l2_" not in item + } + for item in sorted(rmse_val.keys()): + results[item] = rmse_val[item] + return results + + def log_loss_valid(_task_key: str = "Default") -> dict: single_results = {} sum_natoms = 0 if not self.multi_task: @@ -882,24 +987,31 @@ def log_loss_valid(_task_key="Default"): else: train_results = {_key: {} for _key in self.model_keys} valid_results = {_key: {} for _key in self.model_keys} - train_results[task_key] = log_loss_train( - loss, more_loss, _task_key=task_key - ) - for _key in self.model_keys: - if _key != task_key: - self.optimizer.zero_grad() - input_dict, label_dict, _ = self.get_data( - is_train=True, task_key=_key - ) - _, loss, more_loss = self.wrapper( - **input_dict, - cur_lr=pref_lr, - label=label_dict, - task_key=_key, - ) + if self.disp_avg: + # For multi-task, use accumulated average loss for all tasks + for _key in self.model_keys: train_results[_key] = log_loss_train( loss, more_loss, _task_key=_key ) + else: + train_results[task_key] = log_loss_train( + loss, more_loss, _task_key=task_key + ) + for _key in self.model_keys: + if _key != task_key: + self.optimizer.zero_grad() + input_dict, label_dict, _ = self.get_data( + is_train=True, task_key=_key + ) + _, loss, more_loss = self.wrapper( + **input_dict, + cur_lr=pref_lr, + label=label_dict, + task_key=_key, + ) + train_results[_key] = log_loss_train( + loss, more_loss, _task_key=_key + ) valid_results[_key] = log_loss_valid(_task_key=_key) if self.rank == 0: log.info( @@ -921,12 +1033,29 @@ def log_loss_valid(_task_key="Default"): ) self.wrapper.train() + if self.disp_avg: + # Reset loss accumulators after display + if not self.multi_task: + for item in self.train_loss_accu: + self.train_loss_accu[item] = 0.0 + else: + for task_key in self.model_keys: + if task_key in self.train_loss_accu: + for item in self.train_loss_accu[task_key]: + self.train_loss_accu[task_key][item] = 0.0 + if task_key in self.step_count_per_task: + self.step_count_per_task[task_key] = 0 + self.step_count_in_interval = 0 + self.last_display_step = display_step_id + current_time = time.time() train_time = current_time - self.t0 self.t0 = current_time if self.rank == 0 and self.timing_in_training: eta = int( - (self.num_steps - display_step_id) / self.disp_freq * train_time + (self.num_steps - display_step_id) + / min(self.disp_freq, display_step_id - self.start_step) + * train_time ) log.info( format_training_message( @@ -935,12 +1064,19 @@ def log_loss_valid(_task_key="Default"): eta=eta, ) ) - # the first training time is not accurate if ( - (_step_id + 1 - self.start_step) > self.disp_freq - or self.num_steps - self.start_step < 2 * self.disp_freq + (self.num_steps - self.start_step) + <= 2 * self.disp_freq # not enough steps + or (_step_id - self.start_step) + >= self.disp_freq # skip first disp_freq steps ): self.total_train_time += train_time + if display_step_id == 1: + self.timed_steps += 1 + else: + self.timed_steps += min( + self.disp_freq, _step_id - self.start_step + ) if fout: if self.lcurve_should_print_header: @@ -951,11 +1087,14 @@ def log_loss_valid(_task_key="Default"): ) if ( - ((_step_id + 1) % self.save_freq == 0 and _step_id != self.start_step) - or (_step_id + 1) == self.num_steps + ( + (display_step_id) % self.save_freq == 0 + and _step_id != self.start_step + ) + or (display_step_id) == self.num_steps ) and (self.rank == 0 or dist.get_rank() == 0): # Handle the case if rank 0 aborted and re-assigned - self.latest_model = Path(self.save_ckpt + f"-{_step_id + 1}.pt") + self.latest_model = Path(self.save_ckpt + f"-{display_step_id}.pt") module = ( self.wrapper.module @@ -982,6 +1121,18 @@ def log_loss_valid(_task_key="Default"): self.wrapper.train() self.t0 = time.time() self.total_train_time = 0.0 + self.timed_steps = 0 + + if self.disp_avg: + # Initialize loss accumulators + if not self.multi_task: + self.train_loss_accu = {} + else: + self.train_loss_accu = {key: {} for key in self.model_keys} + self.step_count_per_task = dict.fromkeys(self.model_keys, 0) + self.step_count_in_interval = 0 + self.last_display_step = 0 + for step_id in range(self.start_step, self.num_steps): step(step_id) if JIT: @@ -1021,24 +1172,12 @@ def log_loss_valid(_task_key="Default"): with open("checkpoint", "w") as f: f.write(str(self.latest_model)) - elapsed_batch = self.num_steps - self.start_step - if self.timing_in_training and elapsed_batch // self.disp_freq > 0: - if self.start_step >= 2 * self.disp_freq: - log.info( - "average training time: %.4f s/batch (exclude first %d batches)", - self.total_train_time - / ( - elapsed_batch // self.disp_freq * self.disp_freq - - self.disp_freq - ), - self.disp_freq, - ) - else: - log.info( - "average training time: %.4f s/batch", - self.total_train_time - / (elapsed_batch // self.disp_freq * self.disp_freq), - ) + if self.timing_in_training and self.timed_steps: + msg = f"average training time: {self.total_train_time / self.timed_steps:.4f} s/batch" + excluded_steps = self.num_steps - self.start_step - self.timed_steps + if excluded_steps > 0: + msg += f" ({excluded_steps} batches excluded)" + log.info(msg) if JIT: pth_model_path = ( @@ -1062,13 +1201,13 @@ def log_loss_valid(_task_key="Default"): log.info( f"The profiling trace has been saved under {self.tensorboard_log_dir}" ) - if self.profiling: + if not self.enable_profiler and self.profiling: prof.export_chrome_trace(self.profiling_file) log.info( f"The profiling trace has been saved to: {self.profiling_file}" ) - def save_model(self, save_path, lr=0.0, step=0) -> None: + def save_model(self, save_path: str, lr: float = 0.0, step: int = 0) -> None: module = ( self.wrapper.module if dist.is_available() and dist.is_initialized() @@ -1093,7 +1232,9 @@ def save_model(self, save_path, lr=0.0, step=0) -> None: checkpoint_files.sort(key=lambda x: x.stat().st_mtime) checkpoint_files[0].unlink() - def get_data(self, is_train=True, task_key="Default"): + def get_data( + self, is_train: bool = True, task_key: str = "Default" + ) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: if is_train: iterator = self.training_data else: @@ -1127,7 +1268,8 @@ def get_data(self, is_train=True, task_key="Default"): label_dict = {} for item_key in batch_data: if item_key in input_keys: - input_dict[item_key] = batch_data[item_key] + if item_key != "fparam" or batch_data["find_fparam"] != 0.0: + input_dict[item_key] = batch_data[item_key] else: if item_key not in ["sid", "fid"]: label_dict[item_key] = batch_data[item_key] @@ -1137,7 +1279,9 @@ def get_data(self, is_train=True, task_key="Default"): log_dict["sid"] = batch_data["sid"] return input_dict, label_dict, log_dict - def print_header(self, fout, train_results, valid_results) -> None: + def print_header( + self, fout: Any, train_results: dict[str, Any], valid_results: dict[str, Any] + ) -> None: train_keys = sorted(train_results.keys()) print_str = "" print_str += "# {:5s}".format("step") @@ -1169,7 +1313,12 @@ def print_header(self, fout, train_results, valid_results) -> None: fout.flush() def print_on_training( - self, fout, step_id, cur_lr, train_results, valid_results + self, + fout: Any, + step_id: int, + cur_lr: float, + train_results: dict, + valid_results: dict, ) -> None: train_keys = sorted(train_results.keys()) print_str = "" @@ -1201,12 +1350,21 @@ def print_on_training( fout.flush() -def get_additional_data_requirement(_model): +def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]: additional_data_requirement = [] if _model.get_dim_fparam() > 0: + _fparam_default = ( + _model.get_default_fparam().cpu().numpy() + if _model.has_default_fparam() + else 0.0 + ) fparam_requirement_items = [ DataRequirementItem( - "fparam", _model.get_dim_fparam(), atomic=False, must=True + "fparam", + _model.get_dim_fparam(), + atomic=False, + must=not _model.has_default_fparam(), + default=_fparam_default, ) ] additional_data_requirement += fparam_requirement_items @@ -1228,12 +1386,14 @@ def get_additional_data_requirement(_model): return additional_data_requirement -def whether_hessian(loss_params): +def whether_hessian(loss_params: dict[str, Any]) -> bool: loss_type = loss_params.get("type", "ener") return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0 -def get_loss(loss_params, start_lr, _ntypes, _model): +def get_loss( + loss_params: dict[str, Any], start_lr: float, _ntypes: int, _model: Any +) -> TaskLoss: loss_type = loss_params.get("type", "ener") if whether_hessian(loss_params): loss_params["starter_learning_rate"] = start_lr @@ -1276,8 +1436,8 @@ def get_loss(loss_params, start_lr, _ntypes, _model): def get_single_model( - _model_params, -): + _model_params: dict[str, Any], +) -> Any: if "use_srtab" in _model_params: model = get_zbl_model(deepcopy(_model_params)).to(DEVICE) else: @@ -1286,10 +1446,10 @@ def get_single_model( def get_model_for_wrapper( - _model_params, - resuming=False, - _loss_params=None, -): + _model_params: dict[str, Any], + resuming: bool = False, + _loss_params: dict[str, Any] | None = None, +) -> Any: if "model_dict" not in _model_params: if _loss_params is not None and whether_hessian(_loss_params): _model_params["hessian_mode"] = True @@ -1312,7 +1472,7 @@ def get_model_for_wrapper( return _model -def get_case_embd_config(_model_params): +def get_case_embd_config(_model_params: dict[str, Any]) -> tuple[bool, dict[str, int]]: assert "model_dict" in _model_params, ( "Only support setting case embedding for multi-task model!" ) @@ -1337,10 +1497,10 @@ def get_case_embd_config(_model_params): def model_change_out_bias( - _model, - _sample_func, - _bias_adjust_mode="change-by-statistic", -): + _model: Any, + _sample_func: Callable[[], Any], + _bias_adjust_mode: str = "change-by-statistic", +) -> Any: old_bias = deepcopy(_model.get_out_bias()) _model.change_out_bias( _sample_func, diff --git a/deepmd/pt/train/wrapper.py b/deepmd/pt/train/wrapper.py index 9a2cbff295..2669e3d832 100644 --- a/deepmd/pt/train/wrapper.py +++ b/deepmd/pt/train/wrapper.py @@ -1,8 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, - Union, + Any, ) import torch @@ -17,10 +16,10 @@ class ModelWrapper(torch.nn.Module): def __init__( self, - model: Union[torch.nn.Module, dict], - loss: Union[torch.nn.Module, dict] = None, - model_params=None, - shared_links=None, + model: torch.nn.Module | dict, + loss: torch.nn.Module | dict = None, + model_params: dict[str, Any] | None = None, + shared_links: dict[str, Any] | None = None, ) -> None: """Construct a DeePMD model wrapper. @@ -59,7 +58,13 @@ def __init__( self.loss[task_key] = loss[task_key] self.inference_only = self.loss is None - def share_params(self, shared_links, resume=False) -> None: + def share_params( + self, + shared_links: dict[str, Any], + model_key_prob_map: dict, + data_stat_protect: float = 1e-2, + resume: bool = False, + ) -> None: """ Share the parameters of classes following rules defined in shared_links during multitask training. If not start from checkpoint (resume is False), @@ -129,8 +134,16 @@ def share_params(self, shared_links, resume=False) -> None: link_class = self.model[ model_key_link ].atomic_model.__getattr__(class_type_link) + frac_prob = ( + model_key_prob_map[model_key_link] + / model_key_prob_map[model_key_base] + ) link_class.share_params( - base_class, shared_level_link, resume=resume + base_class, + shared_level_link, + model_prob=frac_prob, + protection=data_stat_protect, + resume=resume, ) log.warning( f"Shared params of {model_key_base}.{class_type_base} and {model_key_link}.{class_type_link}!" @@ -138,18 +151,18 @@ def share_params(self, shared_links, resume=False) -> None: def forward( self, - coord, - atype, - spin: Optional[torch.Tensor] = None, - box: Optional[torch.Tensor] = None, - cur_lr: Optional[torch.Tensor] = None, - label: Optional[torch.Tensor] = None, - task_key: Optional[torch.Tensor] = None, - inference_only=False, - do_atomic_virial=False, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): + coord: torch.Tensor, + atype: torch.Tensor, + spin: torch.Tensor | None = None, + box: torch.Tensor | None = None, + cur_lr: torch.Tensor | None = None, + label: torch.Tensor | None = None, + task_key: torch.Tensor | None = None, + inference_only: bool = False, + do_atomic_virial: bool = False, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, + ) -> tuple[Any, Any, Any]: if not self.multi_task: task_key = "Default" else: diff --git a/deepmd/pt/utils/dataloader.py b/deepmd/pt/utils/dataloader.py index bc771b41d4..c991e59daa 100644 --- a/deepmd/pt/utils/dataloader.py +++ b/deepmd/pt/utils/dataloader.py @@ -4,6 +4,9 @@ from multiprocessing.dummy import ( Pool, ) +from typing import ( + Any, +) import h5py import numpy as np @@ -45,7 +48,7 @@ torch.multiprocessing.set_sharing_strategy("file_system") -def setup_seed(seed) -> None: +def setup_seed(seed: int | list[int] | tuple[int, ...]) -> None: if isinstance(seed, (list, tuple)): mixed_seed = mix_entropy(seed) else: @@ -75,11 +78,11 @@ class DpLoaderSet(Dataset): def __init__( self, - systems, - batch_size, - type_map, - seed=None, - shuffle=True, + systems: str | list[str], + batch_size: int, + type_map: list[str] | None, + seed: int | None = None, + shuffle: bool = True, ) -> None: if seed is not None: setup_seed(seed) @@ -87,7 +90,7 @@ def __init__( with h5py.File(systems) as file: systems = [os.path.join(systems, item) for item in file.keys()] - def construct_dataset(system): + def construct_dataset(system: str) -> DeepmdDataSetForLoader: return DeepmdDataSetForLoader( system=system, type_map=type_map, @@ -180,7 +183,7 @@ def construct_dataset(system): for item in self.dataloaders: self.iters.append(iter(item)) - def set_noise(self, noise_settings) -> None: + def set_noise(self, noise_settings: dict[str, Any]) -> None: # noise_settings['noise_type'] # "trunc_normal", "normal", "uniform" # noise_settings['noise'] # float, default 1.0 # noise_settings['noise_mode'] # "prob", "fix_num" @@ -193,7 +196,7 @@ def set_noise(self, noise_settings) -> None: def __len__(self) -> int: return len(self.dataloaders) - def __getitem__(self, idx): + def __getitem__(self, idx: int) -> dict[str, torch.Tensor]: # log.warning(str(torch.distributed.get_rank())+" idx: "+str(idx)+" index: "+str(self.index[idx])) with torch.device("cpu"): try: @@ -231,7 +234,7 @@ def print_summary( ) -def collate_batch(batch): +def collate_batch(batch: list[dict[str, Any]]) -> dict[str, Any]: example = batch[0] result = {} for key in example.keys(): @@ -251,7 +254,9 @@ def collate_batch(batch): return result -def get_weighted_sampler(training_data, prob_style, sys_prob=False): +def get_weighted_sampler( + training_data: Any, prob_style: str, sys_prob: bool = False +) -> WeightedRandomSampler: if sys_prob is False: if prob_style == "prob_uniform": prob_v = 1.0 / float(training_data.__len__()) @@ -276,7 +281,7 @@ def get_weighted_sampler(training_data, prob_style, sys_prob=False): return sampler -def get_sampler_from_params(_data, _params): +def get_sampler_from_params(_data: Any, _params: dict[str, Any]) -> Any: if ( "sys_probs" in _params and _params["sys_probs"] is not None ): # use sys_probs first diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index 3043839308..481fa04497 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -2,7 +2,7 @@ from typing import ( - Optional, + Any, ) from torch.utils.data import ( @@ -16,7 +16,7 @@ class DeepmdDataSetForLoader(Dataset): - def __init__(self, system: str, type_map: Optional[list[str]] = None) -> None: + def __init__(self, system: str, type_map: list[str] | None = None) -> None: """Construct DeePMD-style dataset containing frames cross different systems. Args: @@ -34,7 +34,7 @@ def __init__(self, system: str, type_map: Optional[list[str]] = None) -> None: def __len__(self) -> int: return self._data_system.nframes - def __getitem__(self, index): + def __getitem__(self, index: int) -> dict[str, Any]: """Get a frame from the selected system.""" b_data = self._data_system.get_item_torch(index) b_data["natoms"] = self._natoms_vec diff --git a/deepmd/pt/utils/env.py b/deepmd/pt/utils/env.py index e490e786f0..90d0d536c1 100644 --- a/deepmd/pt/utils/env.py +++ b/deepmd/pt/utils/env.py @@ -2,6 +2,7 @@ import logging import multiprocessing import os +import sys import numpy as np import torch @@ -16,6 +17,17 @@ set_default_nthreads, ) +log = logging.getLogger(__name__) + +if sys.platform != "win32": + try: + multiprocessing.set_start_method("fork", force=True) + log.debug("Successfully set multiprocessing start method to 'fork'.") + except (RuntimeError, ValueError) as err: + log.warning(f"Could not set multiprocessing start method: {err}") +else: + log.debug("Skipping fork start method on Windows (not supported).") + SAMPLER_RECORD = os.environ.get("SAMPLER_RECORD", False) DP_DTYPE_PROMOTION_STRICT = os.environ.get("DP_DTYPE_PROMOTION_STRICT", "0") == "1" try: @@ -26,7 +38,6 @@ NUM_WORKERS = int(os.environ.get("NUM_WORKERS", min(4, ncpus))) if multiprocessing.get_start_method() != "fork": # spawn or forkserver does not support NUM_WORKERS > 0 for DataLoader - log = logging.getLogger(__name__) log.warning( "NUM_WORKERS > 0 is not supported with spawn or forkserver start method. " "Setting NUM_WORKERS to 0." diff --git a/deepmd/pt/utils/env_mat_stat.py b/deepmd/pt/utils/env_mat_stat.py index 23e8627bcd..01822c7f3f 100644 --- a/deepmd/pt/utils/env_mat_stat.py +++ b/deepmd/pt/utils/env_mat_stat.py @@ -4,7 +4,6 @@ ) from typing import ( TYPE_CHECKING, - Union, ) import numpy as np @@ -77,7 +76,7 @@ def __init__(self, descriptor: "DescriptorBlock") -> None: ) # se_r=1, se_a=4 def iter( - self, data: list[dict[str, Union[torch.Tensor, list[tuple[int, int]]]]] + self, data: list[dict[str, torch.Tensor | list[tuple[int, int]]]] ) -> Iterator[dict[str, StatItem]]: """Get the iterator of the environment matrix. @@ -200,7 +199,7 @@ def get_hash(self) -> str: } ) - def __call__(self): + def __call__(self) -> tuple[np.ndarray, np.ndarray]: avgs = self.get_avg() stds = self.get_std() diff --git a/deepmd/pt/utils/exclude_mask.py b/deepmd/pt/utils/exclude_mask.py index 0a99c0777f..cf39220f1b 100644 --- a/deepmd/pt/utils/exclude_mask.py +++ b/deepmd/pt/utils/exclude_mask.py @@ -32,10 +32,10 @@ def reinit( ) self.type_mask = to_torch_tensor(self.type_mask).view([-1]) - def get_exclude_types(self): + def get_exclude_types(self) -> list[int]: return self.exclude_types - def get_type_mask(self): + def get_type_mask(self) -> torch.Tensor: return self.type_mask def forward( @@ -98,7 +98,7 @@ def reinit( self.type_mask = to_torch_tensor(self.type_mask).view([-1]) self.no_exclusion = len(self._exclude_types) == 0 - def get_exclude_types(self): + def get_exclude_types(self) -> set[tuple[int, int]]: return self._exclude_types # may have a better place for this method... diff --git a/deepmd/pt/utils/finetune.py b/deepmd/pt/utils/finetune.py index 96a420bf6a..0e86c9aa6c 100644 --- a/deepmd/pt/utils/finetune.py +++ b/deepmd/pt/utils/finetune.py @@ -3,6 +3,9 @@ from copy import ( deepcopy, ) +from typing import ( + Any, +) import torch @@ -12,18 +15,21 @@ from deepmd.utils.finetune import ( FinetuneRuleItem, ) +from deepmd.utils.model_branch_dict import ( + get_model_dict, +) log = logging.getLogger(__name__) def get_finetune_rule_single( - _single_param_target, - _model_param_pretrained, - from_multitask=False, - model_branch="Default", - model_branch_from="", - change_model_params=False, -): + _single_param_target: dict[str, Any], + _model_param_pretrained: dict[str, Any], + from_multitask: bool = False, + model_branch: str = "Default", + model_branch_from: str = "", + change_model_params: bool = False, +) -> tuple[dict[str, Any], FinetuneRuleItem]: single_config = deepcopy(_single_param_target) new_fitting = False model_branch_chosen = "Default" @@ -44,10 +50,13 @@ def get_finetune_rule_single( ) else: model_branch_chosen = model_branch_from - assert model_branch_chosen in model_dict_params, ( - f"No model branch named '{model_branch_chosen}'! " + model_alias_dict, model_branch_dict = get_model_dict(model_dict_params) + assert model_branch_chosen in model_alias_dict, ( + f"No model branch or alias named '{model_branch_chosen}'! " f"Available ones are {list(model_dict_params.keys())}." + f"Use `dp --pt show your_model.pt model-branch` to show detail information." ) + model_branch_chosen = model_alias_dict[model_branch_chosen] single_config_chosen = deepcopy(model_dict_params[model_branch_chosen]) old_type_map, new_type_map = ( single_config_chosen["type_map"], @@ -80,8 +89,11 @@ def get_finetune_rule_single( def get_finetune_rules( - finetune_model, model_config, model_branch="", change_model_params=True -): + finetune_model: str, + model_config: dict[str, Any], + model_branch: str = "", + change_model_params: bool = True, +) -> tuple[dict[str, Any], dict[str, FinetuneRuleItem]]: """ Get fine-tuning rules and (optionally) change the model_params according to the pretrained one. diff --git a/deepmd/pt/utils/multi_task.py b/deepmd/pt/utils/multi_task.py index 6c397400bf..f5e7799df7 100644 --- a/deepmd/pt/utils/multi_task.py +++ b/deepmd/pt/utils/multi_task.py @@ -2,6 +2,9 @@ from copy import ( deepcopy, ) +from typing import ( + Any, +) from deepmd.pt.model.descriptor import ( BaseDescriptor, @@ -11,7 +14,9 @@ ) -def preprocess_shared_params(model_config): +def preprocess_shared_params( + model_config: dict[str, Any], +) -> tuple[dict[str, Any], dict[str, Any]]: """Preprocess the model params for multitask model, and generate the links dict for further sharing. Args: @@ -97,7 +102,11 @@ def preprocess_shared_params(model_config): type_map_keys = [] def replace_one_item( - params_dict, key_type, key_in_dict, suffix="", index=None + params_dict: dict[str, Any], + key_type: str, + key_in_dict: str, + suffix: str = "", + index: int | None = None, ) -> None: shared_type = key_type shared_key = key_in_dict @@ -155,7 +164,7 @@ def replace_one_item( return model_config, shared_links -def get_class_name(item_key, item_params): +def get_class_name(item_key: str, item_params: dict[str, Any]) -> type: if item_key == "descriptor": return BaseDescriptor.get_class_by_type(item_params.get("type", "se_e2_a")) elif item_key == "fitting_net": diff --git a/deepmd/pt/utils/neighbor_stat.py b/deepmd/pt/utils/neighbor_stat.py index 64ad695827..292a27080b 100644 --- a/deepmd/pt/utils/neighbor_stat.py +++ b/deepmd/pt/utils/neighbor_stat.py @@ -2,9 +2,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np import torch @@ -52,7 +49,7 @@ def forward( self, coord: torch.Tensor, atype: torch.Tensor, - cell: Optional[torch.Tensor], + cell: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor]: """Calculate the neareest neighbor distance between atoms, maximum nbor size of atoms and the output data range of the environment matrix. @@ -170,8 +167,8 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - cell: Optional[np.ndarray], - ): + cell: np.ndarray | None, + ) -> tuple[np.ndarray, np.ndarray]: """Execute the operation. Parameters diff --git a/deepmd/pt/utils/nlist.py b/deepmd/pt/utils/nlist.py index af84151829..7f74e65f26 100644 --- a/deepmd/pt/utils/nlist.py +++ b/deepmd/pt/utils/nlist.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import torch @@ -16,13 +12,13 @@ def extend_input_and_build_neighbor_list( - coord, - atype, + coord: torch.Tensor, + atype: torch.Tensor, rcut: float, sel: list[int], mixed_types: bool = False, - box: Optional[torch.Tensor] = None, -): + box: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: nframes, nloc = atype.shape[:2] if box is not None: box_gpu = box.to(coord.device, non_blocking=True) @@ -53,7 +49,7 @@ def build_neighbor_list( atype: torch.Tensor, nloc: int, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> torch.Tensor: """Build neighbor list for a single frame. keeps nsel neighbors. @@ -192,7 +188,7 @@ def build_directional_neighbor_list( coord_neig: torch.Tensor, atype_neig: torch.Tensor, rcut: float, - sel: Union[int, list[int]], + sel: int | list[int], distinguish_types: bool = True, ) -> torch.Tensor: """Build directional neighbor list. @@ -292,7 +288,7 @@ def nlist_distinguish_types( nlist: torch.Tensor, atype: torch.Tensor, sel: list[int], -): +) -> torch.Tensor: """Given a nlist that does not distinguish atom types, return a nlist that distinguish atom types. @@ -411,10 +407,10 @@ def build_multiple_neighbor_list( def extend_coord_with_ghosts( coord: torch.Tensor, atype: torch.Tensor, - cell: Optional[torch.Tensor], + cell: torch.Tensor | None, rcut: float, - cell_cpu: Optional[torch.Tensor] = None, -): + cell_cpu: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Extend the coordinates of the atoms by appending peridoc images. The number of images is large enough to ensure all the neighbors within rcut are appended. diff --git a/deepmd/pt/utils/preprocess.py b/deepmd/pt/utils/preprocess.py index 7161bac692..0cc31b5d7a 100644 --- a/deepmd/pt/utils/preprocess.py +++ b/deepmd/pt/utils/preprocess.py @@ -6,7 +6,9 @@ log = logging.getLogger(__name__) -def compute_smooth_weight(distance, rmin: float, rmax: float): +def compute_smooth_weight( + distance: torch.Tensor, rmin: float, rmax: float +) -> torch.Tensor: """Compute smooth weight for descriptor elements.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") @@ -17,7 +19,7 @@ def compute_smooth_weight(distance, rmin: float, rmax: float): return vv -def compute_exp_sw(distance, rmin: float, rmax: float): +def compute_exp_sw(distance: torch.Tensor, rmin: float, rmax: float) -> torch.Tensor: """Compute the exponential switch function for neighbor update.""" if rmin >= rmax: raise ValueError("rmin should be less than rmax.") diff --git a/deepmd/pt/utils/region.py b/deepmd/pt/utils/region.py index 3272434995..21af694c2c 100644 --- a/deepmd/pt/utils/region.py +++ b/deepmd/pt/utils/region.py @@ -68,7 +68,7 @@ def to_face_distance( return dist.view(list(cshape[:-2]) + [3]) # noqa:RUF005 -def b_to_face_distance(cell): +def b_to_face_distance(cell: torch.Tensor) -> torch.Tensor: volume = torch.linalg.det(cell) c_yz = torch.cross(cell[:, 1], cell[:, 2], dim=-1) _h2yz = volume / torch.linalg.norm(c_yz, dim=-1) diff --git a/deepmd/pt/utils/spin.py b/deepmd/pt/utils/spin.py index 285dcaf93e..74ddb5ca13 100644 --- a/deepmd/pt/utils/spin.py +++ b/deepmd/pt/utils/spin.py @@ -4,10 +4,10 @@ def concat_switch_virtual( - extended_tensor, - extended_tensor_virtual, + extended_tensor: torch.Tensor, + extended_tensor_virtual: torch.Tensor, nloc: int, -): +) -> torch.Tensor: """ Concat real and virtual extended tensors, and switch all the local ones to the first nloc * 2 atoms. - [:, :nloc]: original nloc real atoms. diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py index cf6892b49d..cf82461a7e 100644 --- a/deepmd/pt/utils/stat.py +++ b/deepmd/pt/utils/stat.py @@ -3,10 +3,11 @@ from collections import ( defaultdict, ) -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) import numpy as np @@ -35,7 +36,9 @@ log = logging.getLogger(__name__) -def make_stat_input(datasets, dataloaders, nbatches): +def make_stat_input( + datasets: list[Any], dataloaders: list[Any], nbatches: int +) -> dict[str, Any]: """Pack data for statistics. Args: @@ -59,6 +62,14 @@ def make_stat_input(datasets, dataloaders, nbatches): except StopIteration: iterator = iter(dataloaders[i]) stat_data = next(iterator) + if ( + "find_fparam" in stat_data + and "fparam" in stat_data + and stat_data["find_fparam"] == 0.0 + ): + # for model using default fparam + stat_data.pop("fparam") + stat_data.pop("find_fparam") for dd in stat_data: if stat_data[dd] is None: sys_stat[dd] = None @@ -86,7 +97,7 @@ def make_stat_input(datasets, dataloaders, nbatches): def _restore_from_file( stat_file_path: DPPath, keys: list[str] = ["energy"], -) -> Optional[dict]: +) -> dict | None: if stat_file_path is None: return None, None stat_files = [stat_file_path / f"bias_atom_{kk}" for kk in keys] @@ -127,9 +138,9 @@ def _save_to_file( def _post_process_stat( - out_bias, - out_std, -): + out_bias: torch.Tensor, + out_std: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: """Post process the statistics. For global statistics, we do not have the std for each type of atoms, @@ -148,10 +159,10 @@ def _post_process_stat( def _compute_model_predict( - sampled: Union[Callable[[], list[dict]], list[dict]], + sampled: Callable[[], list[dict]] | list[dict], keys: list[str], model_forward: Callable[..., torch.Tensor], -): +) -> dict[str, list[torch.Tensor]]: auto_batch_size = AutoBatchSize() model_predict = {kk: [] for kk in keys} for system in sampled: @@ -165,7 +176,7 @@ def _compute_model_predict( fparam = system.get("fparam", None) aparam = system.get("aparam", None) - def model_forward_auto_batch_size(*args, **kwargs): + def model_forward_auto_batch_size(*args: Any, **kwargs: Any) -> Any: return auto_batch_size.execute_all( model_forward, nframes, @@ -188,8 +199,8 @@ def model_forward_auto_batch_size(*args, **kwargs): def _make_preset_out_bias( ntypes: int, - ibias: list[Optional[np.ndarray]], -) -> Optional[np.ndarray]: + ibias: list[np.ndarray | None], +) -> np.ndarray | None: """Make preset out bias. output: @@ -212,9 +223,9 @@ def _make_preset_out_bias( def _fill_stat_with_global( - atomic_stat: Union[np.ndarray, None], + atomic_stat: np.ndarray | None, global_stat: np.ndarray, -): +) -> np.ndarray | None: """This function is used to fill atomic stat with global stat. Parameters @@ -238,16 +249,16 @@ def _fill_stat_with_global( def compute_output_stats( - merged: Union[Callable[[], list[dict]], list[dict]], + merged: Callable[[], list[dict]] | list[dict], ntypes: int, - keys: Union[str, list[str]] = ["energy"], - stat_file_path: Optional[DPPath] = None, - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_forward: Optional[Callable[..., torch.Tensor]] = None, + keys: str | list[str] = ["energy"], + stat_file_path: DPPath | None = None, + rcond: float | None = None, + preset_bias: dict[str, list[np.ndarray | None]] | None = None, + model_forward: Callable[..., torch.Tensor] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, -): +) -> dict[str, Any]: """ Compute the output statistics (e.g. energy bias) for the fitting net from packed data. @@ -409,12 +420,12 @@ def compute_output_stats_global( sampled: list[dict], ntypes: int, keys: list[str], - rcond: Optional[float] = None, - preset_bias: Optional[dict[str, list[Optional[np.ndarray]]]] = None, - model_pred: Optional[dict[str, np.ndarray]] = None, + rcond: float | None = None, + preset_bias: dict[str, list[np.ndarray | None]] | None = None, + model_pred: dict[str, np.ndarray] | None = None, stats_distinguish_types: bool = True, intensive: bool = False, -): +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: """This function only handle stat computation from reduced global labels.""" # return directly if model predict is empty for global if model_pred == {}: @@ -522,7 +533,7 @@ def compute_output_stats_global( } atom_numbs = {kk: merged_natoms[kk].sum(-1) for kk in bias_atom_e.keys()} - def rmse(x): + def rmse(x: np.ndarray) -> float: return np.sqrt(np.mean(np.square(x))) for kk in bias_atom_e.keys(): @@ -540,8 +551,8 @@ def compute_output_stats_atomic( sampled: list[dict], ntypes: int, keys: list[str], - model_pred: Optional[dict[str, np.ndarray]] = None, -): + model_pred: dict[str, np.ndarray] | None = None, +) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: # get label dict from sample; for each key, only picking the system with atomic labels. outputs = { kk: [ diff --git a/deepmd/pt/utils/tabulate.py b/deepmd/pt/utils/tabulate.py index db743ff98c..a308f2d36b 100644 --- a/deepmd/pt/utils/tabulate.py +++ b/deepmd/pt/utils/tabulate.py @@ -3,6 +3,9 @@ from functools import ( cached_property, ) +from typing import ( + Any, +) import numpy as np import torch @@ -43,12 +46,13 @@ class DPTabulate(BaseTabulate): The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1. activation_function - The activation function in the embedding net. Supported options are {"tanh","gelu"} in common.ActivationFn. + The activation function in the embedding net. See :class:`ActivationFn` + for supported options (e.g. "tanh", "gelu", "relu", "silu"). """ def __init__( self, - descrpt, + descrpt: Any, neuron: list[int], type_one_side: bool = False, exclude_types: list[list[int]] = [], @@ -63,12 +67,7 @@ def __init__( ) self.descrpt_type = self._get_descrpt_type() - supported_descrpt_type = ( - "Atten", - "A", - "T", - "R", - ) + supported_descrpt_type = ("Atten", "A", "T", "T_TEBD", "R") if self.descrpt_type in supported_descrpt_type: self.sel_a = self.descrpt.get_sel() @@ -86,6 +85,7 @@ def __init__( "relu6": 4, "softplus": 5, "sigmoid": 6, + "silu": 7, } activation = activation_fn.activation @@ -113,7 +113,7 @@ def __init__( self.data_type = self._get_data_type() self.last_layer_size = self._get_last_layer_size() - def _make_data(self, xx, idx): + def _make_data(self, xx: np.ndarray, idx: int) -> Any: """Generate tabulation data for the given input. Parameters @@ -153,7 +153,7 @@ def _make_data(self, xx, idx): self.matrix["layer_" + str(layer + 1)][idx], xbar, self.functype, - ) + torch.ones((1, 1), dtype=yy.dtype) # pylint: disable=no-explicit-device + ) + torch.ones((1, 1), dtype=yy.dtype, device=yy.device) dy2 = unaggregated_dy2_dx_s( yy - xx, dy, @@ -172,7 +172,7 @@ def _make_data(self, xx, idx): self.matrix["layer_" + str(layer + 1)][idx], xbar, self.functype, - ) + torch.ones((1, 2), dtype=yy.dtype) # pylint: disable=no-explicit-device + ) + torch.ones((1, 2), dtype=yy.dtype, device=yy.device) dy2 = unaggregated_dy2_dx_s( yy - tt, dy, @@ -282,12 +282,12 @@ def _make_data(self, xx, idx): d2 = dy2.detach().cpu().numpy().astype(self.data_type) return vv, dd, d2 - def _layer_0(self, x, w, b): + def _layer_0(self, x: torch.Tensor, w: np.ndarray, b: np.ndarray) -> torch.Tensor: w = torch.from_numpy(w).to(env.DEVICE) b = torch.from_numpy(b).to(env.DEVICE) return self.activation_fn(torch.matmul(x, w) + b) - def _layer_1(self, x, w, b): + def _layer_1(self, x: torch.Tensor, w: np.ndarray, b: np.ndarray) -> torch.Tensor: w = torch.from_numpy(w).to(env.DEVICE) b = torch.from_numpy(b).to(env.DEVICE) t = torch.cat([x, x], dim=1) @@ -308,9 +308,11 @@ def _get_descrpt_type(self) -> str: return "R" elif isinstance(self.descrpt, deepmd.pt.model.descriptor.DescrptSeT): return "T" + elif isinstance(self.descrpt, deepmd.pt.model.descriptor.DescrptSeTTebd): + return "T_TEBD" raise RuntimeError(f"Unsupported descriptor {self.descrpt}") - def _get_layer_size(self): + def _get_layer_size(self) -> int: # get the number of layers in EmbeddingNet layer_size = 0 basic_size = 0 @@ -322,7 +324,7 @@ def _get_layer_size(self): * len(self.embedding_net_nodes[0]) * len(self.neuron) ) - if self.descrpt_type == "Atten": + if self.descrpt_type in ("Atten", "T_TEBD"): layer_size = len(self.embedding_net_nodes[0]["layers"]) elif self.descrpt_type == "A": layer_size = len(self.embedding_net_nodes[0]["layers"]) @@ -391,6 +393,13 @@ def _get_network_variable(self, var_name: str) -> dict: "layers" ][layer - 1]["@variables"][var_name] result["layer_" + str(layer)].append(node) + elif self.descrpt_type == "T_TEBD": + # For the se_e3_tebd descriptor, a single, + # shared embedding network is used for all type pairs + node = self.embedding_net_nodes[0]["layers"][layer - 1]["@variables"][ + var_name + ] + result["layer_" + str(layer)].append(node) elif self.descrpt_type == "R": if self.type_one_side: for ii in range(0, self.ntypes): @@ -417,10 +426,10 @@ def _get_network_variable(self, var_name: str) -> dict: raise RuntimeError("Unsupported descriptor") return result - def _get_bias(self): + def _get_bias(self) -> Any: return self._get_network_variable("b") - def _get_matrix(self): + def _get_matrix(self) -> Any: return self._get_network_variable("w") def _convert_numpy_to_tensor(self) -> None: @@ -435,7 +444,7 @@ def _n_all_excluded(self) -> int: # customized op -def grad(xbar: torch.Tensor, y: torch.Tensor, functype: int): +def grad(xbar: torch.Tensor, y: torch.Tensor, functype: int) -> torch.Tensor: if functype == 1: return 1 - y * y @@ -461,11 +470,16 @@ def grad(xbar: torch.Tensor, y: torch.Tensor, functype: int): elif functype == 6: return y * (1 - y) + elif functype == 7: + # silu'(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))) + sig = torch.sigmoid(xbar) + return sig + xbar * sig * (1 - sig) + else: raise ValueError(f"Unsupported function type: {functype}") -def grad_grad(xbar: torch.Tensor, y: torch.Tensor, functype: int): +def grad_grad(xbar: torch.Tensor, y: torch.Tensor, functype: int) -> torch.Tensor: if functype == 1: return -2 * y * (1 - y * y) @@ -488,13 +502,19 @@ def grad_grad(xbar: torch.Tensor, y: torch.Tensor, functype: int): elif functype == 6: return y * (1 - y) * (1 - 2 * y) + elif functype == 7: + sig = torch.sigmoid(xbar) + d_sig = sig * (1 - sig) + # silu''(x) = 2 * d_sig + x * d_sig * (1 - 2 * sig) + return 2 * d_sig + xbar * d_sig * (1 - 2 * sig) + else: return -torch.ones_like(xbar) def unaggregated_dy_dx_s( y: torch.Tensor, w_np: np.ndarray, xbar: torch.Tensor, functype: int -): +) -> torch.Tensor: w = torch.from_numpy(w_np).to(env.DEVICE) y = y.to(env.DEVICE) xbar = xbar.to(env.DEVICE) @@ -520,7 +540,7 @@ def unaggregated_dy2_dx_s( w_np: np.ndarray, xbar: torch.Tensor, functype: int, -): +) -> torch.Tensor: w = torch.from_numpy(w_np).to(env.DEVICE) y = y.to(env.DEVICE) dy = dy.to(env.DEVICE) @@ -549,7 +569,7 @@ def unaggregated_dy_dx( dy_dx: torch.Tensor, ybar: torch.Tensor, functype: int, -): +) -> torch.Tensor: w = torch.from_numpy(w_np).to(env.DEVICE) if z.dim() != 2: raise ValueError("z tensor must have 2 dimensions") @@ -587,7 +607,7 @@ def unaggregated_dy2_dx( dy2_dx: torch.Tensor, ybar: torch.Tensor, functype: int, -): +) -> torch.Tensor: w = torch.from_numpy(w_np).to(env.DEVICE) if z.dim() != 2: raise ValueError("z tensor must have 2 dimensions") diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index 054dc3c80b..ab066bdf93 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( - Optional, - Union, + Any, overload, ) @@ -69,7 +68,7 @@ def silut_double_backward( class SiLUTScript(torch.nn.Module): - def __init__(self, threshold: float = 3.0): + def __init__(self, threshold: float = 3.0) -> None: super().__init__() self.threshold = threshold @@ -81,14 +80,20 @@ def __init__(self, threshold: float = 3.0): self.const_val = float(threshold * sigmoid_threshold) self.get_script_code() - def get_script_code(self): + def get_script_code(self) -> None: silut_forward_script = torch.jit.script(silut_forward) silut_backward_script = torch.jit.script(silut_backward) silut_double_backward_script = torch.jit.script(silut_double_backward) class SiLUTFunction(torch.autograd.Function): @staticmethod - def forward(ctx, x, threshold, slope, const_val): + def forward( + ctx: Any, + x: torch.Tensor, + threshold: float, + slope: float, + const_val: float, + ) -> torch.Tensor: ctx.save_for_backward(x) ctx.threshold = threshold ctx.slope = slope @@ -96,7 +101,9 @@ def forward(ctx, x, threshold, slope, const_val): return silut_forward_script(x, threshold, slope, const_val) @staticmethod - def backward(ctx, grad_output): + def backward( + ctx: Any, grad_output: torch.Tensor + ) -> tuple[torch.Tensor, None, None, None]: (x,) = ctx.saved_tensors threshold = ctx.threshold slope = ctx.slope @@ -106,7 +113,13 @@ def backward(ctx, grad_output): class SiLUTGradFunction(torch.autograd.Function): @staticmethod - def forward(ctx, x, grad_output, threshold, slope): + def forward( + ctx: Any, + x: torch.Tensor, + grad_output: torch.Tensor, + threshold: float, + slope: float, + ) -> torch.Tensor: ctx.threshold = threshold ctx.slope = slope grad_input = silut_backward_script(x, grad_output, threshold, slope) @@ -114,7 +127,9 @@ def forward(ctx, x, grad_output, threshold, slope): return grad_input @staticmethod - def backward(ctx, grad_grad_output): + def backward( + ctx: Any, grad_grad_output: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: (x, grad_output) = ctx.saved_tensors threshold = ctx.threshold slope = ctx.slope @@ -126,21 +141,21 @@ def backward(ctx, grad_grad_output): self.SiLUTFunction = SiLUTFunction - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: return self.SiLUTFunction.apply(x, self.threshold, self.slope, self.const_val) class SiLUT(torch.nn.Module): - def __init__(self, threshold=3.0): + def __init__(self, threshold: float = 3.0) -> None: super().__init__() - def sigmoid(x): + def sigmoid(x: float) -> float: return 1 / (1 + np.exp(-x)) - def silu(x): + def silu(x: float) -> float: return x * sigmoid(x) - def silu_grad(x): + def silu_grad(x: float) -> float: sig = sigmoid(x) return sig + x * sig * (1 - sig) @@ -156,7 +171,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class ActivationFn(torch.nn.Module): - def __init__(self, activation: Optional[str]) -> None: + def __init__(self, activation: str | None) -> None: super().__init__() self.activation: str = activation if activation is not None else "linear" if self.activation.lower().startswith( @@ -212,8 +227,8 @@ def to_numpy_array(xx: None) -> None: ... def to_numpy_array( - xx, -): + xx: torch.Tensor | None, +) -> np.ndarray | None: if xx is None: return None assert xx is not None @@ -239,8 +254,8 @@ def to_torch_tensor(xx: None) -> None: ... def to_torch_tensor( - xx, -): + xx: np.ndarray | None, +) -> torch.Tensor | None: if xx is None: return None assert xx is not None @@ -259,7 +274,7 @@ def to_torch_tensor( return torch.tensor(xx, dtype=prec, device=DEVICE) -def dict_to_device(sample_dict) -> None: +def dict_to_device(sample_dict: dict[str, Any]) -> None: for key in sample_dict: if isinstance(sample_dict[key], list): sample_dict[key] = [item.to(DEVICE) for item in sample_dict[key]] @@ -280,7 +295,7 @@ def dict_to_device(sample_dict) -> None: XSHIFT = 16 -def hashmix(value: int, hash_const: list[int]): +def hashmix(value: int, hash_const: list[int]) -> int: value ^= INIT_A hash_const[0] *= MULT_A value *= INIT_A @@ -291,7 +306,7 @@ def hashmix(value: int, hash_const: list[int]): return value -def mix(x: int, y: int): +def mix(x: int, y: int) -> int: result = MIX_MULT_L * x - MIX_MULT_R * y # prevent overflow result &= 0xFFFF_FFFF_FFFF_FFFF @@ -309,8 +324,8 @@ def mix_entropy(entropy_array: list[int]) -> int: def get_generator( - seed: Optional[Union[int, list[int]]] = None, -) -> Optional[torch.Generator]: + seed: int | list[int] | None = None, +) -> torch.Generator | None: if seed is not None: if isinstance(seed, list): seed = mix_entropy(seed) diff --git a/deepmd/tf/cluster/__init__.py b/deepmd/tf/cluster/__init__.py index 0f8916038d..a394b30f12 100644 --- a/deepmd/tf/cluster/__init__.py +++ b/deepmd/tf/cluster/__init__.py @@ -1,16 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Module that reads node resources, auto detects if running local or on SLURM.""" -from typing import ( - Optional, -) - from .local import get_resource as get_local_res __all__ = ["get_resource"] -def get_resource() -> tuple[str, list[str], Optional[list[int]]]: +def get_resource() -> tuple[str, list[str], list[int] | None]: """Get local or slurm resources: nodename, nodelist, and gpus. Returns diff --git a/deepmd/tf/cluster/local.py b/deepmd/tf/cluster/local.py index 25fb1cc645..0d1af07d68 100644 --- a/deepmd/tf/cluster/local.py +++ b/deepmd/tf/cluster/local.py @@ -3,9 +3,6 @@ import subprocess as sp import sys -from typing import ( - Optional, -) from deepmd.tf.env import ( tf, @@ -49,7 +46,7 @@ def get_gpus(): return list(range(num_gpus)) if num_gpus > 0 else None -def get_resource() -> tuple[str, list[str], Optional[list[int]]]: +def get_resource() -> tuple[str, list[str], list[int] | None]: """Get local resources: nodename, nodelist, and gpus. Returns diff --git a/deepmd/tf/common.py b/deepmd/tf/common.py index 985c36c686..bf2681594f 100644 --- a/deepmd/tf/common.py +++ b/deepmd/tf/common.py @@ -2,13 +2,15 @@ """Collection of functions and classes used throughout the whole package.""" import warnings +from collections.abc import ( + Callable, +) from functools import ( wraps, ) from typing import ( TYPE_CHECKING, Any, - Callable, Union, ) diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py index bd1af8c72e..e0b61ebb7e 100644 --- a/deepmd/tf/descriptor/descriptor.py +++ b/deepmd/tf/descriptor/descriptor.py @@ -4,7 +4,6 @@ ) from typing import ( Any, - Optional, ) import numpy as np @@ -173,7 +172,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict[str, Any], - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -253,7 +252,7 @@ def enable_compression( f"Descriptor {type(self).__name__} doesn't support compression!" ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -466,9 +465,9 @@ def explicit_ntypes(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py index 57c21f0ee6..1611344c9a 100644 --- a/deepmd/tf/descriptor/hybrid.py +++ b/deepmd/tf/descriptor/hybrid.py @@ -1,8 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -45,9 +43,9 @@ class DescrptHybrid(Descriptor): def __init__( self, - list: list[Union[Descriptor, dict[str, Any]]], - ntypes: Optional[int] = None, - spin: Optional[Spin] = None, + list: list[Descriptor | dict[str, Any]], + ntypes: int | None = None, + spin: Spin | None = None, **kwargs, ) -> None: """Constructor.""" @@ -143,7 +141,7 @@ def compute_input_stats( mesh: list, input_dict: dict, mixed_type: bool = False, - real_natoms_vec: Optional[list] = None, + real_natoms_vec: list | None = None, **kwargs, ) -> None: """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. @@ -213,7 +211,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -351,7 +349,7 @@ def enable_compression( suffix=f"{suffix}_{idx}", ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -426,9 +424,9 @@ def explicit_ntypes(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py index 9b338a5d25..87db875038 100644 --- a/deepmd/tf/descriptor/loc_frame.py +++ b/deepmd/tf/descriptor/loc_frame.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -218,7 +215,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -435,9 +432,9 @@ def init_variables( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py index 2863704143..e62885eb62 100644 --- a/deepmd/tf/descriptor/se.py +++ b/deepmd/tf/descriptor/se.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import re -from typing import ( - Optional, -) from deepmd.dpmodel.utils.network import ( EmbeddingNet, @@ -154,9 +151,9 @@ def precision(self) -> tf.DType: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -192,6 +189,7 @@ def serialize_network( resnet_dt: bool, variables: dict, excluded_types: set[tuple[int, int]] = set(), + trainable: bool = True, suffix: str = "", ) -> dict: """Serialize network. @@ -214,6 +212,8 @@ def serialize_network( The input variables excluded_types : set[tuple[int, int]], optional The excluded types + trainable : bool + Whether the network is trainable suffix : str, optional The suffix of the scope @@ -236,6 +236,7 @@ def serialize_network( activation_function=activation_function, resnet_dt=resnet_dt, precision=self.precision.name, + trainable=trainable, ) embeddings[(type_j, type_i)] = EmbeddingNet( in_dim=in_dim, @@ -243,6 +244,7 @@ def serialize_network( activation_function=activation_function, resnet_dt=resnet_dt, precision=self.precision.name, + trainable=trainable, ) embeddings[(type_i, type_j)].clear() embeddings[(type_j, type_i)].clear() @@ -278,6 +280,7 @@ def serialize_network( activation_function=activation_function, resnet_dt=resnet_dt, precision=self.precision.name, + trainable=trainable, ) assert embeddings[network_idx] is not None if weight_name == "idt": diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py index e3ae7bf99c..3d2ab39c9f 100644 --- a/deepmd/tf/descriptor/se_a.py +++ b/deepmd/tf/descriptor/se_a.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -172,16 +169,16 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - spin: Optional[Spin] = None, + spin: Spin | None = None, tebd_input_mode: str = "concat", - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: @@ -547,7 +544,7 @@ def enable_compression( self.davg = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_avg") self.dstd = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_std") - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -566,7 +563,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -1462,6 +1459,7 @@ def serialize(self, suffix: str = "") -> dict: resnet_dt=self.filter_resnet_dt, variables=self.embedding_net_variables, excluded_types=self.exclude_types, + trainable=self.trainable, suffix=suffix, ), "env_mat": EnvMat(self.rcut_r, self.rcut_r_smth).serialize(), diff --git a/deepmd/tf/descriptor/se_a_ebd.py b/deepmd/tf/descriptor/se_a_ebd.py index ae76308e69..189a013c13 100644 --- a/deepmd/tf/descriptor/se_a_ebd.py +++ b/deepmd/tf/descriptor/se_a_ebd.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -78,7 +75,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, type_nchanl: int = 2, type_nlayer: int = 1, @@ -118,7 +115,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. diff --git a/deepmd/tf/descriptor/se_a_ebd_v2.py b/deepmd/tf/descriptor/se_a_ebd_v2.py index 035fc6509c..65e28e99d9 100644 --- a/deepmd/tf/descriptor/se_a_ebd_v2.py +++ b/deepmd/tf/descriptor/se_a_ebd_v2.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) from deepmd.tf.utils.spin import ( Spin, @@ -35,14 +32,14 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - spin: Optional[Spin] = None, + spin: Spin | None = None, **kwargs, ) -> None: DescrptSeA.__init__( diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py index d5476cbe89..37aa830431 100644 --- a/deepmd/tf/descriptor/se_a_ef.py +++ b/deepmd/tf/descriptor/se_a_ef.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -77,7 +74,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, @@ -207,7 +204,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -308,7 +305,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, diff --git a/deepmd/tf/descriptor/se_a_mask.py b/deepmd/tf/descriptor/se_a_mask.py index 5667122809..a0c5adaf76 100644 --- a/deepmd/tf/descriptor/se_a_mask.py +++ b/deepmd/tf/descriptor/se_a_mask.py @@ -2,7 +2,6 @@ import warnings from typing import ( Any, - Optional, ) import numpy as np @@ -124,7 +123,7 @@ def __init__( trainable: bool = True, type_one_side: bool = False, exclude_types: list[list[int]] = [], - seed: Optional[int] = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, @@ -269,7 +268,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict[str, Any], - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -427,9 +426,9 @@ def prod_force_virial( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 3a9b86a0d6..3978709715 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -4,8 +4,6 @@ import warnings from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -172,13 +170,13 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, set_davg_zero: bool = True, exclude_types: list[list[int]] = [], @@ -196,11 +194,11 @@ def __init__( normalize=True, temperature=None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-3, + ln_eps: float | None = 1e-3, concat_output_tebd: bool = True, env_protection: float = 0.0, # not implement!! - stripped_type_embedding: Optional[bool] = None, - type_map: Optional[list[str]] = None, # to be compat with input + stripped_type_embedding: bool | None = None, + type_map: list[str] | None = None, # to be compat with input **kwargs, ) -> None: # Ensure compatibility with the deprecated stripped_type_embedding option. @@ -343,7 +341,7 @@ def compute_input_stats( mesh: list, input_dict: dict, mixed_type: bool = False, - real_natoms_vec: Optional[list] = None, + real_natoms_vec: list | None = None, **kwargs, ) -> None: """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. @@ -522,7 +520,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -1513,9 +1511,9 @@ def explicit_ntypes(self) -> bool: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -1593,6 +1591,7 @@ def serialize_attention_layers( bias=bias, use_timestep=False, precision=self.precision.name, + trainable=self.trainable, ) matrix_list = [ attention_layer_params[layer_idx][key]["matrix"] @@ -1611,6 +1610,7 @@ def serialize_attention_layers( bias=bias, use_timestep=False, precision=self.precision.name, + trainable=self.trainable, ) out_proj["matrix"] = attention_layer_params[layer_idx]["c_out"]["matrix"] if bias: @@ -1654,6 +1654,7 @@ def serialize_network_strip( variables: dict, suffix: str = "", type_one_side: bool = False, + trainable: bool = True, ) -> dict: """Serialize network. @@ -1679,6 +1680,8 @@ def serialize_network_strip( If 'False', type embeddings of both neighbor and central atoms are considered. If 'True', only type embeddings of neighbor atoms are considered. Default is 'False'. + trainable : bool + Whether the network is trainable Returns ------- @@ -1719,6 +1722,7 @@ def serialize_network_strip( activation_function=activation_function, resnet_dt=resnet_dt, precision=self.precision.name, + trainable=trainable, ) assert embeddings[network_idx] is not None if weight_name == "idt": @@ -1983,6 +1987,7 @@ def serialize(self, suffix: str = "") -> dict: resnet_dt=self.filter_resnet_dt, variables=self.embedding_net_variables, excluded_types=self.exclude_types, + trainable=self.trainable, suffix=suffix, ), "attention_layers": self.serialize_attention_layers( @@ -2032,6 +2037,7 @@ def serialize(self, suffix: str = "") -> dict: variables=self.two_side_embeeding_net_variables, suffix=suffix, type_one_side=self.type_one_side, + trainable=self.trainable, ) } ) @@ -2176,7 +2182,7 @@ def __init__( self, rcut: float, rcut_smth: float, - sel: Union[list[int], int], + sel: list[int] | int, ntypes: int, neuron: list[int] = [25, 50, 100], axis_neuron: int = 8, @@ -2196,17 +2202,17 @@ def __init__( precision: str = "default", scaling_factor=1.0, normalize: bool = True, - temperature: Optional[float] = None, + temperature: float | None = None, trainable_ln: bool = True, - ln_eps: Optional[float] = 1e-3, + ln_eps: float | None = 1e-3, smooth_type_embedding: bool = True, concat_output_tebd: bool = True, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, - spin: Optional[Any] = None, + type_map: list[str] | None = None, + spin: Any | None = None, # consistent with argcheck, not used though - seed: Optional[int] = None, + seed: int | None = None, uniform_seed: bool = False, ) -> None: if not normalize: @@ -2294,7 +2300,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: type_embedding = self.type_embedding.build(self.ntypes, suffix=suffix) diff --git a/deepmd/tf/descriptor/se_atten_v2.py b/deepmd/tf/descriptor/se_atten_v2.py index 69efe004c4..269fa1c083 100644 --- a/deepmd/tf/descriptor/se_atten_v2.py +++ b/deepmd/tf/descriptor/se_atten_v2.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) from deepmd.tf.utils.type_embed import ( TypeEmbedNet, @@ -77,7 +74,7 @@ def __init__( axis_neuron: int = 8, resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, set_davg_zero: bool = False, exclude_types: list[list[int]] = [], diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py index ed66d6ad25..c38a13d35a 100644 --- a/deepmd/tf/descriptor/se_r.py +++ b/deepmd/tf/descriptor/se_r.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -95,15 +92,15 @@ def __init__( neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, type_one_side: bool = True, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - spin: Optional[Spin] = None, - type_map: Optional[list[str]] = None, # to be compat with input + spin: Spin | None = None, + type_map: list[str] | None = None, # to be compat with input env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: @@ -382,7 +379,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -795,6 +792,7 @@ def serialize(self, suffix: str = "") -> dict: resnet_dt=self.filter_resnet_dt, variables=self.embedding_net_variables, excluded_types=self.exclude_types, + trainable=self.trainable, suffix=suffix, ), "env_mat": EnvMat(self.rcut, self.rcut_smth).serialize(), diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py index c5d50744af..b03746a9c6 100644 --- a/deepmd/tf/descriptor/se_t.py +++ b/deepmd/tf/descriptor/se_t.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import re -from typing import ( - Optional, -) import numpy as np @@ -99,13 +96,13 @@ def __init__( neuron: list[int] = [24, 48, 96], resnet_dt: bool = False, trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, exclude_types: list[list[int]] = [], set_davg_zero: bool = False, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input env_protection: float = 0.0, # not implement!! **kwargs, ) -> None: @@ -387,7 +384,7 @@ def build( box_: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - reuse: Optional[bool] = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for the descriptor. @@ -726,6 +723,7 @@ def serialize_network( resnet_dt: bool, variables: dict, excluded_types: set[tuple[int, int]] = set(), + trainable: bool = True, suffix: str = "", ) -> dict: """Serialize network. @@ -748,6 +746,8 @@ def serialize_network( The input variables excluded_types : set[tuple[int, int]], optional The excluded types + trainable : bool, optional + Whether the network is trainable suffix : str, optional The suffix of the scope @@ -771,6 +771,7 @@ def clear_ij(type_i, type_j) -> None: activation_function=activation_function, resnet_dt=resnet_dt, precision=self.precision.name, + trainable=trainable, ) embeddings[(type_i, type_j)].clear() @@ -805,6 +806,7 @@ def clear_ij(type_i, type_j) -> None: activation_function=activation_function, resnet_dt=resnet_dt, precision=self.precision.name, + trainable=trainable, ) assert embeddings[network_idx] is not None if weight_name == "idt": @@ -941,6 +943,7 @@ def serialize(self, suffix: str = "") -> dict: resnet_dt=self.filter_resnet_dt, variables=self.embedding_net_variables, excluded_types=self.exclude_types, + trainable=self.trainable, suffix=suffix, ), "env_mat": EnvMat(self.rcut_r, self.rcut_r_smth).serialize(), diff --git a/deepmd/tf/entrypoints/__init__.py b/deepmd/tf/entrypoints/__init__.py index bf8c51067e..a33dc5b983 100644 --- a/deepmd/tf/entrypoints/__init__.py +++ b/deepmd/tf/entrypoints/__init__.py @@ -4,6 +4,9 @@ from ..infer.model_devi import ( make_model_devi, ) +from .change_bias import ( + change_bias, +) from .compress import ( compress, ) @@ -34,6 +37,7 @@ ) __all__ = [ + "change_bias", "compress", "convert", "doc_train_input", diff --git a/deepmd/tf/entrypoints/change_bias.py b/deepmd/tf/entrypoints/change_bias.py new file mode 100644 index 0000000000..adf65c0e2b --- /dev/null +++ b/deepmd/tf/entrypoints/change_bias.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""DeePMD change bias entrypoint script.""" + +import logging +import os +import shutil +import tempfile +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.common import ( + expand_sys_str, + j_loader, +) +from deepmd.tf.entrypoints.freeze import ( + freeze, +) +from deepmd.tf.env import ( + tf, +) +from deepmd.tf.infer import ( + DeepPotential, +) +from deepmd.tf.train.run_options import ( + RunOptions, +) +from deepmd.tf.train.trainer import ( + DPTrainer, +) +from deepmd.tf.utils.argcheck import ( + normalize, +) +from deepmd.tf.utils.compat import ( + update_deepmd_input, +) +from deepmd.tf.utils.sess import ( + run_sess, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + +__all__ = ["change_bias"] + +log = logging.getLogger(__name__) + + +def change_bias( + INPUT: str, + mode: str = "change", + bias_value: list | None = None, + datafile: str | None = None, + system: str = ".", + numb_batch: int = 0, + model_branch: str | None = None, + output: str | None = None, + log_level: int = 0, + **kwargs, +) -> None: + """Change model out bias according to the input data. + + Parameters + ---------- + INPUT : str + The input checkpoint file or frozen model file + mode : str, optional + The mode for changing energy bias, by default "change" + bias_value : Optional[list], optional + The user defined value for each type, by default None + datafile : Optional[str], optional + The path to the datafile, by default None + system : str, optional + The system dir, by default "." + numb_batch : int, optional + The number of frames for bias changing, by default 0 + model_branch : Optional[str], optional + Model branch chosen for changing bias if multi-task model, by default None + output : Optional[str], optional + The model after changing bias, by default None + log_level : int, optional + The log level for output, by default 0 + """ + # Determine input type and handle accordingly + if INPUT.endswith(".pb"): + # Frozen model (.pb) + return _change_bias_frozen_model( + INPUT, + mode, + bias_value, + datafile, + system, + numb_batch, + model_branch, + output, + log_level, + ) + elif INPUT.endswith(".pbtxt"): + # Text format frozen model (.pbtxt) - not supported + raise NotImplementedError( + "Bias changing for .pbtxt models is not supported. " + "Please convert to .pb format first using: dp convert-from pbtxt -i model.pbtxt -o model.pb" + ) + elif INPUT.endswith((".ckpt", ".meta", ".data", ".index")): + # Individual checkpoint files + checkpoint_prefix = INPUT + if INPUT.endswith((".meta", ".data", ".index")): + checkpoint_prefix = INPUT.rsplit(".", 1)[0] + return _change_bias_checkpoint_file( + checkpoint_prefix, + mode, + bias_value, + datafile, + system, + numb_batch, + model_branch, + output, + log_level, + ) + else: + raise RuntimeError( + "The model provided must be a checkpoint file or frozen model file (.pb)" + ) + + +def _change_bias_checkpoint_file( + checkpoint_prefix: str, + mode: str, + bias_value: list | None, + datafile: str | None, + system: str, + numb_batch: int, + model_branch: str | None, + output: str | None, + log_level: int, +) -> None: + """Change bias for individual checkpoint files.""" + # Reset the default graph to avoid variable conflicts + tf.reset_default_graph() + + checkpoint_path = Path(checkpoint_prefix) + checkpoint_dir = checkpoint_path.parent + + # Check for valid checkpoint and find the actual checkpoint path + checkpoint_state_file = checkpoint_dir / "checkpoint" + if not checkpoint_state_file.exists(): + raise RuntimeError(f"No valid checkpoint found in {checkpoint_dir}") + + # Get the latest checkpoint path from the checkpoint state file + checkpoint_state = tf.train.get_checkpoint_state(str(checkpoint_dir)) + if checkpoint_state is None or checkpoint_state.model_checkpoint_path is None: + raise RuntimeError(f"No valid checkpoint state found in {checkpoint_dir}") + + # The model_checkpoint_path from get_checkpoint_state is the full path to the checkpoint + actual_checkpoint_path = checkpoint_state.model_checkpoint_path + + bias_adjust_mode = "change-by-statistic" if mode == "change" else "set-by-statistic" + + # Read the checkpoint to get the model configuration + input_json_path = _find_input_json(checkpoint_dir) + jdata = j_loader(input_json_path) + + # Update and normalize the configuration + jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + jdata = normalize(jdata) + + # Determine output path - should be a single model file + if output is None: + output = str(checkpoint_path.with_suffix(".pb")) + elif not output.endswith(".pb"): + output = output + ".pb" + + # Create trainer to access model methods + run_opt = RunOptions( + init_model=actual_checkpoint_path, # Use the actual checkpoint file path + restart=None, + finetune=None, + init_frz_model=None, + log_level=log_level, + ) + + trainer = DPTrainer(jdata, run_opt) + + # Load data for bias calculation using trainer data requirements + data = _load_data_systems(datafile, system, trainer) + + # Get stop_batch and origin_type_map like in train.py + stop_batch = jdata.get("training", {}).get("numb_steps", 0) + origin_type_map = jdata["model"].get("origin_type_map", None) + if origin_type_map is not None and not origin_type_map: + # get the type_map from data if not provided + origin_type_map = data.get_type_map() + + try: + # Build the model graph first with proper parameters, then initialize session + # and restore variables from checkpoint - following train.py pattern + trainer.build(data, stop_batch, origin_type_map=origin_type_map) + trainer._init_session() + + if bias_value is not None: + # Use user-defined bias + _apply_user_defined_bias(trainer, bias_value) + else: + # Use data-based bias calculation + type_map = data.get_type_map() + if len(type_map) == 0: + # If data doesn't have type_map, get from model + type_map = trainer.model.get_type_map() + + log.info(f"Changing bias for model with type_map: {type_map}") + log.info(f"Using bias adjustment mode: {bias_adjust_mode}") + + # Read current bias values from the session (after variables are restored) + _apply_data_based_bias(trainer, data, type_map, bias_adjust_mode) + + # Save the updated variables back to checkpoint format first + # Create a separate directory for updated checkpoint to avoid polluting original + updated_checkpoint_dir = checkpoint_dir / f"{checkpoint_path.name}_updated" + updated_checkpoint_dir.mkdir(exist_ok=True) + + # Copy the input.json file to the new directory + updated_input_json_path = updated_checkpoint_dir / "input.json" + shutil.copy2(input_json_path, updated_input_json_path) + + updated_checkpoint_prefix = str(updated_checkpoint_dir / checkpoint_path.name) + if hasattr(trainer, "saver") and trainer.saver is not None: + log.info(f"Saving updated checkpoint to {updated_checkpoint_prefix}") + trainer.saver.save(trainer.sess, updated_checkpoint_prefix) + + # Create a new checkpoint state file in the updated directory + updated_checkpoint_state_file = updated_checkpoint_dir / "checkpoint" + with open(updated_checkpoint_state_file, "w") as f: + f.write(f'model_checkpoint_path: "{checkpoint_path.name}"\n') + f.write(f'all_model_checkpoint_paths: "{checkpoint_path.name}"\n') + + # Then save the updated model as a frozen model using the updated checkpoint directory + freeze( + checkpoint_folder=str(updated_checkpoint_dir), + output=output, + ) + + log.info(f"Bias changing complete. Model saved to {output}") + + finally: + # Ensure session is properly closed + if hasattr(trainer, "sess") and trainer.sess is not None: + trainer.sess.close() + + +def _change_bias_frozen_model( + frozen_model_path: str, + mode: str, + bias_value: list | None, + datafile: str | None, + system: str, + numb_batch: int, + model_branch: str | None, + output: str | None, + log_level: int, +) -> None: + """Change bias for frozen model (.pb file).""" + if bias_value is None: + raise NotImplementedError( + "Data-based bias changing for frozen models is not yet implemented. " + "Please provide user-defined bias values using the -b/--bias-value option, " + "or use a checkpoint directory instead." + ) + + # For frozen models, we need to modify the graph and save a new frozen model + # This is complex and requires graph manipulation + # For now, provide a clear error message with workaround + raise NotImplementedError( + "Bias modification for frozen models (.pb) is not yet fully implemented. " + "Recommended workaround:\n" + "1. Use a checkpoint directory instead of a frozen model\n" + "2. Or load the model, modify bias in training, then freeze again\n" + f" dp --tf change-bias -b {' '.join(map(str, bias_value)) if bias_value else ''} -o \n" + " dp freeze -c -o modified_model.pb" + ) + + +def _load_data_systems( + datafile: str | None, system: str, trainer: DPTrainer +) -> DeepmdDataSystem: + """Load data systems for bias calculation.""" + if datafile is not None: + with open(datafile) as datalist: + all_sys = datalist.read().splitlines() + else: + all_sys = expand_sys_str(system) + + # Load the data systems with proper data requirements + data = DeepmdDataSystem( + systems=all_sys, + batch_size=1, + test_size=1, + rcut=None, + set_prefix="set", + ) + # Use the data requirements from the trainer model instead of hardcoding them + data.add_data_requirements(trainer.data_requirements) + return data + + +def _find_input_json(checkpoint_dir: Path) -> Path: + """Find the input.json file for the checkpoint.""" + input_json_path = checkpoint_dir / "input.json" + if not input_json_path.exists(): + # Look for input.json in parent directories or common locations + for parent in checkpoint_dir.parents: + potential_input = parent / "input.json" + if potential_input.exists(): + input_json_path = potential_input + break + else: + raise RuntimeError( + f"Cannot find input.json configuration file needed to load the model. " + f"Please ensure input.json is available in {checkpoint_dir} or its parent directories." + ) + return input_json_path + + +def _apply_data_based_bias( + trainer: DPTrainer, data: DeepmdDataSystem, type_map: list, bias_adjust_mode: str +) -> None: + """Apply data-based bias calculation by reading current bias from session.""" + from deepmd.tf.env import ( + tf, + ) + from deepmd.tf.fit.ener import ( + change_energy_bias_lower, + ) + + # Get the fitting object which contains the bias tensor + fitting = trainer.model.get_fitting() + if not hasattr(fitting, "t_bias_atom_e"): + raise RuntimeError( + "Model does not have t_bias_atom_e tensor for bias modification" + ) + + # Read current bias values from the session (these are the restored values) + current_bias = run_sess(trainer.sess, fitting.t_bias_atom_e) + + log.info(f"Current bias values from session: {current_bias.flatten()}") + + # Create a temporary frozen model to use with change_energy_bias_lower + with tempfile.NamedTemporaryFile(suffix=".pb", delete=False) as temp_frozen: + freeze( + checkpoint_folder=str(Path(trainer.run_opt.init_model).parent), + output=temp_frozen.name, + ) + + try: + # Create DeepPotential object for evaluation + dp = DeepPotential(temp_frozen.name) + + # Use change_energy_bias_lower with the current bias values from session + new_bias = change_energy_bias_lower( + data, + dp, + type_map, # origin_type_map + type_map, # full_type_map + current_bias, # Use the restored bias values + bias_adjust_mode=bias_adjust_mode, + ntest=1, + ) + + # Update the bias in the session + if len(new_bias.shape) == 1: + # 1D tensor, keep bias as 1D + new_bias_tensor = new_bias.flatten() + else: + # 2D tensor, reshape to match + new_bias_tensor = new_bias.reshape(-1, 1) + + assign_op = tf.assign(fitting.t_bias_atom_e, new_bias_tensor) + run_sess(trainer.sess, assign_op) + + # Also update the numpy array in the fitting object for consistency + fitting.bias_atom_e = new_bias + + finally: + # Clean up temporary file + os.unlink(temp_frozen.name) + + +def _apply_user_defined_bias(trainer: DPTrainer, bias_value: list) -> None: + """Apply user-defined bias values to the model.""" + # Get the type map from the model + type_map = trainer.model.get_type_map() + + # Validate bias_value length + if len(bias_value) != len(type_map): + raise ValueError( + f"The number of elements in the bias ({len(bias_value)}) should be the same as " + f"that in the type_map ({len(type_map)}): {type_map}" + ) + + # Check model type + if trainer.model.model_type != "ener": + raise RuntimeError( + f"User-defined bias is only supported for energy models, got: {trainer.model.model_type}" + ) + + # Get current bias + fitting = trainer.model.get_fitting() + if not hasattr(fitting, "bias_atom_e"): + raise RuntimeError( + "Model does not have bias_atom_e attribute for bias modification" + ) + + # Convert user bias to numpy array with proper shape matching the tensor + new_bias = np.array(bias_value, dtype=np.float64) + + # Check the shape of the existing bias tensor to match it + if hasattr(fitting, "t_bias_atom_e"): + existing_shape = fitting.t_bias_atom_e.get_shape().as_list() + if len(existing_shape) == 1: + # 1D tensor, keep bias as 1D + new_bias = new_bias.flatten() + else: + # 2D tensor, reshape to match + new_bias = new_bias.reshape(-1, 1) + else: + # If no tensor, use the fitting.bias_atom_e shape + new_bias = new_bias.reshape(fitting.bias_atom_e.shape) + + log.info(f"Changing bias from user-defined values for type_map: {type_map}") + log.info(f"Old bias: {fitting.bias_atom_e.flatten()}") + log.info(f"New bias: {new_bias.flatten()}") + + # Update the bias in the model + fitting.bias_atom_e = new_bias + + # Update the tensor in the session if needed + if hasattr(fitting, "t_bias_atom_e"): + assign_op = tf.assign(fitting.t_bias_atom_e, new_bias) + run_sess(trainer.sess, assign_op) diff --git a/deepmd/tf/entrypoints/compress.py b/deepmd/tf/entrypoints/compress.py index e8ceec7e9c..96bbd2c416 100644 --- a/deepmd/tf/entrypoints/compress.py +++ b/deepmd/tf/entrypoints/compress.py @@ -4,9 +4,6 @@ import json import logging import os -from typing import ( - Optional, -) from deepmd.tf.common import ( j_loader, @@ -58,7 +55,7 @@ def compress( checkpoint_folder: str, training_script: str, mpi_log: str, - log_path: Optional[str], + log_path: str | None, log_level: int, **kwargs, ) -> None: diff --git a/deepmd/tf/entrypoints/freeze.py b/deepmd/tf/entrypoints/freeze.py index 83d2ff33f4..366dba9c31 100755 --- a/deepmd/tf/entrypoints/freeze.py +++ b/deepmd/tf/entrypoints/freeze.py @@ -14,10 +14,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) import google.protobuf.message @@ -76,9 +72,9 @@ def _transfer_fitting_net_trainable_variables(sess, old_graph_def, raw_graph_def def _make_node_names( model_type: str, - modifier_type: Optional[str] = None, + modifier_type: str | None = None, out_suffix: str = "", - node_names: Optional[Union[str, list]] = None, + node_names: str | list | None = None, ) -> list[str]: """Get node names based on model type. @@ -297,8 +293,8 @@ def freeze( *, checkpoint_folder: str, output: str, - node_names: Optional[str] = None, - nvnmd_weight: Optional[str] = None, + node_names: str | None = None, + nvnmd_weight: str | None = None, **kwargs, ) -> None: """Freeze the graph in supplied folder. diff --git a/deepmd/tf/entrypoints/main.py b/deepmd/tf/entrypoints/main.py index 5058c51c17..8d83eac826 100644 --- a/deepmd/tf/entrypoints/main.py +++ b/deepmd/tf/entrypoints/main.py @@ -5,10 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, - Union, -) from deepmd.backend.suffix import ( format_model_suffix, @@ -22,6 +18,7 @@ clear_session, ) from deepmd.tf.entrypoints import ( + change_bias, compress, convert, freeze, @@ -38,7 +35,7 @@ __all__ = ["get_ll", "main", "main_parser", "parse_args"] -def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: +def main(args: list[str] | argparse.Namespace | None = None) -> None: """DeePMD-Kit entry point. Parameters @@ -86,6 +83,8 @@ def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None: compress(**dict_args) elif args.command == "convert-from": convert(**dict_args) + elif args.command == "change-bias": + change_bias(**dict_args) elif args.command == "train-nvnmd": # nvnmd train_nvnmd(**dict_args) elif args.command is None: diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index b12e4fe1af..3ab55e190c 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -10,10 +10,9 @@ import time from typing import ( Any, - Optional, ) -from deepmd.tf.common import ( +from deepmd.common import ( j_loader, ) from deepmd.tf.env import ( @@ -55,16 +54,16 @@ def train( *, INPUT: str, - init_model: Optional[str], - restart: Optional[str], + init_model: str | None, + restart: str | None, output: str, init_frz_model: str, mpi_log: str, log_level: int, - log_path: Optional[str], + log_path: str | None, is_compress: bool = False, skip_neighbor_stat: bool = False, - finetune: Optional[str] = None, + finetune: str | None = None, use_pretrain_script: bool = False, **kwargs, ) -> None: @@ -276,7 +275,7 @@ def _do_work( def get_modifier(modi_data=None): - modifier: Optional[BaseModifier] + modifier: BaseModifier | None if modi_data is not None: modifier_params = copy.deepcopy(modi_data) modifier_type = modifier_params.pop("type") diff --git a/deepmd/tf/entrypoints/transfer.py b/deepmd/tf/entrypoints/transfer.py index 0ad022776a..e441322764 100644 --- a/deepmd/tf/entrypoints/transfer.py +++ b/deepmd/tf/entrypoints/transfer.py @@ -6,9 +6,6 @@ from collections.abc import ( Sequence, ) -from typing import ( - Optional, -) import numpy as np @@ -32,7 +29,7 @@ def convert_number(number: int) -> float: def convert_matrix( - matrix: np.ndarray, shape: Sequence[int], dtype: Optional[type] = None + matrix: np.ndarray, shape: Sequence[int], dtype: type | None = None ) -> np.ndarray: """Convert matrix of integers to self defined binary format. @@ -188,7 +185,7 @@ def __init__(self, node) -> None: self.node = node def from_array( - self, tensor: np.ndarray, dtype: type, shape: Optional[Sequence[int]] = None + self, tensor: np.ndarray, dtype: type, shape: Sequence[int] | None = None ) -> None: if shape is None: shape = tensor.shape diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 4428d06536..961198b8e7 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.tf.common import ( cast_precision, get_activation_func, @@ -75,6 +75,13 @@ class DipoleFittingSeA(Fitting): different fitting nets for different atom types. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. + trainable : list[bool], Optional + If the weights of fitting net are trainable. + Suppose that we have :math:`N_l` hidden layers in the fitting net, + this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. """ def __init__( @@ -87,13 +94,15 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - sel_type: Optional[list[int]] = None, - seed: Optional[int] = None, + sel_type: list[int] | None = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input + trainable: list[bool] | None = None, **kwargs, ) -> None: """Constructor.""" @@ -123,18 +132,30 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam if numb_fparam > 0: raise ValueError("numb_fparam is not supported in the dipole fitting") if numb_aparam > 0: raise ValueError("numb_aparam is not supported in the dipole fitting") if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + if default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None self.aparam_avg = None self.aparam_std = None self.aparam_inv_std = None + if trainable is None: + self.trainable = [True for _ in range(len(self.n_neuron) + 1)] + elif isinstance(trainable, bool): + self.trainable = [trainable] * (len(self.n_neuron) + 1) + else: + self.trainable = trainable + assert len(self.trainable) == len(self.n_neuron) + 1, ( + "length of trainable should be that of n_neuron + 1" + ) def get_sel_type(self) -> int: """Get selected type.""" @@ -166,6 +187,7 @@ def _build_lower(self, start_index, natoms, inputs, rot_mat, suffix="", reuse=No uniform_seed=self.uniform_seed, initial_variables=self.fitting_net_variables, mixed_prec=self.mixed_prec, + trainable=self.trainable[ii], ) else: layer = one_layer( @@ -179,6 +201,7 @@ def _build_lower(self, start_index, natoms, inputs, rot_mat, suffix="", reuse=No uniform_seed=self.uniform_seed, initial_variables=self.fitting_net_variables, mixed_prec=self.mixed_prec, + trainable=self.trainable[ii], ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift @@ -195,6 +218,7 @@ def _build_lower(self, start_index, natoms, inputs, rot_mat, suffix="", reuse=No initial_variables=self.fitting_net_variables, mixed_prec=self.mixed_prec, final_layer=True, + trainable=self.trainable[-1], ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift @@ -213,8 +237,8 @@ def build( input_d: tf.Tensor, rot_mat: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for fitting net. @@ -346,7 +370,7 @@ def init_variables( graph_def, suffix=suffix ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -391,20 +415,22 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "dipole", - "@version": 3, + "@version": 4, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, "mixed_types": self.mixed_types, - "dim_out": 3, "neuron": self.n_neuron, "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, - "exclude_types": [], + "exclude_types": [] + if self.sel_type is None + else [ii for ii in range(self.ntypes) if ii not in self.sel_type], "nets": self.serialize_network( ntypes=self.ntypes, ndim=0 if self.mixed_types else 1, @@ -414,9 +440,29 @@ def serialize(self, suffix: str) -> dict: activation_function=self.activation_function_name, resnet_dt=self.resnet_dt, variables=self.fitting_net_variables, + trainable=self.trainable, suffix=suffix, ), + "@variables": { + "fparam_avg": self.fparam_avg, + "fparam_inv_std": self.fparam_inv_std, + "aparam_avg": self.aparam_avg, + "aparam_inv_std": self.aparam_inv_std, + "case_embd": None, + "bias_atom_e": np.zeros( + (self.ntypes, self.dim_rot_mat_1), dtype=GLOBAL_NP_FLOAT_PRECISION + ), + }, "type_map": self.type_map, + "var_name": "dipole", + "rcond": None, + "tot_ener_zero": False, + "trainable": self.trainable, + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + "r_differentiable": True, + "c_differentiable": True, } return data @@ -435,7 +481,12 @@ def deserialize(cls, data: dict, suffix: str): The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) + exclude_types = data.pop("exclude_types", []) + if len(exclude_types) > 0: + data["sel_type"] = [ + ii for ii in range(data["ntypes"]) if ii not in exclude_types + ] fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index 8fa3167bfc..250d803d8f 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -101,6 +98,9 @@ class DOSFitting(Fitting): mixed_types : bool If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. """ @@ -115,16 +115,17 @@ def __init__( numb_aparam: int = 0, dim_case_embd: int = 0, numb_dos: int = 300, - rcond: Optional[float] = None, - trainable: Optional[list[bool]] = None, - seed: Optional[int] = None, + rcond: float | None = None, + trainable: list[bool] | None = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input **kwargs, ) -> None: """Constructor.""" @@ -136,8 +137,11 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + if default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.numb_dos = numb_dos @@ -399,8 +403,8 @@ def build( self, inputs: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for fitting net. @@ -633,7 +637,7 @@ def init_variables( # for compatibility, old models has no t_bias_dos pass - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -678,7 +682,7 @@ def deserialize(cls, data: dict, suffix: str = ""): The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) data["numb_dos"] = data.pop("dim_out") fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( @@ -705,7 +709,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "dos", - "@version": 3, + "@version": 4, "var_name": "dos", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, @@ -716,6 +720,7 @@ def serialize(self, suffix: str = "") -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "rcond": self.rcond, "trainable": self.trainable, "activation_function": self.activation_function, @@ -730,6 +735,7 @@ def serialize(self, suffix: str = "") -> dict: activation_function=self.activation_function, resnet_dt=self.resnet_dt, variables=self.fitting_net_variables, + trainable=self.trainable, suffix=suffix, ), "@variables": { @@ -741,6 +747,11 @@ def serialize(self, suffix: str = "") -> dict: "case_embd": None, }, "type_map": self.type_map, + "tot_ener_zero": False, + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + "atom_ener": None, } return data diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index e10468df32..2b8b1b906e 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -119,6 +116,8 @@ class EnerFitting(Fitting): Number of atomic parameter dim_case_embd Dimension of case specific embedding. + default_fparam + The default frame parameter. This parameter is not supported in TensorFlow. rcond The condition number for the regression of atomic energy. tot_ener_zero @@ -146,6 +145,9 @@ class EnerFitting(Fitting): mixed_types : bool If true, use a uniform fitting net for all atom types, otherwise use different fitting nets for different atom types. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. """ @@ -159,19 +161,20 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - rcond: Optional[float] = None, + rcond: float | None = None, tot_ener_zero: bool = False, - trainable: Optional[list[bool]] = None, - seed: Optional[int] = None, + trainable: list[bool] | None = None, + seed: int | None = None, atom_ener: list[float] = [], activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, - layer_name: Optional[list[Optional[str]]] = None, + layer_name: list[str | None] | None = None, use_aparam_as_mask: bool = False, - spin: Optional[Spin] = None, + spin: Spin | None = None, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input **kwargs, ) -> None: """Constructor.""" @@ -196,6 +199,9 @@ def __init__( self.dim_case_embd = dim_case_embd if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + self.default_fparam = default_fparam + if self.default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.n_neuron = neuron self.resnet_dt = resnet_dt self.rcond = rcond @@ -469,8 +475,8 @@ def build( self, inputs: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ) -> tf.Tensor: """Build the computational graph for fitting net. @@ -832,7 +838,7 @@ def change_energy_bias( ntest=ntest, ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -884,7 +890,7 @@ def deserialize(cls, data: dict, suffix: str = ""): The deserialized model """ data = data.copy() - check_version_compatibility(data.pop("@version", 1), 3, 1) + check_version_compatibility(data.pop("@version", 1), 4, 1) fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( data["nets"], @@ -910,7 +916,7 @@ def serialize(self, suffix: str = "") -> dict: data = { "@class": "Fitting", "type": "ener", - "@version": 3, + "@version": 4, "var_name": "energy", "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt + self.tebd_dim, @@ -921,6 +927,7 @@ def serialize(self, suffix: str = "") -> dict: "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "rcond": self.rcond, "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, @@ -944,6 +951,7 @@ def serialize(self, suffix: str = "") -> dict: activation_function=self.activation_function_name, resnet_dt=self.resnet_dt, variables=self.fitting_net_variables, + trainable=self.trainable, suffix=suffix, ), "@variables": { diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index f159de1628..b33559f12f 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -3,9 +3,6 @@ from abc import ( abstractmethod, ) -from typing import ( - Optional, -) from deepmd.common import ( j_get_type, @@ -134,7 +131,8 @@ def serialize_network( activation_function: str, resnet_dt: bool, variables: dict, - out_dim: Optional[int] = 1, + out_dim: int | None = 1, + trainable: list[bool] | None = None, suffix: str = "", ) -> dict: """Serialize network. @@ -155,6 +153,8 @@ def serialize_network( Whether to use resnet variables : dict The input variables + trainable : list[bool] + Whether the network is trainable suffix : str, optional The suffix of the scope out_dim : int, optional @@ -191,6 +191,8 @@ def serialize_network( raise ValueError(f"Invalid ndim: {ndim}") if fittings[network_idx] is None: # initialize the network if it is not initialized + if trainable is None: + trainable = [True for _ in range(len(neuron) + 1)] fittings[network_idx] = FittingNet( in_dim=in_dim, out_dim=out_dim, @@ -199,6 +201,7 @@ def serialize_network( resnet_dt=resnet_dt, precision=self.precision.name, bias_out=True, + trainable=trainable, ) assert fittings[network_idx] is not None if weight_name == "idt": @@ -238,7 +241,9 @@ def deserialize_network(cls, data: dict, suffix: str = "") -> dict: else: raise ValueError(f"Invalid ndim: {fittings.ndim}") network = fittings[net_idx] - assert network is not None + if network is None: + # Skip types that are not selected (when sel_type is used) + continue for layer_idx, layer in enumerate(network.layers): if layer_idx == len(network.layers) - 1: layer_name = "final_layer" diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index c8fd4e86e8..1e48a5fa59 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import warnings -from typing import ( - Optional, -) import numpy as np +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.tf.common import ( cast_precision, get_activation_func, @@ -90,6 +90,13 @@ class PolarFittingSeA(Fitting): different fitting nets for different atom types. type_map: list[str], Optional A list of strings. Give the name to each type of atoms. + default_fparam: list[float], optional + The default frame parameter. If set, when `fparam.npy` files are not included in the data system, + this value will be used as the default value for the frame parameter in the fitting net. + trainable : list[bool], Optional + If the weights of fitting net are trainable. + Suppose that we have :math:`N_l` hidden layers in the fitting net, + this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. """ def __init__( @@ -102,17 +109,19 @@ def __init__( numb_fparam: int = 0, numb_aparam: int = 0, dim_case_embd: int = 0, - sel_type: Optional[list[int]] = None, + sel_type: list[int] | None = None, fit_diag: bool = True, - scale: Optional[list[float]] = None, + scale: list[float] | None = None, shift_diag: bool = True, # YWolfeee: will support the user to decide whether to use this function # diag_shift : list[float] = None, YWolfeee: will not support the user to assign a shift - seed: Optional[int] = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, mixed_types: bool = False, - type_map: Optional[list[str]] = None, # to be compat with input + type_map: list[str] | None = None, # to be compat with input + default_fparam: list[float] | None = None, # to be compat with input + trainable: list[bool] | None = None, **kwargs, ) -> None: """Constructor.""" @@ -170,18 +179,30 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam self.dim_case_embd = dim_case_embd + self.default_fparam = default_fparam if numb_fparam > 0: raise ValueError("numb_fparam is not supported in the dipole fitting") if numb_aparam > 0: raise ValueError("numb_aparam is not supported in the dipole fitting") if dim_case_embd > 0: raise ValueError("dim_case_embd is not supported in TensorFlow.") + if default_fparam is not None: + raise ValueError("default_fparam is not supported in TensorFlow.") self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None self.aparam_avg = None self.aparam_std = None self.aparam_inv_std = None + if trainable is None: + self.trainable = [True for _ in range(len(self.n_neuron) + 1)] + elif isinstance(trainable, bool): + self.trainable = [trainable] * (len(self.n_neuron) + 1) + else: + self.trainable = trainable + assert len(self.trainable) == len(self.n_neuron) + 1, ( + "length of trainable should be that of n_neuron + 1" + ) def get_sel_type(self) -> list[int]: """Get selected atom types.""" @@ -312,6 +333,7 @@ def _build_lower(self, start_index, natoms, inputs, rot_mat, suffix="", reuse=No uniform_seed=self.uniform_seed, initial_variables=self.fitting_net_variables, mixed_prec=self.mixed_prec, + trainable=self.trainable[ii], ) else: layer = one_layer( @@ -325,6 +347,7 @@ def _build_lower(self, start_index, natoms, inputs, rot_mat, suffix="", reuse=No uniform_seed=self.uniform_seed, initial_variables=self.fitting_net_variables, mixed_prec=self.mixed_prec, + trainable=self.trainable[ii], ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift @@ -347,6 +370,7 @@ def _build_lower(self, start_index, natoms, inputs, rot_mat, suffix="", reuse=No initial_variables=self.fitting_net_variables, mixed_prec=self.mixed_prec, final_layer=True, + trainable=self.trainable[-1], ) if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift @@ -398,8 +422,8 @@ def build( input_d: tf.Tensor, rot_mat: tf.Tensor, natoms: tf.Tensor, - input_dict: Optional[dict] = None, - reuse: Optional[bool] = None, + input_dict: dict | None = None, + reuse: bool | None = None, suffix: str = "", ): """Build the computational graph for fitting net. @@ -580,7 +604,7 @@ def init_variables( stacklevel=2, ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters @@ -612,22 +636,21 @@ def serialize(self, suffix: str) -> dict: data = { "@class": "Fitting", "type": "polar", - "@version": 4, + "@version": 5, "ntypes": self.ntypes, "dim_descrpt": self.dim_descrpt, "embedding_width": self.dim_rot_mat_1, "mixed_types": self.mixed_types, - "dim_out": 3, "neuron": self.n_neuron, "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, "dim_case_embd": self.dim_case_embd, + "default_fparam": self.default_fparam, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, "exclude_types": [], "fit_diag": self.fit_diag, - "scale": list(self.scale), "shift_diag": self.shift_diag, "nets": self.serialize_network( ntypes=self.ntypes, @@ -638,6 +661,7 @@ def serialize(self, suffix: str) -> dict: activation_function=self.activation_function_name, resnet_dt=self.resnet_dt, variables=self.fitting_net_variables, + trainable=self.trainable, suffix=suffix, ), "@variables": { @@ -648,8 +672,18 @@ def serialize(self, suffix: str) -> dict: "case_embd": None, "scale": self.scale.reshape(-1, 1), "constant_matrix": self.constant_matrix.reshape(-1), + "bias_atom_e": np.zeros( + (self.ntypes, self.dim_rot_mat_1), dtype=GLOBAL_NP_FLOAT_PRECISION + ), }, "type_map": self.type_map, + "var_name": "polar", + "rcond": None, + "tot_ener_zero": False, + "trainable": self.trainable, + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, } return data @@ -669,7 +703,7 @@ def deserialize(cls, data: dict, suffix: str): """ data = data.copy() check_version_compatibility( - data.pop("@version", 1), 4, 1 + data.pop("@version", 1), 5, 1 ) # to allow PT version. fitting = cls(**data) fitting.fitting_net_variables = cls.deserialize_network( @@ -713,11 +747,11 @@ def __init__( descrpt: tf.Tensor, neuron: list[int] = [120, 120, 120], resnet_dt: bool = True, - sel_type: Optional[list[int]] = None, + sel_type: list[int] | None = None, fit_diag: bool = True, - scale: Optional[list[float]] = None, - diag_shift: Optional[list[float]] = None, - seed: Optional[int] = None, + scale: list[float] | None = None, + diag_shift: list[float] | None = None, + seed: int | None = None, activation_function: str = "tanh", precision: str = "default", ) -> None: @@ -752,7 +786,7 @@ def build( input_d, rot_mat, natoms, - input_dict: Optional[dict] = None, + input_dict: dict | None = None, reuse=None, suffix="", ) -> tf.Tensor: @@ -812,7 +846,7 @@ def init_variables( graph=graph, graph_def=graph_def, suffix=suffix ) - def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: + def enable_mixed_precision(self, mixed_prec: dict | None = None) -> None: """Receive the mixed precision setting. Parameters diff --git a/deepmd/tf/infer/deep_dipole.py b/deepmd/tf/infer/deep_dipole.py index 09252d1d6c..834c0b7d8f 100644 --- a/deepmd/tf/infer/deep_dipole.py +++ b/deepmd/tf/infer/deep_dipole.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( TYPE_CHECKING, - Optional, ) from deepmd.infer.deep_dipole import ( @@ -50,7 +49,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, neighbor_list=None, ) -> None: # use this in favor of dict update to move attribute from class to diff --git a/deepmd/tf/infer/deep_eval.py b/deepmd/tf/infer/deep_eval.py index a7682d2e58..d1eca83055 100644 --- a/deepmd/tf/infer/deep_eval.py +++ b/deepmd/tf/infer/deep_eval.py @@ -1,14 +1,14 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import json +from collections.abc import ( + Callable, +) from functools import ( cached_property, ) from typing import ( TYPE_CHECKING, Any, - Callable, - Optional, - Union, ) import numpy as np @@ -93,8 +93,8 @@ def __init__( *args: list, load_prefix: str = "load", default_tf_graph: bool = False, - auto_batch_size: Union[bool, int, AutoBatchSize] = False, - input_map: Optional[dict] = None, + auto_batch_size: bool | int | AutoBatchSize = False, + input_map: dict | None = None, neighbor_list=None, **kwargs: dict, ) -> None: @@ -358,7 +358,7 @@ def _load_graph( frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, ): # We load the protobuf file from the disk and parse it to retrieve the # unserialized graph_def @@ -393,7 +393,7 @@ def _load_graph( def sort_input( coord: np.ndarray, atom_type: np.ndarray, - sel_atoms: Optional[list[int]] = None, + sel_atoms: list[int] | None = None, ): """Sort atoms in the system according their types. @@ -532,7 +532,7 @@ def eval_typeebd(self) -> np.ndarray: def build_neighbor_list( self, coords: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, atype: np.ndarray, imap: np.ndarray, neighbor_list, @@ -634,7 +634,7 @@ def get_type_map(self) -> list[str]: """Get the type map (element name of the atom types) of this model.""" return self.tmap - def get_sel_type(self) -> Optional[np.ndarray]: + def get_sel_type(self) -> np.ndarray | None: """Get the selected atom types of this model. Only atoms with selected atom types have atomic contribution @@ -682,7 +682,7 @@ def eval_func(*args, **kwargs): def _get_natoms_and_nframes( self, coords: np.ndarray, - atom_types: Union[list[int], np.ndarray], + atom_types: list[int] | np.ndarray, ) -> tuple[int, int]: natoms = len(atom_types[0]) if natoms == 0: @@ -695,12 +695,12 @@ def _get_natoms_and_nframes( def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, atomic: bool = False, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, **kwargs: Any, ) -> dict[str, np.ndarray]: """Evaluate the energy, force and virial by using this DP. @@ -1025,11 +1025,11 @@ def _get_output_shape(self, odef, nframes, natoms): def eval_descriptor( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, ) -> np.ndarray: """Evaluate descriptors by using this DP. @@ -1082,11 +1082,11 @@ def eval_descriptor( def _eval_descriptor_inner( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: np.ndarray, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, ) -> np.ndarray: natoms, nframes = self._get_natoms_and_nframes( coords, @@ -1126,6 +1126,16 @@ def get_model_def_script(self) -> dict: model_def_script = script.decode("utf-8") return json.loads(model_def_script)["model"] + def get_model(self) -> "tf.Graph": + """Get the TensorFlow graph. + + Returns + ------- + tf.Graph + The TensorFlow graph. + """ + return self.graph + class DeepEvalOld: # old class for DipoleChargeModifier only @@ -1156,8 +1166,8 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, - auto_batch_size: Union[bool, int, AutoBatchSize] = False, - input_map: Optional[dict] = None, + auto_batch_size: bool | int | AutoBatchSize = False, + input_map: dict | None = None, neighbor_list=None, ) -> None: self.graph = self._load_graph( @@ -1244,9 +1254,7 @@ def _graph_compatable(self) -> bool: else: return True - def _get_tensor( - self, tensor_name: str, attr_name: Optional[str] = None - ) -> tf.Tensor: + def _get_tensor(self, tensor_name: str, attr_name: str | None = None) -> tf.Tensor: """Get TF graph tensor and assign it to class namespace. Parameters @@ -1276,7 +1284,7 @@ def _load_graph( frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, ): # We load the protobuf file from the disk and parse it to retrieve the # unserialized graph_def @@ -1311,7 +1319,7 @@ def _load_graph( def sort_input( coord: np.ndarray, atom_type: np.ndarray, - sel_atoms: Optional[list[int]] = None, + sel_atoms: list[int] | None = None, mixed_type: bool = False, ): """Sort atoms in the system according their types. @@ -1468,7 +1476,7 @@ def eval_typeebd(self) -> np.ndarray: def build_neighbor_list( self, coords: np.ndarray, - cell: Optional[np.ndarray], + cell: np.ndarray | None, atype: np.ndarray, imap: np.ndarray, neighbor_list, diff --git a/deepmd/tf/infer/deep_tensor.py b/deepmd/tf/infer/deep_tensor.py index a4d30c4f2f..48fa811ff9 100644 --- a/deepmd/tf/infer/deep_tensor.py +++ b/deepmd/tf/infer/deep_tensor.py @@ -2,7 +2,6 @@ from typing import ( TYPE_CHECKING, ClassVar, - Optional, ) import numpy as np @@ -59,7 +58,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, - input_map: Optional[dict] = None, + input_map: dict | None = None, neighbor_list=None, ) -> None: """Constructor.""" @@ -143,12 +142,12 @@ def get_dim_aparam(self) -> int: def eval( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: list[int], atomic: bool = True, - fparam: Optional[np.ndarray] = None, - aparam: Optional[np.ndarray] = None, - efield: Optional[np.ndarray] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, mixed_type: bool = False, ) -> np.ndarray: """Evaluate the model. @@ -273,12 +272,12 @@ def eval( def eval_full( self, coords: np.ndarray, - cells: Optional[np.ndarray], + cells: np.ndarray | None, atom_types: list[int], atomic: bool = False, - fparam: Optional[np.array] = None, - aparam: Optional[np.array] = None, - efield: Optional[np.array] = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + efield: np.ndarray | None = None, mixed_type: bool = False, ) -> tuple[np.ndarray, ...]: """Evaluate the model with interface similar to the energy model. diff --git a/deepmd/tf/loss/ener.py b/deepmd/tf/loss/ener.py index ba65450613..cba60bb7f0 100644 --- a/deepmd/tf/loss/ener.py +++ b/deepmd/tf/loss/ener.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -98,7 +95,7 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, start_pref_gf: float = 0.0, limit_pref_gf: float = 0.0, @@ -556,9 +553,9 @@ def __init__( limit_pref_ae: float = 0.0, start_pref_pf: float = 0.0, limit_pref_pf: float = 0.0, - relative_f: Optional[float] = None, + relative_f: float | None = None, enable_atom_ener_coeff: bool = False, - use_spin: Optional[list] = None, + use_spin: list | None = None, ) -> None: self.starter_learning_rate = starter_learning_rate self.start_pref_e = start_pref_e diff --git a/deepmd/tf/model/dos.py b/deepmd/tf/model/dos.py index 1bebb4b971..548756f3e9 100644 --- a/deepmd/tf/model/dos.py +++ b/deepmd/tf/model/dos.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) from deepmd.tf.env import ( MODEL_VERSION, @@ -49,8 +45,8 @@ def __init__( self, descriptor: dict, fitting_net: dict, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, **kwargs, @@ -137,7 +133,7 @@ def build( mesh, input_dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix="", reuse=None, ): @@ -149,6 +145,9 @@ def build( t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string) t_od = tf.constant(self.numb_dos, name="output_dim", dtype=tf.int32) + # Initialize out_bias and out_std for DOS models + self.init_out_stat(suffix=suffix) + coord = tf.reshape(coord_, [-1, natoms[1] * 3]) atype = tf.reshape(atype_, [-1, natoms[1]]) input_dict["nframes"] = tf.shape(coord)[0] @@ -181,6 +180,10 @@ def build( atom_dos = self.fitting.build( dout, natoms, input_dict, reuse=reuse, suffix=suffix ) + + # Apply out_bias and out_std directly to DOS output + atom_dos = self._apply_out_bias_std(atom_dos, atype, natoms, coord) + self.atom_dos = atom_dos dos_raw = atom_dos diff --git a/deepmd/tf/model/ener.py b/deepmd/tf/model/ener.py index 6d2ff4615f..7109ad9d73 100644 --- a/deepmd/tf/model/ener.py +++ b/deepmd/tf/model/ener.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np @@ -75,16 +71,16 @@ def __init__( self, descriptor: dict, fitting_net: dict, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, - use_srtab: Optional[str] = None, - smin_alpha: Optional[float] = None, - sw_rmin: Optional[float] = None, - sw_rmax: Optional[float] = None, + use_srtab: str | None = None, + smin_alpha: float | None = None, + sw_rmin: float | None = None, + sw_rmax: float | None = None, srtab_add_bias: bool = True, - spin: Optional[Spin] = None, + spin: Spin | None = None, data_bias_nsample: int = 10, **kwargs, ) -> None: @@ -182,7 +178,7 @@ def build( mesh, input_dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix="", reuse=None, ): @@ -193,6 +189,9 @@ def build( t_mt = tf.constant(self.model_type, name="model_type", dtype=tf.string) t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string) + # Initialize out_bias and out_std for energy models + self.init_out_stat(suffix=suffix) + if self.srtab is not None: tab_info, tab_data = self.srtab.get() self.tab_info = tf.get_variable( @@ -253,6 +252,10 @@ def build( atom_ener = self.fitting.build( dout, natoms, input_dict, reuse=reuse, suffix=suffix ) + + # Apply out_bias and out_std directly to atom energy + atom_ener = self._apply_out_bias_std(atom_ener, atype, natoms, coord) + self.atom_ener = atom_ener if self.srtab is not None: diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 6ca18ed7bd..b63fe80a87 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -7,8 +7,6 @@ ) from typing import ( NoReturn, - Optional, - Union, ) from deepmd.entrypoints.convert_backend import ( @@ -83,10 +81,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ) -> dict: """Build the model. @@ -185,11 +183,11 @@ def build( "Contribution is welcome!" ) - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" return {} - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" # loss should be never used for a frozen model return @@ -243,9 +241,9 @@ def get_type_map(self) -> list: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index 63f55eae9e..f036edbfc3 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -7,10 +7,6 @@ lru_cache, reduce, ) -from typing import ( - Optional, - Union, -) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, @@ -64,13 +60,13 @@ def __init__(self, models: list[dict], weights: list[float], **kwargs) -> None: else: raise ValueError(f"Invalid weights {weights}") - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" return { f"model{ii}": model.get_fitting() for ii, model in enumerate(self.models) } - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" # the first model that is not None, or None if all models are None for model in self.models: @@ -138,9 +134,9 @@ def get_type_map(self) -> list: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -194,10 +190,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ) -> dict: """Build the model. diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 3377ed2d51..26a3f7bf73 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import logging from abc import ( ABC, abstractmethod, @@ -8,8 +9,6 @@ ) from typing import ( NoReturn, - Optional, - Union, ) import numpy as np @@ -70,6 +69,8 @@ check_version_compatibility, ) +log = logging.getLogger(__name__) + class Model(ABC, make_plugin_registry("model")): """Abstract base model. @@ -112,13 +113,13 @@ def __new__(cls, *args, **kwargs): def __init__( self, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_bias_nsample: int = 10, data_stat_protect: float = 1e-2, - spin: Optional[Spin] = None, - compress: Optional[dict] = None, + spin: Spin | None = None, + compress: dict | None = None, **kwargs, ) -> None: super().__init__() @@ -152,10 +153,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ): """Build the model. @@ -220,10 +221,10 @@ def build_descrpt( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ): """Build the descriptor part of the model. @@ -308,10 +309,10 @@ def build_descrpt( def build_type_embedding( self, ntypes: int, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ) -> tf.Tensor: """Build the type embedding part of the model. @@ -429,24 +430,24 @@ def enable_compression(self, suffix: str = "") -> NoReturn: """ raise RuntimeError("Not supported") - def get_numb_fparam(self) -> Union[int, dict]: + def get_numb_fparam(self) -> int | dict: """Get the number of frame parameters.""" return 0 - def get_numb_aparam(self) -> Union[int, dict]: + def get_numb_aparam(self) -> int | dict: """Get the number of atomic parameters.""" return 0 - def get_numb_dos(self) -> Union[int, dict]: + def get_numb_dos(self) -> int | dict: """Get the number of gridpoints in energy space.""" return 0 @abstractmethod - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" @abstractmethod - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" @abstractmethod @@ -515,9 +516,9 @@ def get_feed_dict( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Notes @@ -644,10 +645,10 @@ def __new__(cls, *args, **kwargs): def __init__( self, - descriptor: Union[dict, Descriptor], - fitting_net: Union[dict, Fitting], - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + descriptor: dict | Descriptor, + fitting_net: dict | Fitting, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, **kwargs, ) -> None: super().__init__( @@ -708,6 +709,63 @@ def __init__( else: self.typeebd = None + # Initialize out_bias and out_std storage + self.out_bias = None + self.out_std = None + + def init_variables( + self, + graph: tf.Graph, + graph_def: tf.GraphDef, + model_type: str = "original_model", + suffix: str = "", + ) -> None: + """Init the model variables with the given frozen model. + + Parameters + ---------- + graph : tf.Graph + The input frozen model graph + graph_def : tf.GraphDef + The input frozen model graph_def + model_type : str + the type of the model + suffix : str + suffix to name scope + """ + from deepmd.tf.utils.errors import ( + GraphWithoutTensorError, + ) + from deepmd.tf.utils.graph import ( + get_tensor_by_name_from_graph, + ) + + # Initialize descriptor and fitting variables + self.descrpt.init_variables(graph, graph_def, suffix=suffix) + self.fitting.init_variables(graph, graph_def, suffix=suffix) + if ( + self.typeebd is not None + and self.typeebd.type_embedding_net_variables is None + ): + self.typeebd.init_variables(graph, graph_def, suffix=suffix) + + # Try to load out_bias and out_std from the graph + try: + self.out_bias = get_tensor_by_name_from_graph( + graph, f"model_attr{suffix}/t_out_bias" + ) + except GraphWithoutTensorError: + # For compatibility, create default out_bias if not found + log.debug("out_bias not found in graph, falling back to default value") + + try: + self.out_std = get_tensor_by_name_from_graph( + graph, f"model_attr{suffix}/t_out_std" + ) + except GraphWithoutTensorError: + # For compatibility, create default out_std if not found + log.debug("out_std not found in graph, falling back to default value") + def enable_mixed_precision(self, mixed_prec: dict) -> None: """Enable mixed precision for the model. @@ -746,11 +804,11 @@ def enable_compression(self, suffix: str = "") -> None: ): self.typeebd.init_variables(graph, graph_def, suffix=suffix) - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" return self.fitting - def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + def get_loss(self, loss: dict, lr) -> Loss | dict: """Get the loss function(s).""" return self.fitting.get_loss(loss, lr) @@ -762,13 +820,137 @@ def get_ntypes(self) -> int: """Get the number of types.""" return self.ntypes + def _get_dim_out(self): + """Get output dimension based on model type. + + Returns + ------- + int + Output dimension + """ + if self.model_type == "ener": + return 1 + elif self.model_type == "dipole": + return 3 + elif self.model_type == "polar": + return 9 + elif self.model_type == "dos": + return self.numb_dos + else: + raise ValueError(f"Unknown model type '{self.model_type}' in _get_dim_out") + + def init_out_stat(self, suffix: str = "") -> None: + """Initialize the output bias and std variables.""" + ntypes = self.get_ntypes() + dim_out = self._get_dim_out() + + # Initialize out_bias and out_std as numpy arrays, preserving existing values if set + if self.out_bias is not None: + out_bias_data = self.out_bias.copy() + else: + out_bias_data = np.zeros( + [1, ntypes, dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION + ) + + if self.out_std is not None: + out_std_data = self.out_std.copy() + else: + out_std_data = np.ones( + [1, ntypes, dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION + ) + + # Create TensorFlow variables + with tf.variable_scope("model_attr" + suffix, reuse=tf.AUTO_REUSE): + self.t_out_bias = tf.get_variable( + "t_out_bias", + out_bias_data.shape, + dtype=GLOBAL_TF_FLOAT_PRECISION, + trainable=False, + initializer=tf.constant_initializer(out_bias_data), + ) + self.t_out_std = tf.get_variable( + "t_out_std", + out_std_data.shape, + dtype=GLOBAL_TF_FLOAT_PRECISION, + trainable=False, + initializer=tf.constant_initializer(out_std_data), + ) + + # Store as instance variables for access + self.out_bias = out_bias_data + self.out_std = out_std_data + + def _apply_out_bias_std(self, output, atype, natoms, coord, selected_atype=None): + """Apply output bias and standard deviation to the model output. + + Parameters + ---------- + output : tf.Tensor + The model output tensor + atype : tf.Tensor + Atom types with shape [nframes, nloc] + natoms : list[int] + Number of atoms [nloc, ntypes, ...] + coord : tf.Tensor + Coordinates for getting nframes + selected_atype : tf.Tensor, optional + Selected atom types for tensor models. If None, uses all atoms. + + Returns + ------- + tf.Tensor + Output with bias and std applied + """ + if self.spin is not None: + # spin is not supported yet; also, it's incompatible with dpmodel + return output + nframes = tf.shape(coord)[0] + + # Get output dimension consistently + nout = self._get_dim_out() + + if selected_atype is not None: + natomsel = tf.shape(selected_atype)[1] + output_reshaped = tf.reshape(output, [nframes, natomsel, nout]) + atype_for_gather = selected_atype + else: + nloc = natoms[0] + nall = natoms[1] + output_reshaped = tf.reshape(output, [nframes, nloc, nout]) + atype_for_gather = tf.reshape(atype, [nframes, nall]) + # slice to local atoms + atype_for_gather = atype_for_gather[:, :nloc] + + # Handle invalid atom types (e.g., -1 for padding/invalid atoms) + # Create a mask for valid atom types (>= 0) + valid_mask = tf.greater_equal(atype_for_gather, 0) + # Replace invalid types with 0 for gathering (will be masked out later) + safe_atype = tf.where( + valid_mask, atype_for_gather, tf.zeros_like(atype_for_gather) + ) + + # Get bias and std for each atom type + bias_per_atom = tf.gather(self.t_out_bias[0], safe_atype) + std_per_atom = tf.gather(self.t_out_std[0], safe_atype) + + # Apply bias and std: output = output * std + bias + adjusted_output = output_reshaped * std_per_atom + bias_per_atom + + # expand axis 2 of valid_mask to nout + valid_mask = tf.tile(tf.expand_dims(valid_mask, -1), [1, 1, nout]) + + # Only apply bias/std to valid atoms, keep original values for invalid atoms + output_reshaped = tf.where(valid_mask, adjusted_output, output_reshaped) + + return tf.reshape(output_reshaped, tf.shape(output)) + @classmethod def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters @@ -820,26 +1002,17 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": data = data.copy() check_version_compatibility(data.pop("@version", 2), 2, 1) descriptor = Descriptor.deserialize(data.pop("descriptor"), suffix=suffix) - if data["fitting"].get("@variables", {}).get("bias_atom_e") is not None: - # careful: copy each level and don't modify the input array, - # otherwise it will affect the original data - # deepcopy is not used for performance reasons - data["fitting"] = data["fitting"].copy() - data["fitting"]["@variables"] = data["fitting"]["@variables"].copy() - if ( - int(np.any(data["fitting"]["@variables"]["bias_atom_e"])) - + int(np.any(data["@variables"]["out_bias"])) - > 1 - ): - raise ValueError( - "fitting/@variables/bias_atom_e and @variables/out_bias should not be both non-zero" - ) - data["fitting"]["@variables"]["bias_atom_e"] = data["fitting"][ - "@variables" - ]["bias_atom_e"] + data["@variables"]["out_bias"].reshape( - data["fitting"]["@variables"]["bias_atom_e"].shape + # bias_atom_e and out_bias are now completely independent - no conversion needed + fitting_dict = data.pop("fitting", {}) + atom_exclude_types = data.pop("atom_exclude_types", []) + if len(atom_exclude_types) > 0: + # get sel_type from complement of atom_exclude_types + full_type_list = np.arange(len(data["type_map"]), dtype=int) + sel_type = np.setdiff1d( + full_type_list, atom_exclude_types, assume_unique=True ) - fitting = Fitting.deserialize(data.pop("fitting"), suffix=suffix) + fitting_dict["sel_type"] = sel_type.tolist() + fitting = Fitting.deserialize(fitting_dict, suffix=suffix) # pass descriptor type embedding to model if descriptor.explicit_ntypes: type_embedding = descriptor.type_embedding @@ -847,20 +1020,27 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": else: type_embedding = None # BEGINE not supported keys - if len(data.pop("atom_exclude_types")) > 0: - raise NotImplementedError("atom_exclude_types is not supported") if len(data.pop("pair_exclude_types")) > 0: raise NotImplementedError("pair_exclude_types is not supported") data.pop("rcond", None) data.pop("preset_out_bias", None) - data.pop("@variables", None) + # Extract out_bias and out_std from variables before removing them + variables = data.pop("@variables", {}) + out_bias = variables.get("out_bias", None) + out_std = variables.get("out_std", None) # END not supported keys - return cls( + model = cls( descriptor=descriptor, fitting_net=fitting, type_embedding=type_embedding, **data, ) + # Restore out_bias and out_std if they exist + if out_bias is not None: + model.out_bias = out_bias + if out_std is not None: + model.out_std = out_std + return model def serialize(self, suffix: str = "") -> dict: """Serialize the model. @@ -886,18 +1066,23 @@ def serialize(self, suffix: str = "") -> dict: raise NotImplementedError("spin is not supported") ntypes = len(self.get_type_map()) + + # Get output dimension + dim_out = self._get_dim_out() + + # Serialize fitting dict_fit = self.fitting.serialize(suffix=suffix) - if dict_fit.get("@variables", {}).get("bias_atom_e") is not None: - out_bias = dict_fit["@variables"]["bias_atom_e"].reshape( - [1, ntypes, dict_fit["dim_out"]] - ) - dict_fit["@variables"]["bias_atom_e"] = np.zeros_like( - dict_fit["@variables"]["bias_atom_e"] - ) + + # Use the actual out_bias and out_std if they exist, otherwise create defaults + if self.out_bias is not None: + out_bias = self.out_bias.copy() else: - out_bias = np.zeros( - [1, ntypes, dict_fit["dim_out"]], dtype=GLOBAL_NP_FLOAT_PRECISION - ) + out_bias = np.zeros([1, ntypes, dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION) + + if self.out_std is not None: + out_std = self.out_std.copy() + else: + out_std = np.ones([1, ntypes, dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION) return { "@class": "Model", "type": "standard", @@ -912,7 +1097,7 @@ def serialize(self, suffix: str = "") -> dict: "preset_out_bias": None, "@variables": { "out_bias": out_bias, - "out_std": np.ones([1, ntypes, dict_fit["dim_out"]]), # pylint: disable=no-explicit-dtype + "out_std": out_std, }, } diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py index d91554b98d..8bdf48accb 100644 --- a/deepmd/tf/model/pairtab.py +++ b/deepmd/tf/model/pairtab.py @@ -2,10 +2,6 @@ from enum import ( Enum, ) -from typing import ( - Optional, - Union, -) import numpy as np @@ -67,7 +63,7 @@ class PairTabModel(Model): model_type = "ener" def __init__( - self, tab_file: str, rcut: float, sel: Union[int, list[int]], **kwargs + self, tab_file: str, rcut: float, sel: int | list[int], **kwargs ) -> None: super().__init__() self.tab_file = tab_file @@ -89,10 +85,10 @@ def build( box: tf.Tensor, mesh: tf.Tensor, input_dict: dict, - frz_model: Optional[str] = None, - ckpt_meta: Optional[str] = None, + frz_model: str | None = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[Union[bool, Enum]] = None, + reuse: bool | Enum | None = None, ): """Build the model. @@ -237,12 +233,12 @@ def init_variables( """ # skip. table can be initialized from the file - def get_fitting(self) -> Union[Fitting, dict]: + def get_fitting(self) -> Fitting | dict: """Get the fitting(s).""" # nothing needs to do return {} - def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + def get_loss(self, loss: dict, lr) -> Loss | dict | None: """Get the loss function(s).""" # nothing needs to do return @@ -273,9 +269,9 @@ def enable_compression(self, suffix: str = "") -> None: def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Notes diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index 5ed98f0c49..45d0c754b9 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) from deepmd.tf.common import ( make_default_mesh, @@ -49,17 +45,17 @@ def __init__( self, qm_model: dict, qmmm_model: dict, - type_embedding: Union[dict, TypeEmbedNet], + type_embedding: dict | TypeEmbedNet, type_map: list[str], data_stat_nbatch: int = 10, data_stat_nsample: int = 10, data_stat_protect: float = 1e-2, - use_srtab: Optional[str] = None, - smin_alpha: Optional[float] = None, - sw_rmin: Optional[float] = None, - sw_rmax: Optional[float] = None, - spin: Optional[Spin] = None, - compress: Optional[dict] = None, + use_srtab: str | None = None, + smin_alpha: float | None = None, + sw_rmin: float | None = None, + sw_rmax: float | None = None, + spin: Spin | None = None, + compress: dict | None = None, **kwargs, ) -> None: # internal variable to compare old and new behavior @@ -116,9 +112,9 @@ def build( mesh: tf.Tensor, input_dict: dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix: str = "", - reuse: Optional[bool] = None, + reuse: bool | None = None, ): feed_dict = self.get_feed_dict( coord_, atype_, natoms, box_, mesh, aparam=input_dict["aparam"] @@ -300,14 +296,14 @@ def build( model_dict["atype"] = atype_ return model_dict - def get_fitting(self) -> Union[str, dict]: + def get_fitting(self) -> str | dict: """Get the fitting(s).""" return { "qm": self.qm_model.get_fitting(), "qmmm": self.qmmm_model.get_fitting(), } - def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + def get_loss(self, loss: dict, lr) -> Loss | dict: """Get the loss function(s).""" return self.qm_model.get_loss(loss, lr) @@ -413,9 +409,9 @@ def get_feed_dict( def update_sel( cls, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, local_jdata: dict, - ) -> tuple[dict, Optional[float]]: + ) -> tuple[dict, float | None]: """Update the selection and perform neighbor statistics. Parameters diff --git a/deepmd/tf/model/tensor.py b/deepmd/tf/model/tensor.py index 1e960907ef..70d9a3d26c 100644 --- a/deepmd/tf/model/tensor.py +++ b/deepmd/tf/model/tensor.py @@ -1,10 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) +import numpy as np + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) from deepmd.tf.env import ( + GLOBAL_TF_FLOAT_PRECISION, MODEL_VERSION, global_cvt_2_ener_float, tf, @@ -49,8 +51,8 @@ def __init__( tensor_name: str, descriptor: dict, fitting_net: dict, - type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, - type_map: Optional[list[str]] = None, + type_embedding: dict | TypeEmbedNet | None = None, + type_map: list[str] | None = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, **kwargs, @@ -113,7 +115,7 @@ def build( mesh, input_dict, frz_model=None, - ckpt_meta: Optional[str] = None, + ckpt_meta: str | None = None, suffix="", reuse=None, ): @@ -126,6 +128,9 @@ def build( t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string) t_od = tf.constant(self.get_out_size(), name="output_dim", dtype=tf.int32) + # Initialize out_bias and out_std for tensor models (dipole/polar) + self.init_out_stat(suffix=suffix) + natomsel = sum(natoms[2 + type_i] for type_i in self.get_sel_type()) nout = self.get_out_size() @@ -164,6 +169,38 @@ def build( output = self.fitting.build( dout, rot_mat, natoms, input_dict, reuse=reuse, suffix=suffix ) + + # Apply out_bias and out_std directly to tensor output + # dipole not applying bias but polar does, per dpmodel + if self.model_type == "polar" and self.fitting.shift_diag: + v_constant_matrix = np.zeros( + self.ntypes, + dtype=GLOBAL_NP_FLOAT_PRECISION, + ) + sel_type = self.get_sel_type() + for itype in range(len(sel_type)): + v_constant_matrix[sel_type[itype]] = np.mean( + np.diagonal(self.out_bias[0, itype].reshape((3, 3))) + ) + nframes = input_dict["nframes"] + nloc_mask = tf.reshape( + tf.tile(tf.repeat(self.fitting.sel_mask, natoms[2:]), [nframes]), + [nframes, -1], + ) + constant_matrix = tf.reshape( + tf.reshape( + tf.tile(tf.repeat(v_constant_matrix, natoms[2:]), [nframes]), + [nframes, -1], + )[nloc_mask], + [nframes, -1], + ) + + # nf x nloc x odims, out_bias: ntypes x odims + output = output + tf.reshape( + tf.expand_dims(tf.expand_dims(constant_matrix, -1), -1) + * tf.eye(3, batch_shape=[1, 1], dtype=GLOBAL_TF_FLOAT_PRECISION), + tf.shape(output), + ) framesize = nout if "global" in self.model_type else natomsel * nout output = tf.reshape( output, [-1, framesize], name="o_" + self.model_type + suffix diff --git a/deepmd/tf/nvnmd/entrypoints/mapt.py b/deepmd/tf/nvnmd/entrypoints/mapt.py index 2e6e56bf51..121263974a 100644 --- a/deepmd/tf/nvnmd/entrypoints/mapt.py +++ b/deepmd/tf/nvnmd/entrypoints/mapt.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -649,9 +646,9 @@ def build_davg_dstd(self): def mapt( *, - nvnmd_config: Optional[str] = "nvnmd/config.npy", - nvnmd_weight: Optional[str] = "nvnmd/weight.npy", - nvnmd_map: Optional[str] = "nvnmd/map.npy", + nvnmd_config: str | None = "nvnmd/config.npy", + nvnmd_weight: str | None = "nvnmd/weight.npy", + nvnmd_map: str | None = "nvnmd/map.npy", **kwargs, ) -> None: # build mapping table diff --git a/deepmd/tf/nvnmd/entrypoints/train.py b/deepmd/tf/nvnmd/entrypoints/train.py index c690190c0d..1650061c94 100644 --- a/deepmd/tf/nvnmd/entrypoints/train.py +++ b/deepmd/tf/nvnmd/entrypoints/train.py @@ -1,9 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import os -from typing import ( - Optional, -) from deepmd.tf.entrypoints.freeze import ( freeze, @@ -118,8 +115,8 @@ def normalized_input_qnn(jdata, PATH_QNN, CONFIG_CNN, WEIGHT_CNN, MAP_CNN): def train_nvnmd( *, INPUT: str, - init_model: Optional[str], - restart: Optional[str], + init_model: str | None, + restart: str | None, step: str, skip_neighbor_stat: bool = False, **kwargs, diff --git a/deepmd/tf/nvnmd/entrypoints/wrap.py b/deepmd/tf/nvnmd/entrypoints/wrap.py index ced97bdbf1..e1878aa254 100755 --- a/deepmd/tf/nvnmd/entrypoints/wrap.py +++ b/deepmd/tf/nvnmd/entrypoints/wrap.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - Optional, -) import numpy as np @@ -608,10 +605,10 @@ def wrap_lut(self): def wrap( *, - nvnmd_config: Optional[str] = "nvnmd/config.npy", - nvnmd_weight: Optional[str] = "nvnmd/weight.npy", - nvnmd_map: Optional[str] = "nvnmd/map.npy", - nvnmd_model: Optional[str] = "nvnmd/model.pb", + nvnmd_config: str | None = "nvnmd/config.npy", + nvnmd_weight: str | None = "nvnmd/weight.npy", + nvnmd_map: str | None = "nvnmd/map.npy", + nvnmd_model: str | None = "nvnmd/model.pb", **kwargs, ) -> None: wrapObj = Wrap(nvnmd_config, nvnmd_weight, nvnmd_map, nvnmd_model) diff --git a/deepmd/tf/train/run_options.py b/deepmd/tf/train/run_options.py index fb6e09e8a9..0b5c3b1b43 100644 --- a/deepmd/tf/train/run_options.py +++ b/deepmd/tf/train/run_options.py @@ -96,7 +96,7 @@ class RunOptions: device type - gpu or cpu """ - gpus: Optional[list[int]] + gpus: list[int] | None world_size: int my_rank: int nodename: str @@ -108,11 +108,11 @@ class RunOptions: def __init__( self, - init_model: Optional[str] = None, - init_frz_model: Optional[str] = None, - finetune: Optional[str] = None, - restart: Optional[str] = None, - log_path: Optional[str] = None, + init_model: str | None = None, + init_frz_model: str | None = None, + finetune: str | None = None, + restart: str | None = None, + log_path: str | None = None, log_level: int = 0, mpi_log: str = "master", ) -> None: @@ -151,9 +151,9 @@ def print_resource_summary(self) -> None: def _setup_logger( self, - log_path: Optional[Path], + log_path: Path | None, log_level: int, - mpi_log: Optional[str], + mpi_log: str | None, ) -> None: """Set up package loggers. diff --git a/deepmd/tf/utils/convert.py b/deepmd/tf/utils/convert.py index 461d870f80..5575770257 100644 --- a/deepmd/tf/utils/convert.py +++ b/deepmd/tf/utils/convert.py @@ -2,9 +2,6 @@ import logging import os import textwrap -from typing import ( - Optional, -) from google.protobuf import ( text_format, @@ -60,7 +57,7 @@ def detect_model_version(input_model: str): def convert_to_21( - input_model: str, output_model: str, version: Optional[str] = None + input_model: str, output_model: str, version: str | None = None ) -> None: """Convert DP graph to 2.1 graph. diff --git a/deepmd/tf/utils/learning_rate.py b/deepmd/tf/utils/learning_rate.py index fee73ca9a3..64427e185d 100644 --- a/deepmd/tf/utils/learning_rate.py +++ b/deepmd/tf/utils/learning_rate.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import numpy as np @@ -50,9 +47,7 @@ def __init__( self.cd["decay_rate"] = decay_rate self.start_lr_ = self.cd["start_lr"] - def build( - self, global_step: tf.Tensor, stop_step: Optional[int] = None - ) -> tf.Tensor: + def build(self, global_step: tf.Tensor, stop_step: int | None = None) -> tf.Tensor: """Build the learning rate. Parameters diff --git a/deepmd/tf/utils/neighbor_stat.py b/deepmd/tf/utils/neighbor_stat.py index 37028b23bc..798c40621a 100644 --- a/deepmd/tf/utils/neighbor_stat.py +++ b/deepmd/tf/utils/neighbor_stat.py @@ -3,9 +3,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np @@ -252,7 +249,7 @@ def _execute( self, coord: np.ndarray, atype: np.ndarray, - box: Optional[np.ndarray], + box: np.ndarray | None, pbc: bool, ): """Execute the operation. diff --git a/deepmd/tf/utils/parallel_op.py b/deepmd/tf/utils/parallel_op.py index c23f347dc1..46f9fb87b7 100644 --- a/deepmd/tf/utils/parallel_op.py +++ b/deepmd/tf/utils/parallel_op.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from collections.abc import ( + Callable, Generator, ) from typing import ( Any, - Callable, - Optional, ) from deepmd.tf.env import ( @@ -46,8 +45,8 @@ class ParallelOp: def __init__( self, builder: Callable[..., tuple[dict[str, tf.Tensor], tuple[tf.Tensor]]], - nthreads: Optional[int] = None, - config: Optional[tf.ConfigProto] = None, + nthreads: int | None = None, + config: tf.ConfigProto | None = None, ) -> None: if nthreads is not None: self.nthreads = nthreads diff --git a/deepmd/tf/utils/spin.py b/deepmd/tf/utils/spin.py index 8919bbd16a..c36a1ee21b 100644 --- a/deepmd/tf/utils/spin.py +++ b/deepmd/tf/utils/spin.py @@ -1,7 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, @@ -24,9 +21,9 @@ class Spin: def __init__( self, - use_spin: Optional[list[bool]] = None, - spin_norm: Optional[list[float]] = None, - virtual_len: Optional[list[float]] = None, + use_spin: list[bool] | None = None, + spin_norm: list[float] | None = None, + virtual_len: list[float] | None = None, ) -> None: """Constructor.""" self.use_spin = use_spin diff --git a/deepmd/tf/utils/tabulate.py b/deepmd/tf/utils/tabulate.py index 67ae2f96bd..4e6c027f1b 100644 --- a/deepmd/tf/utils/tabulate.py +++ b/deepmd/tf/utils/tabulate.py @@ -1,11 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging +from collections.abc import ( + Callable, +) from functools import ( cached_property, ) -from typing import ( - Callable, -) import numpy as np from scipy.special import ( diff --git a/deepmd/tf/utils/type_embed.py b/deepmd/tf/utils/type_embed.py index 9b7b17528d..34f3ae225b 100644 --- a/deepmd/tf/utils/type_embed.py +++ b/deepmd/tf/utils/type_embed.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging import re -from typing import ( - Optional, - Union, -) import numpy as np @@ -116,15 +112,15 @@ def __init__( ntypes: int, neuron: list[int], resnet_dt: bool = False, - activation_function: Union[str, None] = "tanh", + activation_function: str | None = "tanh", precision: str = "default", trainable: bool = True, - seed: Optional[int] = None, + seed: int | None = None, uniform_seed: bool = False, padding: bool = False, use_econf_tebd: bool = False, use_tebd_bias: bool = False, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, **kwargs, ) -> None: """Constructor.""" diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index fb911550dd..7fcc117ab5 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -2,10 +2,11 @@ import json import logging import warnings -from typing import ( +from collections.abc import ( Callable, - Optional, - Union, +) +from typing import ( + Any, ) from dargs import ( @@ -40,6 +41,7 @@ doc_only_tf_supported = "(Supported Backend: TensorFlow) " doc_only_pt_supported = "(Supported Backend: PyTorch) " +doc_only_pd_supported = "(Supported Backend: Paddle) " # descriptors doc_loc_frame = "Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame." doc_se_e2_a = "Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor." @@ -59,7 +61,7 @@ doc_dipole_charge = "Use WFCC to model the electronic structure of the system. Correct the long-range interaction." -def list_to_doc(xx): +def list_to_doc(xx: list[Any]) -> str: items = [] for ii in xx: if len(items) == 0: @@ -70,7 +72,7 @@ def list_to_doc(xx): return "".join(items) -def make_link(content, ref_key) -> str: +def make_link(content: str, ref_key: str) -> str: return ( f"`{content} <{ref_key}_>`_" if not dargs.RAW_ANCHOR @@ -87,7 +89,7 @@ def deprecate_argument_extra_check(key: str) -> Callable[[dict], bool]: The name of the deprecated argument. """ - def deprecate_something(data: Optional[dict]) -> bool: + def deprecate_something(data: dict | None) -> bool: if data is not None and key in data: warnings.warn(f"{key} has been removed and takes no effect.", FutureWarning) data.pop(key) @@ -96,7 +98,7 @@ def deprecate_something(data: Optional[dict]) -> bool: return deprecate_something -def type_embedding_args(): +def type_embedding_args() -> list[Argument]: doc_neuron = "Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built." doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' doc_seed = "Random seed for parameter initialization" @@ -132,7 +134,7 @@ def type_embedding_args(): ] -def spin_args(): +def spin_args() -> list[Argument]: doc_use_spin = ( "Whether to use atomic spin model for each atom type. " "List of boolean values with the shape of [ntypes] to specify which types use spin, " @@ -181,10 +183,10 @@ def __init__(self) -> None: self.__plugin = Plugin() def register( - self, name: str, alias: Optional[list[str]] = None, doc: str = "" + self, name: str, alias: list[str] | None = None, doc: str = "" ) -> Callable[ - [Union[Callable[[], Argument], Callable[[], list[Argument]]]], - Union[Callable[[], Argument], Callable[[], list[Argument]]], + [Callable[[], Argument] | Callable[[], list[Argument]]], + Callable[[], Argument] | Callable[[], list[Argument]], ]: """Register a descriptor argument plugin. @@ -247,7 +249,7 @@ def get_all_argument(self, exclude_hybrid: bool = False) -> list[Argument]: @descrpt_args_plugin.register("loc_frame", doc=doc_only_tf_supported + doc_loc_frame) -def descrpt_local_frame_args(): +def descrpt_local_frame_args() -> list[Argument]: doc_sel_a = "A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_a[i]` gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor." doc_sel_r = "A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_r[i]` gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius." doc_rcut = "The cut-off radius. The default value is 6.0" @@ -268,7 +270,7 @@ def descrpt_local_frame_args(): @descrpt_args_plugin.register("se_e2_a", alias=["se_a"], doc=doc_se_e2_a) -def descrpt_se_a_args(): +def descrpt_se_a_args() -> list[Argument]: doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' @@ -338,7 +340,7 @@ def descrpt_se_a_args(): @descrpt_args_plugin.register( "se_e3", alias=["se_at", "se_a_3be", "se_t"], doc=doc_se_e3 ) -def descrpt_se_t_args(): +def descrpt_se_t_args() -> list[Argument]: doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' @@ -395,7 +397,7 @@ def descrpt_se_t_args(): @descrpt_args_plugin.register( "se_a_tpe", alias=["se_a_ebd"], doc=doc_only_tf_supported + doc_se_a_tpe ) -def descrpt_se_a_tpe_args(): +def descrpt_se_a_tpe_args() -> list[Argument]: doc_type_nchanl = "number of channels for type embedding" doc_type_nlayer = "number of hidden layers of type embedding net" doc_numb_aparam = "dimension of atomic parameter. if set to a value > 0, the atomic parameters are embedded." @@ -409,7 +411,7 @@ def descrpt_se_a_tpe_args(): @descrpt_args_plugin.register("se_e2_r", alias=["se_r"], doc=doc_se_e2_r) -def descrpt_se_r_args(): +def descrpt_se_r_args() -> list[Argument]: doc_sel = 'This parameter set the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' @@ -468,7 +470,7 @@ def descrpt_se_r_args(): @descrpt_args_plugin.register("hybrid", doc=doc_hybrid) -def descrpt_hybrid_args(): +def descrpt_hybrid_args() -> list[Argument]: doc_list = "A list of descriptor definitions" return [ @@ -485,7 +487,7 @@ def descrpt_hybrid_args(): ] -def descrpt_se_atten_common_args(): +def descrpt_se_atten_common_args() -> list[Argument]: doc_sel = 'This parameter set the number of selected neighbors. Note that this parameter is a little different from that in other descriptors. Instead of separating each type of atoms, only the summation matters. And this number is highly related with the efficiency, thus one should not make it too large. Usually 200 or less is enough, far away from the GPU limitation 4096. It can be:\n\n\ - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.\ @@ -503,7 +505,7 @@ def descrpt_se_atten_common_args(): doc_exclude_types = "The excluded pairs of types which have no interaction with each other. For example, `[[0, 1]]` means no interaction between type 0 and type 1." doc_env_protection = "Protection parameter to prevent division by zero errors during environment matrix calculations. For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection." doc_attn = "The length of hidden vectors in attention layers" - doc_attn_layer = "The number of attention layers. Note that model compression of `se_atten` is only enabled when attn_layer==0 and tebd_input_mode=='strip'" + doc_attn_layer = "The number of attention layers. Note that model compression of `se_atten` works for any attn_layer value (for pytorch backend only, for other backends, attn_layer=0 is still needed to compress) when tebd_input_mode=='strip'. When attn_layer!=0, only type embedding is compressed, geometric parts are not compressed." doc_attn_dotr = "Whether to do dot product with the normalized relative coordinates" doc_attn_mask = "Whether to do mask on the diagonal in the attention matrix" @@ -560,7 +562,7 @@ def descrpt_se_atten_common_args(): @descrpt_args_plugin.register("se_atten", alias=["dpa1"], doc=doc_se_atten) -def descrpt_se_atten_args(): +def descrpt_se_atten_args() -> list[Argument]: doc_smooth_type_embedding = f"Whether to use smooth process in attention weights calculation. {doc_only_tf_supported} When using stripped type embedding, whether to dot smooth factor on the network output of type embedding to keep the network smooth, instead of setting `set_davg_zero` to be True." doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `se_atten` descriptor or `atom_ener` in the energy fitting is used" doc_trainable_ln = ( @@ -682,7 +684,7 @@ def descrpt_se_atten_args(): @descrpt_args_plugin.register("se_e3_tebd", doc=doc_only_pt_supported) -def descrpt_se_e3_tebd_args(): +def descrpt_se_e3_tebd_args() -> list[Argument]: doc_sel = 'This parameter set the number of selected neighbors. Note that this parameter is a little different from that in other descriptors. Instead of separating each type of atoms, only the summation matters. And this number is highly related with the efficiency, thus one should not make it too large. Usually 200 or less is enough, far away from the GPU limitation 4096. It can be:\n\n\ - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. Only the summation of `sel[i]` matters, and it is recommended to be less than 200.\ @@ -796,7 +798,7 @@ def descrpt_se_e3_tebd_args(): @descrpt_args_plugin.register("se_atten_v2", doc=doc_se_atten_v2) -def descrpt_se_atten_v2_args(): +def descrpt_se_atten_v2_args() -> list[Argument]: doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `se_atten` descriptor or `atom_ener` in the energy fitting is used" doc_trainable_ln = ( "Whether to use trainable shift and scale weights in layer normalization." @@ -880,7 +882,7 @@ def descrpt_se_atten_v2_args(): @descrpt_args_plugin.register("dpa2", doc=doc_only_pt_supported) -def descrpt_dpa2_args(): +def descrpt_dpa2_args() -> list[Argument]: # repinit args doc_repinit = "The arguments used to initialize the repinit block." # repformer args @@ -957,7 +959,7 @@ def descrpt_dpa2_args(): # repinit for dpa2 -def dpa2_repinit_args(): +def dpa2_repinit_args() -> list[Argument]: # repinit args doc_rcut = "The cut-off radius." doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`." @@ -1101,7 +1103,7 @@ def dpa2_repinit_args(): # repformer for dpa2 -def dpa2_repformer_args(): +def dpa2_repformer_args() -> list[Argument]: # repformer args doc_rcut = "The cut-off radius." doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`." @@ -1357,7 +1359,7 @@ def dpa2_repformer_args(): @descrpt_args_plugin.register("dpa3", doc=doc_only_pt_supported) -def descrpt_dpa3_args(): +def descrpt_dpa3_args() -> list[Argument]: # repflow args doc_repflow = "The arguments used to initialize the repflow block." # descriptor args @@ -1436,7 +1438,7 @@ def descrpt_dpa3_args(): # repflow for dpa3 -def dpa3_repflow_args(): +def dpa3_repflow_args() -> list[Argument]: # repflow args doc_n_dim = "The dimension of node representation." doc_e_dim = "The dimension of edge representation." @@ -1671,12 +1673,12 @@ def dpa3_repflow_args(): @descrpt_args_plugin.register( "se_a_ebd_v2", alias=["se_a_tpe_v2"], doc=doc_only_tf_supported ) -def descrpt_se_a_ebd_v2_args(): +def descrpt_se_a_ebd_v2_args() -> list[Argument]: return descrpt_se_a_args() @descrpt_args_plugin.register("se_a_mask", doc=doc_only_tf_supported + doc_se_a_mask) -def descrpt_se_a_mask_args(): +def descrpt_se_a_mask_args() -> list[Argument]: doc_sel = 'This parameter sets the number of selected neighbors for each type of atom. It can be:\n\n\ - `list[int]`. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius. It is noted that the total sel value must be less than 4096 in a GPU environment.\n\n\ - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' @@ -1743,9 +1745,10 @@ def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: @fitting_args_plugin.register("ener", doc=doc_ener) -def fitting_ener(): +def fitting_ener() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1773,6 +1776,13 @@ def fitting_ener(): return [ Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -1810,7 +1820,7 @@ def fitting_ener(): Argument("seed", [int, None], optional=True, doc=doc_seed), Argument( "atom_ener", - list[Optional[float]], + list[float | None], optional=True, default=[], doc=doc_atom_ener, @@ -1827,9 +1837,10 @@ def fitting_ener(): @fitting_args_plugin.register("dos", doc=doc_dos) -def fitting_dos(): +def fitting_dos() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1847,6 +1858,13 @@ def fitting_dos(): return [ Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -1882,9 +1900,10 @@ def fitting_dos(): @fitting_args_plugin.register("property", doc=doc_only_pt_supported) -def fitting_property(): +def fitting_property() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built" doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1900,6 +1919,13 @@ def fitting_property(): return [ Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -1944,9 +1970,10 @@ def fitting_property(): @fitting_args_plugin.register("polar", doc=doc_polar) -def fitting_polar(): +def fitting_polar() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -1976,6 +2003,13 @@ def fitting_polar(): default=0, doc=doc_only_pt_supported + doc_numb_aparam, ), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -2022,9 +2056,10 @@ def fitting_polar(): @fitting_args_plugin.register("dipole", doc=doc_dipole) -def fitting_dipole(): +def fitting_dipole() -> list[Argument]: doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." + doc_default_fparam = "The default frame parameter. If set, when `fparam.npy` files are not included in the data system, this value will be used as the default value for the frame parameter in the fitting net." doc_dim_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' @@ -2047,6 +2082,13 @@ def fitting_dipole(): default=0, doc=doc_only_pt_supported + doc_numb_aparam, ), + Argument( + "default_fparam", + list[float], + optional=True, + default=None, + doc=doc_only_pt_supported + doc_default_fparam, + ), Argument( "dim_case_embd", int, @@ -2083,7 +2125,7 @@ def fitting_dipole(): # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. -def fitting_variant_type_args(): +def fitting_variant_type_args() -> Variant: doc_descrpt_type = "The type of the fitting." return Variant( @@ -2100,7 +2142,7 @@ def fitting_variant_type_args(): @modifier_args_plugin.register("dipole_charge", doc=doc_dipole_charge) -def modifier_dipole_charge(): +def modifier_dipole_charge() -> list[Argument]: doc_model_name = "The name of the frozen dipole model file." doc_model_charge_map = f"The charge of the WFCC. The list length should be the same as the {make_link('sel_type', 'model[standard]/fitting_net[dipole]/sel_type')}. " doc_sys_charge_map = f"The charge of real atoms. The list length should be the same as the {make_link('type_map', 'model/type_map')}" @@ -2118,7 +2160,7 @@ def modifier_dipole_charge(): ] -def modifier_variant_type_args(): +def modifier_variant_type_args() -> Variant: doc_modifier_type = "The type of modifier." return Variant( "type", @@ -2129,7 +2171,7 @@ def modifier_variant_type_args(): # --- model compression configurations: --- # -def model_compression(): +def model_compression() -> list[Argument]: doc_model_file = "The input model file, which will be compressed by the DeePMD-kit." doc_table_config = "The arguments of model compression, including extrapolate(scale of model extrapolation), stride(uniform stride of tabulation's first and second table), and frequency(frequency of tabulation overflow check)." doc_min_nbor_dist = ( @@ -2144,7 +2186,7 @@ def model_compression(): # --- model compression configurations: --- # -def model_compression_type_args(): +def model_compression_type_args() -> Variant: doc_compress_type = "The type of model compression, which should be consistent with the descriptor type." return Variant( @@ -2161,7 +2203,7 @@ def model_compression_type_args(): hybrid_model_args_plugin = ArgsPlugin() -def model_args(exclude_hybrid=False): +def model_args(exclude_hybrid: bool = False) -> list[Argument]: doc_type_map = "A list of strings. Give the name to each type of atoms. It is noted that the number of atom type of training system must be less than 128 in a GPU environment. If not given, type.raw in each system should use the same type indexes, and type_map.raw will take no effect." doc_data_stat_nbatch = "The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics." doc_data_stat_protect = "Protect parameter for atomic energy regression." @@ -2242,7 +2284,7 @@ def model_args(exclude_hybrid=False): ), Argument( "preset_out_bias", - dict[str, list[Optional[Union[float, list[float]]]]], + dict[str, list[float | list[float] | None]], optional=True, default=None, doc=doc_only_pt_supported + doc_preset_out_bias, @@ -2305,6 +2347,16 @@ def model_args(exclude_hybrid=False): def standard_model_args() -> Argument: doc_descrpt = "The descriptor of atomic environment." doc_fitting = "The fitting of physical properties." + doc_model_branch_alias = ( + "List of aliases for this model branch. " + "Multiple aliases can be defined, and any alias can reference this branch throughout the model usage. " + "Used only in multitask models." + ) + doc_info = ( + "Dictionary of metadata for this model branch. " + "Store arbitrary key-value pairs with branch-specific information. " + "Used only in multitask models." + ) ca = Argument( "standard", @@ -2320,6 +2372,20 @@ def standard_model_args() -> Argument: [fitting_variant_type_args()], doc=doc_fitting, ), + Argument( + "model_branch_alias", + list[str], + optional=True, + default=[], + doc=doc_only_pt_supported + doc_model_branch_alias, + ), + Argument( + "info", + dict, + optional=True, + default={}, + doc=doc_only_pt_supported + doc_info, + ), ], doc="Standard model, which contains a descriptor and a fitting.", ) @@ -2411,7 +2477,7 @@ def linear_ener_model_args() -> Argument: # --- Learning rate configurations: --- # -def learning_rate_exp(): +def learning_rate_exp() -> list[Argument]: doc_start_lr = "The learning rate at the start of the training." doc_stop_lr = ( "The desired learning rate at the end of the training. " @@ -2443,7 +2509,7 @@ def learning_rate_exp(): return args -def learning_rate_variant_type_args(): +def learning_rate_variant_type_args() -> Variant: doc_lr = "The type of the learning rate." return Variant( @@ -2478,7 +2544,7 @@ def learning_rate_args(fold_subdoc: bool = False) -> Argument: # --- Loss configurations: --- # -def start_pref(item, label=None, abbr=None) -> str: +def start_pref(item: str, label: str | None = None, abbr: str | None = None) -> str: if label is None: label = item if abbr is None: @@ -2486,7 +2552,7 @@ def start_pref(item, label=None, abbr=None) -> str: return f"The prefactor of {item} loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the {label} label should be provided by file {label}.npy in each data system. If both start_pref_{abbr} and limit_pref_{abbr} are set to 0, then the {item} will be ignored." -def limit_pref(item) -> str: +def limit_pref(item: str) -> str: return f"The prefactor of {item} loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity." @@ -2494,7 +2560,7 @@ def limit_pref(item) -> str: @loss_args_plugin.register("ener") -def loss_ener(): +def loss_ener() -> list[Argument]: doc_start_pref_e = start_pref("energy", abbr="e") doc_limit_pref_e = limit_pref("energy") doc_start_pref_f = start_pref("force", abbr="f") @@ -2654,7 +2720,7 @@ def loss_ener(): @loss_args_plugin.register("ener_spin") -def loss_ener_spin(): +def loss_ener_spin() -> list[Argument]: doc_start_pref_e = start_pref("energy") doc_limit_pref_e = limit_pref("energy") doc_start_pref_fr = start_pref("force_real_atom") @@ -2766,7 +2832,7 @@ def loss_ener_spin(): @loss_args_plugin.register("dos") -def loss_dos(): +def loss_dos() -> list[Argument]: doc_start_pref_dos = start_pref("Density of State (DOS)") doc_limit_pref_dos = limit_pref("Density of State (DOS)") doc_start_pref_cdf = start_pref( @@ -2840,7 +2906,7 @@ def loss_dos(): @loss_args_plugin.register("property") -def loss_property(): +def loss_property() -> list[Argument]: doc_loss_func = "The loss function to minimize, such as 'mae','smooth_mae'." doc_metric = "The metric for display. This list can include 'smooth_mae', 'mae', 'mse' and 'rmse'." doc_beta = "The 'beta' parameter in 'smooth_mae' loss." @@ -2871,7 +2937,7 @@ def loss_property(): # YWolfeee: Modified to support tensor type of loss args. @loss_args_plugin.register("tensor") -def loss_tensor(): +def loss_tensor() -> list[Argument]: # doc_global_weight = "The prefactor of the weight of global loss. It should be larger than or equal to 0. If only `pref` is provided or both are not provided, training will be global mode, i.e. the shape of 'polarizability.npy` or `dipole.npy` should be #frams x [9 or 3]." # doc_local_weight = "The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If only `pref_atomic` is provided, training will be atomic mode, i.e. the shape of `polarizability.npy` or `dipole.npy` should be #frames x ([9 or 3] x #selected atoms). If both `pref` and `pref_atomic` are provided, training will be combined mode, and atomic label should be provided as well." doc_global_weight = "The prefactor of the weight of global loss. It should be larger than or equal to 0. It controls the weight of loss corresponding to global label, i.e. 'polarizability.npy` or `dipole.npy`, whose shape should be #frames x [9 or 3]. If it's larger than 0.0, this npy should be included." @@ -2898,7 +2964,7 @@ def loss_tensor(): ] -def loss_variant_type_args(): +def loss_variant_type_args() -> Variant: doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener` or left unset. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`." return Variant( @@ -2910,7 +2976,7 @@ def loss_variant_type_args(): ) -def loss_args(): +def loss_args() -> list[Argument]: doc_loss = "The definition of loss function. The loss type should be set to `tensor`, `ener` or left unset." ca = Argument( "loss", dict, [], [loss_variant_type_args()], optional=True, doc=doc_loss @@ -2919,12 +2985,15 @@ def loss_args(): # --- Training configurations: --- # -def training_data_args(): # ! added by Ziyao: new specification style for data systems. +def training_data_args() -> list[ + Argument +]: # ! added by Ziyao: new specification style for data systems. link_sys = make_link("systems", "training/training_data/systems") doc_systems = ( "The data systems for training. " - "This key can be provided with a list that specifies the systems, or be provided with a string " - "by which the prefix of all systems are given and the list of the systems is automatically generated." + "This key can be a list or a str. " + "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " + "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." ) doc_patterns = ( "The customized patterns used in `rglob` to collect all training systems. " @@ -2998,12 +3067,15 @@ def training_data_args(): # ! added by Ziyao: new specification style for data ) -def validation_data_args(): # ! added by Ziyao: new specification style for data systems. +def validation_data_args() -> list[ + Argument +]: # ! added by Ziyao: new specification style for data systems. link_sys = make_link("systems", "training/validation_data/systems") doc_systems = ( "The data systems for validation. " - "This key can be provided with a list that specifies the systems, or be provided with a string " - "by which the prefix of all systems are given and the list of the systems is automatically generated." + "This key can be a list or a str. " + "When provided as a string, it can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. " + "When provided as a list, each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories." ) doc_patterns = ( "The customized patterns used in `rglob` to collect all validation systems. " @@ -3088,7 +3160,7 @@ def validation_data_args(): # ! added by Ziyao: new specification style for dat ) -def mixed_precision_args(): # ! added by Denghui. +def mixed_precision_args() -> list[Argument]: # ! added by Denghui. doc_output_prec = 'The precision for mixed precision params. " \ "The trainable variables precision during the mixed precision training process, " \ "supported options are float32 only currently.' @@ -3117,8 +3189,8 @@ def mixed_precision_args(): # ! added by Denghui. def training_args( - multi_task=False, -): # ! modified by Ziyao: data configuration isolated. + multi_task: bool = False, +) -> list[Argument]: # ! modified by Ziyao: data configuration isolated. doc_numb_steps = "Number of training batch. Each training uses one batch of data." doc_seed = "The random seed for getting frames from the training data set." doc_disp_file = "The file for printing learning curve." @@ -3137,7 +3209,10 @@ def training_args( ) doc_disp_training = "Displaying verbose information during training." doc_time_training = "Timing during training." - doc_profiling = "Export the profiling results to the Chrome JSON file for performance analysis, driven by the legacy TensorFlow profiling API or PyTorch Profiler. The output file will be saved to `profiling_file`." + doc_disp_avg = ( + "Display the average loss over the display interval for training sets." + ) + doc_profiling = "Export the profiling results to the Chrome JSON file for performance analysis, driven by the legacy TensorFlow profiling API or PyTorch Profiler. The output file will be saved to `profiling_file`. In the PyTorch backend, when enable_profiler is True, this option is ignored, since the profiling results will be saved to the TensorBoard log." doc_profiling_file = "Output file for profiling." doc_enable_profiler = "Export the profiling results to the TensorBoard log for performance analysis, driven by TensorFlow Profiler (available in TensorFlow 2.3) or PyTorch Profiler. The log will be saved to `tensorboard_log_dir`." doc_tensorboard = "Enable tensorboard" @@ -3164,6 +3239,7 @@ def training_args( doc_kf_blocksize = "The blocksize for the Kalman filter." doc_model_prob = "The visiting probability of each model for each training step in the multi-task mode." doc_data_dict = "The multiple definition of the data, used in the multi-task mode." + doc_acc_freq = "Gradient accumulation steps (number of steps to accumulate gradients before performing an update)." arg_training_data = training_data_args() arg_validation_data = validation_data_args() @@ -3213,6 +3289,13 @@ def training_args( Argument( "time_training", bool, optional=True, default=True, doc=doc_time_training ), + Argument( + "disp_avg", + bool, + optional=True, + default=False, + doc=doc_only_pt_supported + doc_disp_avg, + ), Argument( "profiling", bool, @@ -3259,6 +3342,13 @@ def training_args( optional=True, doc=doc_only_pt_supported + doc_gradient_max_norm, ), + Argument( + "acc_freq", + int, + optional=True, + default=1, + doc=doc_only_pd_supported + doc_acc_freq, + ), ] variants = [ Variant( @@ -3291,7 +3381,7 @@ def training_args( return Argument("training", dict, args, variants, doc=doc_training) -def multi_model_args(): +def multi_model_args() -> list[Argument]: model_dict = model_args() model_dict.name = "model_dict" model_dict.repeat = True @@ -3312,7 +3402,7 @@ def multi_model_args(): ) -def multi_loss_args(): +def multi_loss_args() -> list[Argument]: loss_dict = loss_args() loss_dict.name = "loss_dict" loss_dict.repeat = True @@ -3321,14 +3411,20 @@ def multi_loss_args(): return loss_dict -def make_index(keys): +def make_index(keys: list[str]) -> str: ret = [] for ii in keys: ret.append(make_link(ii, ii)) return ", ".join(ret) -def gen_doc(*, make_anchor=True, make_link=True, multi_task=False, **kwargs) -> str: +def gen_doc( + *, + make_anchor: bool = True, + make_link: bool = True, + multi_task: bool = False, + **kwargs: Any, +) -> str: if make_link: make_anchor = True ptr = [] @@ -3344,7 +3440,7 @@ def gen_doc(*, make_anchor=True, make_link=True, multi_task=False, **kwargs) -> return "\n\n".join(ptr) -def gen_json(multi_task: bool = False, **kwargs) -> str: +def gen_json(multi_task: bool = False, **kwargs: Any) -> str: return json.dumps( tuple(gen_args(multi_task=multi_task)), cls=ArgumentEncoder, @@ -3397,7 +3493,7 @@ def gen_json_schema(multi_task: bool = False) -> str: return json.dumps(generate_json_schema(arg)) -def normalize(data, multi_task: bool = False): +def normalize(data: dict[str, Any], multi_task: bool = False) -> dict[str, Any]: base = Argument("base", dict, gen_args(multi_task=multi_task)) data = base.normalize_value(data, trim_pattern="_*") base.check_value(data, strict=True) diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 516c4d2ead..012ffb4260 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -5,9 +5,12 @@ ABC, abstractmethod, ) -from typing import ( +from collections.abc import ( Callable, ) +from typing import ( + Any, +) import array_api_compat import numpy as np @@ -145,7 +148,12 @@ def _adjust_batch_size(self, factor: float) -> None: ) def execute_all( - self, callable: Callable, total_size: int, natoms: int, *args, **kwargs + self, + callable: Callable, + total_size: int, + natoms: int, + *args: Any, + **kwargs: Any, ) -> tuple[np.ndarray]: """Excuate a method with all given data. @@ -209,7 +217,7 @@ def execute_with_batch_size( result = (result,) if not isinstance(result, tuple) else result index += n_batch - def append_to_list(res_list, res): + def append_to_list(res_list: list[Any], res: Any) -> list[Any]: if n_batch: res_list.append(res) return res_list @@ -223,7 +231,7 @@ def append_to_list(res_list, res): assert results is not None assert returned_dict is not None - def concate_result(r): + def concate_result(r: list[Any]) -> Any: if array_api_compat.is_array_api_obj(r[0]): xp = array_api_compat.array_namespace(r[0]) ret = xp.concat(r, axis=0) diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index 72948c96f4..f15c2b4909 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -11,8 +11,6 @@ ) from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -23,7 +21,7 @@ def convert_input_v0_v1( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: """Convert input from v0 format to v1. @@ -54,7 +52,7 @@ def convert_input_v0_v1( return output -def _warning_input_v0_v1(fname: Optional[Union[str, Path]]) -> None: +def _warning_input_v0_v1(fname: str | Path | None) -> None: msg = ( "It seems that you are using a deepmd-kit input of version 0.x.x, " "which is deprecated. we have converted the input to >2.0.0 compatible" @@ -288,7 +286,7 @@ def remove_decay_rate(jdata: dict[str, Any]) -> None: def convert_input_v1_v2( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: tr_cfg = jdata["training"] tr_data_keys = { @@ -324,7 +322,7 @@ def convert_input_v1_v2( return jdata -def _warning_input_v1_v2(fname: Optional[Union[str, Path]]) -> None: +def _warning_input_v1_v2(fname: str | Path | None) -> None: msg = ( "It seems that you are using a deepmd-kit input of version 1.x.x, " "which is deprecated. we have converted the input to >2.0.0 compatible" @@ -335,7 +333,7 @@ def _warning_input_v1_v2(fname: Optional[Union[str, Path]]) -> None: def deprecate_numb_test( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: """Deprecate `numb_test` since v2.1. It has taken no effect since v2.0. @@ -373,12 +371,12 @@ def deprecate_numb_test( def update_deepmd_input( - jdata: dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None + jdata: dict[str, Any], warning: bool = True, dump: str | Path | None = None ) -> dict[str, Any]: - def is_deepmd_v0_input(jdata): + def is_deepmd_v0_input(jdata: dict[str, Any]) -> bool: return "model" not in jdata.keys() - def is_deepmd_v1_input(jdata): + def is_deepmd_v1_input(jdata: dict[str, Any]) -> bool: return "systems" in jdata["training"].keys() if is_deepmd_v0_input(jdata): diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index ccf65d3243..287107a7ff 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -2,9 +2,17 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import bisect +import functools import logging +from concurrent.futures import ( + ThreadPoolExecutor, + as_completed, +) +from pathlib import ( + Path, +) from typing import ( - Optional, + Any, ) import numpy as np @@ -12,9 +20,11 @@ from deepmd.env import ( GLOBAL_ENER_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION, + LRU_CACHE_SIZE, ) from deepmd.utils import random as dp_random from deepmd.utils.path import ( + DPH5Path, DPPath, ) @@ -52,9 +62,9 @@ def __init__( sys_path: str, set_prefix: str = "set", shuffle_test: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, optional_type_map: bool = True, - modifier=None, + modifier: Any | None = None, trn_all_set: bool = False, sort_atoms: bool = True, ) -> None: @@ -67,10 +77,7 @@ def __init__( raise FileNotFoundError(f"No {set_prefix}.* is found in {sys_path}") self.dirs.sort() # check mix_type format - error_format_msg = ( - "if one of the set is of mixed_type format, " - "then all of the sets in this system should be of mixed_type format!" - ) + error_format_msg = "if one of the set is of mixed_type format, then all of the sets in this system should be of mixed_type format!" self.mixed_type = self._check_mode(self.dirs[0]) for set_item in self.dirs[1:]: assert self._check_mode(set_item) == self.mixed_type, error_format_msg @@ -95,10 +102,10 @@ def __init__( f"Elements {missing_elements} are not present in the provided `type_map`." ) if not self.mixed_type: - atom_type_ = [ - type_map.index(self.type_map[ii]) for ii in self.atom_type - ] - self.atom_type = np.array(atom_type_, dtype=np.int32) + old_to_new_type_idx = np.array( + [type_map.index(name) for name in self.type_map], dtype=np.int32 + ) + self.atom_type = old_to_new_type_idx[self.atom_type].astype(np.int32) else: self.enforce_type_map = True sorter = np.argsort(type_map) @@ -127,8 +134,7 @@ def __init__( self.shuffle_test = shuffle_test # set modifier self.modifier = modifier - # calculate prefix sum for get_item method - frames_list = [self._get_nframes(item) for item in self.dirs] + frames_list = [self._get_nframes(set_name) for set_name in self.dirs] self.nframes = np.sum(frames_list) # The prefix sum stores the range of indices contained in each directory, which is needed by get_item method self.prefix_sum = np.cumsum(frames_list).tolist() @@ -140,12 +146,12 @@ def add( atomic: bool = False, must: bool = False, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, repeat: int = 1, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, - ): + ) -> "DeepmdData": """Add a data item that to be loaded. Parameters @@ -188,7 +194,7 @@ def add( } return self - def reduce(self, key_out: str, key_in: str): + def reduce(self, key_out: str, key_in: str) -> "DeepmdData": """Generate a new item from the reduction of another atom. Parameters @@ -220,7 +226,7 @@ def get_data_dict(self) -> dict: """Get the `data_dict`.""" return self.data_dict - def check_batch_size(self, batch_size): + def check_batch_size(self, batch_size: int) -> bool: """Check if the system can get a batch of data with `batch_size` frames.""" for ii in self.dirs: if self.data_dict["coord"]["high_prec"]: @@ -235,7 +241,7 @@ def check_batch_size(self, batch_size): return ii, tmpe.shape[0] return None - def check_test_size(self, test_size): + def check_test_size(self, test_size: int) -> bool: """Check if the system can get a test dataset with `test_size` frames.""" return self.check_batch_size(test_size) @@ -247,27 +253,18 @@ def get_item_torch(self, index: int) -> dict: index index of the frame """ - i = bisect.bisect_right(self.prefix_sum, index) - frames = self._load_set(self.dirs[i]) - frame = self._get_subdata(frames, index - self.prefix_sum[i]) - frame = self.reformat_data_torch(frame) - frame["fid"] = index - return frame + return self.get_single_frame(index) def get_item_paddle(self, index: int) -> dict: """Get a single frame data . The frame is picked from the data system by index. The index is coded across all the sets. + Same with PyTorch backend. Parameters ---------- index index of the frame """ - i = bisect.bisect_right(self.prefix_sum, index) - frames = self._load_set(self.dirs[i]) - frame = self._get_subdata(frames, index - self.prefix_sum[i]) - frame = self.reformat_data_torch(frame) - frame["fid"] = index - return frame + return self.get_single_frame(index) def get_batch(self, batch_size: int) -> dict: """Get a batch of data with `batch_size` frames. The frames are randomly picked from the data system. @@ -339,8 +336,10 @@ def get_numb_set(self) -> int: def get_numb_batch(self, batch_size: int, set_idx: int) -> int: """Get the number of batches in a set.""" - data = self._load_set(self.dirs[set_idx]) - ret = data["coord"].shape[0] // batch_size + set_name = self.dirs[set_idx] + # Directly obtain the number of frames to avoid loading the entire dataset + nframes = self._get_nframes(set_name) + ret = nframes // batch_size if ret == 0: ret = 1 return ret @@ -352,11 +351,11 @@ def get_sys_numb_batch(self, batch_size: int) -> int: ret += self.get_numb_batch(batch_size, ii) return ret - def get_natoms(self): + def get_natoms(self) -> int: """Get number of atoms.""" return len(self.atom_type) - def get_natoms_vec(self, ntypes: int): + def get_natoms_vec(self, ntypes: int) -> np.ndarray: """Get number of atoms and number of atoms in different types. Parameters @@ -376,7 +375,104 @@ def get_natoms_vec(self, ntypes: int): tmp = np.append(tmp, natoms_vec) return tmp.astype(np.int32) - def avg(self, key): + def get_single_frame(self, index: int) -> dict: + """Orchestrates loading a single frame efficiently using memmap.""" + if index < 0 or index >= self.nframes: + raise IndexError(f"Frame index {index} out of range [0, {self.nframes})") + # 1. Find the correct set directory and local frame index + set_idx = bisect.bisect_right(self.prefix_sum, index) + set_dir = self.dirs[set_idx] + if not isinstance(set_dir, DPPath): + set_dir = DPPath(set_dir) + # Calculate local index within the set.* directory + local_idx = index - (0 if set_idx == 0 else self.prefix_sum[set_idx - 1]) + # Calculate the number of frames in this set to avoid redundant _get_nframes calls + set_nframes = ( + self.prefix_sum[set_idx] + if set_idx == 0 + else self.prefix_sum[set_idx] - self.prefix_sum[set_idx - 1] + ) + + frame_data = {} + # 2. Concurrently load all non-reduced items + non_reduced_keys = [k for k, v in self.data_dict.items() if v["reduce"] is None] + reduced_keys = [k for k, v in self.data_dict.items() if v["reduce"] is not None] + # Use a thread pool to parallelize loading + if non_reduced_keys: + with ThreadPoolExecutor(max_workers=len(non_reduced_keys)) as executor: + future_to_key = { + executor.submit( + self._load_single_data, set_dir, key, local_idx, set_nframes + ): key + for key in non_reduced_keys + } + for future in as_completed(future_to_key): + key = future_to_key[future] + frame_data["find_" + key], frame_data[key] = future.result() + + # 3. Compute reduced items from already loaded data + for key in reduced_keys: + vv = self.data_dict[key] + k_in = vv["reduce"] + ndof = vv["ndof"] + frame_data["find_" + key] = frame_data["find_" + k_in] + # Reshape to (natoms, ndof) and sum over atom axis + tmp_in = ( + frame_data[k_in].reshape(-1, ndof).astype(GLOBAL_ENER_FLOAT_PRECISION) + ) + frame_data[key] = np.sum(tmp_in, axis=0) + + # 4. Handle atom types (mixed or standard) + if self.mixed_type: + type_path = set_dir / "real_atom_types.npy" + # For HDF5 files, use load_numpy; for filesystem, use memmap + if isinstance(type_path, DPH5Path): + mmap_types = type_path.load_numpy() + else: + mmap_types = self._get_memmap(type_path) + real_type = mmap_types[local_idx].copy().astype(np.int32) + + if self.enforce_type_map: + try: + real_type = self.type_idx_map[real_type].astype(np.int32) + except IndexError as e: + raise IndexError( + f"some types in 'real_atom_types.npy' of set {set_dir} are not contained in {self.get_ntypes()} types!" + ) from e + + frame_data["type"] = real_type + ntypes = self.get_ntypes() + natoms = len(real_type) + # Use bincount for efficient counting of each type + natoms_vec = np.bincount( + real_type[real_type >= 0], minlength=ntypes + ).astype(np.int32) + frame_data["real_natoms_vec"] = np.concatenate( + (np.array([natoms, natoms], dtype=np.int32), natoms_vec) + ) + else: + frame_data["type"] = self.atom_type[self.idx_map] + + # 5. Standardize keys + frame_data = {kk.replace("atomic", "atom"): vv for kk, vv in frame_data.items()} + + # 6. Reshape atomic data to match expected format [natoms, ndof] + for kk in self.data_dict.keys(): + if ( + "find_" not in kk + and kk in frame_data + and not self.data_dict[kk]["atomic"] + ): + frame_data[kk] = frame_data[kk].reshape(-1) + frame_data["atype"] = frame_data["type"] + + if not self.pbc: + frame_data["box"] = None + + frame_data["fid"] = index + return frame_data + + def avg(self, key: str) -> float: """Return the average value of an item.""" if key not in self.data_dict.keys(): raise RuntimeError(f"key {key} has not been added") @@ -393,18 +489,16 @@ def avg(self, key): else: return np.average(eners, axis=0) - def _idx_map_sel(self, atom_type, type_sel): - new_types = [] - for ii in atom_type: - if ii in type_sel: - new_types.append(ii) - new_types = np.array(new_types, dtype=int) + def _idx_map_sel(self, atom_type: np.ndarray, type_sel: list[int]) -> np.ndarray: + # Use vectorized operations instead of Python loop + sel_mask = np.isin(atom_type, type_sel) + new_types = atom_type[sel_mask] natoms = new_types.shape[0] idx = np.arange(natoms, dtype=np.int64) idx_map = np.lexsort((idx, new_types)) return idx_map - def _get_natoms_2(self, ntypes): + def _get_natoms_2(self, ntypes: int) -> tuple[int, np.ndarray]: sample_type = self.atom_type natoms = len(sample_type) natoms_vec = np.zeros(ntypes, dtype=np.int64) @@ -412,7 +506,18 @@ def _get_natoms_2(self, ntypes): natoms_vec[ii] = np.count_nonzero(sample_type == ii) return natoms, natoms_vec - def _get_subdata(self, data, idx=None): + def _get_memmap(self, path: DPPath) -> np.memmap: + """Get or create a memory-mapped object for a given npy file. + Uses file path and modification time as cache keys to detect file changes + and invalidate cache when files are modified. + """ + abs_path = Path(str(path)).absolute() + file_mtime = abs_path.stat().st_mtime + return self._create_memmap(str(abs_path), str(file_mtime)) + + def _get_subdata( + self, data: dict[str, Any], idx: np.ndarray | None = None + ) -> dict[str, Any]: new_data = {} for ii in data: dd = data[ii] @@ -454,7 +559,7 @@ def _load_test_set(self, shuffle_test: bool) -> None: if shuffle_test: self.test_set, _ = self._shuffle_data(self.test_set) - def _shuffle_data(self, data): + def _shuffle_data(self, data: dict[str, Any]) -> dict[str, Any]: ret = {} nframes = data["coord"].shape[0] idx = np.arange(nframes, dtype=np.int64) @@ -473,21 +578,30 @@ def _shuffle_data(self, data): ret[kk] = data[kk] return ret, idx - def _get_nframes(self, set_name: DPPath): - # get nframes + def _get_nframes(self, set_name: DPPath | str) -> int: if not isinstance(set_name, DPPath): set_name = DPPath(set_name) path = set_name / "coord.npy" - if self.data_dict["coord"]["high_prec"]: - coord = path.load_numpy().astype(GLOBAL_ENER_FLOAT_PRECISION) + if isinstance(set_name, DPH5Path): + nframes = path.root[path._name].shape[0] else: - coord = path.load_numpy().astype(GLOBAL_NP_FLOAT_PRECISION) - if coord.ndim == 1: - coord = coord.reshape([1, -1]) - nframes = coord.shape[0] + # Read only the header to get shape + with open(str(path), "rb") as f: + version = np.lib.format.read_magic(f) + if version[0] == 1: + shape, _fortran_order, _dtype = np.lib.format.read_array_header_1_0( + f + ) + elif version[0] in [2, 3]: + shape, _fortran_order, _dtype = np.lib.format.read_array_header_2_0( + f + ) + else: + raise ValueError(f"Unsupported .npy file version: {version}") + nframes = shape[0] if len(shape) > 1 else 1 return nframes - def reformat_data_torch(self, data): + def reformat_data_torch(self, data: dict[str, Any]) -> dict[str, Any]: """Modify the data format for the requirements of Torch backend. Parameters @@ -506,7 +620,7 @@ def reformat_data_torch(self, data): data["box"] = None return data - def _load_set(self, set_name: DPPath): + def _load_set(self, set_name: DPPath) -> dict[str, Any]: # get nframes if not isinstance(set_name, DPPath): set_name = DPPath(set_name) @@ -593,27 +707,27 @@ def _load_set(self, set_name: DPPath): def _load_data( self, - set_name, - key, - nframes, - ndof_, - atomic=False, - must=True, - repeat=1, - high_prec=False, - type_sel=None, + set_name: str, + key: str, + nframes: int, + ndof_: int, + atomic: bool = False, + must: bool = True, + repeat: int = 1, + high_prec: bool = False, + type_sel: list[int] | None = None, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, - ): + ) -> np.ndarray: if atomic: natoms = self.natoms idx_map = self.idx_map # if type_sel, then revise natoms and idx_map if type_sel is not None: - natoms_sel = 0 - for jj in type_sel: - natoms_sel += np.sum(self.atom_type == jj) + # Use vectorized operations for better performance + sel_mask = np.isin(self.atom_type, type_sel) + natoms_sel = np.sum(sel_mask) idx_map_sel = self._idx_map_sel(self.atom_type, type_sel) else: natoms_sel = natoms @@ -641,7 +755,6 @@ def _load_data( tmp = np.zeros( [nframes, natoms, ndof_], dtype=data.dtype ) - sel_mask = np.isin(self.atom_type, type_sel) tmp[:, sel_mask] = data.reshape( [nframes, natoms_sel, ndof_] ) @@ -654,7 +767,6 @@ def _load_data( if output_natoms_for_type_sel: pass else: - sel_mask = np.isin(self.atom_type, type_sel) data = data.reshape([nframes, natoms, ndof_]) data = data[:, sel_mask] natoms = natoms_sel @@ -687,7 +799,7 @@ def _load_data( data = data.reshape([nframes, -1]) data = np.reshape(data, [nframes, ndof]) except ValueError as err_message: - explanation = "This error may occur when your label mismatch it's name, i.e. you might store global tensor in `atomic_tensor.npy` or atomic tensor in `tensor.npy`." + explanation = "This error may occur when your label mismatch its name, i.e. you might store global tensor in `atomic_tensor.npy` or atomic tensor in `tensor.npy`." log.error(str(err_message)) log.error(explanation) raise ValueError(str(err_message) + ". " + explanation) from err_message @@ -704,16 +816,157 @@ def _load_data( data = np.repeat(data, repeat).reshape([nframes, -1]) return np.float32(0.0), data - def _load_type(self, sys_path: DPPath): + def _load_single_data( + self, set_dir: DPPath, key: str, frame_idx: int, set_nframes: int + ) -> tuple[np.float32, np.ndarray]: + """ + Loads and processes data for a SINGLE frame from a SINGLE key, + fully replicating the logic from the original _load_data method. + + Parameters + ---------- + set_dir : DPPath + The directory path of the set + key : str + The key name of the data to load + frame_idx : int + The local frame index within the set + set_nframes : int + The total number of frames in this set (to avoid redundant _get_nframes calls) + """ + vv = self.data_dict[key] + path = set_dir / (key + ".npy") + + if vv["atomic"]: + natoms = self.natoms + idx_map = self.idx_map + # if type_sel, then revise natoms and idx_map + if vv["type_sel"] is not None: + # Use vectorized operations for better performance + sel_mask = np.isin(self.atom_type, vv["type_sel"]) + natoms_sel = np.sum(sel_mask) + idx_map_sel = self._idx_map_sel(self.atom_type, vv["type_sel"]) + else: + natoms_sel = natoms + idx_map_sel = idx_map + else: + natoms = 1 + natoms_sel = 0 + idx_map_sel = None + ndof = vv["ndof"] + + # Determine target data type from requirements + dtype = vv.get("dtype") + if dtype is None: + dtype = ( + GLOBAL_ENER_FLOAT_PRECISION + if vv.get("high_prec") + else GLOBAL_NP_FLOAT_PRECISION + ) + + # Branch 1: File does not exist + if not path.is_file(): + if vv.get("must"): + raise RuntimeError(f"{path} not found!") + + # Create a default array based on requirements + if vv["atomic"]: + if vv["type_sel"] is not None and not vv["output_natoms_for_type_sel"]: + natoms = natoms_sel + data = np.full([natoms, ndof], vv["default"], dtype=dtype) + else: + # For non-atomic data, shape should be [ndof] + data = np.full([ndof], vv["default"], dtype=dtype) + return np.float32(0.0), data + + # Branch 2: Data loading + if isinstance(path, DPH5Path): + # For HDF5 files, use load_numpy which handles HDF5 datasets + mmap_obj = path.load_numpy().astype(dtype) + else: + # For filesystem paths, use memmap for better performance + mmap_obj = self._get_memmap(path) + + # corner case: single frame + if set_nframes == 1: + mmap_obj = mmap_obj[None, ...] + # Slice the single frame and make an in-memory copy for modification + data = mmap_obj[frame_idx].copy().astype(dtype, copy=False) + + try: + if vv["atomic"]: + # Handle type_sel logic + if vv["type_sel"] is not None: + if mmap_obj.shape[1] == natoms_sel * ndof: + if vv["output_natoms_for_type_sel"]: + tmp = np.zeros([natoms, ndof], dtype=data.dtype) + # sel_mask needs to be applied to the original atom layout + tmp[sel_mask] = data.reshape([natoms_sel, ndof]) + data = tmp + else: # output is natoms_sel + natoms = natoms_sel + idx_map = idx_map_sel + elif mmap_obj.shape[1] == natoms * ndof: + data = data.reshape([natoms, ndof]) + if vv["output_natoms_for_type_sel"]: + pass + else: + data = data[sel_mask] + idx_map = idx_map_sel + natoms = natoms_sel + else: # Shape mismatch error + raise ValueError( + f"The shape of the data {key} in {set_dir} has width {mmap_obj.shape[1]}, which doesn't match either ({natoms_sel * ndof}) or ({natoms * ndof})" + ) + + # Handle special case for Hessian + if key == "hessian": + data = data.reshape(3 * natoms, 3 * natoms) + num_chunks, chunk_size = len(idx_map), 3 + idx_map_hess = np.arange( + num_chunks * chunk_size, dtype=int + ).reshape(num_chunks, chunk_size) + idx_map_hess = idx_map_hess[idx_map].flatten() + data = data[idx_map_hess, :] + data = data[:, idx_map_hess] + data = data.reshape(-1) + # size of hessian is 3Natoms * 3Natoms + # ndof = 3 * ndof * 3 * ndof + else: + # data should be 2D here: [natoms, ndof] + data = data.reshape([natoms, -1]) + data = data[idx_map, :] + else: + data = data.reshape([ndof]) + + # Atomic: return [natoms, ndof] or flattened hessian above + # Non-atomic: return [ndof] + return np.float32(1.0), data + + except ValueError as err_message: + explanation = ( + "This error may occur when your label mismatches its name, " + "e.g., global tensor stored in `atomic_tensor.npy` or atomic tensor in `tensor.npy`." + ) + log.exception( + "Single-frame load failed for key=%s, set=%s, frame=%d. %s", + key, + set_dir, + frame_idx, + explanation, + ) + raise ValueError(f"{err_message}. {explanation}") from err_message + + def _load_type(self, sys_path: DPPath) -> np.ndarray: atom_type = (sys_path / "type.raw").load_txt(ndmin=1).astype(np.int32) return atom_type - def _load_type_mix(self, set_name: DPPath): + def _load_type_mix(self, set_name: DPPath) -> np.ndarray: type_path = set_name / "real_atom_types.npy" real_type = type_path.load_numpy().astype(np.int32).reshape([-1, self.natoms]) return real_type - def _make_idx_map(self, atom_type): + def _make_idx_map(self, atom_type: np.ndarray) -> np.ndarray: natoms = atom_type.shape[0] idx = np.arange(natoms, dtype=np.int64) if self.sort_atoms: @@ -722,22 +975,50 @@ def _make_idx_map(self, atom_type): idx_map = idx return idx_map - def _load_type_map(self, sys_path: DPPath): + def _load_type_map(self, sys_path: DPPath) -> list[str] | None: fname = sys_path / "type_map.raw" if fname.is_file(): return fname.load_txt(dtype=str, ndmin=1).tolist() else: return None - def _check_pbc(self, sys_path: DPPath): + def _check_pbc(self, sys_path: DPPath) -> bool: pbc = True if (sys_path / "nopbc").is_file(): pbc = False return pbc - def _check_mode(self, set_path: DPPath): + def _check_mode(self, set_path: DPPath) -> bool: return (set_path / "real_atom_types.npy").is_file() + @staticmethod + @functools.lru_cache(maxsize=LRU_CACHE_SIZE) + def _create_memmap(path_str: str, mtime_str: str) -> np.memmap: + """A cached helper function to create memmap objects. + Using lru_cache to limit the number of open file handles. + + Parameters + ---------- + path_str + The file path as a string. + mtime_str + The modification time as a string, used for cache invalidation. + """ + with open(path_str, "rb") as f: + version = np.lib.format.read_magic(f) + if version[0] == 1: + shape, fortran_order, dtype = np.lib.format.read_array_header_1_0(f) + elif version[0] in [2, 3]: + shape, fortran_order, dtype = np.lib.format.read_array_header_2_0(f) + else: + raise ValueError(f"Unsupported .npy file version: {version}") + offset = f.tell() + order = "F" if fortran_order else "C" + # Create a read-only memmap + return np.memmap( + path_str, dtype=dtype, mode="r", shape=shape, order=order, offset=offset + ) + class DataRequirementItem: """A class to store the data requirement for data systems. @@ -776,10 +1057,10 @@ def __init__( atomic: bool = False, must: bool = False, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, repeat: int = 1, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, ) -> None: self.key = key @@ -808,7 +1089,7 @@ def to_dict(self) -> dict: "output_natoms_for_type_sel": self.output_natoms_for_type_sel, } - def __getitem__(self, key: str): + def __getitem__(self, key: str) -> np.ndarray: if key not in self.dict: raise KeyError(key) return self.dict[key] diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 07dab35a90..82ef5ec25d 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -7,8 +7,6 @@ ) from typing import ( Any, - Optional, - Union, ) import numpy as np @@ -44,15 +42,15 @@ def __init__( systems: list[str], batch_size: int, test_size: int, - rcut: Optional[float] = None, + rcut: float | None = None, set_prefix: str = "set", shuffle_test: bool = True, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, optional_type_map: bool = True, - modifier=None, - trn_all_set=False, - sys_probs=None, - auto_prob_style="prob_sys_size", + modifier: Any | None = None, + trn_all_set: bool = False, + sys_probs: list[float] | None = None, + auto_prob_style: str = "prob_sys_size", sort_atoms: bool = True, ) -> None: """Constructor. @@ -152,6 +150,51 @@ def __init__( else: raise RuntimeError("batch size must be specified for mixed systems") self.batch_size = rule * np.ones(self.nsystems, dtype=int) + elif "max" == words[0]: + # Determine batch size so that batch_size * natoms <= rule, at least 1 + if len(words) != 2: + raise RuntimeError("batch size must be specified for max systems") + rule = int(words[1]) + bs = [] + for ii in self.data_systems: + ni = ii.get_natoms() + bsi = rule // ni + if bsi == 0: + bsi = 1 + bs.append(bsi) + self.batch_size = bs + elif "filter" == words[0]: + # Remove systems with natoms > rule, then set batch size like "max:rule" + if len(words) != 2: + raise RuntimeError( + "batch size must be specified for filter systems" + ) + rule = int(words[1]) + filtered_data_systems = [] + filtered_system_dirs = [] + for sys_dir, data_sys in zip(self.system_dirs, self.data_systems): + if data_sys.get_natoms() <= rule: + filtered_data_systems.append(data_sys) + filtered_system_dirs.append(sys_dir) + if len(filtered_data_systems) == 0: + raise RuntimeError( + f"No system left after removing systems with more than {rule} atoms" + ) + if len(filtered_data_systems) != len(self.data_systems): + warnings.warn( + f"Remove {len(self.data_systems) - len(filtered_data_systems)} systems with more than {rule} atoms" + ) + self.data_systems = filtered_data_systems + self.system_dirs = filtered_system_dirs + self.nsystems = len(self.data_systems) + bs = [] + for ii in self.data_systems: + ni = ii.get_natoms() + bsi = rule // ni + if bsi == 0: + bsi = 1 + bs.append(bsi) + self.batch_size = bs else: raise RuntimeError("unknown batch_size rule " + words[0]) elif isinstance(self.batch_size, list): @@ -224,7 +267,7 @@ def __init__( f"system {self.system_dirs[ii]} required test size is larger than the size of the dataset {chk_ret[0]} ({self.test_size[ii]} > {chk_ret[1]})" ) - def _load_test(self, ntests=-1) -> None: + def _load_test(self, ntests: int = -1) -> None: self.test_data = collections.defaultdict(list) for ii in range(self.nsystems): test_system_data = self.data_systems[ii].get_test(ntests=ntests) @@ -241,7 +284,9 @@ def default_mesh(self) -> list[np.ndarray]: for ii in range(self.nsystems) ] - def compute_energy_shift(self, rcond=None, key="energy"): + def compute_energy_shift( + self, rcond: float | None = None, key: str = "energy" + ) -> tuple[np.ndarray, np.ndarray]: sys_ener = [] for ss in self.data_systems: sys_ener.append(ss.avg(key)) @@ -301,10 +346,10 @@ def add( atomic: bool = False, must: bool = False, high_prec: bool = False, - type_sel: Optional[list[int]] = None, + type_sel: list[int] | None = None, repeat: int = 1, default: float = 0.0, - dtype: Optional[np.dtype] = None, + dtype: np.dtype | None = None, output_natoms_for_type_sel: bool = False, ) -> None: """Add a data item that to be loaded. @@ -349,7 +394,7 @@ def add( output_natoms_for_type_sel=output_natoms_for_type_sel, ) - def reduce(self, key_out, key_in) -> None: + def reduce(self, key_out: str, key_in: str) -> None: """Generate a new item from the reduction of another atom. Parameters @@ -366,7 +411,9 @@ def get_data_dict(self, ii: int = 0) -> dict: return self.data_systems[ii].get_data_dict() def set_sys_probs( - self, sys_probs=None, auto_prob_style: str = "prob_sys_size" + self, + sys_probs: list[float] | None = None, + auto_prob_style: str = "prob_sys_size", ) -> None: if sys_probs is None: if auto_prob_style == "prob_uniform": @@ -386,7 +433,7 @@ def set_sys_probs( probs = process_sys_probs(sys_probs, self.nbatches) self.sys_probs = probs - def get_batch(self, sys_idx: Optional[int] = None) -> dict: + def get_batch(self, sys_idx: int | None = None) -> dict: # batch generation style altered by Ziyao Li: # one should specify the "sys_prob" and "auto_prob_style" params # via set_sys_prob() function. The sys_probs this function uses is @@ -413,7 +460,7 @@ def get_batch(self, sys_idx: Optional[int] = None) -> dict: b_data = self.get_batch_mixed() return b_data - def get_batch_standard(self, sys_idx: Optional[int] = None) -> dict: + def get_batch_standard(self, sys_idx: int | None = None) -> dict: """Get a batch of data from the data systems in the standard way. Parameters @@ -512,7 +559,9 @@ def _merge_batch_data(self, batch_data: list[dict]) -> dict: return b_data # ! altered by Marián Rynik - def get_test(self, sys_idx: Optional[int] = None, n_test: int = -1): # depreciated + def get_test( + self, sys_idx: int | None = None, n_test: int = -1 + ) -> dict[str, np.ndarray]: # depreciated """Get test data from the the data systems. Parameters @@ -537,7 +586,7 @@ def get_test(self, sys_idx: Optional[int] = None, n_test: int = -1): # deprecia test_system_data["default_mesh"] = self.default_mesh[idx] return test_system_data - def get_sys_ntest(self, sys_idx=None): + def get_sys_ntest(self, sys_idx: int | None = None) -> int: """Get number of tests for the currently selected system, or one defined by sys_idx. """ @@ -582,7 +631,7 @@ def print_summary(self, name: str) -> None: [ii.pbc for ii in self.data_systems], ) - def _make_auto_bs(self, rule): + def _make_auto_bs(self, rule: int) -> list[int]: bs = [] for ii in self.data_systems: ni = ii.get_natoms() @@ -593,7 +642,7 @@ def _make_auto_bs(self, rule): return bs # ! added by Marián Rynik - def _make_auto_ts(self, percent): + def _make_auto_ts(self, percent: float) -> list[int]: ts = [] for ii in range(self.nsystems): ni = self.batch_size[ii] * self.nbatches[ii] @@ -602,7 +651,9 @@ def _make_auto_ts(self, percent): return ts - def _check_type_map_consistency(self, type_map_list): + def _check_type_map_consistency( + self, type_map_list: list[list[str] | None] + ) -> list[str]: ret = [] for ii in type_map_list: if ii is not None: @@ -619,7 +670,7 @@ def _check_type_map_consistency(self, type_map_list): return ret -def _format_name_length(name, width): +def _format_name_length(name: str, width: int) -> str: if len(name) <= width: return "{: >{}}".format(name, width) else: @@ -689,7 +740,7 @@ def print_summary( ) -def process_sys_probs(sys_probs, nbatch): +def process_sys_probs(sys_probs: list[float], nbatch: int) -> np.ndarray: sys_probs = np.array(sys_probs) type_filter = sys_probs >= 0 assigned_sum_prob = np.sum(type_filter * sys_probs) @@ -708,7 +759,7 @@ def process_sys_probs(sys_probs, nbatch): return ret_prob -def prob_sys_size_ext(keywords, nsystems, nbatch): +def prob_sys_size_ext(keywords: str, nsystems: int, nbatch: int) -> list[float]: block_str = keywords.split(";")[1:] block_stt = [] block_end = [] @@ -732,11 +783,12 @@ def prob_sys_size_ext(keywords, nsystems, nbatch): def process_systems( - systems: Union[str, list[str]], patterns: Optional[list[str]] = None + systems: str | list[str], patterns: list[str] | None = None ) -> list[str]: """Process the user-input systems. If it is a single directory, search for all the systems in the directory. + If it is a list, each item in the list is treated as a directory to search. Check if the systems are valid. Parameters @@ -748,21 +800,39 @@ def process_systems( Returns ------- - list of str + result_systems: list of str The valid systems """ + # Normalize input to a list of paths to search if isinstance(systems, str): + search_paths = [systems] + elif isinstance(systems, list): + search_paths = systems + else: + # Handle unsupported input types + raise ValueError( + f"Invalid systems type: {type(systems)}. Must be str or list[str]." + ) + + # Iterate over the search_paths list and apply expansion logic to each path + result_systems = [] + for path in search_paths: if patterns is None: - systems = expand_sys_str(systems) + expanded_paths = expand_sys_str(path) else: - systems = rglob_sys_str(systems, patterns) - elif isinstance(systems, list): - systems = systems.copy() - return systems + expanded_paths = rglob_sys_str(path, patterns) + + result_systems.extend(expanded_paths) + + return result_systems def get_data( - jdata: dict[str, Any], rcut, type_map, modifier, multi_task_mode=False + jdata: dict[str, Any], + rcut: float, + type_map: list[str] | None, + modifier: Any | None, + multi_task_mode: bool = False, ) -> DeepmdDataSystem: """Get the data system. diff --git a/deepmd/utils/econf_embd.py b/deepmd/utils/econf_embd.py index 5ff136b373..191e804506 100644 --- a/deepmd/utils/econf_embd.py +++ b/deepmd/utils/econf_embd.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later + import numpy as np from mendeleev import ( element, @@ -9,6 +10,7 @@ "electronic_configuration_embedding", "make_econf_embedding", "normalized_electronic_configuration_embedding", + "sort_element_type", "transform_to_spin_rep", ] @@ -181,7 +183,7 @@ ECONF_DIM = electronic_configuration_embedding[type_map[0]].shape[0] -def normalize_vec_length(res): +def normalize_vec_length(res: dict[str, np.ndarray]) -> dict[str, np.ndarray]: scale = 1.0 / np.sqrt(ECONF_DIM) return {kk: scale * vv for kk, vv in res.items()} @@ -191,7 +193,7 @@ def normalize_vec_length(res): ) -def make_empty_list_vec(): +def make_empty_list_vec() -> dict[str, np.ndarray]: ret = {} for kk in conf_keys: ll = lett_to_ln[kk[1]] @@ -199,7 +201,7 @@ def make_empty_list_vec(): return ret -def flatten_list_vec(lv): +def flatten_list_vec(lv: dict[str, np.ndarray]) -> np.ndarray: ret = np.array([], dtype=np.int32) for kk in conf_keys: ret = np.append(ret, lv[kk]) @@ -240,7 +242,7 @@ def transform_to_spin_rep(res: dict[str, np.ndarray]) -> dict[str, np.ndarray]: """Transform electron occupation of 0/1/2 to -1,-1/-1,1/1,1.""" ret = {} - def transform(ii): + def transform(ii: int) -> list[int]: if ii == 0: return [-1, -1] elif ii == 1: @@ -263,3 +265,16 @@ def print_econf_embedding(res: dict[str, np.ndarray]) -> None: vvstr = ",".join([str(ii) for ii in vv]) space = " " * (2 - len(kk)) print(f'"{kk}"{space} : [{vvstr}],') # noqa: T201 + + +def sort_element_type(elements: list[str]) -> list[str]: + """Sort element types based on their atomic number.""" + + def get_atomic_number(symbol: str) -> int | float: + try: + return element(symbol).atomic_number + except ValueError: + return float("inf") + + sorted_elements = sorted(elements, key=lambda x: get_atomic_number(x)) + return sorted_elements diff --git a/deepmd/utils/env_mat_stat.py b/deepmd/utils/env_mat_stat.py index ecc0b7b62f..b10842fb41 100644 --- a/deepmd/utils/env_mat_stat.py +++ b/deepmd/utils/env_mat_stat.py @@ -10,9 +10,6 @@ from collections.abc import ( Iterator, ) -from typing import ( - Optional, -) import numpy as np @@ -28,7 +25,7 @@ class StatItem: Parameters ---------- - number : int + number : float The total size of given array. sum : float The sum value of the matrix. @@ -36,7 +33,9 @@ class StatItem: The sum squared value of the matrix. """ - def __init__(self, number: int = 0, sum: float = 0, squared_sum: float = 0) -> None: + def __init__( + self, number: float = 0, sum: float = 0, squared_sum: float = 0 + ) -> None: self.number = number self.sum = sum self.squared_sum = squared_sum @@ -48,6 +47,13 @@ def __add__(self, other: "StatItem") -> "StatItem": squared_sum=self.squared_sum + other.squared_sum, ) + def __mul__(self, scalar: float) -> "StatItem": + return StatItem( + number=self.number * scalar, + sum=self.sum * scalar, + squared_sum=self.squared_sum * scalar, + ) + def compute_avg(self, default: float = 0) -> float: """Compute the average of the environment matrix. @@ -160,7 +166,7 @@ def load_stats(self, path: DPPath) -> None: ) def load_or_compute_stats( - self, data: list[dict[str, np.ndarray]], path: Optional[DPPath] = None + self, data: list[dict[str, np.ndarray]], path: DPPath | None = None ) -> None: """Load the statistics of the environment matrix if it exists, otherwise compute and save it. diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 644da3649d..c019cc68ab 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -36,35 +36,35 @@ def __init__( self.resuming = resuming self.update_type = self.p_type_map != self.type_map - def get_index_mapping(self): + def get_index_mapping(self) -> list[int]: """Returns the mapping index of newly defined types to those in the pretrained model.""" return get_index_between_two_maps(self.p_type_map, self.type_map)[0] - def get_has_new_type(self): + def get_has_new_type(self) -> bool: """Returns whether there are unseen types in the new type_map.""" return get_index_between_two_maps(self.p_type_map, self.type_map)[1] - def get_model_branch(self): + def get_model_branch(self) -> str: """Returns the chosen model branch.""" return self.model_branch - def get_random_fitting(self): + def get_random_fitting(self) -> bool: """Returns whether to use random fitting.""" return self.random_fitting - def get_resuming(self): + def get_resuming(self) -> bool: """Returns whether to only do resuming.""" return self.resuming - def get_update_type(self): + def get_update_type(self) -> bool: """Returns whether to update the type related params when loading from pretrained model with redundant types.""" return self.update_type - def get_pretrained_tmap(self): + def get_pretrained_tmap(self) -> list[str]: """Returns the type map in the pretrained model.""" return self.p_type_map - def get_finetune_tmap(self): + def get_finetune_tmap(self) -> list[str]: """Returns the type map in the fine-tuned model.""" return self.type_map @@ -72,7 +72,7 @@ def get_finetune_tmap(self): def get_index_between_two_maps( old_map: list[str], new_map: list[str], -): +) -> tuple[list[int], bool]: """Returns the mapping index of types in new_map to those in the old_map. Parameters @@ -110,7 +110,7 @@ def get_index_between_two_maps( def map_atom_exclude_types( atom_exclude_types: list[int], remap_index: list[int], -): +) -> list[int]: """Return the remapped atom_exclude_types according to remap_index. Parameters @@ -135,7 +135,7 @@ def map_atom_exclude_types( def map_pair_exclude_types( pair_exclude_types: list[tuple[int, int]], remap_index: list[int], -): +) -> list[tuple[int, int]]: """Return the remapped atom_exclude_types according to remap_index. Parameters diff --git a/deepmd/utils/model_branch_dict.py b/deepmd/utils/model_branch_dict.py new file mode 100644 index 0000000000..501f5287b6 --- /dev/null +++ b/deepmd/utils/model_branch_dict.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import textwrap +from collections import ( + OrderedDict, +) +from typing import ( + Any, +) + + +def get_model_dict(model_dict: dict[str, Any]) -> tuple[dict[str, str], dict[str, str]]: + """ + Get the model branch alias dictionary from the model_dict. + + Returns + ------- + model_alias_dict: dict + A dictionary where the keys are the aliases and the values are the original model branch names. + model_branch_dict: dict + A dictionary where the keys are the original model branch names, and the values are dictionaries with: + - alias + the list of aliases of this model branch. + - info + the info dictionary of this model branch. + """ + model_alias_dict = {} + model_branch_dict = {} + for key in model_dict: + model_branch_dict[key] = {} + model_alias_dict[key] = key + alias_list = model_dict[key].get("model_branch_alias", []) + model_branch_dict[key]["alias"] = alias_list + branch_info = model_dict[key].get("info", {}) + model_branch_dict[key]["info"] = branch_info + for alias in alias_list: + assert alias not in model_alias_dict, ( + f"Alias {alias} for model_branch {key} already exists in model_branch {model_alias_dict[alias]}!" + ) + model_alias_dict[alias] = key + + return model_alias_dict, model_branch_dict + + +# generated with GPT for formatted print +class OrderedDictTableWrapper: + """ + A wrapper for pretty-printing an OrderedDict that has a specific structure. + + Expected structure: + OrderedDict({ + "BranchName1": { + "alias": ["A", "B"], # Required key: alias (list of strings) + "info": { # Optional key: info (dict of arbitrary key-value pairs) + "description": "Some text", + "description2": "Some long text..." + } + }, + "BranchName2": { + "alias": ["C"], + "info": { "owner": "Alice" } + }, + ... + }) + + Features: + - Prints the data as an ASCII table with borders and aligned columns. + - The first two columns are fixed: "Model Branch Name" and "Alias". + - The remaining columns are all unique keys found in `info` across all branches (order preserved by first occurrence). + - Long text in cells is automatically wrapped to fit the column width, except column 1 & 2 auto-expanding to the **maximum content length** in that column. + - Missing info values are shown as empty strings. + """ + + def __init__( + self, data: "OrderedDict[str, dict[str, Any]]", col_width: int = 30 + ) -> None: + """ + Initialize the table wrapper. + + Args: + data: OrderedDict containing the branch data. + col_width: Maximum width of each column (characters). Longer text will wrap. + """ + # Ensure we are working with an OrderedDict to preserve branch order + if not isinstance(data, OrderedDict): + data = OrderedDict(data) + self.data = data + self.col_width = col_width + + # Collect all unique keys from "info" across all branches in order of first appearance + seen = set() + self.info_keys: list[str] = [] + for _, payload in self.data.items(): + info = payload.get("info") or {} + for k in info.keys(): + if k not in seen: + seen.add(k) + self.info_keys.append(k) + + # Construct table header: fixed columns + dynamic info keys + self.headers: list[str] = ["Model Branch", "Alias", *self.info_keys] + + def _wrap_cell(self, text: Any, width: int | None = None) -> list[str]: + """ + Convert a cell value into a list of wrapped text lines. + + Args: + text: Any value that will be converted to a string. + width: Optional custom wrap width. If None, defaults to `self.col_width`. + + Returns + ------- + A list of strings, each representing one wrapped line of the cell. + """ + text = "" if text is None else str(text) + eff_width = self.col_width if width is None else width + # If eff_width is very large, this effectively disables wrapping for that cell. + return textwrap.wrap(text, eff_width) or [""] + + def as_table(self) -> str: + """ + Generate a formatted ASCII table with borders and aligned columns. + + Returns + ------- + A string representation of the table. + """ + # Step 0: Precompute dynamic widths for the first two columns. + # Column 0 (branch): width = max length over header + all branch names + branch_col_width = len(self.headers[0]) # "Model Branch Name" + for branch in self.data.keys(): + branch_col_width = max(branch_col_width, len(str(branch))) + + # Column 1 (alias): width = max length over header + all alias strings (joined by ", \n") + alias_col_width = len(self.headers[1]) # "Alias" + for payload in self.data.values(): + alias_list = payload.get("alias", []) + for alias in alias_list: + alias_col_width = max(alias_col_width, len(str(alias))) + + # Step 1: Create raw rows (without wrapping) + raw_rows: list[list[str]] = [] + # First row: header + raw_rows.append(self.headers) + + # Data rows + for branch, payload in self.data.items(): + alias_str = ", ".join(map(str, payload.get("alias", []))) + info = payload.get("info") or {} + row = [branch, alias_str] + [info.get(k, "") for k in self.info_keys] + raw_rows.append(row) + + # Step 2: Wrap each cell, using dynamic widths for the first two columns, + # and fixed `self.col_width` for info columns. + wrapped_rows: list[list[list[str]]] = [] + for row in raw_rows: + wrapped_row: list[list[str]] = [] + for j, cell in enumerate(row): + if j == 0: + # First column: branch name -> no wrap by using its max width + wrapped_row.append(self._wrap_cell(cell, width=branch_col_width)) + elif j == 1: + # Second column: alias -> no wrap by using its max width + wrapped_row.append(self._wrap_cell(cell, width=alias_col_width)) + else: + # Info columns: keep using fixed col_width (wrapping allowed) + wrapped_row.append(self._wrap_cell(cell)) + wrapped_rows.append(wrapped_row) + + # Step 3: Determine actual width for each column + # For the first two columns, we already decided the exact widths above. + col_widths: list[int] = [] + for idx, col in enumerate(zip(*wrapped_rows)): + if idx == 0: + col_widths.append(branch_col_width) + elif idx == 1: + col_widths.append(alias_col_width) + else: + # Info columns: width is the maximum wrapped line length (<= self.col_width) + col_widths.append(max(len(line) for cell in col for line in cell)) + + # Helper: Draw a horizontal separator line + def draw_separator() -> str: + return "+" + "+".join("-" * (w + 2) for w in col_widths) + "+" + + # Helper: Draw one row of text parts (single lines per cell) + def draw_row_line(cells_parts: list[list[str]]) -> str: + return ( + "| " + + " | ".join( + part.ljust(width) for part, width in zip(cells_parts, col_widths) + ) + + " |" + ) + + # Step 4: Build the table string + table_lines = [] + table_lines.append(draw_separator()) + + for i, row_cells in enumerate(wrapped_rows): + # Determine the maximum number of wrapped lines in this row + max_lines = max(len(cell) for cell in row_cells) + # Draw each wrapped line + for line_idx in range(max_lines): + line_parts = [ + cell[line_idx] if line_idx < len(cell) else "" for cell in row_cells + ] + table_lines.append(draw_row_line(line_parts)) + table_lines.append(draw_separator()) + + return "\n".join(table_lines) + + +# Example usage +if __name__ == "__main__": + data = OrderedDict( + { + "Omat": { + "alias": ["Default", "Materials"], + "info": { + "observed-type": [ + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Ac", + ], + "description": "OMat24 is a large-scale open dataset containing over 110 million DFT calculations " + "spanning diverse structures and compositions. It is designed to support AI-driven " + "materials discovery by providing broad and deep coverage of chemical space.", + }, + }, + } + ) + + wrapper = OrderedDictTableWrapper(data, col_width=20) + print(wrapper.as_table()) # noqa:T201 diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index d2cc918b64..8061c7aa9c 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -2,11 +2,14 @@ from collections import ( defaultdict, ) +from typing import ( + Any, +) import numpy as np -def _make_all_stat_ref(data, nbatches): +def _make_all_stat_ref(data: Any, nbatches: int) -> dict[str, list[Any]]: all_stat = defaultdict(list) for ii in range(data.get_nsystems()): for jj in range(nbatches): @@ -18,7 +21,9 @@ def _make_all_stat_ref(data, nbatches): return all_stat -def make_stat_input(data, nbatches, merge_sys=True): +def make_stat_input( + data: Any, nbatches: int, merge_sys: bool = True +) -> dict[str, list[Any]]: """Pack data for statistics. Parameters @@ -57,7 +62,7 @@ def make_stat_input(data, nbatches, merge_sys=True): return all_stat -def merge_sys_stat(all_stat): +def merge_sys_stat(all_stat: dict[str, list[Any]]) -> dict[str, list[Any]]: first_key = next(iter(all_stat.keys())) nsys = len(all_stat[first_key]) ret = defaultdict(list) diff --git a/deepmd/utils/out_stat.py b/deepmd/utils/out_stat.py index ecbd379e2d..6aabf4a4f4 100644 --- a/deepmd/utils/out_stat.py +++ b/deepmd/utils/out_stat.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Output statistics.""" -from typing import ( - Optional, -) - import numpy as np from deepmd.env import ( @@ -15,8 +11,8 @@ def compute_stats_from_redu( output_redu: np.ndarray, natoms: np.ndarray, - assigned_bias: Optional[np.ndarray] = None, - rcond: Optional[float] = None, + assigned_bias: np.ndarray | None = None, + rcond: float | None = None, ) -> tuple[np.ndarray, np.ndarray]: """Compute the output statistics. @@ -135,7 +131,7 @@ def compute_stats_from_atomic( def compute_stats_do_not_distinguish_types( output_redu: np.ndarray, natoms: np.ndarray, - assigned_bias: Optional[np.ndarray] = None, + assigned_bias: np.ndarray | None = None, intensive: bool = False, ) -> tuple[np.ndarray, np.ndarray]: """Compute element-independent statistics for property fitting. diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 93c8b7a1f9..503f721c98 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from typing import ( - Optional, + Any, ) import numpy as np @@ -34,12 +34,12 @@ class PairTab: cutoff raduis for the tabulated potential """ - def __init__(self, filename: str, rcut: Optional[float] = None) -> None: + def __init__(self, filename: str, rcut: float | None = None) -> None: """Constructor.""" self.data_type = np.float64 self.reinit(filename, rcut) - def reinit(self, filename: str, rcut: Optional[float] = None) -> None: + def reinit(self, filename: str, rcut: float | None = None) -> None: """Initialize the tabulated interaction. Parameters @@ -95,7 +95,7 @@ def serialize(self) -> dict: } @classmethod - def deserialize(cls, data) -> "PairTab": + def deserialize(cls, data: dict[str, Any]) -> "PairTab": data = data.copy() check_version_compatibility(data.pop("@version", 1), 1, 1) data.pop("@class") @@ -257,7 +257,7 @@ def _extrapolate_table(self, pad_extrapolation: np.array) -> np.array: ) return pad_extrapolation - def _make_data(self): + def _make_data(self) -> np.ndarray: data = np.zeros( [self.ntypes * self.ntypes * 4 * self.nspline], dtype=self.data_type ) diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index 87a44aa70d..5f156298e3 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -12,9 +12,8 @@ Path, ) from typing import ( + Any, ClassVar, - Optional, - Union, ) import h5py @@ -35,7 +34,7 @@ class DPPath(ABC): mode, by default "r" """ - def __new__(cls, path: str, mode: str = "r"): + def __new__(cls, path: str, mode: str = "r") -> "DPPath": if cls is DPPath: if os.path.isdir(path): return super().__new__(DPOSPath) @@ -56,7 +55,7 @@ def load_numpy(self) -> np.ndarray: """ @abstractmethod - def load_txt(self, **kwargs) -> np.ndarray: + def load_txt(self, **kwargs: Any) -> np.ndarray: """Load NumPy array from text. Returns @@ -115,7 +114,7 @@ def is_dir(self) -> bool: """Check if self is directory.""" @abstractmethod - def __getnewargs__(self): + def __getnewargs__(self) -> tuple[str, str]: """Return the arguments to be passed to __new__ when unpickling an instance.""" @abstractmethod @@ -133,10 +132,10 @@ def __str__(self) -> str: def __repr__(self) -> str: return f"{type(self)} ({self!s})" - def __eq__(self, other) -> bool: + def __eq__(self, other: object) -> bool: return str(self) == str(other) - def __hash__(self): + def __hash__(self) -> int: return hash(str(self)) @property @@ -168,12 +167,12 @@ class DPOSPath(DPPath): mode, by default "r" """ - def __init__(self, path: Union[str, Path], mode: str = "r") -> None: + def __init__(self, path: str | Path, mode: str = "r") -> None: super().__init__() self.mode = mode self.path = Path(path) - def __getnewargs__(self): + def __getnewargs__(self) -> tuple[str, str]: return (self.path, self.mode) def load_numpy(self) -> np.ndarray: @@ -186,7 +185,7 @@ def load_numpy(self) -> np.ndarray: """ return np.load(str(self.path)) - def load_txt(self, **kwargs) -> np.ndarray: + def load_txt(self, **kwargs: Any) -> np.ndarray: """Load NumPy array from text. Returns @@ -311,7 +310,7 @@ def __init__(self, path: str, mode: str = "r") -> None: # h5 path: default is the root path self._name = s[1] if len(s) > 1 else "/" - def __getnewargs__(self): + def __getnewargs__(self) -> tuple[str, str]: return (self.root_path, self.mode) @classmethod @@ -341,7 +340,7 @@ def load_numpy(self) -> np.ndarray: """ return self.root[self._name][:] - def load_txt(self, dtype: Optional[np.dtype] = None, **kwargs) -> np.ndarray: + def load_txt(self, dtype: np.dtype | None = None, **kwargs: Any) -> np.ndarray: """Load NumPy array from text. Returns @@ -416,7 +415,7 @@ def _keys(self) -> list[str]: __file_new_keys: ClassVar[dict[h5py.File, list[str]]] = {} @property - def _new_keys(self): + def _new_keys(self) -> list[str]: """New keys that haven't been cached.""" self.__file_new_keys.setdefault(self.root, []) return self.__file_new_keys[self.root] diff --git a/deepmd/utils/plugin.py b/deepmd/utils/plugin.py index 37ff784d61..e1817327a3 100644 --- a/deepmd/utils/plugin.py +++ b/deepmd/utils/plugin.py @@ -6,9 +6,11 @@ from abc import ( ABCMeta, ) -from typing import ( +from collections.abc import ( Callable, - Optional, +) +from typing import ( + Any, ) @@ -32,7 +34,7 @@ def xxx(): def __init__(self) -> None: self.plugins = {} - def __add__(self, other) -> "Plugin": + def __add__(self, other: "Plugin") -> "Plugin": self.plugins.update(other.plugins) return self @@ -56,7 +58,7 @@ def decorator(object: object) -> object: return decorator - def get_plugin(self, key) -> object: + def get_plugin(self, key: str) -> object: """Visit a plugin by key. Parameters @@ -73,7 +75,7 @@ def get_plugin(self, key) -> object: class VariantMeta: - def __call__(self, *args, **kwargs): + def __call__(self, *args: Any, **kwargs: Any) -> Any: """Remove `type` and keys that starts with underline.""" obj = self.__new__(self, *args, **kwargs) kwargs.pop("type", None) @@ -97,7 +99,7 @@ class PluginVariant(metaclass=VariantABCMeta): pass -def make_plugin_registry(name: Optional[str] = None) -> type[object]: +def make_plugin_registry(name: str | None = None) -> type[object]: """Make a plugin registry. Parameters diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 10ebdf0790..8710693607 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -1,8 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, - Union, -) import numpy as np @@ -10,11 +6,11 @@ def choice( - a: Union[np.ndarray, int], - size: Optional[Union[int, tuple[int, ...]]] = None, + a: np.ndarray | int, + size: int | tuple[int, ...] | None = None, replace: bool = True, - p: Optional[np.ndarray] = None, -): + p: np.ndarray | None = None, +) -> np.ndarray | int: """Generates a random sample from a given 1-D array. Parameters @@ -40,7 +36,9 @@ def choice( return _RANDOM_GENERATOR.choice(a, size=size, replace=replace, p=p) -def random(size=None): +def random( + size: int | tuple[int, ...] | None = None, +) -> float | np.ndarray: """Return random floats in the half-open interval [0.0, 1.0). Parameters @@ -56,7 +54,7 @@ def random(size=None): return _RANDOM_GENERATOR.random_sample(size) -def seed(val: Optional[Union[int, list[int]]] = None) -> None: +def seed(val: int | list[int] | None = None) -> None: """Seed the generator. Parameters diff --git a/deepmd/utils/spin.py b/deepmd/utils/spin.py index c6cbb32f52..aed82cae8b 100644 --- a/deepmd/utils/spin.py +++ b/deepmd/utils/spin.py @@ -1,8 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy -from typing import ( - Union, -) import numpy as np @@ -38,7 +35,7 @@ class Spin: def __init__( self, use_spin: list[bool], - virtual_scale: Union[list[float], float], + virtual_scale: list[float] | float, ) -> None: type_dtype = np.int32 self.ntypes_real = len(use_spin) @@ -135,7 +132,9 @@ def init_atom_exclude_types_placeholder(self) -> None: """ self.atom_exclude_types_p = self.placeholder_type.tolist() - def get_pair_exclude_types(self, exclude_types=None) -> list[tuple[int, int]]: + def get_pair_exclude_types( + self, exclude_types: list[tuple[int, int]] | None = None + ) -> list[tuple[int, int]]: """ Return the pair-wise exclusion types for descriptor. The placeholder types for those without spin are excluded. @@ -151,7 +150,9 @@ def get_pair_exclude_types(self, exclude_types=None) -> list[tuple[int, int]]: _exclude_types.append((tt[0], tt[1])) return _exclude_types - def get_atom_exclude_types(self, exclude_types=None) -> list[int]: + def get_atom_exclude_types( + self, exclude_types: list[int] | None = None + ) -> list[int]: """ Return the atom-wise exclusion types for fitting before out_def. Both the placeholder types and spin types are excluded. @@ -164,7 +165,9 @@ def get_atom_exclude_types(self, exclude_types=None) -> list[int]: _exclude_types = list(set(_exclude_types)) return _exclude_types - def get_atom_exclude_types_placeholder(self, exclude_types=None) -> list[int]: + def get_atom_exclude_types_placeholder( + self, exclude_types: list[int] | None = None + ) -> list[int]: """ Return the atom-wise exclusion types for fitting after out_def. The placeholder types for those without spin are excluded. @@ -177,14 +180,14 @@ def get_atom_exclude_types_placeholder(self, exclude_types=None) -> list[int]: _exclude_types = list(set(_exclude_types)) return _exclude_types - def get_spin_mask(self): + def get_spin_mask(self) -> np.ndarray: """ Return the spin mask of shape [ntypes], with spin types being 1, and non-spin types being 0. """ return self.spin_mask - def get_virtual_scale_mask(self): + def get_virtual_scale_mask(self) -> np.ndarray: """ Return the virtual scale mask of shape [ntypes], with spin types being its virtual scale, and non-spin types being 0. diff --git a/deepmd/utils/summary.py b/deepmd/utils/summary.py index 51171ccc19..c00e6deb9e 100644 --- a/deepmd/utils/summary.py +++ b/deepmd/utils/summary.py @@ -56,7 +56,7 @@ class SummaryPrinter(ABC): "build variant": GLOBAL_CONFIG["dp_variant"], } - def __call__(self): + def __call__(self) -> None: """Print build and current running cluster configuration summary.""" nodename, nodelist = get_host_names() build_info = self.BUILD.copy() diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index 6e5e988b4a..fb40f798e2 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -7,6 +7,9 @@ from functools import ( lru_cache, ) +from typing import ( + Any, +) import numpy as np from scipy.special import ( @@ -21,11 +24,11 @@ class BaseTabulate(ABC): def __init__( self, - descrpt, - neuron, - type_one_side, - exclude_types, - is_pt, + descrpt: Any, + neuron: list[int], + type_one_side: bool, + exclude_types: list[list[int]], + is_pt: bool, ) -> None: """Constructor.""" super().__init__() @@ -193,6 +196,48 @@ def build( nspline[ii][0] if self.is_pt else nspline[ii], ) idx += 1 + elif self.descrpt_type == "T_TEBD": + # 1. Find the global range [ll, uu] of cos(theta) across all types + uu = np.max(upper) + ll = np.min(lower) + + # 2. Create a unique input grid xx for this shared geometric network based on the global range + xx = np.arange(extrapolate * ll, ll, stride1, dtype=self.data_type) + xx = np.append( + xx, + np.arange(ll, uu, stride0, dtype=self.data_type), + ) + xx = np.append( + xx, + np.arange( + uu, + extrapolate * uu, + stride1, + dtype=self.data_type, + ), + ) + xx = np.append(xx, np.array([extrapolate * uu], dtype=self.data_type)) + + # 3. Calculate the number of spline points + nspline = ( + (uu - ll) / stride0 + + ((extrapolate * uu - uu) / stride1) + + ((ll - extrapolate * ll) / stride1) + ).astype(int) + + # 4. Call _build_lower only once to generate the table for this shared network + geometric_net_name = "filter_net" + self._build_lower( + geometric_net_name, + xx, + 0, + uu, + ll, + stride0, + stride1, + extrapolate, + nspline, + ) elif self.descrpt_type == "R": for ii in range(self.table_size): if (self.type_one_side and not self._all_excluded(ii)) or ( @@ -238,8 +283,18 @@ def build( self._convert_numpy_float_to_int() return self.lower, self.upper + # generate_spline_table def _build_lower( - self, net, xx, idx, upper, lower, stride0, stride1, extrapolate, nspline + self, + net: str, + xx: np.ndarray, + idx: int, + upper: float, + lower: float, + stride0: int, + stride1: int, + extrapolate: bool, + nspline: int, ) -> None: vv, dd, d2 = self._make_data(xx, idx) self.data[net] = np.zeros( @@ -247,21 +302,14 @@ def _build_lower( ) # tt.shape: [nspline, self.last_layer_size] - if self.descrpt_type in ("Atten", "A", "AEbdV2"): + if self.descrpt_type in ("Atten", "A", "AEbdV2", "R"): tt = np.full((nspline, self.last_layer_size), stride1) # pylint: disable=no-explicit-dtype tt[: int((upper - lower) / stride0), :] = stride0 - elif self.descrpt_type == "T": - tt = np.full((nspline, self.last_layer_size), stride1) # pylint: disable=no-explicit-dtype - tt[ - int((lower - extrapolate * lower) / stride1) + 1 : ( - int((lower - extrapolate * lower) / stride1) - + int((upper - lower) / stride0) - ), - :, - ] = stride0 - elif self.descrpt_type == "R": + elif self.descrpt_type in ("T", "T_TEBD"): tt = np.full((nspline, self.last_layer_size), stride1) # pylint: disable=no-explicit-dtype - tt[: int((upper - lower) / stride0), :] = stride0 + start_index = int((lower - extrapolate * lower) / stride1) + 1 + end_index = start_index + int((upper - lower) / stride0) + tt[start_index:end_index, :] = stride0 else: raise RuntimeError("Unsupported descriptor") @@ -334,7 +382,9 @@ def _build_lower( self.lower[net] = lower @abstractmethod - def _make_data(self, xx, idx) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + def _make_data( + self, xx: np.ndarray, idx: int + ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Generate tabulation data for the given input. Parameters @@ -368,18 +418,18 @@ def _all_excluded(self, ii: int) -> bool: return all((ii, type_i) in self.exclude_types for type_i in range(self.ntypes)) @abstractmethod - def _get_descrpt_type(self): + def _get_descrpt_type(self) -> str: """Get the descrpt type.""" pass @abstractmethod - def _get_layer_size(self): + def _get_layer_size(self) -> int: """Get the number of embedding layer.""" pass - def _get_table_size(self): + def _get_table_size(self) -> int: table_size = 0 - if self.descrpt_type in ("Atten", "AEbdV2"): + if self.descrpt_type in ("Atten", "AEbdV2", "T_TEBD"): table_size = 1 elif self.descrpt_type == "A": table_size = self.ntypes * self.ntypes @@ -395,30 +445,30 @@ def _get_table_size(self): raise RuntimeError("Unsupported descriptor") return table_size - def _get_data_type(self): + def _get_data_type(self) -> type | None: for item in self.matrix["layer_" + str(self.layer_size)]: if len(item) != 0: return type(item[0][0]) return None - def _get_last_layer_size(self): + def _get_last_layer_size(self) -> int: for item in self.matrix["layer_" + str(self.layer_size)]: if len(item) != 0: return item.shape[1] return 0 @abstractmethod - def _get_bias(self): + def _get_bias(self) -> dict[str, Any]: """Get bias of embedding net.""" pass @abstractmethod - def _get_matrix(self): + def _get_matrix(self) -> dict[str, Any]: """Get weight matrx of embedding net.""" pass @abstractmethod - def _convert_numpy_to_tensor(self): + def _convert_numpy_to_tensor(self) -> None: """Convert self.data from np.ndarray to torch.Tensor.""" pass @@ -427,13 +477,13 @@ def _convert_numpy_float_to_int(self) -> None: self.lower = {k: int(v) for k, v in self.lower.items()} self.upper = {k: int(v) for k, v in self.upper.items()} - def _get_env_mat_range(self, min_nbor_dist): + def _get_env_mat_range(self, min_nbor_dist: float) -> tuple[np.ndarray, np.ndarray]: """Change the embedding net range to sw / min_nbor_dist.""" sw = self._spline5_switch(min_nbor_dist, self.rcut_smth, self.rcut) if self.descrpt_type in ("Atten", "A", "AEbdV2"): lower = -self.davg[:, 0] / self.dstd[:, 0] upper = ((1 / min_nbor_dist) * sw - self.davg[:, 0]) / self.dstd[:, 0] - elif self.descrpt_type == "T": + elif self.descrpt_type in ("T", "T_TEBD"): var = np.square(sw / (min_nbor_dist * self.dstd[:, 1:4])) lower = np.min(-var, axis=1) upper = np.max(var, axis=1) @@ -447,7 +497,7 @@ def _get_env_mat_range(self, min_nbor_dist): # returns element-wise lower and upper return np.floor(lower), np.ceil(upper) - def _spline5_switch(self, xx, rmin, rmax): + def _spline5_switch(self, xx: float, rmin: float, rmax: float) -> float: if xx < rmin: vv = 1 elif xx < rmax: diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index c9213de699..616b69c2ca 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -5,8 +5,7 @@ abstractmethod, ) from typing import ( - Optional, - Union, + Any, ) from deepmd.utils.data_system import ( @@ -25,9 +24,9 @@ class BaseUpdateSel(ABC): def update_one_sel( self, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, rcut: float, - sel: Union[int, list[int], str], + sel: int | list[int] | str, mixed_type: bool = False, ) -> tuple[float, list[int]]: min_nbor_dist, tmp_sel = self.get_nbor_stat( @@ -55,7 +54,7 @@ def update_one_sel( ) return min_nbor_dist, sel - def parse_auto_sel(self, sel) -> bool: + def parse_auto_sel(self, sel: Any) -> bool: if not isinstance(sel, str): return False words = sel.split(":") @@ -64,7 +63,7 @@ def parse_auto_sel(self, sel) -> bool: else: return False - def parse_auto_sel_ratio(self, sel): + def parse_auto_sel_ratio(self, sel: Any) -> float: if not self.parse_auto_sel(sel): raise RuntimeError(f"invalid auto sel format {sel}") else: @@ -77,16 +76,16 @@ def parse_auto_sel_ratio(self, sel): raise RuntimeError(f"invalid auto sel format {sel}") return ratio - def wrap_up_4(self, xx): + def wrap_up_4(self, xx: int) -> int: return 4 * ((int(xx) + 3) // 4) def get_nbor_stat( self, train_data: DeepmdDataSystem, - type_map: Optional[list[str]], + type_map: list[str] | None, rcut: float, mixed_type: bool = False, - ) -> tuple[float, Union[int, list[int]]]: + ) -> tuple[float, int | list[int]]: """Get the neighbor statistics of the data. Parameters @@ -131,7 +130,7 @@ def neighbor_stat(self) -> type[NeighborStat]: def get_min_nbor_dist( self, train_data: DeepmdDataSystem, - ): + ) -> float: min_nbor_dist, _ = self.get_nbor_stat( train_data, None, # type_map doesn't affect min_nbor_dist diff --git a/doc/backend.md b/doc/backend.md index cab8dc2128..b2f7dc4826 100644 --- a/doc/backend.md +++ b/doc/backend.md @@ -12,7 +12,7 @@ In the documentation, TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon - Model filename extension: `.pb` - Checkpoint filename extension: `.meta`, `.index`, `.data-00000-of-00001` -[TensorFlow](https://tensorflow.org) 2.7 or above is required, since NumPy 1.21 or above is required. +[TensorFlow](https://tensorflow.org) 2.8 is the first version to support Python 3.10. DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`tf.compat.v1`) in the graph mode. ### PyTorch {{ pytorch_icon }} @@ -28,7 +28,7 @@ While `.pth` and `.pt` are the same in the PyTorch package, they have different - Model filename extension: `.xlo`, `.savedmodel` - Checkpoint filename extension: `.jax` -[JAX](https://jax.readthedocs.io/) 0.4.33 (which requires Python 3.10 or above) or above is required. +[JAX](https://jax.readthedocs.io/) 0.4.33 or above is required. Both `.xlo` and `.jax` are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions. `.savedmodel` is the TensorFlow [SavedModel format](https://www.tensorflow.org/guide/saved_model) generated by [JAX2TF](https://www.tensorflow.org/guide/jax2tf), which needs the installation of TensorFlow. Only the `.savedmodel` format supports C++ inference, which needs the TensorFlow C++ interface. diff --git a/doc/development/coding-conventions.rst b/doc/development/coding-conventions.rst index 4f82b34a60..80d8f915d1 100644 --- a/doc/development/coding-conventions.rst +++ b/doc/development/coding-conventions.rst @@ -30,7 +30,7 @@ Rules ----- The code must be compatible with the oldest supported version of python -which is 3.9. +which is 3.10. The project follows the generic coding conventions as specified in the `Style Guide for Python Code`_, `Docstring diff --git a/doc/env.md b/doc/env.md index 4ca7101236..1688e0af9c 100644 --- a/doc/env.md +++ b/doc/env.md @@ -88,5 +88,37 @@ These environment variables also apply to third-party programs using the C++ int **Type**: List of paths, split by `:` on Unix and `;` on Windows List of customized OP plugin libraries to load, such as `/path/to/plugin1.so:/path/to/plugin2.so` on Linux and `/path/to/plugin1.dll;/path/to/plugin2.dll` on Windows. +::: + +:::{envvar} DP_PROFILER + +{{ pytorch_icon }} Enable the built-in PyTorch Kineto profiler for the PyTorch C++ (inference) backend. + +**Type**: string (output file stem) + +**Default**: unset (disabled) + +When set to a non-empty value, profiling is enabled for the lifetime of the loaded PyTorch model (e.g. during LAMMPS runs). A JSON trace file is created on finish. The final file name is constructed as: + +- `_gpu.json` if running on GPU +- `.json` if running on CPU + +The trace can be examined with [Chrome trace viewer](https://ui.perfetto.dev/) (alternatively chrome://tracing). It includes: + +- CPU operator activities +- CUDA activities (if available) + +Example: + +```bash +export DP_PROFILER=result +mpirun -np 4 lmp -in in.lammps +# Produces result_gpuX.json, where X is the GPU id used by each MPI rank. +``` + +Tips: + +- Large runs can generate sizable JSON files; consider limiting numbers of MD steps, like 20. +- Currently this feature only supports single process, or multi-process runs where each process uses a distinct GPU on the same node. ::: diff --git a/doc/environment.yml b/doc/environment.yml index 85d5a97c5b..fc96b56485 100644 --- a/doc/environment.yml +++ b/doc/environment.yml @@ -4,7 +4,7 @@ channels: - defaults dependencies: - doxygen>=1.9.1 - - python=3.9 + - python=3.10 - pip>=20.1 - pip: - ..[docs,cpu,torch] diff --git a/doc/getting-started/quick_start.ipynb b/doc/getting-started/quick_start.ipynb index 2ddb6a195e..9af2d456bf 100644 --- a/doc/getting-started/quick_start.ipynb +++ b/doc/getting-started/quick_start.ipynb @@ -89,7 +89,7 @@ "\n", "\"Fig1\"\n", "\n", - "For more detailed usage, you can refer to the [DeePMD-kit’s documentation](https://docs.deepmodeling.org/projects/deepmd/en/master/index.html) as a comprehensive reference.\n", + "For more detailed usage, you can refer to the [DeePMD-kit’s documentation](https://docs.deepmodeling.com/projects/deepmd/en/master/index.html) as a comprehensive reference.\n", "\n", "In this case, the Deep Potential (DP) model was generated using the **DeePMD-kit package**.\n", "\n" @@ -638,7 +638,7 @@ "id": "7b0edb0f-df47-4e6c-8c37-5f32c4bd6b39", "metadata": {}, "source": [ - "More detailed docs about Data conversion can be found [here](https://docs.deepmodeling.org/projects/deepmd/en/master/data/data-conv.html)" + "More detailed docs about Data conversion can be found [here](https://docs.deepmodeling.com/projects/deepmd/en/master/data/data-conv.html)" ] }, { diff --git a/doc/index.rst b/doc/index.rst index 38ff2fe97e..238dc0d25d 100644 --- a/doc/index.rst +++ b/doc/index.rst @@ -16,6 +16,7 @@ DeePMD-kit is a package written in Python/C++, designed to minimize the effort r - Han Wang, Linfeng Zhang, Jiequn Han, and Weinan E. "DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics." Computer Physics Communications 228 (2018): 178-184. - Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang. "DeePMD-kit v2: A software package for Deep Potential models." J. Chem. Phys., 159, 054801 (2023). + - Jinzhe Zeng, Duo Zhang, Anyang Peng, Xiangyu Zhang, Sensen He, Yan Wang, Xinzijian Liu, Hangrui Bi, Yifan Li, Chun Cai, Chengqian Zhang, Yiming Du, Jia-Xin Zhu, Pinghui Mo, Zhengtao Huang, Qiyu Zeng, Shaochen Shi, Xuejian Qin, Zhaoxi Yu, Chenxing Luo, Ye Ding, Yun-Pei Liu, Ruosong Shi, Zhenyu Wang, Sigbjørn Løland Bore, Junhan Chang, Zhe Deng, Zhaohan Ding, Siyuan Han, Wanrun Jiang, Guolin Ke, Zhaoqing Liu, Denghui Lu, Koki Muraoka, Hananeh Oliaei, Anurag Kumar Singh, Haohui Que, Weihong Xu, Zhangmancang Xu, Yong-Bin Zhuang, Jiayu Dai, Timothy J. Giese, Weile Jia, Ben Xu, Darrin M. York, Linfeng Zhang, Han Wang. "DeePMD-kit v3: A Multiple-Backend Framework for Machine Learning Potentials." J. Chem. Theory Comput. 21 (2025): 4375-4385. In addition, please follow :ref:`this page ` to cite the methods you used. diff --git a/doc/inference/python.md b/doc/inference/python.md index b2603c85f8..361db7b64f 100644 --- a/doc/inference/python.md +++ b/doc/inference/python.md @@ -19,6 +19,21 @@ e, f, v = dp.eval(coord, cell, atype) where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively. +One can also evaluate the descriptors of the model: + +```python +from deepmd.infer import DeepPot +import numpy as np + +dp = DeepPot("graph.pb") +coord = np.array([[1, 0, 0], [0, 0, 1.5], [1, 0, 3]]).reshape([1, -1]) +cell = np.diag(10 * np.ones(3)).reshape([1, -1]) +atype = [1, 0, 1] +descriptors = dp.eval_descriptor(coord, cell, atype) +``` + +where `descriptors` is the descriptor matrix of the system. This can also be done using the command line interface `dp eval-desc` as described in the [test documentation](../test/test.md). + Furthermore, one can use the python interface to calculate model deviation. ```python diff --git a/doc/install/easy-install-dev.md b/doc/install/easy-install-dev.md index 54309a8582..9a0154320e 100644 --- a/doc/install/easy-install-dev.md +++ b/doc/install/easy-install-dev.md @@ -1,19 +1,17 @@ # Easy install the latest development version -DeePMD-kit is actively developed in the `devel` branch. The documentation of the [`latest`](https://docs.deepmodeling.com/projects/deepmd/en/latest/) version matches the `devel` branch. +DeePMD-kit is actively developed in the `master` branch. The documentation of the [`latest`](https://docs.deepmodeling.com/projects/deepmd/en/latest/) version matches the `master` branch. The following is the way to install the pre-compiled packages without [building from source](./install-from-source.md). All of them are built with [GitHub Actions](../development/cicd.md). ## Install with docker -The [`devel` tag](https://github.com/deepmodeling/deepmd-kit/pkgs/container/deepmd-kit/131827568?tag=devel) is used to mark the latest development version with CUDA 12.2 support: +The `master` tag is used to mark the latest development version with CUDA 12.2 support: ```bash -docker pull ghcr.io/deepmodeling/deepmd-kit:devel +docker pull ghcr.io/deepmodeling/deepmd-kit:master ``` -For CUDA 11.8 support, use the `devel_cu11` tag. - ## Install with pip Follow [the documentation for the stable version](easy-install.md#install-python-interface-with-pip), but add `--pre` and `--extra-index-url` options like below: @@ -28,10 +26,10 @@ pip install -U --pre deepmd-kit[gpu,cu12,lmp,torch] --extra-index-url https://de **Supported backends**: TensorFlow {{ tensorflow_icon }} ::: -The [pre-comiled C library](./install-from-c-library.md) can be downloaded from [here](https://nightly.link/deepmodeling/deepmd-kit/workflows/package_c/devel/libdeepmd_c-0-libdeepmd_c.tar.gz.zip), or via a shell command: +The [pre-compiled C library](./install-from-c-library.md) can be downloaded from [here](https://nightly.link/deepmodeling/deepmd-kit/workflows/package_c/master/libdeepmd_c-0-libdeepmd_c.tar.gz.zip), or via a shell command: ```sh -wget https://nightly.link/deepmodeling/deepmd-kit/workflows/package_c/devel/libdeepmd_c-0-libdeepmd_c.tar.gz.zip && unzip libdeepmd_c-0-libdeepmd_c.tar.gz.zip +wget https://nightly.link/deepmodeling/deepmd-kit/workflows/package_c/master/libdeepmd_c-0-libdeepmd_c.tar.gz.zip && unzip libdeepmd_c-0-libdeepmd_c.tar.gz.zip ``` ## Pre-release conda-forge packages diff --git a/doc/install/easy-install.md b/doc/install/easy-install.md index d28c0d0773..119342add9 100644 --- a/doc/install/easy-install.md +++ b/doc/install/easy-install.md @@ -10,7 +10,7 @@ You can refer to [DeepModeling conda FAQ](https://docs.deepmodeling.com/faq/cond ::: :::{note} -Python 3.9 or above is required for Python interface. +Python 3.10 or above is required for Python interface. ::: - [Install off-line packages](#install-off-line-packages) @@ -96,14 +96,6 @@ pip install deepmd-kit[gpu,cu12] :::: -::::{tab-item} CUDA 11 - -```bash -pip install deepmd-kit-cu11[gpu,cu11] -``` - -:::: - ::::{tab-item} CPU ```bash @@ -128,15 +120,6 @@ pip install deepmd-kit[torch] :::: -::::{tab-item} CUDA 11.8 - -```bash -pip install torch --index-url https://download.pytorch.org/whl/cu118 -pip install deepmd-kit-cu11 -``` - -:::: - ::::{tab-item} CPU ```bash @@ -185,16 +168,10 @@ Switch to the TensorFlow {{ tensorflow_icon }} tab for more information. ::::{tab-item} CUDA 12.6 ```bash -pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/ -pip install deepmd-kit -``` - -:::: - -::::{tab-item} CUDA 11.8 - -```bash -pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/ +# release version +pip install paddlepaddle-gpu==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/ +# nightly-build version +# pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/ pip install deepmd-kit ``` @@ -203,7 +180,10 @@ pip install deepmd-kit ::::{tab-item} CPU ```bash -pip install paddlepaddle==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/ +# release version +pip install paddlepaddle==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/ +# nightly-build version +# pip install --pre paddlepaddle -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/ pip install deepmd-kit ``` @@ -227,4 +207,4 @@ If your platform is not supported, or you want to build against the installed ba pip install deepmd-kit[gpu,cu12,lmp,ipi] ``` -MPICH is required for parallel running. +MPICH will be installed automatically - you do not need to install a MPI library by yourself. diff --git a/doc/install/install-from-c-library.md b/doc/install/install-from-c-library.md index 806be51ca9..4568cdb6c9 100644 --- a/doc/install/install-from-c-library.md +++ b/doc/install/install-from-c-library.md @@ -12,7 +12,7 @@ wget https://github.com/deepmodeling/deepmd-kit/releases/latest/download/libdeep tar xzf libdeepmd_c.tar.gz ``` -The library is built in Linux (GLIBC 2.17) with CUDA 12.2 (`libdeepmd_c.tar.gz`) or 11.8 (`libdeepmd_c_cu11.tar.gz`). It's noted that this package does not contain CUDA Toolkit and cuDNN, so one needs to download them from the NVIDIA website. +The library is built in Linux (GLIBC 2.17) with CUDA 12.2 (`libdeepmd_c.tar.gz`). It's noted that this package does not contain CUDA Toolkit and cuDNN, so one needs to download them from the NVIDIA website. ## Use Pre-compiled C Library to build the LAMMPS plugin, i-PI driver, and GROMACS patch diff --git a/doc/install/install-from-source.md b/doc/install/install-from-source.md index 1dc72c51fa..1e03563c66 100644 --- a/doc/install/install-from-source.md +++ b/doc/install/install-from-source.md @@ -1,6 +1,6 @@ # Install from source code -Please follow our [GitHub](https://github.com/deepmodeling/deepmd-kit) webpage to download the [latest released version](https://github.com/deepmodeling/deepmd-kit/tree/master) and [development version](https://github.com/deepmodeling/deepmd-kit/tree/devel). +Please follow our [GitHub](https://github.com/deepmodeling/deepmd-kit) webpage to download the source code of a specific version or the [development version](https://github.com/deepmodeling/deepmd-kit/tree/master). Or get the DeePMD-kit source code by `git clone` @@ -21,7 +21,7 @@ deepmd_source_dir=`pwd` ### Install Backend's Python interface First, check the Python version on your machine. -Python 3.9 or above is required. +Python 3.10 or above is required. ```bash python --version @@ -52,7 +52,7 @@ If one does not need the GPU support of DeePMD-kit and is concerned about packag pip install --upgrade tensorflow-cpu ``` -One can also [use conda](https://docs.deepmodeling.org/faq/conda.html) to install TensorFlow from [conda-forge](https://conda-forge.org). +One can also [use conda](https://docs.deepmodeling.com/faq/conda.html) to install TensorFlow from [conda-forge](https://conda-forge.org). To verify the installation, run @@ -74,7 +74,7 @@ pip install torch Follow [PyTorch documentation](https://pytorch.org/get-started/locally/) to install PyTorch built against different CUDA versions or without CUDA. -One can also [use conda](https://docs.deepmodeling.org/faq/conda.html) to install PyTorch from [conda-forge](https://conda-forge.org). +One can also [use conda](https://docs.deepmodeling.com/faq/conda.html) to install PyTorch from [conda-forge](https://conda-forge.org). ::: @@ -89,7 +89,7 @@ pip install jax-ai-stack One can also install packages in JAX AI Stack manually. Follow [JAX documentation](https://jax.readthedocs.io/en/latest/installation.html) to install JAX built against different CUDA versions or without CUDA. -One can also [use conda](https://docs.deepmodeling.org/faq/conda.html) to install JAX from [conda-forge](https://conda-forge.org). +One can also [use conda](https://docs.deepmodeling.com/faq/conda.html) to install JAX from [conda-forge](https://conda-forge.org). ::: @@ -99,11 +99,16 @@ To install Paddle, run ```sh # cu126 -pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/ -# cu118 -pip install paddlepaddle-gpu==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/ +# release version +pip install paddlepaddle-gpu==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cu126/ +# nightly-build version +# pip install --pre paddlepaddle-gpu -i https://www.paddlepaddle.org.cn/packages/nightly/cu126/ + # cpu -pip install paddlepaddle==3.0.0 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/ +# release version +pip install paddlepaddle==3.1.1 -i https://www.paddlepaddle.org.cn/packages/stable/cpu/ +# nightly-build version +# pip install --pre paddlepaddle -i https://www.paddlepaddle.org.cn/packages/nightly/cpu/ ``` ::: @@ -125,7 +130,7 @@ deactivate If one has multiple python interpreters named something like python3.x, it can be specified by, for example ```bash -virtualenv -p python3.9 $deepmd_venv +virtualenv -p python3.12 $deepmd_venv ``` One should remember to activate the virtual environment every time he/she uses DeePMD-kit. @@ -262,6 +267,12 @@ It will print the help information like ### Install horovod and mpi4py {{ tensorflow_icon }} +:::{warning} +Horovod has not released a new version for a long time. +As of December 2025, the latest Horovod release does not support the latest TensorFlow versions. +You can check the patches required to support the latest TensorFlow at [conda-forge/horovod-feedstock](https://github.com/conda-forge/horovod-feedstock/blob/main/recipe/meta.yaml). +::: + [Horovod](https://github.com/horovod/horovod) and [mpi4py](https://github.com/mpi4py/mpi4py) are used for parallel training. For better performance on GPU, please follow the tuning steps in [Horovod on GPU](https://github.com/horovod/horovod/blob/master/docs/gpus.rst). ```bash @@ -344,9 +355,7 @@ download the TensorFlow C library from [this page](https://www.tensorflow.org/in If you want to use C++ interface of Paddle, you need to compile the Paddle inference library(C++ interface) manually from the [linux-compile-by-make](https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/install/compile/linux-compile-by-make.html), then use the `.so` and `.a` files in `Paddle/build/paddle_inference_install_dir/`. -We also provide a weekly-build Paddle C++ inference library for Linux x86_64 with CUDA 11.8/12.3/CPU below: - -CUDA 11.8: [Cuda118_cudnn860_Trt8531_D1/latest/paddle_inference.tgz](https://paddle-qa.bj.bcebos.com/paddle-pipeline/GITHUB_Docker_Compile_Test_Cuda118_cudnn860_Trt8531_D1/latest/paddle_inference.tgz) +We also provide a weekly-build Paddle C++ inference library for Linux x86_64 with CUDA 12.3/CPU below: CUDA 12.3: [Cuda123_cudnn900_Trt8616_D1/latest/paddle_inference.tgz](https://paddle-qa.bj.bcebos.com/paddle-pipeline/GITHUB_Docker_Compile_Test_Cuda123_cudnn900_Trt8616_D1/latest/paddle_inference.tgz) @@ -366,7 +375,7 @@ mkdir build cd build ``` -The installation requires CMake 3.16 or later for the CPU version, CMake 3.23 or later for the CUDA support, and CMake 3.21 or later for the ROCM support. One can install CMake via `pip` if it is not installed or the installed version does not satisfy the requirement: +The installation requires CMake 3.25.2 or later for all platforms (CPU, CUDA, and ROCM). One can install CMake via `pip` if it is not installed or the installed version does not satisfy the requirement: ```sh pip install -U cmake @@ -374,7 +383,7 @@ pip install -U cmake You must enable at least one backend. If you enable two or more backends, these backend libraries must be built in a compatible way, e.g. using the same `_GLIBCXX_USE_CXX11_ABI` flag. -We recommend using [conda packages](https://docs.deepmodeling.org/faq/conda.html) from [conda-forge](https://conda-forge.org), which are usually compatible to each other. +We recommend using [conda packages](https://docs.deepmodeling.com/faq/conda.html) from [conda-forge](https://conda-forge.org), which are usually compatible to each other. ::::{tab-set} diff --git a/doc/install/install-lammps.md b/doc/install/install-lammps.md index 00b887e9c3..cb65188002 100644 --- a/doc/install/install-lammps.md +++ b/doc/install/install-lammps.md @@ -17,11 +17,11 @@ DeePMD-kit will generate a module called `USER-DEEPMD` in the `build` directory, ```bash cd /some/workspace -wget https://github.com/lammps/lammps/archive/stable_29Aug2024_update1.tar.gz -tar xf stable_29Aug2024_update1.tar.gz +wget https://github.com/lammps/lammps/archive/stable_22Jul2025_update2.tar.gz +tar xf stable_22Jul2025_update2.tar.gz ``` -The source code of LAMMPS is stored in the directory `lammps-stable_29Aug2024_update1`. +The source code of LAMMPS is stored in the directory `lammps-stable_22Jul2025_update2`. Then, you can [build LAMMPS](https://docs.lammps.org/Build.html) with either make or CMake. @@ -30,7 +30,7 @@ Then, you can [build LAMMPS](https://docs.lammps.org/Build.html) with either mak Now go into the LAMMPS code and copy the DeePMD-kit module like this ```bash -cd lammps-stable_29Aug2024_update1/src/ +cd lammps-stable_22Jul2025_update2/src/ cp -r $deepmd_source_dir/source/build/USER-DEEPMD . make yes-kspace make yes-extra-fix @@ -60,8 +60,8 @@ make no-user-deepmd Now go into the LAMMPS directory and create a directory called `build`: ```bash -mkdir -p lammps-stable_29Aug2024_update1/build/ -cd lammps-stable_29Aug2024_update1/build/ +mkdir -p lammps-stable_22Jul2025_update2/build/ +cd lammps-stable_22Jul2025_update2/build/ ``` Patch the LAMMPS `CMakeLists.txt` file: @@ -94,15 +94,15 @@ Now download the LAMMPS code (`8Apr2021` or later), and uncompress it: ```bash cd /some/workspace -wget https://github.com/lammps/lammps/archive/stable_29Aug2024_update1.tar.gz -tar xf stable_29Aug2024_update1.tar.gz +wget https://github.com/lammps/lammps/archive/stable_22Jul2025_update2.tar.gz +tar xf stable_22Jul2025_update2.tar.gz ``` -The source code of LAMMPS is stored in the directory `lammps-stable_29Aug2024_update1`. The directory of the source code should be specified as the CMAKE argument `LAMMPS_SOURCE_ROOT` during installation of the DeePMD-kit C++ interface. Now go into the LAMMPS directory and create a directory called `build` +The source code of LAMMPS is stored in the directory `lammps-stable_22Jul2025_update2`. The directory of the source code should be specified as the CMAKE argument `LAMMPS_SOURCE_ROOT` during installation of the DeePMD-kit C++ interface. Now go into the LAMMPS directory and create a directory called `build` ```bash -mkdir -p lammps-stable_29Aug2024_update1/build/ -cd lammps-stable_29Aug2024_update1/build/ +mkdir -p lammps-stable_22Jul2025_update2/build/ +cd lammps-stable_22Jul2025_update2/build/ ``` Now build LAMMPS. Note that `PLUGIN` must be enabled, and `BUILD_SHARED_LIBS` must be set to `yes`. You can install any other package you want. diff --git a/doc/install/install-tf.1.12.md b/doc/install/install-tf.1.12.md deleted file mode 100644 index 13abd8f7a7..0000000000 --- a/doc/install/install-tf.1.12.md +++ /dev/null @@ -1,131 +0,0 @@ -# Install TensorFlow's C++ interface - -The TensorFlow's C++ interface will be compiled from the source code. Firstly one installs bazel. It is highly recommended that the bazel version 0.15.0 is used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.15.0/bazel-0.15.0-dist.zip -mkdir bazel-0.15.0 -cd bazel-0.15.0 -unzip ../bazel-0.15.0-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.12.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by CMake, so we need to compile and integrate TensorFlow with CMake projects. The rest of this section follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -```bash -./configure -``` - -You will answer a list of questions that help configure the building of TensorFlow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` with the virtual environment directory): - -```bash -Please specify the location of python. [Default is $tensorflow_venv/bin/python]: -``` - -The library path for Python should be set accordingly. - -Now build the shared library of TensorFlow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue for your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Before moving on, we need to compile the dependencies of TensorFlow, including Protobuf, Eigen, nsync and absl. Firstly, protobuf - -```bash -mkdir /tmp/proto -sed -i 's;PROTOBUF_URL=.*;PROTOBUF_URL=\"https://mirror.bazel.build/github.com/google/protobuf/archive/v3.6.0.tar.gz\";g' tensorflow/contrib/makefile/download_dependencies.sh -tensorflow/contrib/makefile/download_dependencies.sh -cd tensorflow/contrib/makefile/downloads/protobuf/ -./autogen.sh -./configure --prefix=/tmp/proto/ -make -make install -``` - -Then Eigen - -```bash -mkdir /tmp/eigen -cd ../eigen -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/eigen/ ../ -make install -``` - -nsync - -```bash -mkdir /tmp/nsync -cd ../../nsync -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/nsync/ ../ -make -make install -``` - -And absl - -```bash -cd ../../absl -bazel build -mkdir -p $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --exclude '*' absl $tensorflow_root/include/ -cd ../../../../.. -``` - -Now, copy the libraries to the tensorflow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp bazel-bin/tensorflow/libtensorflow_cc.so $tensorflow_root/lib/ -cp bazel-bin/tensorflow/libtensorflow_framework.so $tensorflow_root/lib/ -cp /tmp/proto/lib/libprotobuf.a $tensorflow_root/lib/ -cp /tmp/nsync/lib64/libnsync.a $tensorflow_root/lib/ -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r /tmp/proto/include/* $tensorflow_root/include -cp -r /tmp/eigen/include/eigen3/* $tensorflow_root/include -cp -r /tmp/nsync/include/*h $tensorflow_root/include -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` - -The temporary installation directories for the dependencies can be removed: - -```bash -rm -fr /tmp/proto /tmp/eigen /tmp/nsync -``` diff --git a/doc/install/install-tf.1.14-gpu.md b/doc/install/install-tf.1.14-gpu.md deleted file mode 100644 index 5850af24ba..0000000000 --- a/doc/install/install-tf.1.14-gpu.md +++ /dev/null @@ -1,164 +0,0 @@ -# Install TensorFlow-GPU's C++ interface - -TensorFlow's C++ interface will be compiled from the source code. Firstly one installs Bazel. It is highly recommended that the Bazel version 0.24.1 is used. Full instructions on Bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.24.1/bazel-0.24.1-dist.zip -mkdir bazel-0.24.1 -cd bazel-0.24.1 -unzip ../bazel-0.24.1-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.14.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by CMake, so we need to compile and integrate TensorFlow with CMake projects. The rest of this section follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -You will answer a list of questions that help configure the building of TensorFlow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` with the virtual environment directory): - -```bash -./configure -Please specify the location of python. [Default is xxx]: - -Traceback (most recent call last): - File "", line 1, in -AttributeError: module 'site' has no attribute 'getsitepackages' -Found possible Python library paths: - /xxx/deepmd_gpu/tensorflow_venv/lib/python3.7/site-packages -Please input the desired Python library path to use. Default is [xxx] - -Do you wish to build TensorFlow with XLA JIT support? [Y/n]: -XLA JIT support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with OpenCL SYCL support? [y/N]: -No OpenCL SYCL support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with ROCm support? [y/N]: -No ROCm support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with CUDA support? [y/N]: y -CUDA support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with TensorRT support? [y/N]: -No TensorRT support will be enabled for TensorFlow. - -Found CUDA 10.1 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include -Found cuDNN 7 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include - -Please specify a list of comma-separated CUDA compute capabilities you want to build with. -You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. -Please note that each additional compute capability significantly increases your build time and binary size, and that TensorFlow only supports compute capabilities >= 3.5 [Default is: 6.1,6.1]: - -Do you want to use clang as CUDA compiler? [y/N]: -nvcc will be used as CUDA compiler. - -Please specify which gcc should be used by nvcc as the host compiler. [Default is /usr/bin/gcc]: - - -Do you wish to build TensorFlow with MPI support? [y/N]: -No MPI support will be enabled for TensorFlow. - -Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native -Wno-sign-compare]: - -Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]: -Not configuring the WORKSPACE for Android builds. - -Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details. - --config=mkl # Build with MKL support. - --config=monolithic # Config for mostly static monolithic build. - --config=gdr # Build with GDR support. - --config=verbs # Build with libverbs support. - --config=ngraph # Build with Intel nGraph support. - --config=numa # Build with NUMA support. - --config=dynamic_kernels # (Experimental) Build kernels into separate shared objects. - --config=v2 # Build TensorFlow 2.x instead of 1.x. -Preconfigured Bazel build configs to DISABLE default on features: - --config=noaws # Disable AWS S3 filesystem support. - --config=nogcp # Disable GCP support. - --config=nohdfs # Disable HDFS support. - --config=noignite # Disable Apache Ignite support. - --config=nokafka # Disable Apache Kafka support. - --config=nonccl # Disable NVIDIA NCCL support. -Configuration finished -``` - -The library path for Python should be set accordingly. - -Now build the shared library of TensorFlow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue for your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Now, copy the libraries to the TensorFlow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp -d bazel-bin/tensorflow/libtensorflow_cc.so* $tensorflow_root/lib/ -cp -d bazel-bin/tensorflow/libtensorflow_framework.so* $tensorflow_root/lib/ -cp -d $tensorflow_root/lib/libtensorflow_framework.so.1 $tensorflow_root/lib/libtensorflow_framework.so -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/Eigen/ $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/unsupported/ $tensorflow_root/include -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/protobuf_archive/src/ $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_absl/absl/ $tensorflow_root/include/absl -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` - -# Troubleshooting - -```bash -git: unknown command -C ... -``` - -This may be your git version issue because the low version of Git does not support this command. Upgrading your Git may be helpful. - -```bash -CMake Error: The following variables are used in this project, but they are set to NOTFOUND. -Please set them or make sure they are set and tested correctly in the CMake files: -FFTW_LIB (ADVANCED) - linked by target "FFTW" in directory xxx -``` - -Currently, when building the Eigen package, you can delete the FFTW in the CMake file. - -```bash -fatal error: absl/numeric/int128_have_intrinsic.inc: No such file or directory -``` - -Basically, you could build an empty file named "int128_have_intrinsic.inc" in the same directory of "int128.h". diff --git a/doc/install/install-tf.1.14.md b/doc/install/install-tf.1.14.md deleted file mode 100644 index 6457d484ad..0000000000 --- a/doc/install/install-tf.1.14.md +++ /dev/null @@ -1,79 +0,0 @@ -# Install tensorflow's C++ interface - -The tensorflow's C++ interface will be compiled from the source code. Firstly one installs bazel. It is highly recommended that the bazel version 0.24.1 is used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.24.1/bazel-0.24.1-dist.zip -mkdir bazel-0.24.1 -cd bazel-0.24.1 -unzip ../bazel-0.24.1-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the tensorflow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.14.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by cmake, so we need to compile and integrate tensorflow with cmake projects. The rest of this section basically follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -```bash -./configure -``` - -You will answer a list of questions that help configure the building of tensorflow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` by the virtual environment directory): - -```bash -Please specify the location of python. [Default is $tensorflow_venv/bin/python]: -``` - -The library path for Python should be set accordingly. - -Now build the shared library of tensorflow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue of your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install tensorflow in directory `$tensorflow_root`. Create the directory if it does not exists - -```bash -mkdir -p $tensorflow_root -``` - -Now, copy the libraries to the tensorflow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp -d bazel-bin/tensorflow/libtensorflow_cc.so* $tensorflow_root/lib/ -cp -d bazel-bin/tensorflow/libtensorflow_framework.so* $tensorflow_root/lib/ -cp -d $tensorflow_root/lib/libtensorflow_framework.so.1 $tensorflow_root/lib/libtensorflow_framework.so -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/Eigen/ $tensorflow_root/include -cp -r bazel-tensorflow/external/eigen_archive/unsupported/ $tensorflow_root/include -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/protobuf_archive/src/ $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_absl/absl/ $tensorflow_root/include/absl -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` diff --git a/doc/install/install-tf.1.8.md b/doc/install/install-tf.1.8.md deleted file mode 100644 index f9554f9348..0000000000 --- a/doc/install/install-tf.1.8.md +++ /dev/null @@ -1,121 +0,0 @@ -# Install tensorflow's C++ interface - -The tensorflow's C++ interface will be compiled from the source code. Firstly one installs bazel. It is highly recommended that the bazel version 0.10.0 is used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/0.10.0/bazel-0.10.0-dist.zip -mkdir bazel-0.10.0 -cd bazel-0.10.0 -unzip ../bazel-0.10.0-dist.zip -./compile.sh -export PATH=`pwd`/output:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -cd /some/workspace -git clone https://github.com/tensorflow/tensorflow tensorflow -b v1.8.0 --depth=1 -cd tensorflow -``` - -DeePMD-kit is compiled by CMake, so we need to compile and integrate TensorFlow with CMake projects. The rest of this section basically follows [the instruction provided by Tuatini](http://tuatini.me/building-tensorflow-as-a-standalone-project/). Now execute - -```bash -./configure -``` - -You will answer a list of questions that help configure the building of TensorFlow. It is recommended to build for Python3. You may want to answer the question like this (please replace `$tensorflow_venv` with the virtual environment directory): - -```bash -Please specify the location of python. [Default is $tensorflow_venv/bin/python]: -``` - -The library path for Python should be set accordingly. - -Now build the shared library of TensorFlow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue of your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Before moving on, we need to compile the dependencies of TensorFlow, including Protobuf, Eigen and nsync. Firstly, protobuf - -```bash -mkdir /tmp/proto -tensorflow/contrib/makefile/download_dependencies.sh -cd tensorflow/contrib/makefile/downloads/protobuf/ -./autogen.sh -./configure --prefix=/tmp/proto/ -make -make install -``` - -Then Eigen - -```bash -mkdir /tmp/eigen -cd ../eigen -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/eigen/ ../ -make install -``` - -And nsync - -```bash -mkdir /tmp/nsync -cd ../../nsync -mkdir build_dir -cd build_dir -cmake -DCMAKE_INSTALL_PREFIX=/tmp/nsync/ ../ -make -make install -cd ../../../../../.. -``` - -Now, copy the libraries to the TensorFlow's installation directory: - -```bash -mkdir $tensorflow_root/lib -cp bazel-bin/tensorflow/libtensorflow_cc.so $tensorflow_root/lib/ -cp bazel-bin/tensorflow/libtensorflow_framework.so $tensorflow_root/lib/ -cp /tmp/proto/lib/libprotobuf.a $tensorflow_root/lib/ -cp /tmp/nsync/lib/libnsync.a $tensorflow_root/lib/ -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -cp -r bazel-genfiles/* $tensorflow_root/include/ -cp -r tensorflow/cc $tensorflow_root/include/tensorflow -cp -r tensorflow/core $tensorflow_root/include/tensorflow -cp -r third_party $tensorflow_root/include -cp -r /tmp/proto/include/* $tensorflow_root/include -cp -r /tmp/eigen/include/eigen3/* $tensorflow_root/include -cp -r /tmp/nsync/include/*h $tensorflow_root/include -``` - -Now clean up the source files in the header directories: - -```bash -cd $tensorflow_root/include -find . -name "*.cc" -type f -delete -``` - -The temporary installation directories for the dependencies can be removed: - -```bash -rm -fr /tmp/proto /tmp/eigen /tmp/nsync -``` diff --git a/doc/install/install-tf.2.3.md b/doc/install/install-tf.2.3.md deleted file mode 100644 index 2fc7b35f2c..0000000000 --- a/doc/install/install-tf.2.3.md +++ /dev/null @@ -1,123 +0,0 @@ -# Install TensorFlow's C++ interface - -The tensorflow's C++ interface will be compiled from the source code. Firstly one installs bazel. The bazel version 3.1.0 should be used. A full instruction of bazel installation can be found [here](https://docs.bazel.build/versions/master/install.html). - -```bash -cd /some/workspace -wget https://github.com/bazelbuild/bazel/releases/download/3.1.0/bazel-3.1.0-installer-linux-x86_64.sh -chmod +x bazel-3.1.0-installer-linux-x86_64.sh -./bazel-3.1.0-installer-linux-x86_64.sh --prefix /some/workspace/bazel -export PATH=/some/workspace/bazel/bin:$PATH -``` - -Firstly get the source code of the TensorFlow - -```bash -git clone https://github.com/tensorflow/tensorflow tensorflow -b v2.3.0 --depth=1 -cd tensorflow -./configure -``` - -You will answer a list of questions that help configure the building of TensorFlow. You may want to answer the question like the following. If you do not want to add CUDA support, please answer no. - -``` -Please specify the location of python. [Default is xxx]: - -Found possible Python library paths: - xxx -Please input the desired Python library path to use. Default is [xxx] - -Do you wish to build TensorFlow with OpenCL SYCL support? [y/N]: -No OpenCL SYCL support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with ROCm support? [y/N]: -No ROCm support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with CUDA support? [y/N]: y -CUDA support will be enabled for TensorFlow. - -Do you wish to build TensorFlow with TensorRT support? [y/N]: -No TensorRT support will be enabled for TensorFlow. - -Found CUDA 10.2 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include -Found cuDNN 7 in: - /usr/local/cuda/lib64 - /usr/local/cuda/include - -Please specify a list of comma-separated CUDA compute capabilities you want to build with. -You can find the compute capability of your device at: https://developer.nvidia.com/cuda-gpus. -Please note that each additional compute capability significantly increases your build time and binary size, and that TensorFlow only supports compute capabilities >= 3.5 [Default is: 7.5,7.5]: - -Do you want to use clang as CUDA compiler? [y/N]: -nvcc will be used as CUDA compiler. - -Please specify which gcc should be used by nvcc as the host compiler. [Default is /usr/bin/gcc]: - -Please specify optimization flags to use during compilation when bazel option "--config=opt" is specified [Default is -march=native -Wno-sign-compare]: - -Would you like to interactively configure ./WORKSPACE for Android builds? [y/N]: -Not configuring the WORKSPACE for Android builds. - -Preconfigured Bazel build configs. You can use any of the below by adding "--config=<>" to your build command. See .bazelrc for more details. - --config=mkl # Build with MKL support. - --config=monolithic # Config for mostly static monolithic build. - --config=ngraph # Build with Intel nGraph support. - --config=numa # Build with NUMA support. - --config=dynamic_kernels # (Experimental) Build kernels into separate shared objects. - --config=v2 # Build TensorFlow 2.x instead of 1.x. -Preconfigured Bazel build configs to DISABLE default on features: - --config=noaws # Disable AWS S3 filesystem support. - --config=nogcp # Disable GCP support. - --config=nohdfs # Disable HDFS support. - --config=nonccl # Disable NVIDIA NCCL support. -Configuration finished -``` - -The library path for Python should be set accordingly. - -Now build the shared library of tensorflow: - -```bash -bazel build -c opt --verbose_failures //tensorflow:libtensorflow_cc.so -``` - -You may want to add options `--copt=-msse4.2`, `--copt=-mavx`, `--copt=-mavx2` and `--copt=-mfma` to enable SSE4.2, AVX, AVX2 and FMA SIMD accelerations, respectively. It is noted that these options should be chosen according to the CPU architecture. If the RAM becomes an issue of your machine, you may limit the RAM usage by using `--local_resources 2048,.5,1.0`. - -Now I assume you want to install TensorFlow in directory `$tensorflow_root`. Create the directory if it does not exist - -```bash -mkdir -p $tensorflow_root -``` - -Now, copy the libraries to the tensorflow's installation directory: - -```bash -mkdir -p $tensorflow_root/lib -cp -d bazel-bin/tensorflow/libtensorflow_cc.so* $tensorflow_root/lib/ -cp -d bazel-bin/tensorflow/libtensorflow_framework.so* $tensorflow_root/lib/ -cp -d $tensorflow_root/lib/libtensorflow_framework.so.2 $tensorflow_root/lib/libtensorflow_framework.so -``` - -Then copy the headers - -```bash -mkdir -p $tensorflow_root/include/tensorflow -rsync -avzh --exclude '_virtual_includes/' --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-bin/ $tensorflow_root/include/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' tensorflow/cc $tensorflow_root/include/tensorflow/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' tensorflow/core $tensorflow_root/include/tensorflow/ -rsync -avzh --include '*/' --include '*' --exclude '*.cc' third_party/ $tensorflow_root/include/third_party/ -rsync -avzh --include '*/' --include '*' --exclude '*.txt' bazel-tensorflow/external/eigen_archive/Eigen/ $tensorflow_root/include/Eigen/ -rsync -avzh --include '*/' --include '*' --exclude '*.txt' bazel-tensorflow/external/eigen_archive/unsupported/ $tensorflow_root/include/unsupported/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_protobuf/src/google/ $tensorflow_root/include/google/ -rsync -avzh --include '*/' --include '*.h' --include '*.inc' --exclude '*' bazel-tensorflow/external/com_google_absl/absl/ $tensorflow_root/include/absl/ -``` - -# Troubleshooting - -```bash -git: unknown command -C ... -``` - -This may be an issue with your git version issue. Early versions of git do not support this command, in this case upgrading your git to a newer version may resolve any issues. diff --git a/doc/model/change-bias.md b/doc/model/change-bias.md index ac28201cb6..2a9b098606 100644 --- a/doc/model/change-bias.md +++ b/doc/model/change-bias.md @@ -1,7 +1,7 @@ -# Change the model output bias for trained model {{ pytorch_icon }} +# Change the model output bias for trained model {{ tensorflow_icon }} {{ pytorch_icon }} :::{note} -**Supported backends**: PyTorch {{ pytorch_icon }} +**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }} ::: The output bias of a trained model typically originates from the statistical results of the training dataset. @@ -10,32 +10,45 @@ There are several scenarios where one might want to adjust the output bias after such as zero-shot testing (similar to the procedure before the first step in fine-tuning) or manually setting the output bias. -The `dp --pt change-bias` command supports the following methods for adjusting the bias: +The `dp change-bias` command supports the following methods for adjusting the bias: ::::{tab-set} -:::{tab-item} Changing bias using provided systems for trained `.pt`/`.pth` models: +:::{tab-item} TensorFlow Backend {{ tensorflow_icon }} + +**Changing bias using provided systems for trained checkpoint:** ```sh -dp --pt change-bias model.pt -s data_dir -o model_updated.pt +dp --tf change-bias model.ckpt -s data_dir -o model_updated.pb ``` -For multitask models, where `--model-branch` must be specified: +**Changing bias using user input for energy model:** ```sh -dp --pt change-bias multi_model.pt -s data_dir -o model_updated.pt --model-branch model_1 +dp --tf change-bias model.ckpt -b -92.523 -187.66 -o model_updated.pb ``` ::: -:::{tab-item} Changing bias using user input for **energy model**: +:::{tab-item} PyTorch Backend {{ pytorch_icon }} + +**Changing bias using provided systems for trained `.pt`/`.pth` models:** + +```sh +dp --pt change-bias model.pt -s data_dir -o model_updated.pt +``` + +**Changing bias using user input for energy model:** ```sh dp --pt change-bias model.pt -b -92.523 -187.66 -o model_updated.pt ``` -Here, `-b` specifies user-defined energy bias for each type, separated by space, -in an order consistent with the `type_map` in the model. +For multitask models, where `--model-branch` must be specified: + +```sh +dp --pt change-bias multi_model.pt -s data_dir -o model_updated.pt --model-branch model_1 +``` ::: diff --git a/doc/model/dpa2.md b/doc/model/dpa2.md index c8e60c514a..466a4de4f2 100644 --- a/doc/model/dpa2.md +++ b/doc/model/dpa2.md @@ -38,6 +38,10 @@ Type embedding is within this descriptor with the {ref}`tebd_dim ` is `strip`, but only the `repinit` part is compressed. +Model compression is supported when {ref}`repinit/tebd_input_mode ` is `strip`. + +- If {ref}`repinit/attn_layer ` is `0`, both the type embedding and geometric parts inside `repinit` are compressed. +- If `repinit/attn_layer` is not `0`, only the type embedding tables are compressed and the geometric attention layers remain as neural networks. + An example is given in `examples/water/dpa2/input_torch_compressible.json`. The performance improvement will be limited if other parts are more expensive. diff --git a/doc/model/dpa3.md b/doc/model/dpa3.md index c63b26f90a..0ff46c438f 100644 --- a/doc/model/dpa3.md +++ b/doc/model/dpa3.md @@ -1,4 +1,4 @@ -# Descriptor DPA3 {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }} +# Descriptor DPA3 {{ pytorch_icon }} {{ jax_icon }} {{ paddle_icon }} {{ dpmodel_icon }} :::{note} **Supported backends**: PyTorch {{ pytorch_icon }}, JAX {{ jax_icon }}, DP {{ dpmodel_icon }} @@ -40,7 +40,11 @@ Virial RMSEs were averaged exclusively for systems containing virial labels (`Al Note that we set `float32` in all DPA3 models, while `float64` in other models by default. -## Requirements of installation from source code {{ pytorch_icon }} +## Requirements of installation from source code {{ pytorch_icon }} {{ paddle_icon }} + +::::{tab-set} + +:::{tab-item} PyTorch {{ pytorch_icon }} To run the DPA3 model on LAMMPS via source code installation (users can skip this step if using [easy installation](../install/easy-install.md)), @@ -53,6 +57,25 @@ If one runs LAMMPS with MPI, the customized OP library for the C++ interface sho If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA, otherwise the communication between GPU cards falls back to the slower CPU implementation. +::: + +:::{tab-item} Paddle {{ paddle_icon }} + +The customized OP library for the Python interface can be installed by + +```sh +cd deepmd-kit/source/op/pd +python setup.py install +``` + +If one runs LAMMPS with MPI, the customized OP library for the C++ interface should be compiled against the same MPI library as the runtime MPI. +If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a [CUDA-Aware MPI](https://developer.nvidia.com/mpi-solutions-gpus) library and CUDA, +otherwise the communication between GPU cards falls back to the slower CPU implementation. + +::: + +:::: + ## Limitations of the JAX backend with LAMMPS {{ jax_icon }} When using the JAX backend, 2 or more MPI ranks are not supported. One must set `map` to `yes` using the [`atom_modify`](https://docs.lammps.org/atom_modify.html) command. diff --git a/doc/model/dplr.md b/doc/model/dplr.md index cf071d4029..61327bb55e 100644 --- a/doc/model/dplr.md +++ b/doc/model/dplr.md @@ -69,7 +69,7 @@ The loss section is provided as follows }, ``` -so that the atomic dipole is trained as labels. Note that the NumPy compressed file `atomic_dipole.npy` should be provided in each dataset. +so that the atomic dipole is trained as labels. Note that the NumPy compressed file `atomic_dipole.npy` should be provided in each dataset. In the context of DPLR models, the atomic dipole data represents the displacement vector from each atom to its associated Wannier centroid (WC), which can be calculated as `atomic_dipole = wannier_centroid_position - atom_position` from DFT calculations using tools such as VASP with Wannier90. The training and freezing can be started from the example directory by diff --git a/doc/model/pairtab.md b/doc/model/pairtab.md index f52ad5dae7..57fe23f5e9 100644 --- a/doc/model/pairtab.md +++ b/doc/model/pairtab.md @@ -103,3 +103,10 @@ To combine with a pairwise potential, use the [linear model](./linear.md): ``` The {ref}`rcut ` can be larger than that of the DP model. + +:::{note} +The above example shows a example of combining D3 dispersion. +However, it is more efficient to train a model using plain DFT calculations without the dispersion correction, and add the dispersion correction during the simulation via the LAMMPS [`pair_style dispersion/d3` command](https://docs.lammps.org/pair_dispersion_d3.html#pair-style-dispersion-d3-command). +Training against data with dispersion directly is discouraged. +See the [D3 dispersion section](../third-party/lammps-command.md#d3-dispersion) for details. +::: diff --git a/doc/model/show-model-info.md b/doc/model/show-model-info.md index 48aa41385d..67d82610de 100644 --- a/doc/model/show-model-info.md +++ b/doc/model/show-model-info.md @@ -11,12 +11,12 @@ dp --pt show - ``: Path to the model checkpoint file or frozen model file. - ``: One or more information categories to display. Supported values are: - - `model-branch`: Shows available branches for multi-task models. - `type-map`: Shows the type mapping used by the model. - `descriptor`: Displays the model descriptor parameters. - `fitting-net`: Displays parameters of the fitting network. - `size`: (Supported Backends: PyTorch and PaddlePaddle) Shows the parameter counts for various components. + - `observed-type`: (Supported Backends: PyTorch) Shows the observed types (elements) of the model during data statistics. Only energy models are supported now. ## Example Usage @@ -33,33 +33,32 @@ dp show frozen_model.pth type-map descriptor fitting-net size Depending on the provided attributes and the model type, the output includes: - **Model Type** - - Logs whether the loaded model is a _singletask_ or _multitask_ model. - **model-branch** - - _Only available for multitask models._ - Lists all available model branches and the special `"RANDOM"` branch, which refers to a randomly initialized fitting net. - **type-map** - - For multitask models: Shows the type map for each branch. - For singletask models: Shows the model's type map. - **descriptor** - - For multitask models: Displays the descriptor parameter for each branch. - For singletask models: Displays the descriptor parameter. - **fitting-net** - - For multitask models: Shows the fitting network parameters for each branch. - For singletask models: Shows the fitting network parameters. - **size** - - Prints the number of parameters for each component (`descriptor`, `fitting-net`, etc.), as well as the total parameter count. +- **observed-type** + - Displays the count and list of observed element types of the model during data statistics. + - For multitask models, it shows the observed types for each branch. + - Note: This info shows the types observed during training data statistics, which may differ from the type map. + ## Example Output For a singletask model, the output might look like: @@ -73,6 +72,9 @@ Parameter counts: Parameters in descriptor: 19,350 Parameters in fitting-net: 119,091 Parameters in total: 138,441 +The observed types for this model: +Number of observed types: 2 +Observed types: ['H', 'O'] ``` For a multitask model, if `model-branch` is selected, it will additionally display available branches: diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md index 6c0ca0817c..2e0c236cf6 100644 --- a/doc/model/train-se-atten.md +++ b/doc/model/train-se-atten.md @@ -134,7 +134,9 @@ You can use descriptor `"se_atten_v2"` and is not allowed to set `tebd_input_mod Practical evidence demonstrates that `"se_atten_v2"` offers better and more stable performance compared to `"se_atten"`. -Notice: Model compression for the `se_atten_v2` descriptor is exclusively designed for models with the training parameter {ref}`attn_layer ` set to 0. +:::{note} +Model compression support differs across backends. See [Model compression](#model-compression) for backend-specific requirements. +::: ## Type embedding @@ -182,7 +184,13 @@ DPA-1 supports both the [standard data format](../data/system.md) and the [mixed ## Model compression -Model compression is supported only when there is no attention layer (`attn_layer` is 0) and `tebd_input_mode` is `strip`. +### TensorFlow {{ tensorflow_icon }} + +Model compression is supported only when the descriptor attention depth {ref}`attn_layer ` is 0 and {ref}`tebd_input_mode ` is `"strip"`. Attention layers higher than 0 cannot be compressed in the TensorFlow implementation because the geometric part is tabulated from the static computation graph. + +### PyTorch {{ pytorch_icon }} + +Model compression is supported for any {ref}`attn_layer ` value when {ref}`tebd_input_mode ` is `"strip"`. When `attn_layer` is 0, both the type embedding and geometric parts are compressed. When `attn_layer` is not 0, only the type embedding is compressed while the geometric part keeps the neural network implementation (a warning is emitted during compression). ## Training example diff --git a/doc/model/train-se-e3-tebd.md b/doc/model/train-se-e3-tebd.md index 3e0b6f9c5a..f0001f4e67 100644 --- a/doc/model/train-se-e3-tebd.md +++ b/doc/model/train-se-e3-tebd.md @@ -83,4 +83,4 @@ Type embedding is within this descriptor with the {ref}`tebd_dim None: +def exhale_environment_ready(app: "Sphinx") -> None: default_project = app.config.breathe_default_project default_exhale_args = dict(app.config.exhale_args) diff --git a/doc/test/test.md b/doc/test/test.md index dfd59d8f1f..9d399cb1ed 100644 --- a/doc/test/test.md +++ b/doc/test/test.md @@ -17,3 +17,25 @@ An explanation will be provided ```{program-output} dp test -h ``` + +## Evaluate descriptors + +The descriptors of a model can be evaluated and saved using `dp eval-desc`. A typical usage of `dp eval-desc` is + +```bash +dp eval-desc -m graph.pb -s /path/to/system -o desc +``` + +where `-m` gives the model file, `-s` the path to the system directory (or `-f` for a datafile containing paths to systems), and `-o` the output directory where descriptor files will be saved. The descriptors for each system will be saved as `.npy` files with the format `desc/(system_name).npy`. Each descriptor file contains a 3D array with shape (nframes, natoms, ndesc). + +Several other command line options can be passed to `dp eval-desc`, which can be checked with + +```bash +$ dp eval-desc --help +``` + +An explanation will be provided + +```{program-output} dp eval-desc -h + +``` diff --git a/doc/third-party/ase.md b/doc/third-party/ase.md index 6ede63e2f9..183efa7cbb 100644 --- a/doc/third-party/ase.md +++ b/doc/third-party/ase.md @@ -6,6 +6,10 @@ See [Environment variables](../env.md) for the runtime environment variables. Deep potential can be set up as a calculator with ASE to obtain potential energies and forces. +::::{tab-set} + +:::{tab-item} TensorFlow {{ tensorflow_icon }} + ```python from ase import Atoms from deepmd.calculator import DP @@ -20,6 +24,46 @@ print(water.get_potential_energy()) print(water.get_forces()) ``` +::: + +:::{tab-item} PyTorch {{ pytorch_icon }} + +```python +from ase import Atoms +from deepmd.calculator import DP + +water = Atoms( + "H2O", + positions=[(0.7601, 1.9270, 1), (1.9575, 1, 1), (1.0, 1.0, 1.0)], + cell=[100, 100, 100], + calculator=DP(model="frozen_model.pth"), +) +print(water.get_potential_energy()) +print(water.get_forces()) +``` + +::: + +:::{tab-item} Paddle {{ paddle_icon }} + +```python +from ase import Atoms +from deepmd.calculator import DP + +water = Atoms( + "H2O", + positions=[(0.7601, 1.9270, 1), (1.9575, 1, 1), (1.0, 1.0, 1.0)], + cell=[100, 100, 100], + calculator=DP(model="frozen_model.json"), +) +print(water.get_potential_energy()) +print(water.get_forces()) +``` + +::: + +:::: + Optimization is also available: ```python diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md index 5d0a90c18b..25a77f8670 100644 --- a/doc/third-party/lammps-command.md +++ b/doc/third-party/lammps-command.md @@ -229,7 +229,9 @@ dump 1 all custom 100 water.dump id type c_dipole[1] c_dipole[2] c_di - The `deeptensor/atom` compute is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the [DeePMD-kit website](https://github.com/deepmodeling/deepmd-kit) for more information. - For the issue of using a unit style for `compute deeptensor/atom`, refer to the discussions in [units](#units) of this page. -## Long-range interaction +## Combine with other commands + +### Long-range interaction The reciprocal space part of the long-range interaction can be calculated by LAMMPS command `kspace_style`. To use it with DeePMD-kit, one writes @@ -242,7 +244,7 @@ kspace_modify gewald 0.45 Please notice that the DeePMD does nothing to the direct space part of the electrostatic interaction, because this part is assumed to be fitted in the DeePMD model (the direct space cut-off is thus the cut-off of the DeePMD model). The splitting parameter `gewald` is modified by the `kspace_modify` command. -## Use of the centroid/stress/atom to get the full 3x3 "atomic-virial" +### Use of the centroid/stress/atom to get the full 3x3 "atomic-virial" The [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit) also allows the computation of per-atom stress tensor defined as: @@ -262,7 +264,7 @@ see [LAMMPS doc page](https://docs.lammps.org/compute_stress_atom.html#thompson2 v2.2.2 or previous versions passed per-atom stress (`cvatom`) with the per-atom pressure tensor, which is inconsistent with [LAMMPS's definition](https://docs.lammps.org/compute_stress_atom.html). LAMMPS defines per-atom stress as the negative of the per-atom pressure tensor. Such behavior is corrected in v2.2.3. ::: -### Examples +#### Examples In order of computing the 9-component per-atom stress @@ -274,7 +276,7 @@ Thus `c_stress` is an array with 9 components in the order `xx,yy,zz,xy,xz,yz,yx If you use this feature please cite [D. Tisi, L. Zhang, R. Bertossa, H. Wang, R. Car, S. Baroni - arXiv preprint arXiv:2108.10850, 2021](https://arxiv.org/abs/2108.10850) -## Computation of heat flux +### Computation of heat flux Using a per-atom stress tensor one can, for example, compute the heat flux defined as: @@ -289,7 +291,7 @@ compute stress_ID group-ID centroid/stress/atom NULL virial compute flux_ID all heat/flux ke_ID pe_ID stress_ID ``` -### Examples +#### Examples ```lammps compute ke all ke/atom @@ -305,3 +307,18 @@ If you use these features please cite [D. Tisi, L. Zhang, R. Bertossa, H. Wang, [DP]: https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.120.143001 [DP-SE]: https://dl.acm.org/doi/10.5555/3327345.3327356 [DPSPIN]: https://doi.org/10.1103/PhysRevB.110.064427 + +### D3 dispersion + +:::{note} +Requires LAMMPS version 4Feb2025 or newer. +::: + +The DP model can be trained using plain DFT calculations without the dispersion correction, and the dispersion correction can be added during the simulation via the [`pair_style dispersion/d3` command](https://docs.lammps.org/pair_dispersion_d3.html#pair-style-dispersion-d3-command). +For example, when `water.pb` is trained against the PBE0 functional, the simulation can be performed under the PBE0-D3 level with the following commands: + +```lammps +pair_style hybrid/overlay deepmd water.pb dispersion/d3 original pbe0 30.0 20.0 +pair_coeff * * deepmd O H +pair_coeff * * dispersion/d3 O H +``` diff --git a/doc/train/multi-task-training.md b/doc/train/multi-task-training.md index 16f6c0e05c..115c463cc2 100644 --- a/doc/train/multi-task-training.md +++ b/doc/train/multi-task-training.md @@ -48,7 +48,6 @@ Specifically, there are several parts that need to be modified: - {ref}`model/model_dict `: The core definition of the model part and the explanation of sharing rules, starting with user-defined model name keys `model_key`, such as `my_model_1`. Each model part needs to align with the components of the single-task training {ref}`model `, but with the following sharing rules: - - If you want to share the current model component with other tasks, which should be part of the {ref}`model/shared_dict `, you can directly fill in the corresponding `part_key`, such as `"descriptor": "my_descriptor", ` diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md index 96cfadf4bb..998f1c3bec 100644 --- a/doc/train/parallel-training.md +++ b/doc/train/parallel-training.md @@ -218,6 +218,21 @@ NUM_WORKERS=0 HDF5_USE_FILE_LOCKING=0 python -m paddle.distributed.launch \ dp --pd train input.json ``` +or you can wrapper the training script with `mpirun`: + +```bash +# ----- train_pp.sh ------- +unset CUDA_DEVICE_MAX_CONNECTIONS +python -m paddle.distributed.launch --gpus="0,1,2,3,4,5,6,7" --log_dir logs dp --pd train input_torch.json -l train_pp.log +# ------------------------- +``` + +Then, run the script on the first node with: + +```bash +mpirun run_pp.sh +``` + :::{note} If `NUM_WORKERS` is too large, it may cause the program to be terminated by the system; diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 174c39d6d9..af4b4b31d9 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -76,8 +76,8 @@ Other training parameters are given in the {ref}`training ` section. The sections {ref}`training_data ` and {ref}`validation_data ` give the training dataset and validation dataset, respectively. Taking the training dataset for example, the keys are explained below: - {ref}`systems ` provide paths of the training data systems. DeePMD-kit allows you to provide multiple systems with different numbers of atoms. This key can be a `list` or a `str`. - - `list`: {ref}`systems ` gives the training data systems. - - `str`: {ref}`systems ` should be a valid path. DeePMD-kit will recursively search all data systems in this path. + - `str`: {ref}`systems ` should be a valid path. It can be a system directory path (containing 'type.raw') or a parent directory path to recursively search for all system subdirectories. + - `list`: {ref}`systems ` gives a list of paths. Each string item in the list is processed the same way as individual string inputs, i.e., each path can be a system directory or a parent directory to recursively search for all system subdirectories. - At each training step, DeePMD-kit randomly picks {ref}`batch_size ` frame(s) from one of the systems. The probability of using a system is by default in proportion to the number of batches in the system. More options are available for automatically determining the probability of using systems. One can set the key {ref}`auto_prob ` to - `"prob_uniform"` all systems are used with the same probability. - `"prob_sys_size"` the probability of using a system is proportional to its size (number of frames). diff --git a/doc/train/training.md b/doc/train/training.md index 6c8b7a5549..6ccb43bbd7 100644 --- a/doc/train/training.md +++ b/doc/train/training.md @@ -29,12 +29,15 @@ $ dp --pt train input.json :::{tab-item} Paddle {{ paddle_icon }} ```bash -# training model in eager mode +# training model $ dp --pd train input.json -# [experimental] training model with CINN compiler for better performance, +# [experimental] training models with the CINN compiler (~40%+ speedup) # see: https://www.paddlepaddle.org.cn/documentation/docs/zh/develop/guides/paddle_v3_features/cinn_cn.html +## If the shape(s) of batch input data are dynamic during training(default). $ CINN=1 dp --pd train input.json +## If the shape(s) of batch input data are fixed during training, e.g., examples/water. +$ CINN=1 CINN_ALLOW_DYNAMIC_SHAPE=0 dp --pd train input.json ``` ::: diff --git a/examples/infer_water/CMakeLists.txt b/examples/infer_water/CMakeLists.txt index 27b7541c2d..c82816c912 100644 --- a/examples/infer_water/CMakeLists.txt +++ b/examples/infer_water/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 3.7) +cmake_minimum_required(VERSION 3.25.2) project(infer_water) # find DeePMD-kit diff --git a/examples/water/d3/README.md b/examples/water/d3/README.md index bd75960010..fdbeea1c1f 100644 --- a/examples/water/d3/README.md +++ b/examples/water/d3/README.md @@ -1,5 +1,9 @@ # DPD3 +> [!IMPORTANT] +> It is more efficient to train a model using plain DFT calculations without the dispersion correction, and add the dispersion correction during the simulation via the LAMMPS [`pair_style dispersion/d3` command](https://docs.lammps.org/pair_dispersion_d3.html#pair-style-dispersion-d3-command). +> Training against data with dispersion directly is discouraged. + `dftd3.txt` tabulates D3 dispersion for each pair of types (O-O, O-H, H-H). It can be generated by [simple-dftd3](https://github.com/dftd3/simple-dftd3). diff --git a/examples/water/dpa2/input_torch_compressible.json b/examples/water/dpa2/input_torch_compressible.json index a9bfdd7b69..14ec347b35 100644 --- a/examples/water/dpa2/input_torch_compressible.json +++ b/examples/water/dpa2/input_torch_compressible.json @@ -53,7 +53,8 @@ "g1_out_mlp": true }, "precision": "float64", - "add_tebd_to_repinit_out": false + "add_tebd_to_repinit_out": false, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/dpa2/input_torch_large.json b/examples/water/dpa2/input_torch_large.json index 49edc2721e..4894cc6915 100644 --- a/examples/water/dpa2/input_torch_large.json +++ b/examples/water/dpa2/input_torch_large.json @@ -52,7 +52,8 @@ "g1_out_mlp": true }, "precision": "float64", - "add_tebd_to_repinit_out": false + "add_tebd_to_repinit_out": false, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/dpa2/input_torch_medium.json b/examples/water/dpa2/input_torch_medium.json index 0ba83f2cd5..b752e28f31 100644 --- a/examples/water/dpa2/input_torch_medium.json +++ b/examples/water/dpa2/input_torch_medium.json @@ -52,7 +52,8 @@ "g1_out_mlp": true }, "precision": "float64", - "add_tebd_to_repinit_out": false + "add_tebd_to_repinit_out": false, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/dpa2/input_torch_small.json b/examples/water/dpa2/input_torch_small.json index 86217d980e..bd136a8666 100644 --- a/examples/water/dpa2/input_torch_small.json +++ b/examples/water/dpa2/input_torch_small.json @@ -52,7 +52,8 @@ "g1_out_mlp": true }, "precision": "float64", - "add_tebd_to_repinit_out": false + "add_tebd_to_repinit_out": false, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/dpa3/input_torch.json b/examples/water/dpa3/input_torch.json index 90e81b5403..ec8bba4821 100644 --- a/examples/water/dpa3/input_torch.json +++ b/examples/water/dpa3/input_torch.json @@ -34,7 +34,8 @@ "activation_function": "silut:10.0", "use_tebd_bias": false, "precision": "float32", - "concat_output_tebd": false + "concat_output_tebd": false, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/dpa3/input_torch_dynamic.json b/examples/water/dpa3/input_torch_dynamic.json index acf48e15b1..b3137feffc 100644 --- a/examples/water/dpa3/input_torch_dynamic.json +++ b/examples/water/dpa3/input_torch_dynamic.json @@ -36,7 +36,8 @@ "activation_function": "silut:10.0", "use_tebd_bias": false, "precision": "float32", - "concat_output_tebd": false + "concat_output_tebd": false, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/se_atten/input_torch.json b/examples/water/se_atten/input_torch.json index cdb4b0db49..6366ae9c67 100644 --- a/examples/water/se_atten/input_torch.json +++ b/examples/water/se_atten/input_torch.json @@ -25,7 +25,8 @@ "activation_function": "tanh", "scaling_factor": 1.0, "normalize": true, - "temperature": 1.0 + "temperature": 1.0, + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water/se_e3_tebd/input_torch.json b/examples/water/se_e3_tebd/input_torch.json index 05aee70a73..a65cd2970f 100644 --- a/examples/water/se_e3_tebd/input_torch.json +++ b/examples/water/se_e3_tebd/input_torch.json @@ -17,7 +17,8 @@ ], "tebd_dim": 8, "tebd_input_mode": "concat", - "activation_function": "tanh" + "activation_function": "tanh", + "seed": 1 }, "fitting_net": { "neuron": [ diff --git a/examples/water_multi_task/pytorch_example/input_torch.json b/examples/water_multi_task/pytorch_example/input_torch.json index f0d9a67ea5..d5aadc8de2 100644 --- a/examples/water_multi_task/pytorch_example/input_torch.json +++ b/examples/water_multi_task/pytorch_example/input_torch.json @@ -54,6 +54,7 @@ }, "precision": "float64", "add_tebd_to_repinit_out": false, + "seed": 1, "_comment": " that's all" }, "_comment": "that's all" diff --git a/examples/water_multi_task/pytorch_example/input_torch_sharefit.json b/examples/water_multi_task/pytorch_example/input_torch_sharefit.json index 2fc23007c6..d915de181c 100644 --- a/examples/water_multi_task/pytorch_example/input_torch_sharefit.json +++ b/examples/water_multi_task/pytorch_example/input_torch_sharefit.json @@ -54,6 +54,7 @@ }, "precision": "float64", "add_tebd_to_repinit_out": false, + "seed": 1, "_comment": " that's all" }, "shared_fit_with_id": { diff --git a/examples/water_multi_task/pytorch_example/input_torch_with_alias.json b/examples/water_multi_task/pytorch_example/input_torch_with_alias.json new file mode 100644 index 0000000000..03ee0f7e11 --- /dev/null +++ b/examples/water_multi_task/pytorch_example/input_torch_with_alias.json @@ -0,0 +1,149 @@ +{ + "_comment": "that's all", + "model": { + "shared_dict": { + "type_map_all": [ + "O", + "H" + ], + "dpa3_descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 128, + "e_dim": 64, + "a_dim": 32, + "nlayers": 6, + "e_rcut": 6.0, + "e_rcut_smth": 5.3, + "e_sel": 1200, + "a_rcut": 4.0, + "a_rcut_smth": 3.5, + "a_sel": 300, + "axis_neuron": 4, + "fix_stat_std": 0.3, + "a_compress_rate": 1, + "a_compress_e_rate": 2, + "a_compress_use_split": true, + "update_angle": true, + "smooth_edge_update": true, + "edge_init_use_dist": true, + "use_dynamic_sel": true, + "sel_reduce_factor": 10.0, + "use_exp_switch": true, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const" + }, + "activation_function": "silut:10.0", + "use_tebd_bias": false, + "precision": "float32", + "concat_output_tebd": false, + "seed": 1 + }, + "shared_fit_with_id": { + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1, + "dim_case_embd": 2, + "_comment": " that's all" + }, + "_comment": "that's all" + }, + "model_dict": { + "water_1": { + "type_map": "type_map_all", + "descriptor": "dpa3_descriptor", + "fitting_net": "shared_fit_with_id", + "model_branch_alias": ["Default","Water"], + "info": { + "description": "Water model with DPA3 descriptor and shared fitting net", + "observed_type": ["H", "O"] + } + }, + "water_2": { + "type_map": "type_map_all", + "descriptor": "dpa3_descriptor", + "fitting_net": "shared_fit_with_id", + "model_branch_alias": ["Water2"], + "info": { + "description": "Water duplicated model with DPA3 descriptor and shared fitting net", + "observed_type": ["H", "O"] + } + } + } + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment": "that's all" + }, + "loss_dict": { + "water_1": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0 + }, + "water_2": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0 + } + }, + "training": { + "model_prob": { + "water_1": 0.5, + "water_2": 0.5 + }, + "data_dict": { + "water_1": { + "training_data": { + "systems": [ + "../../water/data/data_0/", + "../../water/data/data_1/", + "../../water/data/data_2/" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../../water/data/data_3/" + ], + "batch_size": 1, + "_comment": "that's all" + } + }, + "water_2": { + "training_data": { + "systems": [ + "../../water/data/data_0/", + "../../water/data/data_1/", + "../../water/data/data_2/" + ], + "batch_size": 1, + "_comment": "that's all" + } + } + }, + "numb_steps": 100000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 100, + "_comment": "that's all" + } +} diff --git a/pyproject.toml b/pyproject.toml index 362bcacd4b..5b089c4558 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,9 +2,10 @@ requires = [ # TODO: unpin the upper bound when scikit-build dynamic metadata API is stable # dynamic metadata API is still unstable - "scikit-build-core>=0.5,<0.11,!=0.6.0", + "scikit-build-core>=0.5,!=0.6.0,<0.12", "packaging", 'tomli >= 1.1.0 ; python_version < "3.11"', + "dependency_groups", ] build-backend = "backend.dp_backend" backend-path = ["."] @@ -28,10 +29,10 @@ classifiers = [ "Programming Language :: Python :: 3 :: Only", "Environment :: GPU :: NVIDIA CUDA :: 12 :: 12.2", "Intended Audience :: Science/Research", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "License :: OSI Approved :: GNU Lesser General Public License v3 (LGPLv3)", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: Scientific/Engineering :: Physics", @@ -54,11 +55,11 @@ dependencies = [ 'mendeleev', 'array-api-compat', ] -requires-python = ">=3.9" +requires-python = ">=3.10" keywords = ["deepmd"] [project.entry-points."lammps.plugins"] -deepmd = "deepmd.lmp:get_op_dir" +deepmd = "deepmd.lmp_check_build:get_op_dir" [project.entry-points."dpgui"] "DeePMD-kit" = "deepmd.utils.argcheck:gen_args" @@ -90,7 +91,7 @@ test = [ docs = [ "sphinx>=3.1.1", "sphinx-book-theme", - "myst-nb>=1.0.0rc0", + "myst-nb>=1.0.0", "myst-parser>=0.19.2", "sphinx-design", "breathe", @@ -108,7 +109,7 @@ docs = [ "sphinx-remove-toctrees", ] lmp = [ - "lammps~=2024.8.29.1.0", + "lammps[mpi]~=2025.7.22.2.0", ] ipi = [ "ipi", @@ -116,16 +117,6 @@ ipi = [ gui = [ "dpgui", ] -cu11 = [ - "nvidia-cuda-runtime-cu11", - "nvidia-cublas-cu11", - "nvidia-cufft-cu11", - "nvidia-curand-cu11", - "nvidia-cusolver-cu11", - "nvidia-cusparse-cu11", - "nvidia-cudnn-cu11<9", - "nvidia-cuda-nvcc-cu11", -] cu12 = [ "nvidia-cuda-runtime-cu12", "nvidia-cublas-cu12", @@ -158,6 +149,28 @@ dev = [ "cmake", "mpich", ] +pin_tensorflow_cpu = [ + # https://github.com/tensorflow/tensorflow/issues/75279 + # macos x86 has been deprecated + "tensorflow-cpu>=2.18,<2.21; platform_machine=='x86_64' and platform_system == 'Linux'", + "tensorflow~=2.18.0; (platform_machine!='x86_64' or platform_system != 'Linux') and (platform_machine!='x86_64' or platform_system != 'Darwin')", + "tensorflow; platform_machine=='x86_64' and platform_system == 'Darwin'", +] +pin_tensorflow_gpu = [ + "tensorflow~=2.18.0", +] +pin_pytorch_cpu = [ + # https://github.com/pytorch/pytorch/issues/114602 + # macos x86 has been deprecated + "torch>=2.8,<2.10; platform_machine!='x86_64' or platform_system != 'Darwin'", + "torch; platform_machine=='x86_64' and platform_system == 'Darwin'", +] +pin_pytorch_gpu = [ + "torch>=2.7,<2.10", +] +pin_jax = [ + "jax==0.5.0;python_version>='3.10'", +] [tool.setuptools_scm] @@ -232,20 +245,15 @@ test-command = [ test-extras = ["cpu", "test", "lmp", "ipi", "torch", "paddle"] build = ["cp311-*"] skip = ["*-win32", "*-manylinux_i686", "*-musllinux*"] -# TODO: uncomment to use the latest image when CUDA 11 is deprecated -# manylinux-x86_64-image = "manylinux_2_28" -manylinux-x86_64-image = "quay.io/pypa/manylinux_2_28_x86_64:2022-11-19-1b19e81" +manylinux-x86_64-image = "manylinux_2_28" manylinux-aarch64-image = "manylinux_2_28" [tool.cibuildwheel.macos] -before-all = [ - '''pip install -i https://pypi.anaconda.org/mpi4py/simple mpich''', -] repair-wheel-command = """delocate-wheel --require-archs {delocate_archs} -w {dest_dir} -v {wheel} --ignore-missing-dependencies""" [tool.cibuildwheel.macos.environment] PIP_PREFER_BINARY = "1" -DP_LAMMPS_VERSION = "stable_29Aug2024_update1" +DP_LAMMPS_VERSION = "stable_22Jul2025_update2" DP_ENABLE_IPI = "1" DP_ENABLE_PYTORCH = "1" DP_ENABLE_PADDLE = "1" @@ -259,7 +267,7 @@ inherit.environment = "append" environment.MACOSX_DEPLOYMENT_TARGET = "11.0" [tool.cibuildwheel.linux] -repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so --exclude _pywrap_tensorflow_internal.so --exclude libtensorflow_cc.so.2 --exclude libc10.so --exclude libtorch.so --exclude libtorch_cpu.so -w {dest_dir} {wheel}" +repair-wheel-command = "auditwheel repair --exclude libtensorflow_framework.so.2 --exclude libtensorflow_framework.so.1 --exclude libtensorflow_framework.so --exclude _pywrap_tensorflow_internal.so --exclude libtensorflow_cc.so.2 --exclude libc10.so --exclude libtorch.so --exclude libtorch_cpu.so --exclude libmpi.so.12 -w {dest_dir} {wheel}" environment-pass = [ "CIBW_BUILD", "DP_VARIANT", @@ -269,31 +277,19 @@ environment-pass = [ ] before-all = [ """if [ ! -z "${DP_PKG_NAME}" ]; then sed -i "s/name = \\"deepmd-kit\\"/name = \\"${DP_PKG_NAME}\\"/g" pyproject.toml; fi""", - # https://almalinux.org/blog/2023-12-20-almalinux-8-key-update/ - """rpm --import https://repo.almalinux.org/almalinux/RPM-GPG-KEY-AlmaLinux""", """{ if [ "$(uname -m)" = "x86_64" ] ; then yum config-manager --add-repo http://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo && yum install -y cuda-nvcc-${CUDA_VERSION/./-} cuda-cudart-devel-${CUDA_VERSION/./-}; fi }""", - '''/opt/python/cp311-cp311/bin/python -m pip install -i https://pypi.anaconda.org/mpi4py/simple mpich''', - # uv is not available in the old manylinux image - """{ if [ "$(uname -m)" = "x86_64" ] ; then pipx install uv; fi }""", ] before-build = [ - # old build doesn't support uv - """{ if [ "$(uname -m)" = "x86_64" ] ; then uv pip install --system -U build; fi }""", ] [tool.cibuildwheel.linux.environment] PIP_PREFER_BINARY = "1" -DP_LAMMPS_VERSION = "stable_29Aug2024_update1" +DP_LAMMPS_VERSION = "stable_22Jul2025_update2" DP_ENABLE_IPI = "1" DP_ENABLE_PYTORCH = "1" DP_ENABLE_PADDLE = "1" -MPI_HOME = "/usr/lib64/mpich" -PATH = "/usr/lib64/mpich/bin:$PATH" # use CPU version of torch for building, which should also work for GPU -# note: uv has different behavior from pip on extra index url -# https://github.com/astral-sh/uv/blob/main/PIP_COMPATIBILITY.md#packages-that-exist-on-multiple-indexes -UV_EXTRA_INDEX_URL = "https://download.pytorch.org/whl/cpu" -# trick to find the correction version of mpich -CMAKE_PREFIX_PATH="/opt/python/cp311-cp311/" +# https://docs.astral.sh/uv/guides/integration/pytorch/#automatic-backend-selection +UV_TORCH_BACKEND = "cpu" [tool.cibuildwheel.windows] test-extras = ["cpu", "torch", "paddle"] @@ -358,6 +354,11 @@ ignore = "D413, D416, D203, D107, D213" profile = "black" force_grid_wrap = 1 +[tool.ruff] +extend-exclude = [ + "source/3rdparty/**", +] + [tool.ruff.format] docstring-code-format = true @@ -380,11 +381,14 @@ select = [ "DTZ", # datetime "TCH", # flake8-type-checking "PYI", # flake8-pyi + "ANN", # type annotations ] ignore = [ + "ANN401", # Allow Any due to too many violations "E501", # line too long "F841", # local variable is assigned to but never used + "RUF059", # unused-unpacked-variable "E741", # ambiguous variable name "E402", # module level import not at top of file "D100", # TODO: missing docstring in public module @@ -397,11 +401,6 @@ ignore = [ "D401", # TODO: first line should be in imperative mood "D404", # TODO: first word of the docstring should not be This ] -ignore-init-module-imports = true - -exclude = [ - "source/3rdparty/**", -] [tool.ruff.lint.pydocstyle] convention = "numpy" @@ -426,19 +425,27 @@ runtime-evaluated-base-classes = ["torch.nn.Module"] [tool.ruff.lint.extend-per-file-ignores] # Also ignore `E402` in all `__init__.py` files. -"deepmd/tf/**" = ["TID253"] +"source/3rdparty/**" = ["ALL"] +"backend/**" = ["ANN"] +"data/**" = ["ANN"] +"deepmd/tf/**" = ["TID253", "ANN"] "deepmd/pt/**" = ["TID253"] "deepmd/jax/**" = ["TID253"] -"deepmd/pd/**" = ["TID253"] -"source/tests/tf/**" = ["TID253"] -"source/tests/pt/**" = ["TID253"] -"source/tests/jax/**" = ["TID253"] -"source/tests/pd/**" = ["TID253"] -"source/tests/universal/pt/**" = ["TID253"] -"source/tests/universal/pd/**" = ["TID253"] -"source/jax2tf_tests/**" = ["TID253"] -"source/ipi/tests/**" = ["TID253"] -"source/lmp/tests/**" = ["TID253"] +"deepmd/pd/**" = ["TID253", "ANN"] + +"source/**" = ["ANN"] +"source/tests/tf/**" = ["TID253", "ANN"] +"source/tests/pt/**" = ["TID253", "ANN"] +"source/tests/jax/**" = ["TID253", "ANN"] +"source/tests/pd/**" = ["TID253", "ANN"] +"source/tests/universal/pt/**" = ["TID253", "ANN"] +"source/tests/universal/pd/**" = ["TID253", "ANN"] +"source/tests/**" = ["ANN"] +"source/jax2tf_tests/**" = ["TID253", "ANN"] +"source/ipi/tests/**" = ["TID253", "ANN"] +"source/lmp/tests/**" = ["TID253", "ANN"] +"**/tests/**/test_*.py" = ["ANN"] +"**/tests/**/*_test.py" = ["ANN"] "**/*.ipynb" = ["T20"] # printing in a nb file is expected [tool.pytest.ini_options] @@ -459,15 +466,6 @@ select = [ "TOR2", ] -[tool.uv.sources] -mpich = { index = "mpi4py" } -openmpi = { index = "mpi4py" } - -[[tool.uv.index]] -name = "mpi4py" -url = "https://pypi.anaconda.org/mpi4py/simple" -explicit = true - [[tool.uv.dependency-metadata]] # Fix https://github.com/deepmodeling/deepmd-kit/issues/4679 name = "tensorflow" diff --git a/source/CMakeLists.txt b/source/CMakeLists.txt index a678802a35..d356d4cba6 100644 --- a/source/CMakeLists.txt +++ b/source/CMakeLists.txt @@ -1,5 +1,5 @@ -# new in 3.16: GET_RUNTIME_DEPENDENCIES, target_precompile_headers -cmake_minimum_required(VERSION 3.16) +# new in 3.25: RUN_OUTPUT_STDOUT_VARIABLE +cmake_minimum_required(VERSION 3.25.2) project(DeePMD) option(ENABLE_TENSORFLOW "Enable TensorFlow interface" OFF) @@ -244,8 +244,14 @@ set(DP_VARIANT "cpu") # define USE_CUDA_TOOLKIT if(USE_CUDA_TOOLKIT) - cmake_minimum_required(VERSION 3.23) + cmake_minimum_required(VERSION 3.25.2) find_package(CUDAToolkit REQUIRED) + # CUDA 13.0+ requires C++17 + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.0") + set_if_higher(CMAKE_CXX_STANDARD 17) + message( + STATUS "CUDA ${CUDAToolkit_VERSION} detected, setting C++ standard to 17") + endif() if(NOT DEFINED CMAKE_CUDA_COMPILER) set(CMAKE_CUDA_COMPILER ${CUDAToolkit_NVCC_EXECUTABLE}) endif() @@ -261,7 +267,7 @@ endif(USE_CUDA_TOOLKIT) # define USE_ROCM_TOOLKIT if(USE_ROCM_TOOLKIT) - cmake_minimum_required(VERSION 3.21) + cmake_minimum_required(VERSION 3.25.2) include(CMakeDetermineHIPCompiler) list(APPEND CMAKE_PREFIX_PATH ${CMAKE_HIP_COMPILER_ROCM_ROOT}) find_package(hip REQUIRED) @@ -373,7 +379,15 @@ if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT) else() # Maybe in macos/windows if(NOT DEFINED OP_CXX_ABI) - set(OP_CXX_ABI 0) + if(UNIX + AND NOT APPLE + AND Torch_VERSION VERSION_GREATER_EQUAL "2.8.0") + # https://github.com/deepmodeling/deepmd-kit/issues/4877 + # torch.compiled_with_cxx11_abi in PyTorch 2.8 always return True + set(OP_CXX_ABI 1) + else() + set(OP_CXX_ABI 0) + endif() endif() set(OP_CXX_ABI_PT "${OP_CXX_ABI}") endif() @@ -520,13 +534,16 @@ if(DEEPMD_C_ROOT) endif() if(NOT DEEPMD_C_ROOT) + add_subdirectory(lib/) if(ENABLE_TENSORFLOW) add_subdirectory(op/tf/) endif() if(ENABLE_PYTORCH) add_subdirectory(op/pt/) endif() - add_subdirectory(lib/) + if(ENABLE_PADDLE) + add_subdirectory(op/pd/) + endif() endif() if(BUILD_PY_IF) add_subdirectory(config/) diff --git a/source/api_c/include/deepmd.hpp b/source/api_c/include/deepmd.hpp index 8a3656bfc2..afa62403e7 100644 --- a/source/api_c/include/deepmd.hpp +++ b/source/api_c/include/deepmd.hpp @@ -26,7 +26,7 @@ namespace hpp { struct deepmd_exception : public std::runtime_error { public: deepmd_exception() : runtime_error("DeePMD-kit C API Error!") {}; - deepmd_exception(const std::string &msg) + deepmd_exception(const std::string& msg) : runtime_error(std::string("DeePMD-kit C API Error: ") + msg) {}; }; } // namespace hpp @@ -36,7 +36,7 @@ struct deepmd_exception : public std::runtime_error { * @brief Check if any exceptions throw in the C++ API. Throw if possible. */ #define DP_CHECK_OK(check_func, dp) \ - const char *err_msg = check_func(dp); \ + const char* err_msg = check_func(dp); \ if (std::strlen(err_msg)) { \ std::string err_msg_str = std::string(err_msg); \ DP_DeleteChar(err_msg); \ @@ -45,173 +45,173 @@ struct deepmd_exception : public std::runtime_error { DP_DeleteChar(err_msg); template -inline void _DP_DeepPotCompute(DP_DeepPot *dp, +inline void _DP_DeepPotCompute(DP_DeepPot* dp, const int nframes, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepPotCompute(DP_DeepPot *dp, +inline void _DP_DeepPotCompute(DP_DeepPot* dp, const int nframes, const int natom, - const double *coord, - const int *atype, - const double *cell, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepPotCompute2(dp, nframes, natom, coord, atype, cell, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepPotCompute(DP_DeepPot *dp, +inline void _DP_DeepPotCompute(DP_DeepPot* dp, const int nframes, const int natom, - const float *coord, - const int *atype, - const float *cell, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepPotComputef2(dp, nframes, natom, coord, atype, cell, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } // support spin template -inline void _DP_DeepSpinCompute(DP_DeepSpin *dp, +inline void _DP_DeepSpinCompute(DP_DeepSpin* dp, const int nframes, const int natom, - const FPTYPE *coord, - const FPTYPE *spin, - const int *atype, - const FPTYPE *cell, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *force_mag, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* coord, + const FPTYPE* spin, + const int* atype, + const FPTYPE* cell, + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* force_mag, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepSpinCompute(DP_DeepSpin *dp, +inline void _DP_DeepSpinCompute(DP_DeepSpin* dp, const int nframes, const int natom, - const double *coord, - const double *spin, - const int *atype, - const double *cell, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *force_mag, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepSpinCompute2(dp, nframes, natom, coord, spin, atype, cell, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepSpinCompute(DP_DeepSpin *dp, +inline void _DP_DeepSpinCompute(DP_DeepSpin* dp, const int nframes, const int natom, - const float *coord, - const float *spin, - const int *atype, - const float *cell, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *force_mag, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepSpinComputef2(dp, nframes, natom, coord, spin, atype, cell, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepPotComputeNList(DP_DeepPot *dp, +inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, const int nframes, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepPotComputeNList(DP_DeepPot *dp, +inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, const int nframes, const int natom, - const double *coord, - const int *atype, - const double *cell, + const double* coord, + const int* atype, + const double* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepPotComputeNList2(dp, nframes, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepPotComputeNList(DP_DeepPot *dp, +inline void _DP_DeepPotComputeNList(DP_DeepPot* dp, const int nframes, const int natom, - const float *coord, - const int *atype, - const float *cell, + const float* coord, + const int* atype, + const float* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepPotComputeNListf2(dp, nframes, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); @@ -219,550 +219,550 @@ inline void _DP_DeepPotComputeNList(DP_DeepPot *dp, // support spin template -inline void _DP_DeepSpinComputeNList(DP_DeepSpin *dp, +inline void _DP_DeepSpinComputeNList(DP_DeepSpin* dp, const int nframes, const int natom, - const FPTYPE *coord, - const FPTYPE *spin, - const int *atype, - const FPTYPE *cell, + const FPTYPE* coord, + const FPTYPE* spin, + const int* atype, + const FPTYPE* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *force_mag, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* force_mag, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepSpinComputeNList(DP_DeepSpin *dp, +inline void _DP_DeepSpinComputeNList(DP_DeepSpin* dp, const int nframes, const int natom, - const double *coord, - const double *spin, - const int *atype, - const double *cell, + const double* coord, + const double* spin, + const int* atype, + const double* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *force_mag, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepSpinComputeNList2(dp, nframes, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepSpinComputeNList(DP_DeepSpin *dp, +inline void _DP_DeepSpinComputeNList(DP_DeepSpin* dp, const int nframes, const int natom, - const float *coord, - const float *spin, - const int *atype, - const float *cell, + const float* coord, + const float* spin, + const int* atype, + const float* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *force_mag, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepSpinComputeNListf2(dp, nframes, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepPotComputeMixedType(DP_DeepPot *dp, +inline void _DP_DeepPotComputeMixedType(DP_DeepPot* dp, const int nframes, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepPotComputeMixedType(DP_DeepPot *dp, +inline void _DP_DeepPotComputeMixedType(DP_DeepPot* dp, const int nframes, const int natom, - const double *coord, - const int *atype, - const double *cell, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepPotComputeMixedType(dp, nframes, natom, coord, atype, cell, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepPotComputeMixedType(DP_DeepPot *dp, +inline void _DP_DeepPotComputeMixedType(DP_DeepPot* dp, const int nframes, const int natom, - const float *coord, - const int *atype, - const float *cell, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepPotComputeMixedTypef(dp, nframes, natom, coord, atype, cell, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi *dp, +inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi *dp, +inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, const int natom, - const double *coord, - const int *atype, - const double *cell, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* coord, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepPotModelDeviCompute2(dp, 1, natom, coord, atype, cell, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi *dp, +inline void _DP_DeepPotModelDeviCompute(DP_DeepPotModelDevi* dp, const int natom, - const float *coord, - const int *atype, - const float *cell, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* coord, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepPotModelDeviComputef2(dp, 1, natom, coord, atype, cell, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi *dp, +inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi* dp, const int natom, - const FPTYPE *coord, - const FPTYPE *spin, - const int *atype, - const FPTYPE *cell, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *force_mag, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* coord, + const FPTYPE* spin, + const int* atype, + const FPTYPE* cell, + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* force_mag, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi *dp, +inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi* dp, const int natom, - const double *coord, - const double *spin, - const int *atype, - const double *cell, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *force_mag, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* coord, + const double* spin, + const int* atype, + const double* cell, + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepSpinModelDeviCompute2(dp, 1, natom, coord, spin, atype, cell, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi *dp, +inline void _DP_DeepSpinModelDeviCompute(DP_DeepSpinModelDevi* dp, const int natom, - const float *coord, - const float *spin, - const int *atype, - const float *cell, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *force_mag, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* coord, + const float* spin, + const int* atype, + const float* cell, + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepSpinModelDeviComputef2(dp, 1, natom, coord, spin, atype, cell, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi *dp, +inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi *dp, +inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, const int natom, - const double *coord, - const int *atype, - const double *cell, + const double* coord, + const int* atype, + const double* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepPotModelDeviComputeNList2(dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi *dp, +inline void _DP_DeepPotModelDeviComputeNList(DP_DeepPotModelDevi* dp, const int natom, - const float *coord, - const int *atype, - const float *cell, + const float* coord, + const int* atype, + const float* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepPotModelDeviComputeNListf2(dp, 1, natom, coord, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi *dp, +inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, const int natom, - const FPTYPE *coord, - const FPTYPE *spin, - const int *atype, - const FPTYPE *cell, + const FPTYPE* coord, + const FPTYPE* spin, + const int* atype, + const FPTYPE* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const FPTYPE *fparam, - const FPTYPE *aparam, - double *energy, - FPTYPE *force, - FPTYPE *force_mag, - FPTYPE *virial, - FPTYPE *atomic_energy, - FPTYPE *atomic_virial); + const FPTYPE* fparam, + const FPTYPE* aparam, + double* energy, + FPTYPE* force, + FPTYPE* force_mag, + FPTYPE* virial, + FPTYPE* atomic_energy, + FPTYPE* atomic_virial); template <> -inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi *dp, +inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, const int natom, - const double *coord, - const double *spin, - const int *atype, - const double *cell, + const double* coord, + const double* spin, + const int* atype, + const double* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const double *fparam, - const double *aparam, - double *energy, - double *force, - double *force_mag, - double *virial, - double *atomic_energy, - double *atomic_virial) { + const double* fparam, + const double* aparam, + double* energy, + double* force, + double* force_mag, + double* virial, + double* atomic_energy, + double* atomic_virial) { DP_DeepSpinModelDeviComputeNList2( dp, 1, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template <> -inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi *dp, +inline void _DP_DeepSpinModelDeviComputeNList(DP_DeepSpinModelDevi* dp, const int natom, - const float *coord, - const float *spin, - const int *atype, - const float *cell, + const float* coord, + const float* spin, + const int* atype, + const float* cell, const int nghost, - const DP_Nlist *nlist, + const DP_Nlist* nlist, const int ago, - const float *fparam, - const float *aparam, - double *energy, - float *force, - float *force_mag, - float *virial, - float *atomic_energy, - float *atomic_virial) { + const float* fparam, + const float* aparam, + double* energy, + float* force, + float* force_mag, + float* virial, + float* atomic_energy, + float* atomic_virial) { DP_DeepSpinModelDeviComputeNListf2( dp, 1, natom, coord, spin, atype, cell, nghost, nlist, ago, fparam, aparam, energy, force, force_mag, virial, atomic_energy, atomic_virial); } template -inline void _DP_DeepTensorComputeTensor(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeTensor(DP_DeepTensor* dt, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, - FPTYPE **tensor, - int *size); + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, + FPTYPE** tensor, + int* size); template <> -inline void _DP_DeepTensorComputeTensor(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeTensor(DP_DeepTensor* dt, const int natom, - const double *coord, - const int *atype, - const double *cell, - double **tensor, - int *size) { + const double* coord, + const int* atype, + const double* cell, + double** tensor, + int* size) { DP_DeepTensorComputeTensor(dt, natom, coord, atype, cell, tensor, size); } template <> -inline void _DP_DeepTensorComputeTensor(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeTensor(DP_DeepTensor* dt, const int natom, - const float *coord, - const int *atype, - const float *cell, - float **tensor, - int *size) { + const float* coord, + const int* atype, + const float* cell, + float** tensor, + int* size) { DP_DeepTensorComputeTensorf(dt, natom, coord, atype, cell, tensor, size); } template -inline void _DP_DeepTensorComputeTensorNList(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeTensorNList(DP_DeepTensor* dt, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, const int nghost, - const DP_Nlist *nlist, - FPTYPE **tensor, - int *size); + const DP_Nlist* nlist, + FPTYPE** tensor, + int* size); template <> -inline void _DP_DeepTensorComputeTensorNList(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeTensorNList(DP_DeepTensor* dt, const int natom, - const double *coord, - const int *atype, - const double *cell, + const double* coord, + const int* atype, + const double* cell, const int nghost, - const DP_Nlist *nlist, - double **tensor, - int *size) { + const DP_Nlist* nlist, + double** tensor, + int* size) { DP_DeepTensorComputeTensorNList(dt, natom, coord, atype, cell, nghost, nlist, tensor, size); } template <> -inline void _DP_DeepTensorComputeTensorNList(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeTensorNList(DP_DeepTensor* dt, const int natom, - const float *coord, - const int *atype, - const float *cell, + const float* coord, + const int* atype, + const float* cell, const int nghost, - const DP_Nlist *nlist, - float **tensor, - int *size) { + const DP_Nlist* nlist, + float** tensor, + int* size) { DP_DeepTensorComputeTensorNListf(dt, natom, coord, atype, cell, nghost, nlist, tensor, size); } template -inline void _DP_DeepTensorCompute(DP_DeepTensor *dt, +inline void _DP_DeepTensorCompute(DP_DeepTensor* dt, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, - FPTYPE *global_tensor, - FPTYPE *force, - FPTYPE *virial, - FPTYPE **atomic_energy, - FPTYPE *atomic_virial, - int *size_at); + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, + FPTYPE* global_tensor, + FPTYPE* force, + FPTYPE* virial, + FPTYPE** atomic_energy, + FPTYPE* atomic_virial, + int* size_at); template <> -inline void _DP_DeepTensorCompute(DP_DeepTensor *dt, +inline void _DP_DeepTensorCompute(DP_DeepTensor* dt, const int natom, - const double *coord, - const int *atype, - const double *cell, - double *global_tensor, - double *force, - double *virial, - double **atomic_tensor, - double *atomic_virial, - int *size_at) { + const double* coord, + const int* atype, + const double* cell, + double* global_tensor, + double* force, + double* virial, + double** atomic_tensor, + double* atomic_virial, + int* size_at) { DP_DeepTensorCompute(dt, natom, coord, atype, cell, global_tensor, force, virial, atomic_tensor, atomic_virial, size_at); } template <> -inline void _DP_DeepTensorCompute(DP_DeepTensor *dt, +inline void _DP_DeepTensorCompute(DP_DeepTensor* dt, const int natom, - const float *coord, - const int *atype, - const float *cell, - float *global_tensor, - float *force, - float *virial, - float **atomic_tensor, - float *atomic_virial, - int *size_at) { + const float* coord, + const int* atype, + const float* cell, + float* global_tensor, + float* force, + float* virial, + float** atomic_tensor, + float* atomic_virial, + int* size_at) { DP_DeepTensorComputef(dt, natom, coord, atype, cell, global_tensor, force, virial, atomic_tensor, atomic_virial, size_at); } template -inline void _DP_DeepTensorComputeNList(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeNList(DP_DeepTensor* dt, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, const int nghost, - const DP_Nlist *nlist, - FPTYPE *global_tensor, - FPTYPE *force, - FPTYPE *virial, - FPTYPE **atomic_energy, - FPTYPE *atomic_virial, - int *size_at); + const DP_Nlist* nlist, + FPTYPE* global_tensor, + FPTYPE* force, + FPTYPE* virial, + FPTYPE** atomic_energy, + FPTYPE* atomic_virial, + int* size_at); template <> -inline void _DP_DeepTensorComputeNList(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeNList(DP_DeepTensor* dt, const int natom, - const double *coord, - const int *atype, - const double *cell, + const double* coord, + const int* atype, + const double* cell, const int nghost, - const DP_Nlist *nlist, - double *global_tensor, - double *force, - double *virial, - double **atomic_tensor, - double *atomic_virial, - int *size_at) { + const DP_Nlist* nlist, + double* global_tensor, + double* force, + double* virial, + double** atomic_tensor, + double* atomic_virial, + int* size_at) { DP_DeepTensorComputeNList(dt, natom, coord, atype, cell, nghost, nlist, global_tensor, force, virial, atomic_tensor, atomic_virial, size_at); } template <> -inline void _DP_DeepTensorComputeNList(DP_DeepTensor *dt, +inline void _DP_DeepTensorComputeNList(DP_DeepTensor* dt, const int natom, - const float *coord, - const int *atype, - const float *cell, + const float* coord, + const int* atype, + const float* cell, const int nghost, - const DP_Nlist *nlist, - float *global_tensor, - float *force, - float *virial, - float **atomic_tensor, - float *atomic_virial, - int *size_at) { + const DP_Nlist* nlist, + float* global_tensor, + float* force, + float* virial, + float** atomic_tensor, + float* atomic_virial, + int* size_at) { DP_DeepTensorComputeNListf(dt, natom, coord, atype, cell, nghost, nlist, global_tensor, force, virial, atomic_tensor, atomic_virial, size_at); } template -inline void _DP_DipoleChargeModifierComputeNList(DP_DipoleChargeModifier *dcm, +inline void _DP_DipoleChargeModifierComputeNList(DP_DipoleChargeModifier* dcm, const int natom, - const FPTYPE *coord, - const int *atype, - const FPTYPE *cell, - const int *pairs, + const FPTYPE* coord, + const int* atype, + const FPTYPE* cell, + const int* pairs, const int npairs, - const FPTYPE *delef_, + const FPTYPE* delef_, const int nghost, - const DP_Nlist *nlist, - FPTYPE *dfcorr_, - FPTYPE *dvcorr_); + const DP_Nlist* nlist, + FPTYPE* dfcorr_, + FPTYPE* dvcorr_); template <> inline void _DP_DipoleChargeModifierComputeNList( - DP_DipoleChargeModifier *dcm, + DP_DipoleChargeModifier* dcm, const int natom, - const double *coord, - const int *atype, - const double *cell, - const int *pairs, + const double* coord, + const int* atype, + const double* cell, + const int* pairs, const int npairs, - const double *delef_, + const double* delef_, const int nghost, - const DP_Nlist *nlist, - double *dfcorr_, - double *dvcorr_) { + const DP_Nlist* nlist, + double* dfcorr_, + double* dvcorr_) { DP_DipoleChargeModifierComputeNList(dcm, natom, coord, atype, cell, pairs, npairs, delef_, nghost, nlist, dfcorr_, dvcorr_); @@ -770,30 +770,30 @@ inline void _DP_DipoleChargeModifierComputeNList( template <> inline void _DP_DipoleChargeModifierComputeNList( - DP_DipoleChargeModifier *dcm, + DP_DipoleChargeModifier* dcm, const int natom, - const float *coord, - const int *atype, - const float *cell, - const int *pairs, + const float* coord, + const int* atype, + const float* cell, + const int* pairs, const int npairs, - const float *delef_, + const float* delef_, const int nghost, - const DP_Nlist *nlist, - float *dfcorr_, - float *dvcorr_) { + const DP_Nlist* nlist, + float* dfcorr_, + float* dvcorr_) { DP_DipoleChargeModifierComputeNListf(dcm, natom, coord, atype, cell, pairs, npairs, delef_, nghost, nlist, dfcorr_, dvcorr_); } -inline double *_DP_Get_Energy_Pointer(std::vector &vec, +inline double* _DP_Get_Energy_Pointer(std::vector& vec, const int nframes) { vec.resize(nframes); return &vec[0]; } -inline double *_DP_Get_Energy_Pointer(double &vec, const int nframes) { +inline double* _DP_Get_Energy_Pointer(double& vec, const int nframes) { assert(nframes == 1); return &vec; } @@ -812,7 +812,7 @@ struct InputNlist { nl(DP_NewNlist(0, nullptr, nullptr, nullptr)) { DP_CHECK_OK(DP_NlistCheckOK, nl); }; - InputNlist(int inum_, int *ilist_, int *numneigh_, int **firstneigh_) + InputNlist(int inum_, int* ilist_, int* numneigh_, int** firstneigh_) : inum(inum_), ilist(ilist_), numneigh(numneigh_), @@ -821,17 +821,17 @@ struct InputNlist { DP_CHECK_OK(DP_NlistCheckOK, nl); }; InputNlist(int inum_, - int *ilist_, - int *numneigh_, - int **firstneigh_, + int* ilist_, + int* numneigh_, + int** firstneigh_, int nswap, - int *sendnum, - int *recvnum, - int *firstrecv, - int **sendlist, - int *sendproc, - int *recvproc, - void *world) + int* sendnum, + int* recvnum, + int* firstrecv, + int** sendlist, + int* sendproc, + int* recvproc, + void* world) : inum(inum_), ilist(ilist_), numneigh(numneigh_), @@ -850,15 +850,15 @@ struct InputNlist { world)) {}; ~InputNlist() { DP_DeleteNlist(nl); }; /// @brief C API neighbor list. - DP_Nlist *nl; + DP_Nlist* nl; /// @brief Number of core region atoms int inum; /// @brief Array stores the core region atom's index - int *ilist; + int* ilist; /// @brief Array stores the core region atom's neighbor atom number - int *numneigh; + int* numneigh; /// @brief Array stores the core region atom's neighbor index - int **firstneigh; + int** firstneigh; /** * @brief Set mask for this neighbor list. */ @@ -867,7 +867,7 @@ struct InputNlist { * @brief Set mapping for this neighbor list. * @param mapping mapping from all atoms to real atoms, in size nall. */ - void set_mapping(int *mapping) { DP_NlistSetMapping(nl, mapping); }; + void set_mapping(int* mapping) { DP_NlistSetMapping(nl, mapping); }; }; /** @@ -884,8 +884,8 @@ void inline convert_pbtxt_to_pb(std::string fn_pb_txt, std::string fn_pb) { * @param[in] from_nlist 2D int vector. The first axis represents the centeral * atoms and the second axis represents the neighbor atoms. */ -void inline convert_nlist(InputNlist &to_nlist, - std::vector> &from_nlist) { +void inline convert_nlist(InputNlist& to_nlist, + std::vector>& from_nlist) { to_nlist.inum = from_nlist.size(); for (int ii = 0; ii < to_nlist.inum; ++ii) { to_nlist.ilist[ii] = ii; @@ -936,8 +936,8 @@ class DeepBaseModel { * @brief Get the type map (element name of the atom types) of this model. * @param[out] type_map The type map of this model. **/ - void get_type_map(std::string &type_map) { - const char *type_map_c = DP_DeepBaseModelGetTypeMap(dpbase); + void get_type_map(std::string& type_map) { + const char* type_map_c = DP_DeepBaseModelGetTypeMap(dpbase); type_map.assign(type_map_c); DP_DeleteChar(type_map_c); }; @@ -946,7 +946,7 @@ class DeepBaseModel { * information. * @param[in] pre The prefix to each line. */ - void print_summary(const std::string &pre) const { + void print_summary(const std::string& pre) const { DP_PrintSummary(pre.c_str()); } /** @@ -967,15 +967,15 @@ class DeepBaseModel { } protected: - DP_DeepBaseModel *dpbase; + DP_DeepBaseModel* dpbase; int dfparam; int daparam; bool aparam_nall; template - void validate_fparam_aparam(const int &nframes, - const int &nloc, - const std::vector &fparam, - const std::vector &aparam) const { + void validate_fparam_aparam(const int& nframes, + const int& nloc, + const std::vector& fparam, + const std::vector& aparam) const { if (fparam.size() != dfparam && fparam.size() != static_cast(nframes) * dfparam) { throw deepmd::hpp::deepmd_exception( @@ -991,10 +991,10 @@ class DeepBaseModel { } } template - void tile_fparam_aparam(std::vector &out_param, - const int &nframes, - const int &dparam, - const std::vector ¶m) const { + void tile_fparam_aparam(std::vector& out_param, + const int& nframes, + const int& dparam, + const std::vector& param) const { if (param.size() == dparam) { out_param.resize(static_cast(nframes) * dparam); for (int ii = 0; ii < nframes; ++ii) { @@ -1023,9 +1023,9 @@ class DeepPot : public DeepBaseModel { * @param[in] gpu_rank The GPU rank. * @param[in] file_content The content of the frozen model file. **/ - DeepPot(const std::string &model, - const int &gpu_rank = 0, - const std::string &file_content = "") + DeepPot(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = "") : dp(nullptr) { try { init(model, gpu_rank, file_content); @@ -1043,9 +1043,9 @@ class DeepPot : public DeepBaseModel { * @param[in] gpu_rank The GPU rank. * @param[in] file_content The content of the frozen model file. **/ - void init(const std::string &model, - const int &gpu_rank = 0, - const std::string &file_content = "") { + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = "") { if (dp) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" @@ -1058,7 +1058,7 @@ class DeepPot : public DeepBaseModel { dfparam = DP_DeepPotGetDimFParam(dp); daparam = DP_DeepPotGetDimAParam(dp); aparam_nall = DP_DeepPotIsAParamNAll(dp); - dpbase = (DP_DeepBaseModel *)dp; + dpbase = (DP_DeepBaseModel*)dp; }; /** @@ -1083,34 +1083,34 @@ class DeepPot : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, virial_, @@ -1142,41 +1142,41 @@ class DeepPot : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &virial, - std::vector &atom_energy, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); atom_energy.resize(static_cast(nframes) * natoms); atom_virial.resize(static_cast(nframes) * natoms * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_ener_ = &atom_energy[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_ener_ = &atom_energy[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotCompute(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, virial_, @@ -1210,31 +1210,31 @@ class DeepPot : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -1242,8 +1242,8 @@ class DeepPot : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotComputeNList( dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, @@ -1278,38 +1278,38 @@ class DeepPot : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &virial, - std::vector &atom_energy, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); atom_energy.resize(static_cast(nframes) * natoms); atom_virial.resize(static_cast(nframes) * natoms * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_ener_ = &atom_energy[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_ener_ = &atom_energy[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -1317,8 +1317,8 @@ class DeepPot : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotComputeNList(dp, nframes, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, @@ -1349,34 +1349,34 @@ class DeepPot : public DeepBaseModel { **/ template void compute_mixed_type( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &virial, - const int &nframes, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size() / nframes; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotComputeMixedType(dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, @@ -1408,41 +1408,41 @@ class DeepPot : public DeepBaseModel { **/ template void compute_mixed_type( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &virial, - std::vector &atom_energy, - std::vector &atom_virial, - const int &nframes, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + ENERGYVTYPE& ener, + std::vector& force, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const int& nframes, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size() / nframes; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); atom_energy.resize(static_cast(nframes) * natoms); atom_virial.resize(static_cast(nframes) * natoms * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_ener_ = &atom_energy[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_ener_ = &atom_energy[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotComputeMixedType( dp, nframes, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, @@ -1451,7 +1451,7 @@ class DeepPot : public DeepBaseModel { }; private: - DP_DeepPot *dp; + DP_DeepPot* dp; }; class DeepSpin : public DeepBaseModel { @@ -1467,9 +1467,9 @@ class DeepSpin : public DeepBaseModel { * @param[in] gpu_rank The GPU rank. * @param[in] file_content The content of the frozen model file. **/ - DeepSpin(const std::string &model, - const int &gpu_rank = 0, - const std::string &file_content = "") + DeepSpin(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = "") : dp(nullptr) { try { init(model, gpu_rank, file_content); @@ -1487,9 +1487,9 @@ class DeepSpin : public DeepBaseModel { * @param[in] gpu_rank The GPU rank. * @param[in] file_content The content of the frozen model file. **/ - void init(const std::string &model, - const int &gpu_rank = 0, - const std::string &file_content = "") { + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& file_content = "") { if (dp) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" @@ -1502,7 +1502,7 @@ class DeepSpin : public DeepBaseModel { dfparam = DP_DeepSpinGetDimFParam(dp); daparam = DP_DeepSpinGetDimAParam(dp); aparam_nall = DP_DeepSpinIsAParamNAll(dp); - dpbase = (DP_DeepBaseModel *)dp; + dpbase = (DP_DeepBaseModel*)dp; }; /** @@ -1531,39 +1531,39 @@ class DeepSpin : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &force_mag, - std::vector &virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); force_mag.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *force_mag_ = &force_mag[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* force_mag_ = &force_mag[0]; + VALUETYPE* virial_ = &virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinCompute(dp, nframes, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, ener_, force_, @@ -1599,46 +1599,46 @@ class DeepSpin : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &force_mag, - std::vector &virial, - std::vector &atom_energy, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); force_mag.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); atom_energy.resize(static_cast(nframes) * natoms); atom_virial.resize(static_cast(nframes) * natoms * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *force_mag_ = &force_mag[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_ener_ = &atom_energy[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* force_mag_ = &force_mag[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_ener_ = &atom_energy[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinCompute( dp, nframes, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, @@ -1675,36 +1675,36 @@ class DeepSpin : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &force_mag, - std::vector &virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, + ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); force_mag.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *force_mag_ = &force_mag[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* force_mag_ = &force_mag[0]; + VALUETYPE* virial_ = &virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -1712,8 +1712,8 @@ class DeepSpin : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinComputeNList(dp, nframes, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, aparam__, ener_, force_, @@ -1752,42 +1752,42 @@ class DeepSpin : public DeepBaseModel { **/ template void compute( - ENERGYVTYPE &ener, - std::vector &force, - std::vector &force_mag, - std::vector &virial, - std::vector &atom_energy, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, + ENERGYVTYPE& ener, + std::vector& force, + std::vector& force_mag, + std::vector& virial, + std::vector& atom_energy, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = natoms > 0 ? coord.size() / natoms / 3 : 1; assert(nframes * natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == nframes * 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; - double *ener_ = _DP_Get_Energy_Pointer(ener, nframes); + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; + double* ener_ = _DP_Get_Energy_Pointer(ener, nframes); force.resize(static_cast(nframes) * natoms * 3); force_mag.resize(static_cast(nframes) * natoms * 3); virial.resize(static_cast(nframes) * 9); atom_energy.resize(static_cast(nframes) * natoms); atom_virial.resize(static_cast(nframes) * natoms * 9); - VALUETYPE *force_ = &force[0]; - VALUETYPE *force_mag_ = &force_mag[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_ener_ = &atom_energy[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* force_mag_ = &force_mag[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_ener_ = &atom_energy[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -1795,8 +1795,8 @@ class DeepSpin : public DeepBaseModel { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinComputeNList( dp, nframes, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, aparam__, ener_, force_, force_mag_, virial_, @@ -1805,7 +1805,7 @@ class DeepSpin : public DeepBaseModel { }; private: - DP_DeepSpin *dp; + DP_DeepSpin* dp; }; /** @@ -1865,8 +1865,8 @@ class DeepBaseModelDevi { * @param[in] xx The vectors of all models. **/ template - void compute_avg(std::vector &avg, - const std::vector> &xx) { + void compute_avg(std::vector& avg, + const std::vector>& xx) { assert(xx.size() == numb_models); if (numb_models == 0) { return; @@ -1893,10 +1893,10 @@ class DeepBaseModelDevi { * @param[in] stride The stride to compute the deviation. **/ template - void compute_std(std::vector &std, - const std::vector &avg, - const std::vector> &xx, - const int &stride) { + void compute_std(std::vector& std, + const std::vector& avg, + const std::vector>& xx, + const int& stride) { assert(xx.size() == numb_models); if (numb_models == 0) { return; @@ -1911,8 +1911,8 @@ class DeepBaseModelDevi { for (unsigned ii = 0; ii < numb_models; ++ii) { for (unsigned jj = 0; jj < nloc; ++jj) { - const VALUETYPE *tmp_f = &(xx[ii][static_cast(jj) * stride]); - const VALUETYPE *tmp_avg = &(avg[static_cast(jj) * stride]); + const VALUETYPE* tmp_f = &(xx[ii][static_cast(jj) * stride]); + const VALUETYPE* tmp_avg = &(avg[static_cast(jj) * stride]); for (unsigned dd = 0; dd < stride; ++dd) { VALUETYPE vdiff = tmp_f[dd] - tmp_avg[dd]; std[jj] += vdiff * vdiff; @@ -1932,16 +1932,16 @@ class DeepBaseModelDevi { * @param[in] stride The stride to compute the deviation. **/ template - void compute_relative_std(std::vector &std, - const std::vector &avg, + void compute_relative_std(std::vector& std, + const std::vector& avg, const VALUETYPE eps, - const int &stride) { + const int& stride) { unsigned ndof = avg.size(); unsigned nloc = std.size(); assert(nloc * stride == ndof); for (unsigned ii = 0; ii < nloc; ++ii) { - const VALUETYPE *tmp_avg = &(avg[static_cast(ii) * stride]); + const VALUETYPE* tmp_avg = &(avg[static_cast(ii) * stride]); VALUETYPE f_norm = 0.0; for (unsigned dd = 0; dd < stride; ++dd) { f_norm += tmp_avg[dd] * tmp_avg[dd]; @@ -1957,9 +1957,9 @@ class DeepBaseModelDevi { * @param[in] xx The vectors of all forces. **/ template - void compute_std_f(std::vector &std, - const std::vector &avg, - const std::vector> &xx) { + void compute_std_f(std::vector& std, + const std::vector& avg, + const std::vector>& xx) { compute_std(std, avg, xx, 3); }; /** @@ -1969,23 +1969,23 @@ class DeepBaseModelDevi { * @param[in] eps The level parameter for computing the deviation. **/ template - void compute_relative_std_f(std::vector &std, - const std::vector &avg, + void compute_relative_std_f(std::vector& std, + const std::vector& avg, const VALUETYPE eps) { compute_relative_std(std, avg, eps, 3); }; protected: - DP_DeepBaseModelDevi *dpbase; + DP_DeepBaseModelDevi* dpbase; int numb_models; int dfparam; int daparam; bool aparam_nall; template - void validate_fparam_aparam(const int &nframes, - const int &nloc, - const std::vector &fparam, - const std::vector &aparam) const { + void validate_fparam_aparam(const int& nframes, + const int& nloc, + const std::vector& fparam, + const std::vector& aparam) const { if (fparam.size() != dfparam && fparam.size() != static_cast(nframes) * dfparam) { throw deepmd::hpp::deepmd_exception( @@ -2001,10 +2001,10 @@ class DeepBaseModelDevi { } } template - void tile_fparam_aparam(std::vector &out_param, - const int &nframes, - const int &dparam, - const std::vector ¶m) const { + void tile_fparam_aparam(std::vector& out_param, + const int& nframes, + const int& dparam, + const std::vector& param) const { if (param.size() == dparam) { out_param.resize(static_cast(nframes) * dparam); for (int ii = 0; ii < nframes; ++ii) { @@ -2031,7 +2031,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { * @brief DP model deviation constructor with initialization. * @param[in] models The names of the frozen model file. **/ - DeepPotModelDevi(const std::vector &models) : dp(nullptr) { + DeepPotModelDevi(const std::vector& models) : dp(nullptr) { try { init(models); } catch (...) { @@ -2048,9 +2048,9 @@ class DeepPotModelDevi : public DeepBaseModelDevi { * @param[in] gpu_rank The GPU rank. * @param[in] file_content The content of the frozen model file. **/ - void init(const std::vector &models, - const int &gpu_rank = 0, - const std::vector &file_content = + void init(const std::vector& models, + const int& gpu_rank = 0, + const std::vector& file_content = std::vector()) { if (dp) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " @@ -2058,17 +2058,17 @@ class DeepPotModelDevi : public DeepBaseModelDevi { << std::endl; return; } - std::vector cstrings; + std::vector cstrings; cstrings.reserve(models.size()); - for (std::string const &str : models) { + for (std::string const& str : models) { cstrings.push_back(str.data()); } - std::vector c_file_contents; + std::vector c_file_contents; std::vector size_file_contents; c_file_contents.reserve(file_content.size()); size_file_contents.reserve(file_content.size()); - for (std::string const &str : file_content) { + for (std::string const& str : file_content) { c_file_contents.push_back(str.data()); size_file_contents.push_back(str.size()); } @@ -2081,7 +2081,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { dfparam = DP_DeepPotModelDeviGetDimFParam(dp); daparam = DP_DeepPotModelDeviGetDimAParam(dp); aparam_nall = DP_DeepPotModelDeviIsAParamNAll(dp); - dpbase = (DP_DeepBaseModelDevi *)dp; + dpbase = (DP_DeepBaseModelDevi*)dp; }; /** @@ -2106,23 +2106,23 @@ class DeepPotModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + std::vector& ener, + std::vector>& force, + std::vector>& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; // memory will be continuous for std::vector but not // std::vector @@ -2130,15 +2130,15 @@ class DeepPotModelDevi : public DeepBaseModelDevi { std::vector force_flat(static_cast(numb_models) * natoms * 3); std::vector virial_flat(numb_models * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotModelDeviCompute(dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, @@ -2185,25 +2185,25 @@ class DeepPotModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &virial, - std::vector> &atom_energy, - std::vector> &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + std::vector& ener, + std::vector>& force, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; std::vector energy_flat(numb_models); std::vector force_flat(static_cast(numb_models) * @@ -2213,17 +2213,17 @@ class DeepPotModelDevi : public DeepBaseModelDevi { natoms); std::vector atom_virial_flat(static_cast(numb_models) * natoms * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; - VALUETYPE *atomic_ener_ = &atom_energy_flat[0]; - VALUETYPE *atomic_virial_ = &atom_virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; + VALUETYPE* atomic_ener_ = &atom_energy_flat[0]; + VALUETYPE* atomic_virial_ = &atom_virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotModelDeviCompute( dp, natoms, coord_, atype_, box_, fparam__, aparam__, ener_, force_, @@ -2282,26 +2282,26 @@ class DeepPotModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + std::vector& ener, + std::vector>& force, + std::vector>& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; // memory will be continuous for std::vector but not // std::vector @@ -2309,9 +2309,9 @@ class DeepPotModelDevi : public DeepBaseModelDevi { std::vector force_flat(static_cast(numb_models) * natoms * 3); std::vector virial_flat(numb_models * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -2319,8 +2319,8 @@ class DeepPotModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, @@ -2370,28 +2370,28 @@ class DeepPotModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &virial, - std::vector> &atom_energy, - std::vector> &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + std::vector& ener, + std::vector>& force, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; std::vector energy_flat(numb_models); std::vector force_flat(static_cast(numb_models) * @@ -2401,11 +2401,11 @@ class DeepPotModelDevi : public DeepBaseModelDevi { natoms); std::vector atom_virial_flat(static_cast(numb_models) * natoms * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; - VALUETYPE *atomic_ener_ = &atom_energy_flat[0]; - VALUETYPE *atomic_virial_ = &atom_virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; + VALUETYPE* atomic_ener_ = &atom_energy_flat[0]; + VALUETYPE* atomic_virial_ = &atom_virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -2413,8 +2413,8 @@ class DeepPotModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepPotModelDeviComputeNList( dp, natoms, coord_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, @@ -2449,7 +2449,7 @@ class DeepPotModelDevi : public DeepBaseModelDevi { }; private: - DP_DeepPotModelDevi *dp; + DP_DeepPotModelDevi* dp; }; class DeepSpinModelDevi : public DeepBaseModelDevi { @@ -2463,7 +2463,7 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * @brief DP model deviation constructor with initialization. * @param[in] models The names of the frozen model file. **/ - DeepSpinModelDevi(const std::vector &models) : dp(nullptr) { + DeepSpinModelDevi(const std::vector& models) : dp(nullptr) { try { init(models); } catch (...) { @@ -2480,9 +2480,9 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { * @param[in] gpu_rank The GPU rank. * @param[in] file_content The content of the frozen model file. **/ - void init(const std::vector &models, - const int &gpu_rank = 0, - const std::vector &file_content = + void init(const std::vector& models, + const int& gpu_rank = 0, + const std::vector& file_content = std::vector()) { if (dp) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " @@ -2490,17 +2490,17 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { << std::endl; return; } - std::vector cstrings; + std::vector cstrings; cstrings.reserve(models.size()); - for (std::string const &str : models) { + for (std::string const& str : models) { cstrings.push_back(str.data()); } - std::vector c_file_contents; + std::vector c_file_contents; std::vector size_file_contents; c_file_contents.reserve(file_content.size()); size_file_contents.reserve(file_content.size()); - for (std::string const &str : file_content) { + for (std::string const& str : file_content) { c_file_contents.push_back(str.data()); size_file_contents.push_back(str.size()); } @@ -2513,7 +2513,7 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { dfparam = DP_DeepSpinModelDeviGetDimFParam(dp); daparam = DP_DeepSpinModelDeviGetDimAParam(dp); aparam_nall = DP_DeepSpinModelDeviIsAParamNAll(dp); - dpbase = (DP_DeepBaseModelDevi *)dp; + dpbase = (DP_DeepBaseModelDevi*)dp; }; /** @@ -2541,26 +2541,26 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &force_mag, - std::vector> &virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + std::vector& ener, + std::vector>& force, + std::vector>& force_mag, + std::vector>& virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; // memory will be continuous for std::vector but not // std::vector @@ -2570,16 +2570,16 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { std::vector force_mag_flat(static_cast(numb_models) * natoms * 3); std::vector virial_flat(numb_models * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *force_mag_ = &force_mag_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* force_mag_ = &force_mag_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinModelDeviCompute( dp, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, ener_, @@ -2634,28 +2634,28 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &force_mag, - std::vector> &virial, - std::vector> &atom_energy, - std::vector> &atom_virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + std::vector& ener, + std::vector>& force, + std::vector>& force_mag, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; std::vector energy_flat(numb_models); std::vector force_flat(static_cast(numb_models) * @@ -2667,18 +2667,18 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { natoms); std::vector atom_virial_flat(static_cast(numb_models) * natoms * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *force_mag_ = &force_mag_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; - VALUETYPE *atomic_ener_ = &atom_energy_flat[0]; - VALUETYPE *atomic_virial_ = &atom_virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* force_mag_ = &force_mag_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; + VALUETYPE* atomic_ener_ = &atom_energy_flat[0]; + VALUETYPE* atomic_virial_ = &atom_virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, natoms, fparam, aparam); tile_fparam_aparam(fparam_, nframes, dfparam, fparam); tile_fparam_aparam(aparam_, nframes, natoms * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinModelDeviCompute( dp, natoms, coord_, spin_, atype_, box_, fparam__, aparam__, ener_, @@ -2745,29 +2745,29 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &force_mag, - std::vector> &virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, + std::vector& ener, + std::vector>& force, + std::vector>& force_mag, + std::vector>& virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; // memory will be continous for std::vector but not std::vector std::vector energy_flat(numb_models); std::vector force_flat(static_cast(numb_models) * @@ -2775,10 +2775,10 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { std::vector force_mag_flat(static_cast(numb_models) * natoms * 3); std::vector virial_flat(numb_models * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *force_mag_ = &force_mag_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* force_mag_ = &force_mag_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -2786,8 +2786,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinModelDeviComputeNList( dp, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, aparam__, ener_, force_, force_mag_, virial_, nullptr, @@ -2845,31 +2845,31 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { **/ template void compute( - std::vector &ener, - std::vector> &force, - std::vector> &force_mag, - std::vector> &virial, - std::vector> &atom_energy, - std::vector> &atom_virial, - const std::vector &coord, - const std::vector &spin, - const std::vector &atype, - const std::vector &box, + std::vector& ener, + std::vector>& force, + std::vector>& force_mag, + std::vector>& virial, + std::vector>& atom_energy, + std::vector>& atom_virial, + const std::vector& coord, + const std::vector& spin, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list, - const int &ago, - const std::vector &fparam = std::vector(), - const std::vector &aparam = std::vector()) { + const InputNlist& lmp_list, + const int& ago, + const std::vector& fparam = std::vector(), + const std::vector& aparam = std::vector()) { unsigned int natoms = atype.size(); unsigned int nframes = 1; assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *spin_ = &spin[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* spin_ = &spin[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; std::vector energy_flat(numb_models); std::vector force_flat(static_cast(numb_models) * natoms * 3); @@ -2880,12 +2880,12 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { natoms); std::vector atom_virial_flat(static_cast(numb_models) * natoms * 9); - double *ener_ = &energy_flat[0]; - VALUETYPE *force_ = &force_flat[0]; - VALUETYPE *force_mag_ = &force_mag_flat[0]; - VALUETYPE *virial_ = &virial_flat[0]; - VALUETYPE *atomic_ener_ = &atom_energy_flat[0]; - VALUETYPE *atomic_virial_ = &atom_virial_flat[0]; + double* ener_ = &energy_flat[0]; + VALUETYPE* force_ = &force_flat[0]; + VALUETYPE* force_mag_ = &force_mag_flat[0]; + VALUETYPE* virial_ = &virial_flat[0]; + VALUETYPE* atomic_ener_ = &atom_energy_flat[0]; + VALUETYPE* atomic_virial_ = &atom_virial_flat[0]; std::vector fparam_, aparam_; validate_fparam_aparam(nframes, (aparam_nall ? natoms : (natoms - nghost)), fparam, aparam); @@ -2893,8 +2893,8 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { tile_fparam_aparam(aparam_, nframes, (aparam_nall ? natoms : (natoms - nghost)) * daparam, aparam); - const VALUETYPE *fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; - const VALUETYPE *aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; + const VALUETYPE* fparam__ = !fparam_.empty() ? &fparam_[0] : nullptr; + const VALUETYPE* aparam__ = !aparam_.empty() ? &aparam_[0] : nullptr; _DP_DeepSpinModelDeviComputeNList( dp, natoms, coord_, spin_, atype_, box_, nghost, lmp_list.nl, ago, fparam__, aparam__, ener_, force_, force_mag_, virial_, atomic_ener_, @@ -2933,7 +2933,7 @@ class DeepSpinModelDevi : public DeepBaseModelDevi { }; private: - DP_DeepSpinModelDevi *dp; + DP_DeepSpinModelDevi* dp; }; /** @@ -2950,9 +2950,9 @@ class DeepTensor { * @brief DeepTensor constructor with initialization. * @param[in] model The name of the frozen model file. **/ - DeepTensor(const std::string &model, - const int &gpu_rank = 0, - const std::string &name_scope = "") + DeepTensor(const std::string& model, + const int& gpu_rank = 0, + const std::string& name_scope = "") : dt(nullptr) { try { init(model, gpu_rank, name_scope); @@ -2968,9 +2968,9 @@ class DeepTensor { * @brief Initialize the DeepTensor. * @param[in] model The name of the frozen model file. **/ - void init(const std::string &model, - const int &gpu_rank = 0, - const std::string &name_scope = "") { + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& name_scope = "") { if (dt) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" @@ -2993,23 +2993,23 @@ class DeepTensor { *x 9 (PBC) or empty (no PBC). **/ template - void compute(std::vector &tensor, - const std::vector &coord, - const std::vector &atype, - const std::vector &box) { + void compute(std::vector& tensor, + const std::vector& coord, + const std::vector& atype, + const std::vector& box) { unsigned int natoms = atype.size(); assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; - VALUETYPE *tensor_; - VALUETYPE **p_tensor = &tensor_; + VALUETYPE* tensor_; + VALUETYPE** p_tensor = &tensor_; int size; - int *p_size = &size; + int* p_size = &size; _DP_DeepTensorComputeTensor(dt, natoms, coord_, atype_, box_, p_tensor, p_size); @@ -3033,25 +3033,25 @@ class DeepTensor { * @param[in] nlist The neighbor list. **/ template - void compute(std::vector &tensor, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + void compute(std::vector& tensor, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { unsigned int natoms = atype.size(); assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; - VALUETYPE *tensor_; - VALUETYPE **p_tensor = &tensor_; + VALUETYPE* tensor_; + VALUETYPE** p_tensor = &tensor_; int size; - int *p_size = &size; + int* p_size = &size; _DP_DeepTensorComputeTensorNList(dt, natoms, coord_, atype_, box_, nghost, lmp_list.nl, @@ -3076,26 +3076,26 @@ class DeepTensor { *x 9 (PBC) or empty (no PBC). **/ template - void compute(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box) { + void compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box) { unsigned int natoms = atype.size(); assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; global_tensor.resize(odim); force.resize(static_cast(odim) * natoms * 3); virial.resize(static_cast(odim) * 9); - VALUETYPE *global_tensor_ = &global_tensor[0]; - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* global_tensor_ = &global_tensor[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; _DP_DeepTensorCompute(dt, natoms, coord_, atype_, box_, global_tensor_, force_, virial_, nullptr, @@ -3117,36 +3117,36 @@ class DeepTensor { *x 9 (PBC) or empty (no PBC). **/ template - void compute(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - std::vector &atom_tensor, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box) { + void compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box) { unsigned int natoms = atype.size(); assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; global_tensor.resize(odim); force.resize(static_cast(odim) * natoms * 3); virial.resize(static_cast(odim) * 9); atom_virial.resize(static_cast(odim) * natoms * 9); - VALUETYPE *global_tensor_ = &global_tensor[0]; - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* global_tensor_ = &global_tensor[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; - VALUETYPE *atomic_tensor_; - VALUETYPE **p_atomic_tensor = &atomic_tensor_; + VALUETYPE* atomic_tensor_; + VALUETYPE** p_atomic_tensor = &atomic_tensor_; int size_at; - int *p_size_at = &size_at; + int* p_size_at = &size_at; _DP_DeepTensorCompute( dt, natoms, coord_, atype_, box_, global_tensor_, force_, virial_, @@ -3173,28 +3173,28 @@ class DeepTensor { * @param[in] nlist The neighbor list. **/ template - void compute(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + void compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { unsigned int natoms = atype.size(); assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; global_tensor.resize(odim); force.resize(static_cast(odim) * natoms * 3); virial.resize(static_cast(odim) * 9); - VALUETYPE *global_tensor_ = &global_tensor[0]; - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; + VALUETYPE* global_tensor_ = &global_tensor[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; _DP_DeepTensorComputeNList( dt, natoms, coord_, atype_, box_, nghost, lmp_list.nl, global_tensor_, @@ -3218,38 +3218,38 @@ class DeepTensor { * @param[in] nlist The neighbor list. **/ template - void compute(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - std::vector &atom_tensor, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, + void compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { unsigned int natoms = atype.size(); assert(natoms * 3 == coord.size()); if (!box.empty()) { assert(box.size() == 9); } - const VALUETYPE *coord_ = &coord[0]; - const VALUETYPE *box_ = !box.empty() ? &box[0] : nullptr; - const int *atype_ = &atype[0]; + const VALUETYPE* coord_ = &coord[0]; + const VALUETYPE* box_ = !box.empty() ? &box[0] : nullptr; + const int* atype_ = &atype[0]; global_tensor.resize(odim); force.resize(static_cast(odim) * natoms * 3); virial.resize(static_cast(odim) * 9); atom_virial.resize(static_cast(odim) * natoms * 9); - VALUETYPE *global_tensor_ = &global_tensor[0]; - VALUETYPE *force_ = &force[0]; - VALUETYPE *virial_ = &virial[0]; - VALUETYPE *atomic_virial_ = &atom_virial[0]; + VALUETYPE* global_tensor_ = &global_tensor[0]; + VALUETYPE* force_ = &force[0]; + VALUETYPE* virial_ = &virial[0]; + VALUETYPE* atomic_virial_ = &atom_virial[0]; - VALUETYPE *atomic_tensor_; - VALUETYPE **p_atomic_tensor = &atomic_tensor_; + VALUETYPE* atomic_tensor_; + VALUETYPE** p_atomic_tensor = &atomic_tensor_; int size_at; - int *p_size_at = &size_at; + int* p_size_at = &size_at; _DP_DeepTensorComputeNList( dt, natoms, coord_, atype_, box_, nghost, lmp_list.nl, global_tensor_, @@ -3286,7 +3286,7 @@ class DeepTensor { } std::vector sel_types() const { - int *sel_types_arr = DP_DeepTensorGetSelTypes(dt); + int* sel_types_arr = DP_DeepTensorGetSelTypes(dt); std::vector sel_types_vec = std::vector(sel_types_arr, sel_types_arr + nsel_types); return sel_types_vec; @@ -3296,21 +3296,21 @@ class DeepTensor { * information. * @param[in] pre The prefix to each line. */ - void print_summary(const std::string &pre) const { + void print_summary(const std::string& pre) const { DP_PrintSummary(pre.c_str()); } /** * @brief Get the type map (element name of the atom types) of this model. * @param[out] type_map The type map of this model. **/ - void get_type_map(std::string &type_map) { - const char *type_map_c = DP_DeepTensorGetTypeMap(dt); + void get_type_map(std::string& type_map) { + const char* type_map_c = DP_DeepTensorGetTypeMap(dt); type_map.assign(type_map_c); DP_DeleteChar(type_map_c); }; private: - DP_DeepTensor *dt; + DP_DeepTensor* dt; int odim; int nsel_types; }; @@ -3328,9 +3328,9 @@ class DipoleChargeModifier { * @param[in] gpu_rank The rank of the GPU to be used. * @param[in] name_scope The name scope of the model. **/ - DipoleChargeModifier(const std::string &model, - const int &gpu_rank = 0, - const std::string &name_scope = "") + DipoleChargeModifier(const std::string& model, + const int& gpu_rank = 0, + const std::string& name_scope = "") : dcm(nullptr) { try { init(model, gpu_rank, name_scope); @@ -3348,9 +3348,9 @@ class DipoleChargeModifier { * @param[in] gpu_rank The rank of the GPU to be used. * @param[in] name_scope The name scope of the model. **/ - void init(const std::string &model, - const int &gpu_rank = 0, - const std::string &name_scope = "") { + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& name_scope = "") { if (dcm) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" @@ -3379,31 +3379,31 @@ class DipoleChargeModifier { * @param[in] lmp_list The neighbor list. **/ template - void compute(std::vector &dfcorr_, - std::vector &dvcorr_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, - const std::vector> &pairs, - const std::vector &delef_, + void compute(std::vector& dfcorr_, + std::vector& dvcorr_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, + const std::vector>& pairs, + const std::vector& delef_, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { unsigned int natoms = datype_.size(); assert(natoms * 3 == dcoord_.size()); if (!dbox.empty()) { assert(dbox.size() == 9); } - const VALUETYPE *dcoord = &dcoord_[0]; - const VALUETYPE *dbox_ = !dbox.empty() ? &dbox[0] : nullptr; - const int *datype = &datype_[0]; + const VALUETYPE* dcoord = &dcoord_[0]; + const VALUETYPE* dbox_ = !dbox.empty() ? &dbox[0] : nullptr; + const int* datype = &datype_[0]; const int npairs = pairs.size(); - const int *dpairs = reinterpret_cast(&pairs[0]); - const VALUETYPE *delef = &delef_[0]; + const int* dpairs = reinterpret_cast(&pairs[0]); + const VALUETYPE* delef = &delef_[0]; dfcorr_.resize(static_cast(natoms) * 3); dvcorr_.resize(9); - VALUETYPE *dfcorr = &dfcorr_[0]; - VALUETYPE *dvcorr = &dvcorr_[0]; + VALUETYPE* dfcorr = &dfcorr_[0]; + VALUETYPE* dvcorr = &dvcorr_[0]; _DP_DipoleChargeModifierComputeNList( dcm, natoms, dcoord, datype, dbox_, dpairs, npairs, delef, nghost, @@ -3428,7 +3428,7 @@ class DipoleChargeModifier { }; std::vector sel_types() const { - int *sel_types_arr = DP_DipoleChargeModifierGetSelTypes(dcm); + int* sel_types_arr = DP_DipoleChargeModifierGetSelTypes(dcm); std::vector sel_types_vec = std::vector(sel_types_arr, sel_types_arr + nsel_types); return sel_types_vec; @@ -3439,12 +3439,12 @@ class DipoleChargeModifier { * information. * @param[in] pre The prefix to each line. */ - void print_summary(const std::string &pre) const { + void print_summary(const std::string& pre) const { DP_PrintSummary(pre.c_str()); } private: - DP_DipoleChargeModifier *dcm; + DP_DipoleChargeModifier* dcm; int nsel_types; }; @@ -3453,9 +3453,9 @@ class DipoleChargeModifier { * @param[in] model Path to the model. * @param[out] file_content Content of the model file. **/ -void inline read_file_to_string(std::string model, std::string &file_content) { +void inline read_file_to_string(std::string model, std::string& file_content) { int size; - const char *c_file_content = DP_ReadFileToChar2(model.c_str(), &size); + const char* c_file_content = DP_ReadFileToChar2(model.c_str(), &size); if (size < 0) { // negative size indicates error std::string error_message = std::string(c_file_content, -size); @@ -3478,13 +3478,13 @@ void inline read_file_to_string(std::string model, std::string &file_content) { * @param[in] sel_type_ The selected atom types. */ template -void select_by_type(std::vector &fwd_map, - std::vector &bkw_map, - int &nghost_real, - const std::vector &dcoord_, - const std::vector &datype_, - const int &nghost, - const std::vector &sel_type_) { +void select_by_type(std::vector& fwd_map, + std::vector& bkw_map, + int& nghost_real, + const std::vector& dcoord_, + const std::vector& datype_, + const int& nghost, + const std::vector& sel_type_) { const int natoms = datype_.size(); const int nsel_type = sel_type_.size(); fwd_map.resize(natoms); @@ -3505,10 +3505,10 @@ void select_by_type(std::vector &fwd_map, * @param[in] stride The stride of the input vector. */ template -void select_map(std::vector &out, - const std::vector &in, - const std::vector &fwd_map, - const int &stride) { +void select_map(std::vector& out, + const std::vector& in, + const std::vector& fwd_map, + const int& stride) { static_assert(std::is_same(), "only support int"); const int nall1 = in.size() / stride; int nall2 = 0; diff --git a/source/api_c/tests/CMakeLists.txt b/source/api_c/tests/CMakeLists.txt index c42055ba6f..ef19305004 100644 --- a/source/api_c/tests/CMakeLists.txt +++ b/source/api_c/tests/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 3.9) +cmake_minimum_required(VERSION 3.25.2) project(deepmd_api_c_test) file(GLOB TEST_SRC test_*.cc) @@ -11,6 +11,19 @@ target_link_libraries(runUnitTests_c PRIVATE GTest::gtest_main ${LIB_DEEPMD_C} coverage_config) target_link_libraries(runUnitTests_c PRIVATE ${LIB_DEEPMD} ${LIB_DEEPMD_CC}) +if(ENABLE_TENSORFLOW) + target_compile_definitions(runUnitTests_c PRIVATE BUILD_TENSORFLOW) +endif() +if(ENABLE_PYTORCH) + target_compile_definitions(runUnitTests_c PRIVATE BUILD_PYTORCH) +endif() +if(ENABLE_JAX) + target_compile_definitions(runUnitTests_c PRIVATE BUILD_JAX) +endif() +if(ENABLE_PADDLE) + target_compile_definitions(runUnitTests_c PRIVATE BUILD_PADDLE) +endif() + if(UNIX AND NOT APPLE) find_library(RT_LIBRARY rt) if(RT_LIBRARY) diff --git a/source/api_c/tests/test_deepdipole_hpp.cc b/source/api_c/tests/test_deepdipole_hpp.cc index f781c34c5b..c4a2b75d09 100644 --- a/source/api_c/tests/test_deepdipole_hpp.cc +++ b/source/api_c/tests/test_deepdipole_hpp.cc @@ -26,6 +26,9 @@ class TestInferDeepDipole : public ::testing::Test { deepmd::hpp::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deepdipole.pbtxt", "deepdipole.pb"); @@ -221,6 +224,9 @@ class TestInferDeepDipoleNew : public ::testing::Test { deepmd::hpp::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deepdipole_new.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deepdipole_new.pbtxt", "deepdipole_new.pb"); @@ -434,6 +440,9 @@ class TestInferDeepDipoleFake : public ::testing::Test { deepmd::hpp::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deepdipole_fake.pbtxt", "deepdipole_fake.pb"); diff --git a/source/api_c/tests/test_deepmd_exception.cc b/source/api_c/tests/test_deepmd_exception.cc index f9f2984588..96f6942a65 100644 --- a/source/api_c/tests/test_deepmd_exception.cc +++ b/source/api_c/tests/test_deepmd_exception.cc @@ -16,7 +16,7 @@ TEST(TestDeepmdException, deepmdexception) { std::string expected_error_message = "DeePMD-kit C API Error: unittest"; try { throw deepmd::hpp::deepmd_exception("unittest"); - } catch (deepmd::hpp::deepmd_exception &ex) { + } catch (deepmd::hpp::deepmd_exception& ex) { EXPECT_STREQ(expected_error_message.c_str(), ex.what()); } } diff --git a/source/api_c/tests/test_deeppolar_hpp.cc b/source/api_c/tests/test_deeppolar_hpp.cc index 63ebf5d760..5e68b23bb4 100644 --- a/source/api_c/tests/test_deeppolar_hpp.cc +++ b/source/api_c/tests/test_deeppolar_hpp.cc @@ -32,6 +32,9 @@ class TestInferDeepPolar : public ::testing::Test { deepmd::hpp::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppolar.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deeppolar.pbtxt", "deeppolar.pb"); @@ -453,6 +456,9 @@ class TestInferDeepPolarNew : public ::testing::Test { deepmd::hpp::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppolar_new.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deeppolar_new.pbtxt", "deeppolar_new.pb"); diff --git a/source/api_c/tests/test_deeppot_a.cc b/source/api_c/tests/test_deeppot_a.cc index b4a9a81f92..d5e2d42bf8 100644 --- a/source/api_c/tests/test_deeppot_a.cc +++ b/source/api_c/tests/test_deeppot_a.cc @@ -61,9 +61,12 @@ class TestInferDeepPotA : public ::testing::Test { double expected_tot_e; std::vector expected_tot_v; - DP_DeepPot* dp; + DP_DeepPot* dp = nullptr; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif const char* file_name = "../../tests/infer/deeppot.pbtxt"; const char* model_file = "deeppot.pb"; DP_ConvertPbtxtToPb(file_name, model_file); @@ -242,9 +245,12 @@ class TestInferDeepPotANoPBC : public ::testing::Test { double expected_tot_e; std::vector expected_tot_v; - DP_DeepPot* dp; + DP_DeepPot* dp = nullptr; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif const char* file_name = "../../tests/infer/deeppot.pbtxt"; const char* model_file = "deeppot.pb"; DP_ConvertPbtxtToPb(file_name, model_file); diff --git a/source/api_c/tests/test_deeppot_a_fparam_aparam.cc b/source/api_c/tests/test_deeppot_a_fparam_aparam.cc index a728ede22d..635b47fe51 100644 --- a/source/api_c/tests/test_deeppot_a_fparam_aparam.cc +++ b/source/api_c/tests/test_deeppot_a_fparam_aparam.cc @@ -72,6 +72,9 @@ class TestInferDeepPotAFParamAParam : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb(file_name, "fparam_aparam.pb"); @@ -394,6 +397,9 @@ class TestInferAParamNAll : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/pairwise_dprc.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb(file_name, "pairwise_dprc.pb"); dp.init("pairwise_dprc.pb"); diff --git a/source/api_c/tests/test_deeppot_a_fparam_aparam_nframes.cc b/source/api_c/tests/test_deeppot_a_fparam_aparam_nframes.cc index 383c8f5fb1..2279ad2f66 100644 --- a/source/api_c/tests/test_deeppot_a_fparam_aparam_nframes.cc +++ b/source/api_c/tests/test_deeppot_a_fparam_aparam_nframes.cc @@ -116,6 +116,9 @@ class TestInferDeepPotAFparamAparamNFrames : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/fparam_aparam.pbtxt", "fparam_aparam.pb"); @@ -757,6 +760,9 @@ class TestInferDeepPotAFparamAparamNFramesSingleParam : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/fparam_aparam.pbtxt", "fparam_aparam.pb"); diff --git a/source/api_c/tests/test_deeppot_a_hpp.cc b/source/api_c/tests/test_deeppot_a_hpp.cc index 388b53fdfd..3adb47bdf9 100644 --- a/source/api_c/tests/test_deeppot_a_hpp.cc +++ b/source/api_c/tests/test_deeppot_a_hpp.cc @@ -66,6 +66,9 @@ class TestInferDeepPotAHPP : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", "deeppot.pb"); @@ -557,6 +560,9 @@ class TestInferDeepPotANoPbcHPP : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb(file_name, "deeppot.pb"); diff --git a/source/api_c/tests/test_deeppot_a_nframes_hpp.cc b/source/api_c/tests/test_deeppot_a_nframes_hpp.cc index af132c0146..bee0546bfc 100644 --- a/source/api_c/tests/test_deeppot_a_nframes_hpp.cc +++ b/source/api_c/tests/test_deeppot_a_nframes_hpp.cc @@ -109,6 +109,9 @@ class TestInferDeepPotANFrames : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", "deeppot.pb"); @@ -719,6 +722,9 @@ class TestInferDeepPotANFramesNoPbc : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb(file_name, "deeppot.pb"); diff --git a/source/api_c/tests/test_deeppot_model_devi_fparam_aparam.cc b/source/api_c/tests/test_deeppot_model_devi_fparam_aparam.cc index 5e77529021..550259b42d 100644 --- a/source/api_c/tests/test_deeppot_model_devi_fparam_aparam.cc +++ b/source/api_c/tests/test_deeppot_model_devi_fparam_aparam.cc @@ -30,6 +30,9 @@ class TestInferDeepPotModeDeviFparamAparam : public ::testing::Test { deepmd::hpp::DeepPotModelDevi dp_md; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif { std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/fparam_aparam.pbtxt", diff --git a/source/api_c/tests/test_deeppot_model_devi_hpp.cc b/source/api_c/tests/test_deeppot_model_devi_hpp.cc index 5751c4f69c..2a8242f99f 100644 --- a/source/api_c/tests/test_deeppot_model_devi_hpp.cc +++ b/source/api_c/tests/test_deeppot_model_devi_hpp.cc @@ -24,6 +24,9 @@ class TestInferDeepPotModeDevi : public ::testing::Test { deepmd::hpp::DeepPotModelDevi dp_md; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif { std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", @@ -74,6 +77,9 @@ class TestInferDeepPotModeDeviPython : public ::testing::Test { deepmd::hpp::DeepPotModelDevi dp_md; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif { std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", diff --git a/source/api_c/tests/test_deepspin_a.cc b/source/api_c/tests/test_deepspin_a.cc index ccffd7221f..8e3236b776 100644 --- a/source/api_c/tests/test_deepspin_a.cc +++ b/source/api_c/tests/test_deepspin_a.cc @@ -53,9 +53,12 @@ class TestInferDeepSpinA : public ::testing::Test { double expected_tot_e; // std::vector expected_tot_v; - DP_DeepSpin* dp; + DP_DeepSpin* dp = nullptr; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp = DP_NewDeepSpin("../../tests/infer/deeppot_dpa_spin.pth"); natoms = expected_e.size(); @@ -233,9 +236,12 @@ class TestInferDeepSpinANoPBC : public ::testing::Test { double expected_tot_e; // std::vector expected_tot_v; - DP_DeepSpin* dp; + DP_DeepSpin* dp = nullptr; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp = DP_NewDeepSpin("../../tests/infer/deeppot_dpa_spin.pth"); natoms = expected_e.size(); diff --git a/source/api_c/tests/test_deepspin_a_hpp.cc b/source/api_c/tests/test_deepspin_a_hpp.cc index 701d517690..c65a89a45f 100644 --- a/source/api_c/tests/test_deepspin_a_hpp.cc +++ b/source/api_c/tests/test_deepspin_a_hpp.cc @@ -56,6 +56,9 @@ class TestInferDeepSpinAHPP : public ::testing::Test { deepmd::hpp::DeepSpin dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa_spin.pth"); natoms = expected_e.size(); @@ -208,6 +211,9 @@ class TestInferDeepSpinANoPbcHPP : public ::testing::Test { deepmd::hpp::DeepSpin dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa_spin.pth"); natoms = expected_e.size(); diff --git a/source/api_c/tests/test_deepspin_model_devi_hpp.cc b/source/api_c/tests/test_deepspin_model_devi_hpp.cc index c6bbdef8a8..19acccc6f7 100644 --- a/source/api_c/tests/test_deepspin_model_devi_hpp.cc +++ b/source/api_c/tests/test_deepspin_model_devi_hpp.cc @@ -26,6 +26,9 @@ class TestInferDeepSpinModeDevi : public ::testing::Test { deepmd::hpp::DeepSpinModelDevi dp_md; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif { dp0.init("../../tests/infer/deeppot_dpa_spin.pth"); } diff --git a/source/api_c/tests/test_dipolecharge.cc b/source/api_c/tests/test_dipolecharge.cc index 03e9e44b34..647fa5c3a5 100644 --- a/source/api_c/tests/test_dipolecharge.cc +++ b/source/api_c/tests/test_dipolecharge.cc @@ -54,6 +54,9 @@ class TestDipoleCharge : public ::testing::Test { deepmd::hpp::DipoleChargeModifier dm; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/dipolecharge_e.pbtxt"; std::string model = "dipolecharge_e.pb"; deepmd::hpp::convert_pbtxt_to_pb(file_name, model); diff --git a/source/api_c/tests/test_mixed_type.cc b/source/api_c/tests/test_mixed_type.cc index 13f9edbad9..84280f4131 100644 --- a/source/api_c/tests/test_mixed_type.cc +++ b/source/api_c/tests/test_mixed_type.cc @@ -29,6 +29,9 @@ class TestInferMixedType : public ::testing::Test { deepmd::hpp::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/virtual_type.pbtxt"; deepmd::hpp::convert_pbtxt_to_pb("../../tests/infer/virtual_type.pbtxt", "virtual_type.pb"); diff --git a/source/api_c/tests/test_read_file_to_string.cc b/source/api_c/tests/test_read_file_to_string.cc index cd26e95928..bef2a28da0 100644 --- a/source/api_c/tests/test_read_file_to_string.cc +++ b/source/api_c/tests/test_read_file_to_string.cc @@ -13,6 +13,9 @@ #include "deepmd.hpp" TEST(TestReadFileToString, readfiletostring) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_content; deepmd::hpp::read_file_to_string("../../tests/infer/deeppot.txt", file_content); @@ -26,6 +29,9 @@ TEST(TestReadFileToString, readfiletostring) { } TEST(TestReadFileToString, readfiletostringerr) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_content; EXPECT_THROW( { diff --git a/source/api_c/tests/test_utils.h b/source/api_c/tests/test_utils.h index 5167732bc8..59c764409a 100644 --- a/source/api_c/tests/test_utils.h +++ b/source/api_c/tests/test_utils.h @@ -14,7 +14,7 @@ typedef testing::Types ValueTypes; template inline void _fold_back(typename std::vector::iterator out, const typename std::vector::const_iterator in, - const std::vector &mapping, + const std::vector& mapping, const int nloc, const int nall, const int ndim, @@ -35,9 +35,9 @@ inline void _fold_back(typename std::vector::iterator out, } template -inline void _fold_back(std::vector &out, - const std::vector &in, - const std::vector &mapping, +inline void _fold_back(std::vector& out, + const std::vector& in, + const std::vector& mapping, const int nloc, const int nall, const int ndim, @@ -48,14 +48,14 @@ inline void _fold_back(std::vector &out, } template -inline void _build_nlist(std::vector> &nlist_data, - std::vector &coord_cpy, - std::vector &atype_cpy, - std::vector &mapping, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const float &rc) { +inline void _build_nlist(std::vector>& nlist_data, + std::vector& coord_cpy, + std::vector& atype_cpy, + std::vector& mapping, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const float& rc) { // convert VALUETYPE to double, it looks like copy_coord only accepts double std::vector coord_cpy_; std::vector coord_(coord.begin(), coord.end()); @@ -90,13 +90,13 @@ class EnergyModelTest { double level = std::is_same::value ? 1e-6 : 1e-2; // expected? public: - virtual void compute(double &ener, - std::vector &force, - std::vector &virial, - const std::vector &coord, - const std::vector &box) = 0; - void test_f(const std::vector &coord, - const std::vector &box) { + virtual void compute(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& box) = 0; + void test_f(const std::vector& coord, + const std::vector& box) { int ndof = coord.size(); double ener; std::vector force, virial; @@ -114,8 +114,8 @@ class EnergyModelTest { EXPECT_LT(fabs(num - ana), level); } } - void test_v(const std::vector &coord, - const std::vector &box) { + void test_v(const std::vector& coord, + const std::vector& box) { std::vector num_diff(9); double ener; std::vector force, virial; diff --git a/source/api_cc/CMakeLists.txt b/source/api_cc/CMakeLists.txt index 0d4bba1047..90b7c08449 100644 --- a/source/api_cc/CMakeLists.txt +++ b/source/api_cc/CMakeLists.txt @@ -6,7 +6,6 @@ file(GLOB LIB_SRC src/*.cc src/*.cpp) file(GLOB INC_SRC include/*.h ${CMAKE_CURRENT_BINARY_DIR}/version.h) set(libname "${LIB_DEEPMD_CC}") - add_library(${libname} SHARED ${LIB_SRC}) # link: libdeepmd libdeepmd_op libtensorflow_cc libtensorflow_framework @@ -48,7 +47,7 @@ set_target_properties( ${libname} PROPERTIES INSTALL_RPATH "$ORIGIN;${BACKEND_LIBRARY_PATH}" INSTALL_RPATH_USE_LINK_PATH TRUE - BUILD_RPATH "$ORIGIN/../op/tf;$ORIGIN/../op/pt") + BUILD_RPATH "$ORIGIN/../op/tf;$ORIGIN/../op/pt;$ORIGIN/../op/pd") target_compile_definitions(${libname} PRIVATE TF_PRIVATE) if(CMAKE_TESTING_ENABLED) target_link_libraries(${libname} PRIVATE coverage_config) diff --git a/source/api_cc/include/DeepPotPD.h b/source/api_cc/include/DeepPotPD.h index ec43300ca0..6fceb19352 100644 --- a/source/api_cc/include/DeepPotPD.h +++ b/source/api_cc/include/DeepPotPD.h @@ -282,12 +282,11 @@ class DeepPotPD : public DeepPotBackend { * @brief Compute the number of elements in a tensor. * @param[in] x Tensor x. **/ - int numel(const paddle_infer::Tensor& x) const { - // TODO: There might be a overflow problem here for multiply int numbers. - int ret = 1; + size_t numel(const paddle_infer::Tensor& x) const { + size_t ret = 1; std::vector x_shape = x.shape(); for (std::size_t i = 0, n = x_shape.size(); i < n; ++i) { - ret *= x_shape[i]; + ret *= static_cast(x_shape[i]); } return ret; }; @@ -392,7 +391,7 @@ class DeepPotPD : public DeepPotBackend { int do_message_passing; // 1:dpa2 model 0:others bool gpu_enabled; std::unique_ptr firstneigh_tensor; - // std::unordered_map comm_dict; # Not used yet + std::unique_ptr mapping_tensor; }; } // namespace deepmd diff --git a/source/api_cc/include/DeepPotPT.h b/source/api_cc/include/DeepPotPT.h index 207a13286c..4a06bf012c 100644 --- a/source/api_cc/include/DeepPotPT.h +++ b/source/api_cc/include/DeepPotPT.h @@ -340,6 +340,8 @@ class DeepPotPT : public DeepPotBackend { at::Tensor firstneigh_tensor; c10::optional mapping_tensor; torch::Dict comm_dict; + bool profiler_enabled{false}; + std::string profiler_file; /** * @brief Translate PyTorch exceptions to the DeePMD-kit exception. * @param[in] f The function to run. diff --git a/source/api_cc/include/DeepTensorPT.h b/source/api_cc/include/DeepTensorPT.h new file mode 100644 index 0000000000..c602fc53e0 --- /dev/null +++ b/source/api_cc/include/DeepTensorPT.h @@ -0,0 +1,247 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#pragma once + +#include +#include + +#include "DeepTensor.h" + +namespace deepmd { +/** + * @brief PyTorch implementation for Deep Tensor. + **/ +class DeepTensorPT : public DeepTensorBase { + public: + /** + * @brief Deep Tensor constructor without initialization. + **/ + DeepTensorPT(); + virtual ~DeepTensorPT(); + /** + * @brief Deep Tensor constructor with initialization. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] name_scope Name scopes of operations. + **/ + DeepTensorPT(const std::string& model, + const int& gpu_rank = 0, + const std::string& name_scope = ""); + /** + * @brief Initialize the Deep Tensor. + * @param[in] model The name of the frozen model file. + * @param[in] gpu_rank The GPU rank. Default is 0. + * @param[in] name_scope Name scopes of operations. + **/ + void init(const std::string& model, + const int& gpu_rank = 0, + const std::string& name_scope = ""); + + private: + /** + * @brief Evaluate the global tensor and component-wise force and virial. + * @param[out] global_tensor The global tensor to evaluate. + * @param[out] force The component-wise force of the global tensor, size odim + *x natoms x 3. + * @param[out] virial The component-wise virial of the global tensor, size + *odim x 9. + * @param[out] atom_tensor The atomic tensor value of the model, size natoms x + *odim. + * @param[out] atom_virial The component-wise atomic virial of the global + *tensor, size odim x natoms x 9. + * @param[in] coord The coordinates of atoms. The array should be of size + *natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size 9. + * @param[in] request_deriv Whether to request the derivative of the global + * tensor, including force and virial. + **/ + template + void compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const bool request_deriv); + /** + * @brief Evaluate the global tensor and component-wise force and virial. + * @param[out] global_tensor The global tensor to evaluate. + * @param[out] force The component-wise force of the global tensor, size odim + *x natoms x 3. + * @param[out] virial The component-wise virial of the global tensor, size + *odim x 9. + * @param[out] atom_tensor The atomic tensor value of the model, size natoms x + *odim. + * @param[out] atom_virial The component-wise atomic virial of the global + *tensor, size odim x natoms x 9. + * @param[in] coord The coordinates of atoms. The array should be of size + *natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] inlist The input neighbour list. + * @param[in] request_deriv Whether to request the derivative of the global + * tensor, including force and virial. + **/ + template + void compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const bool request_deriv); + + public: + /** + * @brief Get the cutoff radius. + * @return The cutoff radius. + **/ + double cutoff() const { + assert(inited); + return rcut; + }; + /** + * @brief Get the number of types. + * @return The number of types. + **/ + int numb_types() const { + assert(inited); + return ntypes; + }; + /** + * @brief Get the output dimension. + * @return The output dimension. + **/ + int output_dim() const { + assert(inited); + return odim; + }; + /** + * @brief Get the list of sel types. + * @return The list of sel types. + */ + const std::vector& sel_types() const { + assert(inited); + return sel_type; + }; + /** + * @brief Get the type map (element name of the atom types) of this model. + * @param[out] type_map The type map of this model. + **/ + void get_type_map(std::string& type_map); + + /** + * @brief Evaluate the global tensor and component-wise force and virial. + * @param[out] global_tensor The global tensor to evaluate. + * @param[out] force The component-wise force of the global tensor, size odim + *x natoms x 3. + * @param[out] virial The component-wise virial of the global tensor, size + *odim x 9. + * @param[out] atom_tensor The atomic tensor value of the model, size natoms x + *odim. + * @param[out] atom_virial The component-wise atomic virial of the global + *tensor, size odim x natoms x 9. + * @param[in] coord The coordinates of atoms. The array should be of size + *natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size 9. + * @param[in] request_deriv Whether to request the derivative of the global + * tensor, including force and virial. + * @{ + **/ + void computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const bool request_deriv); + void computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const bool request_deriv); + /** @} */ + /** + * @brief Evaluate the global tensor and component-wise force and virial. + * @param[out] global_tensor The global tensor to evaluate. + * @param[out] force The component-wise force of the global tensor, size odim + *x natoms x 3. + * @param[out] virial The component-wise virial of the global tensor, size + *odim x 9. + * @param[out] atom_tensor The atomic tensor value of the model, size natoms x + *odim. + * @param[out] atom_virial The component-wise atomic virial of the global + *tensor, size odim x natoms x 9. + * @param[in] coord The coordinates of atoms. The array should be of size + *natoms x 3. + * @param[in] atype The atom types. The list should contain natoms ints. + * @param[in] box The cell of the region. The array should be of size 9. + * @param[in] nghost The number of ghost atoms. + * @param[in] inlist The input neighbour list. + * @param[in] request_deriv Whether to request the derivative of the global + * tensor, including force and virial. + * @{ + **/ + void computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const bool request_deriv); + void computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const bool request_deriv); + /** @} */ + + private: + int num_intra_nthreads, num_inter_nthreads; + bool inited; + double rcut; + int ntypes; + mutable int odim; + std::vector sel_type; + std::string name_scope; + // PyTorch module and device management + mutable torch::jit::script::Module module; + int gpu_id; + bool gpu_enabled; + NeighborListData nlist_data; + // Neighbor list tensors for efficient computation + at::Tensor firstneigh_tensor; + + /** + * @brief Translate PyTorch exceptions to the DeePMD-kit exception. + * @param[in] f The function to run. + * @example translate_error([&](){...}); + */ + void translate_error(std::function f); +}; + +} // namespace deepmd diff --git a/source/api_cc/src/DeepPotPD.cc b/source/api_cc/src/DeepPotPD.cc index 3a3d880c4b..94931a8415 100644 --- a/source/api_cc/src/DeepPotPD.cc +++ b/source/api_cc/src/DeepPotPD.cc @@ -11,15 +11,176 @@ using namespace deepmd; -std::vector createNlistTensorPD( - const std::vector>& data) { - std::vector ret; - for (const auto& row : data) { - ret.insert(ret.end(), row.begin(), row.end()); +#include +#include +#include +#include +#include + +class Logger { + public: + enum Level { DEBUG = 0, INFO = 1, WARNING = 2, ERROR = 3 }; + + private: + static Level minLevel; + static bool colorEnabled; + static bool showTimestamp; + + static const char* getColorCode(Level level) { + if (!colorEnabled) { + return ""; + } + switch (level) { + case DEBUG: + return "\033[1;36m"; + case INFO: + return "\033[1;32m"; + case WARNING: + return "\033[1;33m"; + case ERROR: + return "\033[1;31m"; + default: + return ""; + } + } + + static const char* getResetCode() { return colorEnabled ? "\033[0m" : ""; } + + static const char* getLevelName(Level level) { + switch (level) { + case DEBUG: + return "DEBUG"; + case INFO: + return "INFO"; + case WARNING: + return "WARNING"; + case ERROR: + return "ERROR"; + default: + return "UNKNOWN"; + } } - return ret; + + static std::string getCurrentTime() { + if (!showTimestamp) { + return ""; + } + + std::time_t now = std::time(0); + std::tm* ltm = std::localtime(&now); + + std::ostringstream oss; + oss << std::setfill('0') << std::setw(4) << (1900 + ltm->tm_year) << "-" + << std::setw(2) << (1 + ltm->tm_mon) << "-" << std::setw(2) + << ltm->tm_mday << " " << std::setw(2) << ltm->tm_hour << ":" + << std::setw(2) << ltm->tm_min << ":" << std::setw(2) << ltm->tm_sec; + return oss.str(); + } + + public: + class LogStream { + private: + std::ostringstream oss; + Level level; + bool shouldLog; + + public: + LogStream(Level lvl) : level(lvl), shouldLog(lvl >= minLevel) { + if (shouldLog) { + std::string timestamp = getCurrentTime(); + if (!timestamp.empty()) { + oss << "[" << timestamp << "] "; + } + oss << getColorCode(level) << "[" << getLevelName(level) << "]" + << getResetCode() << " "; + } + } + + ~LogStream() { + if (shouldLog) { + std::cout << oss.str() << std::flush; + } + } + + template + LogStream& operator<<(const T& value) { + if (shouldLog) { + oss << value; + } + return *this; + } + + LogStream& operator<<(std::ostream& (*manip)(std::ostream&)) { + if (shouldLog) { + oss << manip; + } + return *this; + } + + LogStream(const LogStream&) = delete; + LogStream& operator=(const LogStream&) = delete; + LogStream(LogStream&& other) noexcept + : oss(std::move(other.oss)), + level(other.level), + shouldLog(other.shouldLog) {} + + LogStream& operator=(LogStream&& other) noexcept { + if (this != &other) { + oss = std::move(other.oss); + level = other.level; + shouldLog = other.shouldLog; + } + return *this; + } + }; + + static void setLevel(Level level) { minLevel = level; } + static void enableColor(bool enable = true) { colorEnabled = enable; } + static void enableTimestamp(bool enable = true) { showTimestamp = enable; } + static Level getLevel() { return minLevel; } + static bool isColorEnabled() { return colorEnabled; } + static bool isTimestampEnabled() { return showTimestamp; } + + static LogStream debug() { return LogStream(DEBUG); } + static LogStream info() { return LogStream(INFO); } + static LogStream warning() { return LogStream(WARNING); } + static LogStream error() { return LogStream(ERROR); } +}; + +Logger::Level Logger::minLevel = Logger::INFO; +bool Logger::colorEnabled = true; +bool Logger::showTimestamp = true; + +namespace logg { +inline Logger::LogStream debug() { return Logger::debug(); } +inline Logger::LogStream info() { return Logger::info(); } +inline Logger::LogStream warning() { return Logger::warning(); } +inline Logger::LogStream error() { return Logger::error(); } + +inline void setLevel(Logger::Level level) { Logger::setLevel(level); } +inline void enableColor(bool enable = true) { Logger::enableColor(enable); } +inline void enableTimestamp(bool enable = true) { + Logger::enableTimestamp(enable); } +} // namespace logg + +void fillNlistTensor(const std::vector>& data, + std::unique_ptr& flat_tensor) { + size_t total_size = 0; + for (const auto& row : data) { + total_size += row.size(); + } + std::vector flat_data; + flat_data.reserve(total_size); + for (const auto& row : data) { + flat_data.insert(flat_data.end(), row.begin(), row.end()); + } + int nloc = data.size(); + int nnei = nloc > 0 ? total_size / nloc : 0; + flat_tensor->Reshape({1, nloc, nnei}); + flat_tensor->CopyFromCpu(flat_data.data()); +} DeepPotPD::DeepPotPD() : inited(false) {} DeepPotPD::DeepPotPD(const std::string& model, const int& gpu_rank, @@ -41,9 +202,7 @@ void DeepPotPD::init(const std::string& model, << std::endl; return; } - // NOTE: There is no custom operators need to be loaded now. - // deepmd::load_op_library(); - + deepmd::load_op_library(); // NOTE: Only support 1 GPU now. int gpu_num = 1; if (gpu_num > 0) { @@ -59,6 +218,7 @@ void DeepPotPD::init(const std::string& model, config->EnableNewIR(true); config->EnableCustomPasses({"add_shadow_output_after_dead_parameter_pass"}, true); + // config->SwitchIrOptim(false); // initialize inference config_fl config_fl = std::make_shared(); @@ -67,6 +227,7 @@ void DeepPotPD::init(const std::string& model, config_fl->EnableNewIR(true); config_fl->EnableCustomPasses({"add_shadow_output_after_dead_parameter_pass"}, true); + // config_fl->SwitchIrOptim(false); // loading inference model std::string pdmodel_path, fl_pdmodel_path; @@ -113,30 +274,31 @@ void DeepPotPD::init(const std::string& model, if (!gpu_enabled) { config->DisableGpu(); config_fl->DisableGpu(); - std::cout << "load model from: " << model << " to cpu " << std::endl; + logg::info() << "load model from: " << model << " to cpu " << std::endl; } else { config->EnableUseGpu(4096, 0); config_fl->EnableUseGpu(4096, 0); - std::cout << "load model from: " << model << " to gpu:" << gpu_id - << std::endl; + logg::info() << "load model from: " << model << " to gpu:" << gpu_id + << std::endl; } if (config->cinn_enabled()) { - std::cout << "model.forward will be compiled with cinn." << std::endl; + logg::info() << "model.forward will be compiled with cinn." << std::endl; } else { - std::cout << "NOTE: You can try: \n'export FLAGS_prim_all=true" - " FLAGS_enable_pir_in_executor=1" - " FLAGS_prim_enable_dynamic=true FLAGS_use_cinn=true'\n" - "to speed up C++ inference with paddle backend" - << std::endl; + logg::info() << "NOTE: You can try: \n'export FLAGS_prim_all=true" + " FLAGS_enable_pir_in_executor=1" + " FLAGS_prim_enable_dynamic=true FLAGS_use_cinn=true' " + "to speed up C++ inference with paddle backend" + << std::endl; } if (config_fl->cinn_enabled()) { - std::cout << "model.forward_lower will be compiled with cinn." << std::endl; + logg::info() << "model.forward_lower will be compiled with cinn." + << std::endl; } else { - std::cout << "NOTE: You can try: \n'export FLAGS_prim_all=true" - " FLAGS_enable_pir_in_executor=1" - " FLAGS_prim_enable_dynamic=true FLAGS_use_cinn=true'\n" - "to speed up C++ inference with paddle backend" - << std::endl; + logg::info() << "NOTE: You can try: \n'export FLAGS_prim_all=true" + " FLAGS_enable_pir_in_executor=1" + " FLAGS_prim_enable_dynamic=true FLAGS_use_cinn=true' " + "to speed up C++ inference with paddle backend" + << std::endl; } // NOTE: Both set to 1 now. @@ -151,16 +313,42 @@ void DeepPotPD::init(const std::string& model, } predictor = paddle_infer::CreatePredictor(*config); + logg::info() << "Setup model.forward model" << std::endl; predictor_fl = paddle_infer::CreatePredictor(*config_fl); + logg::info() << "Setup model.forward_lower" << std::endl; + auto print_handle_names = [](const std::vector& name_vec) { + int n = name_vec.size(); + std::string ret; + for (int i = 0; i < n; ++i) { + ret += "[" + std::to_string(i) + "]" + name_vec[i] + " \n"[i == n - 1]; + } + logg::debug() << ret; + }; + logg::debug() << "Input names of model.forward below:" << std::endl; + print_handle_names(predictor->GetInputNames()); + logg::debug() << "Output names of model.forward below:" << std::endl; + print_handle_names(predictor->GetOutputNames()); + std::cout << std::endl; + logg::debug() << "Input names of model.forward_lower below:" << std::endl; + print_handle_names(predictor_fl->GetInputNames()); + logg::debug() << "Output names of model.forward_lower below:" << std::endl; + print_handle_names(predictor_fl->GetOutputNames()); // initialize hyper params from model buffers ntypes_spin = 0; DeepPotPD::get_buffer("buffer_has_message_passing", do_message_passing); + logg::debug() << "buffer_has_message_passing = " << this->do_message_passing + << std::endl; DeepPotPD::get_buffer("buffer_rcut", rcut); + logg::debug() << "buffer_rcut = " << this->rcut << std::endl; DeepPotPD::get_buffer("buffer_ntypes", ntypes); + logg::debug() << "buffer_ntypes = " << this->ntypes << std::endl; DeepPotPD::get_buffer("buffer_dfparam", dfparam); + logg::debug() << "buffer_dfparam = " << this->dfparam << std::endl; DeepPotPD::get_buffer("buffer_daparam", daparam); + logg::debug() << "buffer_daparam = " << this->daparam << std::endl; DeepPotPD::get_buffer("buffer_aparam_nall", aparam_nall); + logg::debug() << "buffer_aparam_nall = " << this->aparam_nall << std::endl; inited = true; } DeepPotPD::~DeepPotPD() {} @@ -195,61 +383,86 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, auto coord_wrapped_Tensor = predictor_fl->GetInputHandle("coord"); coord_wrapped_Tensor->Reshape({1, nall_real, 3}); coord_wrapped_Tensor->CopyFromCpu(coord_wrapped.data()); - auto atype_Tensor = predictor_fl->GetInputHandle("atype"); atype_Tensor->Reshape({1, nall_real}); atype_Tensor->CopyFromCpu(datype.data()); - if (ago == 0) { - nlist_data.copy_from_nlist(lmp_list); + nlist_data.copy_from_nlist(lmp_list, nall - nghost); nlist_data.shuffle_exclude_empty(fwd_map); nlist_data.padding(); - if (do_message_passing == 1 && nghost > 0) { - throw deepmd::deepmd_exception( - "(do_message_passing == 1 && nghost > 0) is not supported yet."); - // int nswap = lmp_list.nswap; - // auto sendproc_tensor = predictor_fl->GetInputHandle("sendproc"); - // sendproc_tensor->Reshape({nswap}); - // sendproc_tensor->CopyFromCpu(lmp_list.sendproc); - // auto recvproc_tensor = predictor_fl->GetInputHandle("recvproc"); - // recvproc_tensor->Reshape({nswap}); - // recvproc_tensor->CopyFromCpu(lmp_list.recvproc); - // auto firstrecv_tensor = predictor_fl->GetInputHandle("firstrecv"); - // firstrecv_tensor->Reshape({nswap}); - // firstrecv_tensor->CopyFromCpu(lmp_list.firstrecv); - // auto recvnum_tensor = predictor_fl->GetInputHandle("recvnum"); - // recvnum_tensor->Reshape({nswap}); - // recvnum_tensor->CopyFromCpu(lmp_list.recvnum); - // auto sendnum_tensor = predictor_fl->GetInputHandle("sendnum"); - // sendnum_tensor->Reshape({nswap}); - // sendnum_tensor->CopyFromCpu(lmp_list.sendnum); - // auto communicator_tensor = - // predictor_fl->GetInputHandle("communicator"); - // communicator_tensor->Reshape({1}); - // communicator_tensor->CopyFromCpu(static_cast(lmp_list.world)); - // auto sendlist_tensor = predictor_fl->GetInputHandle("sendlist"); - - // int total_send = - // std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0); + if (do_message_passing) { + auto sendproc_tensor = predictor_fl->GetInputHandle("send_proc"); + auto recvproc_tensor = predictor_fl->GetInputHandle("recv_proc"); + auto recvnum_tensor = predictor_fl->GetInputHandle("recv_num"); + auto sendnum_tensor = predictor_fl->GetInputHandle("send_num"); + auto communicator_tensor = predictor_fl->GetInputHandle("communicator"); + auto sendlist_tensor = predictor_fl->GetInputHandle("send_list"); + + int nswap = lmp_list.nswap; + sendproc_tensor->Reshape({nswap}); + sendproc_tensor->CopyFromCpu(lmp_list.sendproc); + + recvproc_tensor->Reshape({nswap}); + recvproc_tensor->CopyFromCpu(lmp_list.recvproc); + + recvnum_tensor->Reshape({nswap}); + recvnum_tensor->CopyFromCpu(lmp_list.recvnum); + + sendnum_tensor->Reshape({nswap}); + if (sizeof(lmp_list.sendnum[0]) != sizeof(int32_t)) { + std::vector temp_data(nswap); + for (int i = 0; i < nswap; i++) { + temp_data[i] = static_cast(lmp_list.sendnum[i]); + } + sendnum_tensor->CopyFromCpu(temp_data.data()); + } else { + sendnum_tensor->CopyFromCpu(lmp_list.sendnum); + } + communicator_tensor->Reshape({1}); + if (lmp_list.world) { + communicator_tensor->CopyFromCpu(static_cast(lmp_list.world)); + } + + assert(sizeof(std::intptr_t) == 8); + int total_send = + std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0); + sendlist_tensor->Reshape({total_send}); + + /** + ** NOTE: paddle do not support construct a Tensor with from_blob(T**, ...) + ** from a double pointer, so we convert int* pointer to indptr_t for each + ** entry and wrap it into int64 Tensor as a workaround. + */ + std::vector pointer_addresses; + pointer_addresses.reserve(nswap); + for (int iswap = 0; iswap < nswap; ++iswap) { + std::intptr_t addr = + reinterpret_cast(lmp_list.sendlist[iswap]); + pointer_addresses.push_back(addr); + } + sendlist_tensor->CopyFromCpu(pointer_addresses.data()); } - if (do_message_passing == 1 && nghost == 0) { - throw deepmd::deepmd_exception( - "(do_message_passing == 1 && nghost == 0) is not supported yet."); + if (lmp_list.mapping) { + std::vector mapping(nall_real); + for (size_t ii = 0; ii < nall_real; ii++) { + mapping[ii] = lmp_list.mapping[fwd_map[ii]]; + } + this->mapping_tensor = predictor_fl->GetInputHandle("mapping"); + this->mapping_tensor->Reshape({1, nall_real}); + this->mapping_tensor->CopyFromCpu(mapping.data()); } } - std::vector firstneigh = createNlistTensorPD(nlist_data.jlist); - firstneigh_tensor = predictor_fl->GetInputHandle("nlist"); - firstneigh_tensor->Reshape({1, nloc, (int)firstneigh.size() / (int)nloc}); - firstneigh_tensor->CopyFromCpu(firstneigh.data()); + this->firstneigh_tensor = predictor_fl->GetInputHandle("nlist"); + fillNlistTensor(nlist_data.jlist, this->firstneigh_tensor); bool do_atom_virial_tensor = atomic; - std::unique_ptr fparam_tensor; if (!fparam.empty()) { + std::unique_ptr fparam_tensor; fparam_tensor = predictor_fl->GetInputHandle("fparam"); fparam_tensor->Reshape({1, static_cast(fparam.size())}); - fparam_tensor->CopyFromCpu((fparam.data())); + fparam_tensor->CopyFromCpu(fparam.data()); } - std::unique_ptr aparam_tensor; if (!aparam_.empty()) { + std::unique_ptr aparam_tensor; aparam_tensor = predictor_fl->GetInputHandle("aparam"); aparam_tensor->Reshape( {1, lmp_list.inum, static_cast(aparam_.size()) / lmp_list.inum}); @@ -264,9 +477,9 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, auto energy_ = predictor_fl->GetOutputHandle(output_names.at(1)); auto force_ = predictor_fl->GetOutputHandle(output_names.at(2)); auto virial_ = predictor_fl->GetOutputHandle(output_names.at(4)); - int output_energy_size = numel(*energy_); - int output_force_size = numel(*force_); - int output_virial_size = numel(*virial_); + size_t output_energy_size = numel(*energy_); + size_t output_force_size = numel(*force_); + size_t output_virial_size = numel(*virial_); // output energy ener.resize(output_energy_size); energy_->CopyToCpu(ener.data()); @@ -300,7 +513,7 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, } } template void DeepPotPD::compute>( - std::vector& dener, + std::vector& ener, std::vector& force, std::vector& virial, std::vector& atom_energy, @@ -312,11 +525,10 @@ template void DeepPotPD::compute>( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_, + const std::vector& aparam, const bool atomic); - template void DeepPotPD::compute>( - std::vector& dener, + std::vector& ener, std::vector& force, std::vector& virial, std::vector& atom_energy, @@ -328,9 +540,8 @@ template void DeepPotPD::compute>( const InputNlist& lmp_list, const int& ago, const std::vector& fparam, - const std::vector& aparam_, + const std::vector& aparam, const bool atomic); - // ENERGYVTYPE: std::vector or ENERGYTYPE template void DeepPotPD::compute(ENERGYVTYPE& ener, @@ -352,9 +563,9 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, coord_wrapped_Tensor->Reshape({1, natoms, 3}); coord_wrapped_Tensor->CopyFromCpu(coord_wrapped.data()); - std::vector atype_64(atype.begin(), atype.end()); auto atype_Tensor = predictor->GetInputHandle("atype"); atype_Tensor->Reshape({1, natoms}); + std::vector atype_64(atype.begin(), atype.end()); atype_Tensor->CopyFromCpu(atype_64.data()); std::unique_ptr box_Tensor; @@ -363,15 +574,15 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, box_Tensor->Reshape({1, 9}); box_Tensor->CopyFromCpu((box.data())); } - std::unique_ptr fparam_tensor; if (!fparam.empty()) { - fparam_tensor = predictor->GetInputHandle("box"); + std::unique_ptr fparam_tensor; + fparam_tensor = predictor->GetInputHandle("fparam"); fparam_tensor->Reshape({1, static_cast(fparam.size())}); fparam_tensor->CopyFromCpu((fparam.data())); } - std::unique_ptr aparam_tensor; if (!aparam.empty()) { - aparam_tensor = predictor->GetInputHandle("box"); + std::unique_ptr aparam_tensor; + aparam_tensor = predictor->GetInputHandle("aparam"); aparam_tensor->Reshape( {1, natoms, static_cast(aparam.size()) / natoms}); aparam_tensor->CopyFromCpu((aparam.data())); @@ -387,17 +598,17 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, auto force_ = predictor->GetOutputHandle(output_names.at(3)); auto virial_ = predictor->GetOutputHandle(output_names.at(5)); - int enery_numel = numel(*energy_); + size_t enery_numel = numel(*energy_); assert(enery_numel > 0); ener.resize(enery_numel); energy_->CopyToCpu(ener.data()); - int force_numel = numel(*force_); + size_t force_numel = numel(*force_); assert(force_numel > 0); force.resize(force_numel); force_->CopyToCpu(force.data()); - int virial_numel = numel(*virial_); + size_t virial_numel = numel(*virial_); assert(virial_numel > 0); virial.resize(virial_numel); virial_->CopyToCpu(virial.data()); @@ -405,8 +616,8 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, if (atomic) { auto atom_energy_ = predictor->GetOutputHandle(output_names.at(0)); auto atom_virial_ = predictor->GetOutputHandle(output_names.at(1)); - int atom_energy_numel = numel(*atom_energy_); - int atom_virial_numel = numel(*atom_virial_); + size_t atom_energy_numel = numel(*atom_energy_); + size_t atom_virial_numel = numel(*atom_virial_); assert(atom_energy_numel > 0); assert(atom_virial_numel > 0); atom_energy.resize(atom_energy_numel); @@ -418,11 +629,11 @@ void DeepPotPD::compute(ENERGYVTYPE& ener, template void DeepPotPD::compute>( std::vector& ener, - std::vector& dforce, + std::vector& force, std::vector& virial, std::vector& atom_energy, std::vector& atom_virial, - const std::vector& dcoord, + const std::vector& coord, const std::vector& atype, const std::vector& box, const std::vector& fparam, @@ -435,7 +646,7 @@ template void DeepPotPD::compute>( std::vector& virial, std::vector& atom_energy, std::vector& atom_virial, - const std::vector& dcoord, + const std::vector& coord, const std::vector& atype, const std::vector& box, const std::vector& fparam, @@ -446,7 +657,7 @@ template void DeepPotPD::compute>( that need to be postprocessed */ void DeepPotPD::get_type_map(std::string& type_map) { auto type_map_tensor = predictor->GetOutputHandle("buffer_type_map"); - int type_map_size = numel(*type_map_tensor); + size_t type_map_size = numel(*type_map_tensor); std::vector type_map_arr(type_map_size, 0); type_map_tensor->CopyToCpu(type_map_arr.data()); @@ -460,7 +671,7 @@ template void DeepPotPD::get_buffer(const std::string& buffer_name, std::vector& buffer_array) { auto buffer_tensor = predictor->GetOutputHandle(buffer_name); - int buffer_size = numel(*buffer_tensor); + size_t buffer_size = numel(*buffer_tensor); buffer_array.resize(buffer_size); buffer_tensor->CopyToCpu(buffer_array.data()); } diff --git a/source/api_cc/src/DeepPotPT.cc b/source/api_cc/src/DeepPotPT.cc index 5f03f7c5cb..3fdfeeae27 100644 --- a/source/api_cc/src/DeepPotPT.cc +++ b/source/api_cc/src/DeepPotPT.cc @@ -2,6 +2,7 @@ #ifdef BUILD_PYTORCH #include "DeepPotPT.h" +#include #include #include @@ -69,13 +70,9 @@ void DeepPotPT::init(const std::string& model, } deepmd::load_op_library(); int gpu_num = torch::cuda::device_count(); - if (gpu_num > 0) { - gpu_id = gpu_rank % gpu_num; - } else { - gpu_id = 0; - } - torch::Device device(torch::kCUDA, gpu_id); + gpu_id = (gpu_num > 0) ? (gpu_rank % gpu_num) : 0; gpu_enabled = torch::cuda::is_available(); + torch::Device device(torch::kCUDA, gpu_id); if (!gpu_enabled) { device = torch::Device(torch::kCPU); std::cout << "load model from: " << model << " to cpu " << std::endl; @@ -86,6 +83,37 @@ void DeepPotPT::init(const std::string& model, std::cout << "load model from: " << model << " to gpu " << gpu_id << std::endl; } + + // Configure PyTorch profiler + const char* env_profiler = std::getenv("DP_PROFILER"); + if (env_profiler && *env_profiler) { + using torch::profiler::impl::ActivityType; + using torch::profiler::impl::ExperimentalConfig; + using torch::profiler::impl::ProfilerConfig; + using torch::profiler::impl::ProfilerState; + std::set activities{ActivityType::CPU}; + if (gpu_enabled) { + activities.insert(ActivityType::CUDA); + } + profiler_file = std::string(env_profiler); + if (gpu_enabled) { + profiler_file += "_gpu" + std::to_string(gpu_id); + } + profiler_file += ".json"; + ExperimentalConfig exp_cfg; + ProfilerConfig cfg(ProfilerState::KINETO, + false, // report_input_shapes + false, // profile_memory + true, // with_stack + false, // with_flops + true, // with_modules + exp_cfg); + torch::autograd::profiler::prepareProfiler(cfg, activities); + torch::autograd::profiler::enableProfiler(cfg, activities); + std::cout << "PyTorch profiler enabled, output file: " << profiler_file + << std::endl; + profiler_enabled = true; + } std::unordered_map metadata = {{"type", ""}}; module = torch::jit::load(model, device, metadata); module.eval(); @@ -119,7 +147,17 @@ void DeepPotPT::init(const std::string& model, aparam_nall = module.run_method("is_aparam_nall").toBool(); inited = true; } -DeepPotPT::~DeepPotPT() {} + +DeepPotPT::~DeepPotPT() { + if (profiler_enabled) { + auto result = torch::autograd::profiler::disableProfiler(); + if (result) { + result->save(profiler_file); + } + std::cout << "PyTorch profiler result saved to " << profiler_file + << std::endl; + } +} template void DeepPotPT::compute(ENERGYVTYPE& ener, @@ -197,12 +235,12 @@ void DeepPotPT::compute(ENERGYVTYPE& ener, std::accumulate(lmp_list.sendnum, lmp_list.sendnum + nswap, 0); torch::Tensor sendlist_tensor = torch::from_blob(lmp_list.sendlist, {total_send}, int32_option); - comm_dict.insert("send_list", sendlist_tensor); - comm_dict.insert("send_proc", sendproc_tensor); - comm_dict.insert("recv_proc", recvproc_tensor); - comm_dict.insert("send_num", sendnum_tensor); - comm_dict.insert("recv_num", recvnum_tensor); - comm_dict.insert("communicator", communicator_tensor); + comm_dict.insert_or_assign("send_list", sendlist_tensor); + comm_dict.insert_or_assign("send_proc", sendproc_tensor); + comm_dict.insert_or_assign("recv_proc", recvproc_tensor); + comm_dict.insert_or_assign("send_num", sendnum_tensor); + comm_dict.insert_or_assign("recv_num", recvnum_tensor); + comm_dict.insert_or_assign("communicator", communicator_tensor); } if (lmp_list.mapping) { std::vector mapping(nall_real); diff --git a/source/api_cc/src/DeepSpinPT.cc b/source/api_cc/src/DeepSpinPT.cc index 19d5368213..8ccf2fd383 100644 --- a/source/api_cc/src/DeepSpinPT.cc +++ b/source/api_cc/src/DeepSpinPT.cc @@ -205,13 +205,13 @@ void DeepSpinPT::compute(ENERGYVTYPE& ener, torch::Tensor sendlist_tensor = torch::from_blob(lmp_list.sendlist, {total_send}, int32_option); torch::Tensor has_spin = torch::tensor({1}, int32_option); - comm_dict.insert("send_list", sendlist_tensor); - comm_dict.insert("send_proc", sendproc_tensor); - comm_dict.insert("recv_proc", recvproc_tensor); - comm_dict.insert("send_num", sendnum_tensor); - comm_dict.insert("recv_num", recvnum_tensor); - comm_dict.insert("communicator", communicator_tensor); - comm_dict.insert("has_spin", has_spin); + comm_dict.insert_or_assign("send_list", sendlist_tensor); + comm_dict.insert_or_assign("send_proc", sendproc_tensor); + comm_dict.insert_or_assign("recv_proc", recvproc_tensor); + comm_dict.insert_or_assign("send_num", sendnum_tensor); + comm_dict.insert_or_assign("recv_num", recvnum_tensor); + comm_dict.insert_or_assign("communicator", communicator_tensor); + comm_dict.insert_or_assign("has_spin", has_spin); } } at::Tensor firstneigh = createNlistTensor2(nlist_data.jlist); diff --git a/source/api_cc/src/DeepTensor.cc b/source/api_cc/src/DeepTensor.cc index a9031472e6..ce9ca9dea2 100644 --- a/source/api_cc/src/DeepTensor.cc +++ b/source/api_cc/src/DeepTensor.cc @@ -6,24 +6,27 @@ #ifdef BUILD_TENSORFLOW #include "DeepTensorTF.h" #endif +#ifdef BUILD_PYTORCH +#include "DeepTensorPT.h" +#endif #include "common.h" using namespace deepmd; DeepTensor::DeepTensor() : inited(false) {} -DeepTensor::DeepTensor(const std::string &model, - const int &gpu_rank, - const std::string &name_scope_) +DeepTensor::DeepTensor(const std::string& model, + const int& gpu_rank, + const std::string& name_scope_) : inited(false) { init(model, gpu_rank, name_scope_); } DeepTensor::~DeepTensor() {} -void DeepTensor::init(const std::string &model, - const int &gpu_rank, - const std::string &name_scope_) { +void DeepTensor::init(const std::string& model, + const int& gpu_rank, + const std::string& name_scope_) { if (inited) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" @@ -38,7 +41,11 @@ void DeepTensor::init(const std::string &model, throw deepmd::deepmd_exception("TensorFlow backend is not built."); #endif } else if (deepmd::DPBackend::PyTorch == backend) { - throw deepmd::deepmd_exception("PyTorch backend is not supported yet"); +#ifdef BUILD_PYTORCH + dt = std::make_shared(model, gpu_rank, name_scope_); +#else + throw deepmd::deepmd_exception("PyTorch backend is not built."); +#endif } else if (deepmd::DPBackend::Paddle == backend) { throw deepmd::deepmd_exception("PaddlePaddle backend is not supported yet"); } else { @@ -47,183 +54,183 @@ void DeepTensor::init(const std::string &model, inited = true; } -void DeepTensor::print_summary(const std::string &pre) const { +void DeepTensor::print_summary(const std::string& pre) const { deepmd::print_summary(pre); } template -void DeepTensor::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensor::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { std::vector force_, virial_, datom_tensor_, datom_virial_; dt->computew(dtensor_, force_, virial_, datom_tensor_, datom_virial_, dcoord_, datype_, dbox, false); } -template void DeepTensor::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensor::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); -template void DeepTensor::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensor::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensor::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensor::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { std::vector force_, virial_, datom_tensor_, datom_virial_; dt->computew(dtensor_, force_, virial_, datom_tensor_, datom_virial_, dcoord_, datype_, dbox, nghost, lmp_list, false); } -template void DeepTensor::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +template void DeepTensor::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); -template void DeepTensor::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +template void DeepTensor::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); template -void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { std::vector datom_tensor_, datom_virial_; dt->computew(dglobal_tensor_, dforce_, dvirial_, datom_tensor_, datom_virial_, dcoord_, datype_, dbox, true); } -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { std::vector datom_tensor_, datom_virial_; dt->computew(dglobal_tensor_, dforce_, dvirial_, datom_tensor_, datom_virial_, dcoord_, datype_, dbox, nghost, lmp_list, true); } -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); - -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + const InputNlist& lmp_list); + +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); template -void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { dt->computew(dglobal_tensor_, dforce_, dvirial_, datom_tensor_, datom_virial_, dcoord_, datype_, dbox, true); } -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); - -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); + +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { dt->computew(dglobal_tensor_, dforce_, dvirial_, datom_tensor_, datom_virial_, dcoord_, datype_, dbox, nghost, lmp_list, true); } -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); - -template void DeepTensor::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + const InputNlist& lmp_list); + +template void DeepTensor::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); -void DeepTensor::get_type_map(std::string &type_map) { +void DeepTensor::get_type_map(std::string& type_map) { dt->get_type_map(type_map); } @@ -231,7 +238,7 @@ double DeepTensor::cutoff() const { return dt->cutoff(); } int DeepTensor::output_dim() const { return dt->output_dim(); } -const std::vector &DeepTensor::sel_types() const { +const std::vector& DeepTensor::sel_types() const { return dt->sel_types(); } diff --git a/source/api_cc/src/DeepTensorPT.cc b/source/api_cc/src/DeepTensorPT.cc new file mode 100644 index 0000000000..1636f3af95 --- /dev/null +++ b/source/api_cc/src/DeepTensorPT.cc @@ -0,0 +1,482 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#ifdef BUILD_PYTORCH +#include "DeepTensorPT.h" + +#include + +#include +#include // for std::iota +#include + +#include "common.h" +#include "device.h" +#include "errors.h" + +using namespace deepmd; + +static torch::Tensor createNlistTensor( + const std::vector>& data) { + size_t total_size = 0; + for (const auto& row : data) { + total_size += row.size(); + } + std::vector flat_data; + flat_data.reserve(total_size); + for (const auto& row : data) { + flat_data.insert(flat_data.end(), row.begin(), row.end()); + } + + torch::Tensor flat_tensor = torch::tensor(flat_data, torch::kInt32); + int nloc = data.size(); + int nnei = nloc > 0 ? total_size / nloc : 0; + return flat_tensor.view({1, nloc, nnei}); +} + +void DeepTensorPT::translate_error(std::function f) { + try { + f(); + // it seems that libtorch may throw different types of exceptions which are + // inherbited from different base classes + // https://github.com/pytorch/pytorch/blob/13316a8d4642454012d34da0d742f1ba93fc0667/torch/csrc/jit/runtime/interpreter.cpp#L924-L939 + } catch (const c10::Error& e) { + throw deepmd::deepmd_exception("DeePMD-kit PyTorch backend error: " + + std::string(e.what())); + } catch (const torch::jit::JITException& e) { + throw deepmd::deepmd_exception("DeePMD-kit PyTorch backend JIT error: " + + std::string(e.what())); + } catch (const std::runtime_error& e) { + throw deepmd::deepmd_exception("DeePMD-kit PyTorch backend error: " + + std::string(e.what())); + } +} + +DeepTensorPT::DeepTensorPT() : inited(false) {} + +DeepTensorPT::DeepTensorPT(const std::string& model, + const int& gpu_rank, + const std::string& name_scope_) + : inited(false), name_scope(name_scope_) { + try { + translate_error([&] { init(model, gpu_rank, name_scope_); }); + } catch (...) { + // Clean up and rethrow, as the destructor will not be called + throw; + } +} + +void DeepTensorPT::init(const std::string& model, + const int& gpu_rank, + const std::string& name_scope_) { + if (inited) { + std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " + "nothing at the second call of initializer" + << std::endl; + return; + } + name_scope = name_scope_; + deepmd::load_op_library(); + int gpu_num = torch::cuda::device_count(); + if (gpu_num > 0) { + gpu_id = gpu_rank % gpu_num; + } else { + gpu_id = 0; + } + torch::Device device(torch::kCUDA, gpu_id); + gpu_enabled = torch::cuda::is_available(); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + std::cout << "load model from: " << model << " to cpu " << std::endl; + } else { +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + DPErrcheck(DPSetDevice(gpu_id)); +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + std::cout << "load model from: " << model << " to gpu " << gpu_id + << std::endl; + } + std::unordered_map metadata = {{"type", ""}}; + module = torch::jit::load(model, device, metadata); + module.eval(); + + get_env_nthreads(num_intra_nthreads, num_inter_nthreads); + if (num_inter_nthreads) { + try { + at::set_num_interop_threads(num_inter_nthreads); + } catch (...) { + } + } + if (num_intra_nthreads) { + try { + at::set_num_threads(num_intra_nthreads); + } catch (...) { + } + } + + // Get model properties using run_method for C++ interface + auto rcut_result = module.run_method("get_rcut"); + rcut = rcut_result.toDouble(); + + auto ntypes_result = module.run_method("get_ntypes"); + ntypes = ntypes_result.toInt(); + + // Get task dimension from model method + auto task_dim_result = module.run_method("get_task_dim"); + odim = task_dim_result.toInt(); + + // Get type map and set up sel_type + auto type_map_result = module.run_method("get_type_map"); + auto type_map_list = type_map_result.toList(); + sel_type.clear(); + + // For PyTorch models, all types are included (the backend handles exclusions + // internally) The model always outputs all types, but some results may be + // zero + for (size_t i = 0; i < type_map_list.size(); ++i) { + sel_type.push_back(i); + } + inited = true; +} + +DeepTensorPT::~DeepTensorPT() {} + +void DeepTensorPT::get_type_map(std::string& type_map) { + auto type_map_result = module.run_method("get_type_map"); + auto type_map_list = type_map_result.toList(); + type_map.clear(); + for (const torch::IValue& element : type_map_list) { + if (!type_map.empty()) { + type_map += " "; + } + type_map += torch::str(element); + } +} + +template +void DeepTensorPT::compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const bool request_deriv) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + + int natoms = atype.size(); + auto options = torch::TensorOptions().dtype(torch::kFloat64); + torch::ScalarType floatType = torch::kFloat64; + if (std::is_same::value) { + options = torch::TensorOptions().dtype(torch::kFloat32); + floatType = torch::kFloat32; + } + auto int_options = torch::TensorOptions().dtype(torch::kInt64); + + // Convert inputs to tensors + std::vector coord_wrapped = coord; + at::Tensor coord_tensor = + torch::from_blob(coord_wrapped.data(), {1, natoms, 3}, options) + .to(device); + + std::vector atype_64(atype.begin(), atype.end()); + at::Tensor atype_tensor = + torch::from_blob(atype_64.data(), {1, natoms}, int_options).to(device); + + c10::optional box_tensor; + if (!box.empty()) { + box_tensor = + torch::from_blob(const_cast(box.data()), {1, 9}, options) + .to(device); + } + + // Create input vector + std::vector inputs; + inputs.push_back(coord_tensor); + inputs.push_back(atype_tensor); + inputs.push_back(box_tensor); + + // Add None for fparam and aparam (not used by tensor models) + inputs.push_back(torch::jit::IValue()); // fparam = None + inputs.push_back(torch::jit::IValue()); // aparam = None + inputs.push_back(request_deriv); // do_atomic_virial + + // Forward pass through model + c10::Dict outputs = + module.forward(inputs).toGenericDict(); + + // Extract global dipole/polar results + c10::IValue global_out; + if (outputs.contains("global_dipole")) { + global_out = outputs.at("global_dipole"); + } else if (outputs.contains("global_polar")) { + global_out = outputs.at("global_polar"); + } else { + throw deepmd::deepmd_exception( + "Cannot find global tensor output in model results"); + } + torch::Tensor flat_global_ = global_out.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_global_ = flat_global_.to(torch::kCPU); + global_tensor.assign(cpu_global_.data_ptr(), + cpu_global_.data_ptr() + cpu_global_.numel()); + + // Extract atomic dipole/polar results + c10::IValue atom_out; + if (outputs.contains("dipole")) { + atom_out = outputs.at("dipole"); + } else if (outputs.contains("polar")) { + atom_out = outputs.at("polar"); + } else { + throw deepmd::deepmd_exception( + "Cannot find atomic tensor output in model results"); + } + torch::Tensor flat_atom_ = atom_out.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_ = flat_atom_.to(torch::kCPU); + atom_tensor.assign(cpu_atom_.data_ptr(), + cpu_atom_.data_ptr() + cpu_atom_.numel()); + + // Extract force results + c10::IValue force_ = outputs.at("force"); + torch::Tensor flat_force_ = force_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); + force.assign(cpu_force_.data_ptr(), + cpu_force_.data_ptr() + cpu_force_.numel()); + + // Extract virial results + c10::IValue virial_ = outputs.at("virial"); + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); + // Extract atomic virial results if requested + if (request_deriv) { + c10::IValue atom_virial_ = outputs.at("atom_virial"); + torch::Tensor flat_atom_virial_ = + atom_virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); + atom_virial.assign( + cpu_atom_virial_.data_ptr(), + cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); + } else { + atom_virial.clear(); + } +} + +template +void DeepTensorPT::compute(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& lmp_list, + const bool request_deriv) { + torch::Device device(torch::kCUDA, gpu_id); + if (!gpu_enabled) { + device = torch::Device(torch::kCPU); + } + + int natoms = atype.size(); + auto options = torch::TensorOptions().dtype(torch::kFloat64); + torch::ScalarType floatType = torch::kFloat64; + if (std::is_same::value) { + options = torch::TensorOptions().dtype(torch::kFloat32); + floatType = torch::kFloat32; + } + auto int32_option = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt32); + auto int_option = + torch::TensorOptions().device(torch::kCPU).dtype(torch::kInt64); + + // Select real atoms following DeepPotPT pattern + std::vector dcoord, aparam_; + std::vector datype, fwd_map, bkw_map; + int nghost_real, nall_real, nloc_real; + int nall = natoms; + int nframes = 1; + std::vector aparam; // Empty for tensor models + select_real_atoms_coord(dcoord, datype, aparam_, nghost_real, fwd_map, + bkw_map, nall_real, nloc_real, coord, atype, aparam, + nghost, ntypes, nframes, 0, nall, false); + + std::vector coord_wrapped = dcoord; + at::Tensor coord_wrapped_Tensor = + torch::from_blob(coord_wrapped.data(), {1, nall_real, 3}, options) + .to(device); + std::vector atype_64(datype.begin(), datype.end()); + at::Tensor atype_Tensor = + torch::from_blob(atype_64.data(), {1, nall_real}, int_option).to(device); + + // Process neighbor list following DeepPotPT pattern + nlist_data.copy_from_nlist(lmp_list, nall - nghost); + nlist_data.shuffle_exclude_empty(fwd_map); + nlist_data.padding(); + + at::Tensor firstneigh = createNlistTensor(nlist_data.jlist); + firstneigh_tensor = firstneigh.to(torch::kInt64).to(device); + + bool do_atom_virial_tensor = request_deriv; + c10::optional fparam_tensor; + c10::optional aparam_tensor; + c10::optional mapping_tensor; + + // Use forward_lower method following DeepPotPT pattern + c10::Dict outputs = + module + .run_method("forward_lower", coord_wrapped_Tensor, atype_Tensor, + firstneigh_tensor, mapping_tensor, fparam_tensor, + aparam_tensor, do_atom_virial_tensor) + .toGenericDict(); + + // Extract outputs following DeepPotPT pattern + c10::IValue global_dipole_; + if (outputs.contains("global_dipole")) { + global_dipole_ = outputs.at("global_dipole"); + } else if (outputs.contains("global_polar")) { + global_dipole_ = outputs.at("global_polar"); + } else { + throw deepmd::deepmd_exception( + "Cannot find global tensor output in model results"); + } + // in Python, here used double; however, in TF C++, float is used + // for consistency, we use float + torch::Tensor flat_global_ = + global_dipole_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_global_ = flat_global_.to(torch::kCPU); + global_tensor.assign(cpu_global_.data_ptr(), + cpu_global_.data_ptr() + cpu_global_.numel()); + + c10::IValue force_ = outputs.at("extended_force"); + torch::Tensor flat_force_ = force_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_force_ = flat_force_.to(torch::kCPU); + std::vector dforce; + dforce.assign(cpu_force_.data_ptr(), + cpu_force_.data_ptr() + cpu_force_.numel()); + + c10::IValue virial_ = outputs.at("virial"); + torch::Tensor flat_virial_ = virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_virial_ = flat_virial_.to(torch::kCPU); + virial.assign(cpu_virial_.data_ptr(), + cpu_virial_.data_ptr() + cpu_virial_.numel()); + + // bkw map for forces + force.resize(static_cast(nframes) * odim * fwd_map.size() * 3); + for (int kk = 0; kk < odim; ++kk) { + select_map(force.begin() + kk * fwd_map.size() * 3, + dforce.begin() + kk * bkw_map.size() * 3, bkw_map, 3); + } + + // Extract atomic dipoles/polars if available + c10::IValue atom_tensor_output; + int task_dim; + if (outputs.contains("dipole")) { + atom_tensor_output = outputs.at("dipole"); + task_dim = 3; // dipole has 3 components + } else if (outputs.contains("polar")) { + atom_tensor_output = outputs.at("polar"); + task_dim = 9; // polarizability has 9 components typically + } else { + throw deepmd::deepmd_exception( + "Cannot find atomic tensor output in model results"); + } + + torch::Tensor flat_atom_tensor_ = + atom_tensor_output.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_tensor_ = flat_atom_tensor_.to(torch::kCPU); + std::vector datom_tensor; + datom_tensor.assign( + cpu_atom_tensor_.data_ptr(), + cpu_atom_tensor_.data_ptr() + cpu_atom_tensor_.numel()); + atom_tensor.resize(static_cast(nframes) * fwd_map.size() * task_dim); + select_map(atom_tensor, datom_tensor, bkw_map, task_dim, nframes, + fwd_map.size(), nall_real); + + if (request_deriv) { + c10::IValue atom_virial_ = outputs.at("extended_virial"); + torch::Tensor flat_atom_virial_ = + atom_virial_.toTensor().view({-1}).to(floatType); + torch::Tensor cpu_atom_virial_ = flat_atom_virial_.to(torch::kCPU); + std::vector datom_virial; + datom_virial.assign( + cpu_atom_virial_.data_ptr(), + cpu_atom_virial_.data_ptr() + cpu_atom_virial_.numel()); + atom_virial.resize(static_cast(nframes) * odim * fwd_map.size() * + 9); + for (int kk = 0; kk < odim; ++kk) { + select_map(atom_virial.begin() + kk * fwd_map.size() * 9, + datom_virial.begin() + kk * bkw_map.size() * 9, + bkw_map, 9); + } + } +} + +// Public wrapper functions +void DeepTensorPT::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const bool request_deriv) { + translate_error([&] { + compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, + atype, box, request_deriv); + }); +} + +void DeepTensorPT::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const bool request_deriv) { + translate_error([&] { + compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, + atype, box, request_deriv); + }); +} + +void DeepTensorPT::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const bool request_deriv) { + translate_error([&] { + compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, + atype, box, nghost, inlist, request_deriv); + }); +} + +void DeepTensorPT::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const int nghost, + const InputNlist& inlist, + const bool request_deriv) { + translate_error([&] { + compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, + atype, box, nghost, inlist, request_deriv); + }); +} + +#endif // BUILD_PYTORCH diff --git a/source/api_cc/src/DeepTensorTF.cc b/source/api_cc/src/DeepTensorTF.cc index 1081473f25..d17c248f7e 100644 --- a/source/api_cc/src/DeepTensorTF.cc +++ b/source/api_cc/src/DeepTensorTF.cc @@ -7,9 +7,9 @@ using namespace tensorflow; DeepTensorTF::DeepTensorTF() : inited(false), graph_def(new GraphDef()) {} -DeepTensorTF::DeepTensorTF(const std::string &model, - const int &gpu_rank, - const std::string &name_scope_) +DeepTensorTF::DeepTensorTF(const std::string& model, + const int& gpu_rank, + const std::string& name_scope_) : inited(false), name_scope(name_scope_), graph_def(new GraphDef()) { try { init(model, gpu_rank, name_scope_); @@ -22,9 +22,9 @@ DeepTensorTF::DeepTensorTF(const std::string &model, DeepTensorTF::~DeepTensorTF() { delete graph_def; } -void DeepTensorTF::init(const std::string &model, - const int &gpu_rank, - const std::string &name_scope_) { +void DeepTensorTF::init(const std::string& model, + const int& gpu_rank, + const std::string& name_scope_) { if (inited) { std::cerr << "WARNING: deepmd-kit should not be initialized twice, do " "nothing at the second call of initializer" @@ -59,7 +59,7 @@ void DeepTensorTF::init(const std::string &model, deepmd::check_status(session->Create(*graph_def)); try { model_version = get_scalar("model_attr/model_version"); - } catch (deepmd::tf_exception &e) { + } catch (deepmd::tf_exception& e) { // no model version defined in old models model_version = "0.0"; } @@ -85,23 +85,23 @@ void DeepTensorTF::init(const std::string &model, } template -VT DeepTensorTF::get_scalar(const std::string &name) const { +VT DeepTensorTF::get_scalar(const std::string& name) const { return session_get_scalar(session, name, name_scope); } template -void DeepTensorTF::get_vector(std::vector &vec, - const std::string &name) const { +void DeepTensorTF::get_vector(std::vector& vec, + const std::string& name) const { session_get_vector(vec, session, name, name_scope); } template void DeepTensorTF::run_model( - std::vector &d_tensor_, - Session *session, - const std::vector> &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& d_tensor_, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost) { unsigned nloc = atommap.get_type().size(); unsigned nall = nloc + nghost; @@ -139,46 +139,46 @@ void DeepTensorTF::run_model( } template void DeepTensorTF::run_model( - std::vector &d_tensor_, - Session *session, - const std::vector> &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& d_tensor_, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &d_tensor_, - Session *session, - const std::vector> &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& d_tensor_, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &d_tensor_, - Session *session, - const std::vector> &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& d_tensor_, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &d_tensor_, - Session *session, - const std::vector> &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& d_tensor_, + Session* session, + const std::vector>& input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - tensorflow::Session *session, - const std::vector> - &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + tensorflow::Session* session, + const std::vector>& + input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost) { unsigned nloc = atommap.get_type().size(); unsigned nall = nloc + nghost; @@ -282,61 +282,61 @@ void DeepTensorTF::run_model( } template void DeepTensorTF::run_model( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - tensorflow::Session *session, - const std::vector> - &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + tensorflow::Session* session, + const std::vector>& + input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - tensorflow::Session *session, - const std::vector> - &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + tensorflow::Session* session, + const std::vector>& + input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - tensorflow::Session *session, - const std::vector> - &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + tensorflow::Session* session, + const std::vector>& + input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template void DeepTensorTF::run_model( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - tensorflow::Session *session, - const std::vector> - &input_tensors, - const AtomMap &atommap, - const std::vector &sel_fwd, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + tensorflow::Session* session, + const std::vector>& + input_tensors, + const AtomMap& atommap, + const std::vector& sel_fwd, const int nghost); template -void DeepTensorTF::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensorTF::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { int nall = datype_.size(); std::vector dcoord, aparam, aparam_; std::vector datype, fwd_map, bkw_map; @@ -347,23 +347,23 @@ void DeepTensorTF::compute(std::vector &dtensor_, compute_inner(dtensor_, dcoord, datype, dbox); } -template void DeepTensorTF::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensorTF::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); -template void DeepTensorTF::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); +template void DeepTensorTF::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensorTF::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensorTF::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { int nall = datype_.size(); std::vector dcoord, dforce, datom_virial, aparam, aparam_; std::vector datype, fwd_map, bkw_map; @@ -380,29 +380,29 @@ void DeepTensorTF::compute(std::vector &dtensor_, compute_inner(dtensor_, dcoord, datype, dbox, nghost_real, nlist); } -template void DeepTensorTF::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +template void DeepTensorTF::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); -template void DeepTensorTF::compute(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +template void DeepTensorTF::compute(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); template -void DeepTensorTF::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensorTF::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { int nall = datype_.size(); std::vector dcoord, dforce, datom_virial, aparam, aparam_; std::vector datype, fwd_map, bkw_map; @@ -434,35 +434,35 @@ void DeepTensorTF::compute(std::vector &dglobal_tensor_, } template void DeepTensorTF::compute( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); - -template void DeepTensorTF::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); + +template void DeepTensorTF::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensorTF::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensorTF::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list) { + const InputNlist& lmp_list) { int nall = datype_.size(); std::vector dcoord, dforce, datom_virial, aparam, aparam_; std::vector datype, fwd_map, bkw_map; @@ -493,33 +493,33 @@ void DeepTensorTF::compute(std::vector &dglobal_tensor_, } template void DeepTensorTF::compute( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); - -template void DeepTensorTF::compute(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + const InputNlist& lmp_list); + +template void DeepTensorTF::compute(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &lmp_list); + const InputNlist& lmp_list); template -void DeepTensorTF::compute_inner(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensorTF::compute_inner(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { int nall = dcoord_.size() / 3; int nloc = nall; AtomMap atommap(datype_.begin(), datype_.begin() + nloc); @@ -550,24 +550,24 @@ void DeepTensorTF::compute_inner(std::vector &dtensor_, } template void DeepTensorTF::compute_inner( - std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); + std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template void DeepTensorTF::compute_inner( - std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); + std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensorTF::compute_inner(std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensorTF::compute_inner(std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &nlist_) { + const InputNlist& nlist_) { int nall = dcoord_.size() / 3; int nloc = nall - nghost; AtomMap atommap(datype_.begin(), datype_.begin() + nloc); @@ -608,30 +608,30 @@ void DeepTensorTF::compute_inner(std::vector &dtensor_, } template void DeepTensorTF::compute_inner( - std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &nlist_); + const InputNlist& nlist_); template void DeepTensorTF::compute_inner( - std::vector &dtensor_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + std::vector& dtensor_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &nlist_); + const InputNlist& nlist_); template -void DeepTensorTF::compute_inner(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox) { +void DeepTensorTF::compute_inner(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox) { int nall = dcoord_.size() / 3; int nloc = nall; AtomMap atommap(datype_.begin(), datype_.begin() + nloc); @@ -664,36 +664,36 @@ void DeepTensorTF::compute_inner(std::vector &dglobal_tensor_, } template void DeepTensorTF::compute_inner( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template void DeepTensorTF::compute_inner( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox); + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox); template -void DeepTensorTF::compute_inner(std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, +void DeepTensorTF::compute_inner(std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &nlist_) { + const InputNlist& nlist_) { int nall = dcoord_.size() / 3; int nloc = nall - nghost; AtomMap atommap(datype_.begin(), datype_.begin() + nloc); @@ -736,41 +736,41 @@ void DeepTensorTF::compute_inner(std::vector &dglobal_tensor_, } template void DeepTensorTF::compute_inner( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &nlist_); + const InputNlist& nlist_); template void DeepTensorTF::compute_inner( - std::vector &dglobal_tensor_, - std::vector &dforce_, - std::vector &dvirial_, - std::vector &datom_tensor_, - std::vector &datom_virial_, - const std::vector &dcoord_, - const std::vector &datype_, - const std::vector &dbox, + std::vector& dglobal_tensor_, + std::vector& dforce_, + std::vector& dvirial_, + std::vector& datom_tensor_, + std::vector& datom_virial_, + const std::vector& dcoord_, + const std::vector& datype_, + const std::vector& dbox, const int nghost, - const InputNlist &nlist_); + const InputNlist& nlist_); -void DeepTensorTF::get_type_map(std::string &type_map) { +void DeepTensorTF::get_type_map(std::string& type_map) { type_map = get_scalar("model_attr/tmap"); } -void DeepTensorTF::computew(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - std::vector &atom_tensor, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, +void DeepTensorTF::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const bool request_deriv) { if (request_deriv) { compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, @@ -783,14 +783,14 @@ void DeepTensorTF::computew(std::vector &global_tensor, atom_virial.clear(); } } -void DeepTensorTF::computew(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - std::vector &atom_tensor, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, +void DeepTensorTF::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const bool request_deriv) { if (request_deriv) { compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, @@ -804,16 +804,16 @@ void DeepTensorTF::computew(std::vector &global_tensor, } } -void DeepTensorTF::computew(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - std::vector &atom_tensor, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, +void DeepTensorTF::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &inlist, + const InputNlist& inlist, const bool request_deriv) { if (request_deriv) { compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, @@ -826,16 +826,16 @@ void DeepTensorTF::computew(std::vector &global_tensor, atom_virial.clear(); } } -void DeepTensorTF::computew(std::vector &global_tensor, - std::vector &force, - std::vector &virial, - std::vector &atom_tensor, - std::vector &atom_virial, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, +void DeepTensorTF::computew(std::vector& global_tensor, + std::vector& force, + std::vector& virial, + std::vector& atom_tensor, + std::vector& atom_virial, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, const int nghost, - const InputNlist &inlist, + const InputNlist& inlist, const bool request_deriv) { if (request_deriv) { compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc index 70755c901a..eace577f89 100644 --- a/source/api_cc/src/common.cc +++ b/source/api_cc/src/common.cc @@ -415,6 +415,9 @@ void deepmd::load_op_library() { #endif #ifdef BUILD_PYTORCH _load_single_op_library("deepmd_op_pt"); +#endif +#ifdef BUILD_PADDLE + _load_single_op_library("deepmd_op_pd"); #endif // load customized plugins const char* env_customized_plugins = std::getenv("DP_PLUGIN_PATH"); @@ -1419,7 +1422,9 @@ deepmd::DPBackend deepmd::get_backend(const std::string& model) { model.substr(model.length() - 11) == ".savedmodel") { return deepmd::DPBackend::JAX; } else if ((model.length() >= 5 && - model.substr(model.length() - 5) == ".json")) { + model.substr(model.length() - 5) == ".json") || + (model.length() >= 8 && + model.substr(model.length() - 8) == ".pdmodel")) { return deepmd::DPBackend::Paddle; } throw deepmd::deepmd_exception("Unsupported model file format"); diff --git a/source/api_cc/tests/CMakeLists.txt b/source/api_cc/tests/CMakeLists.txt index 5599b63243..8ec3361d3a 100644 --- a/source/api_cc/tests/CMakeLists.txt +++ b/source/api_cc/tests/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 3.9) +cmake_minimum_required(VERSION 3.25.2) project(deepmd_api_test) file(GLOB TEST_SRC test_*.cc) @@ -6,6 +6,18 @@ file(GLOB TEST_SRC test_*.cc) add_executable(runUnitTests_cc ${TEST_SRC}) target_link_libraries(runUnitTests_cc GTest::gtest_main ${LIB_DEEPMD_CC} coverage_config) +if(ENABLE_TENSORFLOW) + target_compile_definitions(runUnitTests_cc PRIVATE BUILD_TENSORFLOW) +endif() +if(ENABLE_PYTORCH) + target_compile_definitions(runUnitTests_cc PRIVATE BUILD_PYTORCH) +endif() +if(ENABLE_JAX) + target_compile_definitions(runUnitTests_cc PRIVATE BUILD_JAX) +endif() +if(ENABLE_PADDLE) + target_compile_definitions(runUnitTests_cc PRIVATE BUILD_PADDLE) +endif() if(UNIX AND NOT APPLE) find_library(RT_LIBRARY rt) diff --git a/source/api_cc/tests/test_deepdipole.cc b/source/api_cc/tests/test_deepdipole.cc index 86a8a4131f..4c5ec345a6 100644 --- a/source/api_cc/tests/test_deepdipole.cc +++ b/source/api_cc/tests/test_deepdipole.cc @@ -30,6 +30,9 @@ class TestInferDeepDipole : public ::testing::Test { deepmd::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif deepmd::convert_pbtxt_to_pb("../../tests/infer/deepdipole.pbtxt", "deepdipole.pb"); @@ -225,6 +228,9 @@ class TestInferDeepDipoleNew : public ::testing::Test { deepmd::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deepdipole_new.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deepdipole_new.pbtxt", "deepdipole_new.pb"); @@ -438,6 +444,9 @@ class TestInferDeepDipoleFake : public ::testing::Test { deepmd::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif deepmd::convert_pbtxt_to_pb("../../tests/infer/deepdipole_fake.pbtxt", "deepdipole_fake.pb"); diff --git a/source/api_cc/tests/test_deepdipole_pt.cc b/source/api_cc/tests/test_deepdipole_pt.cc new file mode 100644 index 0000000000..473f697188 --- /dev/null +++ b/source/api_cc/tests/test_deepdipole_pt.cc @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "DeepTensor.h" +#include "neighbor_list.h" +#include "test_utils.h" + +template +class TestInferDeepTensorPt : public ::testing::Test { + protected: + std::vector coord = {12.83, 2.56, 2.18, 12.09, 2.87, 2.74, + 00.25, 3.32, 1.68, 3.36, 3.00, 1.81, + 3.51, 2.51, 2.60, 4.27, 3.22, 1.56}; + std::vector atype = {0, 1, 1, 0, 1, 1}; + std::vector box = {13., 0., 0., 0., 13., 0., 0., 0., 13.}; + + // Expected global tensor values from Python inference + std::vector expected_global_tensor = {0.2338104, 0.23701073, + 0.2334505}; + + // Expected atomic tensor values from Python inference (flattened) + std::vector expected_atom_tensor = {-0.1808925408386811, + 0.3190798607195795, + 0.04760079958216837, + -0.0, + -0.0, + 0.0, + 0.0, + 0.0, + -0.0, + 0.4147029447879755, + -0.08206913353381971, + 0.1858497008385067, + 0.0, + -0.0, + 0.0, + 0.0, + 0.0, + -0.0}; + + int natoms = 6; + int output_dim = 3; + + deepmd::DeepTensor dt; + + void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif + std::string file_name = "../../tests/infer/deepdipole_pt.pth"; + dt.init(file_name); + }; + + void TearDown() override {}; +}; + +TYPED_TEST_SUITE(TestInferDeepTensorPt, ValueTypes); + +TYPED_TEST(TestInferDeepTensorPt, cpu_build_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_global_tensor = this->expected_global_tensor; + std::vector& expected_atom_tensor = this->expected_atom_tensor; + int& natoms = this->natoms; + int& output_dim = this->output_dim; + deepmd::DeepTensor& dt = this->dt; + // Use reasonable tolerance for minimal trained model + double tensor_tol = 1e-6; + + std::vector global_tensor, force, virial, atom_tensor, atom_virial; + + dt.compute(global_tensor, force, virial, atom_tensor, atom_virial, coord, + atype, box); + + EXPECT_EQ(global_tensor.size(), output_dim); + EXPECT_EQ(atom_tensor.size(), natoms * output_dim); + EXPECT_EQ(force.size(), natoms * output_dim * 3); + EXPECT_EQ(virial.size(), output_dim * 9); + EXPECT_EQ(atom_virial.size(), natoms * output_dim * 9); + + for (int ii = 0; ii < output_dim; ++ii) { + EXPECT_LT(fabs(global_tensor[ii] - expected_global_tensor[ii]), tensor_tol); + } + + for (int ii = 0; ii < natoms * output_dim; ++ii) { + EXPECT_LT(fabs(atom_tensor[ii] - expected_atom_tensor[ii]), tensor_tol); + } +} + +TYPED_TEST(TestInferDeepTensorPt, cpu_lmp_nlist) { + using VALUETYPE = TypeParam; + std::vector& coord = this->coord; + std::vector& atype = this->atype; + std::vector& box = this->box; + std::vector& expected_global_tensor = this->expected_global_tensor; + std::vector& expected_atom_tensor = this->expected_atom_tensor; + int& natoms = this->natoms; + int& output_dim = this->output_dim; + deepmd::DeepTensor& dt = this->dt; + double ener_tol = 1e-6; + + float rc = dt.cutoff(); + int nloc = coord.size() / 3; + std::vector coord_cpy; + std::vector atype_cpy, mapping; + std::vector > nlist_data; + _build_nlist(nlist_data, coord_cpy, atype_cpy, mapping, coord, + atype, box, rc); + int nall = coord_cpy.size() / 3; + std::vector ilist(nloc), numneigh(nloc); + std::vector firstneigh(nloc); + deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); + convert_nlist(inlist, nlist_data); + + std::vector global_tensor, force, virial, atom_tensor, atom_virial; + + dt.compute(global_tensor, force, virial, atom_tensor, atom_virial, coord_cpy, + atype_cpy, box, nall - nloc, inlist); + + EXPECT_EQ(global_tensor.size(), output_dim); + EXPECT_EQ(atom_tensor.size(), nall * output_dim); + + for (int ii = 0; ii < output_dim; ++ii) { + EXPECT_LT(fabs(global_tensor[ii] - expected_global_tensor[ii]), ener_tol); + } + + for (int ii = 0; ii < natoms * output_dim; ++ii) { + EXPECT_LT(fabs(atom_tensor[ii] - expected_atom_tensor[ii]), ener_tol); + } +} + +TYPED_TEST(TestInferDeepTensorPt, print_summary) { + deepmd::DeepTensor& dt = this->dt; + dt.print_summary(""); +} + +TYPED_TEST(TestInferDeepTensorPt, get_type_map) { + deepmd::DeepTensor& dt = this->dt; + std::string type_map_str; + dt.get_type_map(type_map_str); + // Parse the type map string manually + std::vector type_map; + std::istringstream iss(type_map_str); + std::string token; + while (iss >> token) { + type_map.push_back(token); + } + EXPECT_EQ(type_map.size(), 2); + EXPECT_EQ(type_map[0], "O"); + EXPECT_EQ(type_map[1], "H"); +} + +TYPED_TEST(TestInferDeepTensorPt, get_properties) { + deepmd::DeepTensor& dt = this->dt; + + EXPECT_EQ(dt.numb_types(), 2); + EXPECT_EQ(dt.output_dim(), 3); + EXPECT_DOUBLE_EQ(dt.cutoff(), 4.0); + + std::vector sel_types = dt.sel_types(); + EXPECT_EQ(sel_types.size(), 2); // PyTorch models always return all types + EXPECT_EQ(sel_types[0], 0); // Type 0 (O) + EXPECT_EQ(sel_types[1], + 1); // Type 1 (H) - included but may have zero results +} diff --git a/source/api_cc/tests/test_deepmd_exception.cc b/source/api_cc/tests/test_deepmd_exception.cc index 77e399d722..d150004af1 100644 --- a/source/api_cc/tests/test_deepmd_exception.cc +++ b/source/api_cc/tests/test_deepmd_exception.cc @@ -18,31 +18,46 @@ TEST(TestDeepmdException, deepmdexception) { std::string expected_error_message = "DeePMD-kit Error: unittest"; try { throw deepmd::deepmd_exception("unittest"); - } catch (deepmd::deepmd_exception &ex) { + } catch (deepmd::deepmd_exception& ex) { EXPECT_STREQ(expected_error_message.c_str(), ex.what()); } } TEST(TestDeepmdException, deepmdexception_nofile_deeppot) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif ASSERT_THROW(deepmd::DeepPot("_no_such_file.pb"), deepmd::deepmd_exception); } TEST(TestDeepmdException, deepmdexception_nofile_deeppot_pt) { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif ASSERT_THROW(deepmd::DeepPot("_no_such_file.pth"), deepmd::deepmd_exception); } TEST(TestDeepmdException, deepmdexception_nofile_deeppotmodeldevi) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif ASSERT_THROW( deepmd::DeepPotModelDevi({"_no_such_file.pb", "_no_such_file.pb"}), deepmd::deepmd_exception); } TEST(TestDeepmdException, deepmdexception_nofile_deeptensor) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif ASSERT_THROW(deepmd::DeepTensor("_no_such_file.pb"), deepmd::deepmd_exception); } TEST(TestDeepmdException, deepmdexception_nofile_dipolechargemodifier) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif ASSERT_THROW(deepmd::DipoleChargeModifier("_no_such_file.pb"), deepmd::deepmd_exception); } diff --git a/source/api_cc/tests/test_deeppolar.cc b/source/api_cc/tests/test_deeppolar.cc index 89014fd245..dbeb61db9f 100644 --- a/source/api_cc/tests/test_deeppolar.cc +++ b/source/api_cc/tests/test_deeppolar.cc @@ -36,6 +36,9 @@ class TestInferDeepPolar : public ::testing::Test { deepmd::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppolar.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppolar.pbtxt", "deeppolar.pb"); @@ -457,6 +460,9 @@ class TestInferDeepPolarNew : public ::testing::Test { deepmd::DeepTensor dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppolar_new.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppolar_new.pbtxt", "deeppolar_new.pb"); diff --git a/source/api_cc/tests/test_deeppot_a.cc b/source/api_cc/tests/test_deeppot_a.cc index ae1384e8a3..ffd2a25a5d 100644 --- a/source/api_cc/tests/test_deeppot_a.cc +++ b/source/api_cc/tests/test_deeppot_a.cc @@ -70,6 +70,9 @@ class TestInferDeepPotA : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", "deeppot.pb"); @@ -607,6 +610,9 @@ class TestInferDeepPotANoPbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::convert_pbtxt_to_pb(file_name, "deeppot.pb"); diff --git a/source/api_cc/tests/test_deeppot_a_fparam_aparam.cc b/source/api_cc/tests/test_deeppot_a_fparam_aparam.cc index 89e83603a3..93efee28fe 100644 --- a/source/api_cc/tests/test_deeppot_a_fparam_aparam.cc +++ b/source/api_cc/tests/test_deeppot_a_fparam_aparam.cc @@ -73,6 +73,9 @@ class TestInferDeepPotAFParamAParam : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::convert_pbtxt_to_pb(file_name, "fparam_aparam.pb"); diff --git a/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes.cc b/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes.cc index 0851523814..e33437805c 100644 --- a/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes.cc +++ b/source/api_cc/tests/test_deeppot_a_fparam_aparam_nframes.cc @@ -117,6 +117,9 @@ class TestInferDeepPotAFparamAparamNFrames : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/fparam_aparam.pbtxt", "fparam_aparam.pb"); @@ -758,6 +761,9 @@ class TestInferDeepPotAFparamAparamNFramesSingleParam : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/fparam_aparam.pbtxt", "fparam_aparam.pb"); diff --git a/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc b/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc index e394f86b2a..bc7661f671 100644 --- a/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc +++ b/source/api_cc/tests/test_deeppot_a_fparam_aparam_pt.cc @@ -77,6 +77,9 @@ class TestInferDeepPotAFParamAParamPt : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/fparam_aparam.pth"); natoms = expected_e.size(); diff --git a/source/api_cc/tests/test_deeppot_a_nframes.cc b/source/api_cc/tests/test_deeppot_a_nframes.cc index c83a7a0b41..1e7c160437 100644 --- a/source/api_cc/tests/test_deeppot_a_nframes.cc +++ b/source/api_cc/tests/test_deeppot_a_nframes.cc @@ -113,6 +113,9 @@ class TestInferDeepPotANFrames : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", "deeppot.pb"); @@ -723,6 +726,9 @@ class TestInferDeepPotANFramesNoPbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::convert_pbtxt_to_pb(file_name, "deeppot.pb"); diff --git a/source/api_cc/tests/test_deeppot_dpa_jax.cc b/source/api_cc/tests/test_deeppot_dpa_jax.cc index bd7f1ee7af..2e857450b7 100644 --- a/source/api_cc/tests/test_deeppot_dpa_jax.cc +++ b/source/api_cc/tests/test_deeppot_dpa_jax.cc @@ -77,6 +77,9 @@ class TestInferDeepPotDpaJAX : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_JAX + GTEST_SKIP() << "Skip because JAX support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa.savedmodel"); natoms = expected_e.size(); @@ -224,6 +227,9 @@ class TestInferDeepPotDpaJAXNopbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_JAX + GTEST_SKIP() << "Skip because JAX support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa.savedmodel"); natoms = expected_e.size(); diff --git a/source/api_cc/tests/test_deeppot_dpa_pt.cc b/source/api_cc/tests/test_deeppot_dpa_pt.cc index 7bf4cbf376..0de83278a9 100644 --- a/source/api_cc/tests/test_deeppot_dpa_pt.cc +++ b/source/api_cc/tests/test_deeppot_dpa_pt.cc @@ -77,6 +77,9 @@ class TestInferDeepPotDpaPt : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa.pth"); natoms = expected_e.size(); @@ -224,6 +227,9 @@ class TestInferDeepPotDpaPtNopbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa.pth"); natoms = expected_e.size(); diff --git a/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc b/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc index f14aa7a52e..8b569dd707 100644 --- a/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc +++ b/source/api_cc/tests/test_deeppot_dpa_pt_spin.cc @@ -85,6 +85,9 @@ class TestInferDeepSpinDpaPt : public ::testing::Test { deepmd::DeepSpin dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa_spin.pth"); natoms = expected_e.size(); @@ -248,6 +251,9 @@ class TestInferDeepSpinDpaPtNopbc : public ::testing::Test { deepmd::DeepSpin dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif dp.init("../../tests/infer/deeppot_dpa_spin.pth"); natoms = expected_e.size(); diff --git a/source/api_cc/tests/test_deeppot_jax.cc b/source/api_cc/tests/test_deeppot_jax.cc index 439a271015..80ddf00860 100644 --- a/source/api_cc/tests/test_deeppot_jax.cc +++ b/source/api_cc/tests/test_deeppot_jax.cc @@ -69,6 +69,9 @@ class TestInferDeepPotAJAX : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_JAX + GTEST_SKIP() << "Skip because JAX support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot_sea.savedmodel"; dp.init(file_name); diff --git a/source/api_cc/tests/test_deeppot_model_devi.cc b/source/api_cc/tests/test_deeppot_model_devi.cc index 5ce4b11c7f..ea89847935 100644 --- a/source/api_cc/tests/test_deeppot_model_devi.cc +++ b/source/api_cc/tests/test_deeppot_model_devi.cc @@ -28,6 +28,9 @@ class TestInferDeepPotModeDevi : public ::testing::Test { deepmd::DeepPotModelDevi dp_md; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif { std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", @@ -78,6 +81,9 @@ class TestInferDeepPotModeDeviPython : public ::testing::Test { deepmd::DeepPotModelDevi dp_md; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif { std::string file_name = "../../tests/infer/deeppot.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppot.pbtxt", diff --git a/source/api_cc/tests/test_deeppot_model_devi_fparam_aparam.cc b/source/api_cc/tests/test_deeppot_model_devi_fparam_aparam.cc index fb1797961d..3c49a34db5 100644 --- a/source/api_cc/tests/test_deeppot_model_devi_fparam_aparam.cc +++ b/source/api_cc/tests/test_deeppot_model_devi_fparam_aparam.cc @@ -31,6 +31,9 @@ class TestInferDeepPotModeDeviFparamAparam : public ::testing::Test { deepmd::DeepPotModelDevi dp_md; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif { std::string file_name = "../../tests/infer/fparam_aparam.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/fparam_aparam.pbtxt", diff --git a/source/api_cc/tests/test_deeppot_pd.cc b/source/api_cc/tests/test_deeppot_pd.cc index f39aa84640..6da9ee643b 100644 --- a/source/api_cc/tests/test_deeppot_pd.cc +++ b/source/api_cc/tests/test_deeppot_pd.cc @@ -61,6 +61,9 @@ class TestInferDeepPotAPd : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PADDLE + GTEST_SKIP() << "Skip because Paddle support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot_sea.json"; dp.init(file_name); @@ -578,6 +581,9 @@ class TestInferDeepPotAPdNoPbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PADDLE + GTEST_SKIP() << "Skip because Paddle support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot_sea.json"; dp.init(file_name); diff --git a/source/api_cc/tests/test_deeppot_pt.cc b/source/api_cc/tests/test_deeppot_pt.cc index af5b533a5f..144ee8da8c 100644 --- a/source/api_cc/tests/test_deeppot_pt.cc +++ b/source/api_cc/tests/test_deeppot_pt.cc @@ -54,6 +54,9 @@ class TestInferDeepPotAPt : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot_sea.pth"; dp.init(file_name); @@ -570,6 +573,9 @@ class TestInferDeepPotAPtNoPbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot_sea.pth"; dp.init(file_name); diff --git a/source/api_cc/tests/test_deeppot_r.cc b/source/api_cc/tests/test_deeppot_r.cc index 86fa4b091f..99c5eca6e5 100644 --- a/source/api_cc/tests/test_deeppot_r.cc +++ b/source/api_cc/tests/test_deeppot_r.cc @@ -70,6 +70,9 @@ class TestInferDeepPotR : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot-r.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppot-r.pbtxt", "deeppot.pb"); @@ -539,6 +542,9 @@ class TestInferDeepPotRNoPbc : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deeppot-r.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deeppot-r.pbtxt", "deeppot.pb"); diff --git a/source/api_cc/tests/test_deeppot_tf_spin.cc b/source/api_cc/tests/test_deeppot_tf_spin.cc index d03a9b0b57..c7623d8faf 100644 --- a/source/api_cc/tests/test_deeppot_tf_spin.cc +++ b/source/api_cc/tests/test_deeppot_tf_spin.cc @@ -41,6 +41,9 @@ class TestInferDeepSpin : public ::testing::Test { deepmd::DeepSpin dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deepspin_nlist.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deepspin_nlist.pbtxt", "deepspin_nlist.pb"); @@ -168,6 +171,9 @@ class TestInferDeepSpinNopbc : public ::testing::Test { deepmd::DeepSpin dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/deepspin_nlist.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/deepspin_nlist.pbtxt", "deepspin_nlist.pb"); diff --git a/source/api_cc/tests/test_deepspin_model_devi.cc b/source/api_cc/tests/test_deepspin_model_devi.cc index fcc4a4315d..56104a4f99 100644 --- a/source/api_cc/tests/test_deepspin_model_devi.cc +++ b/source/api_cc/tests/test_deepspin_model_devi.cc @@ -30,6 +30,9 @@ class TestInferDeepSpinModeDevi : public ::testing::Test { deepmd::DeepSpinModelDevi dp_md; void SetUp() override { +#ifndef BUILD_PYTORCH + GTEST_SKIP() << "Skip because PyTorch support is not enabled."; +#endif { dp0.init("../../tests/infer/deeppot_dpa_spin.pth"); } diff --git a/source/api_cc/tests/test_dipolecharge.cc b/source/api_cc/tests/test_dipolecharge.cc index 1d92215c5c..1a33d5ed1c 100644 --- a/source/api_cc/tests/test_dipolecharge.cc +++ b/source/api_cc/tests/test_dipolecharge.cc @@ -59,6 +59,9 @@ class TestDipoleCharge : public ::testing::Test { deepmd::DipoleChargeModifier dm; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/dipolecharge_e.pbtxt"; std::string model = "dipolecharge_e.pb"; deepmd::convert_pbtxt_to_pb(file_name, model); diff --git a/source/api_cc/tests/test_mixed_type.cc b/source/api_cc/tests/test_mixed_type.cc index b11bd5123a..0a577e21f2 100644 --- a/source/api_cc/tests/test_mixed_type.cc +++ b/source/api_cc/tests/test_mixed_type.cc @@ -30,6 +30,9 @@ class TestInferMixedType : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/virtual_type.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/virtual_type.pbtxt", "virtual_type.pb"); @@ -137,6 +140,9 @@ class TestInferVirtualType : public ::testing::Test { deepmd::DeepPot dp; void SetUp() override { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_name = "../../tests/infer/virtual_type.pbtxt"; deepmd::convert_pbtxt_to_pb("../../tests/infer/virtual_type.pbtxt", "virtual_type.pb"); diff --git a/source/api_cc/tests/test_read_file_to_string.cc b/source/api_cc/tests/test_read_file_to_string.cc index aaa7e6a83f..dd3249d771 100644 --- a/source/api_cc/tests/test_read_file_to_string.cc +++ b/source/api_cc/tests/test_read_file_to_string.cc @@ -13,6 +13,9 @@ #include "common.h" TEST(TestReadFileToString, readfiletostring) { +#ifndef BUILD_TENSORFLOW + GTEST_SKIP() << "Skip because TensorFlow support is not enabled."; +#endif std::string file_content; deepmd::read_file_to_string("../../tests/infer/deeppot.txt", file_content); diff --git a/source/api_cc/tests/test_utils.h b/source/api_cc/tests/test_utils.h index d06823b4e0..64d8a37ef5 100644 --- a/source/api_cc/tests/test_utils.h +++ b/source/api_cc/tests/test_utils.h @@ -14,7 +14,7 @@ typedef testing::Types ValueTypes; template inline void _fold_back(typename std::vector::iterator out, const typename std::vector::const_iterator in, - const std::vector &mapping, + const std::vector& mapping, const int nloc, const int nall, const int ndim, @@ -35,9 +35,9 @@ inline void _fold_back(typename std::vector::iterator out, } template -inline void _fold_back(std::vector &out, - const std::vector &in, - const std::vector &mapping, +inline void _fold_back(std::vector& out, + const std::vector& in, + const std::vector& mapping, const int nloc, const int nall, const int ndim, @@ -48,14 +48,14 @@ inline void _fold_back(std::vector &out, } template -inline void _build_nlist(std::vector> &nlist_data, - std::vector &coord_cpy, - std::vector &atype_cpy, - std::vector &mapping, - const std::vector &coord, - const std::vector &atype, - const std::vector &box, - const float &rc) { +inline void _build_nlist(std::vector>& nlist_data, + std::vector& coord_cpy, + std::vector& atype_cpy, + std::vector& mapping, + const std::vector& coord, + const std::vector& atype, + const std::vector& box, + const float& rc) { // convert VALUETYPE to double, it looks like copy_coord only accepts double std::vector coord_cpy_; std::vector coord_(coord.begin(), coord.end()); @@ -90,13 +90,13 @@ class EnergyModelTest { double level = std::is_same::value ? 1e-6 : 1e-2; // expected? public: - virtual void compute(double &ener, - std::vector &force, - std::vector &virial, - const std::vector &coord, - const std::vector &box) = 0; - void test_f(const std::vector &coord, - const std::vector &box) { + virtual void compute(double& ener, + std::vector& force, + std::vector& virial, + const std::vector& coord, + const std::vector& box) = 0; + void test_f(const std::vector& coord, + const std::vector& box) { int ndof = coord.size(); double ener; std::vector force, virial; @@ -114,8 +114,8 @@ class EnergyModelTest { EXPECT_LT(fabs(num - ana), level); } } - void test_v(const std::vector &coord, - const std::vector &box) { + void test_v(const std::vector& coord, + const std::vector& box) { std::vector num_diff(9); double ener; std::vector force, virial; diff --git a/source/cmake/Findtensorflow.cmake b/source/cmake/Findtensorflow.cmake index d579af7679..b5b8c92f3d 100644 --- a/source/cmake/Findtensorflow.cmake +++ b/source/cmake/Findtensorflow.cmake @@ -291,8 +291,10 @@ if(NOT DEFINED TENSORFLOW_VERSION) TENSORFLOW_VERSION_RUN_RESULT_VAR TENSORFLOW_VERSION_COMPILE_RESULT_VAR ${CMAKE_CURRENT_BINARY_DIR}/tf_version "${CMAKE_CURRENT_LIST_DIR}/tf_version.cpp" - CMAKE_FLAGS "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}" - RUN_OUTPUT_VARIABLE TENSORFLOW_VERSION + CMAKE_FLAGS + "-DINCLUDE_DIRECTORIES:STRING=${TensorFlow_INCLUDE_DIRS}" LINK_LIBRARIES + ${TensorFlowFramework_LIBRARY} ${TensorFlow_LIBRARY} + RUN_OUTPUT_STDOUT_VARIABLE TENSORFLOW_VERSION COMPILE_OUTPUT_VARIABLE TENSORFLOW_VERSION_COMPILE_OUTPUT_VAR) if(NOT ${TENSORFLOW_VERSION_COMPILE_RESULT_VAR}) message( @@ -304,6 +306,23 @@ if(NOT DEFINED TENSORFLOW_VERSION) endif() endif() +if(TENSORFLOW_VERSION VERSION_GREATER_EQUAL 2.20) + # since TF 2.20, macros like TF_MAJOR_VERSION, TF_MINOR_VERSION, and + # TF_PATCH_VERSION are not defined We manuanlly define them in our CMake files + # first, split TENSORFLOW_VERSION (e.g. 2.20.0rc0) to 2 20 0 rc0 + string(REGEX MATCH "^([0-9]+)\\.([0-9]+)\\.([0-9]+)(.*)$" _match + ${TENSORFLOW_VERSION}) + if(_match) + set(TF_MAJOR_VERSION ${CMAKE_MATCH_1}) + set(TF_MINOR_VERSION ${CMAKE_MATCH_2}) + set(TF_PATCH_VERSION ${CMAKE_MATCH_3}) + # add defines + add_definitions(-DTF_MAJOR_VERSION=${TF_MAJOR_VERSION}) + add_definitions(-DTF_MINOR_VERSION=${TF_MINOR_VERSION}) + add_definitions(-DTF_PATCH_VERSION=${TF_PATCH_VERSION}) + endif() +endif() + # print message if(NOT TensorFlow_FIND_QUIETLY) message( diff --git a/source/cmake/tf_version.cpp b/source/cmake/tf_version.cpp index 6d09e33493..2ad2125291 100644 --- a/source/cmake/tf_version.cpp +++ b/source/cmake/tf_version.cpp @@ -1,12 +1,14 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include -#include "tensorflow/core/public/version.h" +#include "tensorflow/c/c_api.h" int main(int argc, char* argv[]) { // See // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/public/version.h // TF_VERSION_STRING has been available since TensorFlow v0.6 - std::cout << TF_VERSION_STRING; + // Aug 2025: since TF 2.20, TF_VERSION_STRING is no more available; + // try to use the C API TF_Version + std::cout << TF_Version(); return 0; } diff --git a/source/config/run_config.ini b/source/config/run_config.ini index 7bb6041af9..596be911db 100644 --- a/source/config/run_config.ini +++ b/source/config/run_config.ini @@ -18,3 +18,5 @@ PD_VERSION = @PADDLE_VERSION@ PD_INFERENCE_DIR = @PADDLE_INFERENCE_DIR@ MODEL_VERSION=@MODEL_VERSION@ DP_VARIANT=@DP_VARIANT@ +LAMMPS_VERSION = @LAMMPS_VERSION@ +CIBUILDWHEEL = @CIBUILDWHEEL@ diff --git a/source/install/build_cc.sh b/source/install/build_cc.sh index dc66343cb2..56c50e92f5 100755 --- a/source/install/build_cc.sh +++ b/source/install/build_cc.sh @@ -26,7 +26,7 @@ cmake -D ENABLE_TENSORFLOW=ON \ -D USE_TF_PYTHON_LIBS=TRUE \ -D USE_PT_PYTHON_LIBS=TRUE \ ${CUDA_ARGS} \ - -D LAMMPS_VERSION=stable_29Aug2024_update1 \ + -D LAMMPS_VERSION=stable_22Jul2025_update2 \ .. cmake --build . -j${NPROC} cmake --install . diff --git a/source/install/build_from_c.sh b/source/install/build_from_c.sh index 22739ec531..903d50bb72 100755 --- a/source/install/build_from_c.sh +++ b/source/install/build_from_c.sh @@ -13,7 +13,7 @@ NPROC=$(nproc --all) BUILD_TMP_DIR=${SCRIPT_PATH}/../build mkdir -p ${BUILD_TMP_DIR} cd ${BUILD_TMP_DIR} -cmake -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DDEEPMD_C_ROOT=${DEEPMD_C_ROOT} -DLAMMPS_VERSION=stable_29Aug2024_update1 .. +cmake -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DDEEPMD_C_ROOT=${DEEPMD_C_ROOT} -DLAMMPS_VERSION=stable_22Jul2025_update2 .. cmake --build . -j${NPROC} cmake --install . cmake --build . --target=lammps diff --git a/source/install/build_lammps.sh b/source/install/build_lammps.sh index d101714739..3cf981c5f2 100755 --- a/source/install/build_lammps.sh +++ b/source/install/build_lammps.sh @@ -14,7 +14,7 @@ BUILD_TMP_DIR=${SCRIPT_PATH}/../build_lammps mkdir -p ${BUILD_TMP_DIR} cd ${BUILD_TMP_DIR} # download LAMMMPS -LAMMPS_VERSION=stable_29Aug2024_update1 +LAMMPS_VERSION=stable_22Jul2025_update2 if [ ! -d "lammps-${LAMMPS_VERSION}" ]; then curl -L -o lammps.tar.gz https://github.com/lammps/lammps/archive/refs/tags/${LAMMPS_VERSION}.tar.gz tar vxzf lammps.tar.gz diff --git a/source/install/build_tf.py b/source/install/build_tf.py index d4fcc357e7..a5ca2d7411 100755 --- a/source/install/build_tf.py +++ b/source/install/build_tf.py @@ -57,7 +57,6 @@ ) from typing import ( NoReturn, - Optional, ) # default config @@ -133,9 +132,9 @@ def __init__( self, filename: str, url: str, - sha256: Optional[str] = None, + sha256: str | None = None, executable: bool = False, - gzip: Optional[str] = None, + gzip: str | None = None, ) -> None: self.filename = filename self.url = url @@ -856,7 +855,7 @@ class RawTextArgumentDefaultsHelpFormatter( pass -def parse_args(args: Optional[list[str]] = None): +def parse_args(args: list[str] | None = None): """TensorFlow C++ Library Installer commandline options argument parser. Parameters diff --git a/source/install/docker/Dockerfile b/source/install/docker/Dockerfile index 1e0f46eec6..da6936b55d 100644 --- a/source/install/docker/Dockerfile +++ b/source/install/docker/Dockerfile @@ -8,8 +8,7 @@ ENV PATH="/opt/deepmd-kit/bin:$PATH" ENV VIRTUAL_ENV="/opt/deepmd-kit" # Install package COPY dist /dist -RUN if [ "${CUDA_VERSION}" = 11 ]; then uv pip install torch --index-url https://download.pytorch.org/whl/cu118; fi \ - && uv pip install "$(ls /dist/deepmd_kit${VARIANT}-*manylinux*_x86_64.whl)[gpu,cu${CUDA_VERSION},lmp,ipi,torch]" \ +RUN uv pip install "$(ls /dist/deepmd_kit${VARIANT}-*manylinux*_x86_64.whl)[gpu,cu${CUDA_VERSION},lmp,ipi,torch]" \ && dp -h \ && lmp -h \ && dp_ipi \ diff --git a/source/install/test_cc.sh b/source/install/test_cc.sh index 1626f36193..8058edda87 100755 --- a/source/install/test_cc.sh +++ b/source/install/test_cc.sh @@ -17,7 +17,7 @@ INSTALL_PREFIX=${SCRIPT_PATH}/../../dp_test BUILD_TMP_DIR=${SCRIPT_PATH}/../build_tests mkdir -p ${BUILD_TMP_DIR} cd ${BUILD_TMP_DIR} -cmake -DINSTALL_TENSORFLOW=TRUE -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DTENSORFLOW_ROOT=${INSTALL_PREFIX} -DBUILD_TESTING:BOOL=TRUE -DLAMMPS_VERSION=stable_29Aug2024_update1 ${CUDA_ARGS} .. +cmake -DINSTALL_TENSORFLOW=TRUE -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} -DTENSORFLOW_ROOT=${INSTALL_PREFIX} -DBUILD_TESTING:BOOL=TRUE -DLAMMPS_VERSION=stable_22Jul2025_update2 ${CUDA_ARGS} .. cmake --build . -j${NPROC} cmake --install . ctest --output-on-failure diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 8152b6f1a4..86088bc0b3 100755 --- a/source/install/test_cc_local.sh +++ b/source/install/test_cc_local.sh @@ -1,3 +1,4 @@ +#!/bin/bash set -ex if [ "$DP_VARIANT" = "cuda" ]; then @@ -20,18 +21,20 @@ PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir mkdir -p ${BUILD_TMP_DIR} cd ${BUILD_TMP_DIR} cmake \ - -D ENABLE_TENSORFLOW=TRUE \ - -D ENABLE_PYTORCH=TRUE \ - -D ENABLE_PADDLE=TRUE \ + -D ENABLE_TENSORFLOW=${ENABLE_TENSORFLOW:-TRUE} \ + -D ENABLE_PYTORCH=${ENABLE_PYTORCH:-TRUE} \ + -D ENABLE_PADDLE=${ENABLE_PADDLE:-TRUE} \ -D INSTALL_TENSORFLOW=FALSE \ - -D USE_TF_PYTHON_LIBS=TRUE \ - -D USE_PT_PYTHON_LIBS=TRUE \ + -D USE_TF_PYTHON_LIBS=${ENABLE_TENSORFLOW:-TRUE} \ + -D USE_PT_PYTHON_LIBS=${ENABLE_PYTORCH:-TRUE} \ -D CMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \ -D BUILD_TESTING:BOOL=TRUE \ - -D LAMMPS_VERSION=stable_29Aug2024_update1 \ + -D LAMMPS_VERSION=stable_22Jul2025_update2 \ ${CUDA_ARGS} .. cmake --build . -j${NPROC} cmake --install . -PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir -export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PADDLE_INFERENCE_DIR}/third_party/install/onednn/lib:${PADDLE_INFERENCE_DIR}/third_party/install/mklml/lib +if [ "${ENABLE_PADDLE:-TRUE}" == "TRUE" ]; then + PADDLE_INFERENCE_DIR=${BUILD_TMP_DIR}/paddle_inference_install_dir + export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PADDLE_INFERENCE_DIR}/third_party/install/onednn/lib:${PADDLE_INFERENCE_DIR}/third_party/install/mklml/lib +fi ctest --output-on-failure diff --git a/source/ipi/driver.cc b/source/ipi/driver.cc index 9a91a27ad3..879e19c46f 100644 --- a/source/ipi/driver.cc +++ b/source/ipi/driver.cc @@ -29,8 +29,8 @@ const double icvt_ener = 1. / cvt_ener; const double cvt_f = cvt_ener / cvt_len; const double icvt_f = 1. / cvt_f; -char *trimwhitespace(char *str) { - char *end; +char* trimwhitespace(char* str) { + char* end; // Trim leading space while (isspace((unsigned char)*str)) { str++; @@ -48,7 +48,7 @@ char *trimwhitespace(char *str) { return str; } -int main(int argc, char *argv[]) { +int main(int argc, char* argv[]) { if (argc == 1) { std::cerr << "usage " << std::endl; std::cerr << argv[0] << " input_script " << std::endl; @@ -68,7 +68,7 @@ int main(int argc, char *argv[]) { } int port = jdata["port"]; std::string host_str = jdata["host"]; - const char *host = host_str.c_str(); + const char* host = host_str.c_str(); std::string graph_file = jdata["graph_file"]; std::string coord_file = jdata["coord_file"]; std::map name_type_map = jdata["atom_type"]; @@ -102,7 +102,7 @@ int main(int argc, char *argv[]) { std::vector dcoord_tmp; std::vector dtype = cvt.get_type(); std::vector dbox(9, 0); - double *msg_buff = NULL; + double* msg_buff = NULL; double ener; double virial[9]; char msg_needinit[] = "NEEDINIT "; @@ -144,7 +144,7 @@ int main(int argc, char *argv[]) { } } else if (header_str == "INIT") { assert(4 == sizeof(int32_t)); - readbuffer_(&socket, (char *)(&cbuf), sizeof(int32_t)); + readbuffer_(&socket, (char*)(&cbuf), sizeof(int32_t)); readbuffer_(&socket, initbuffer, cbuf); if (b_verb) { std::cout << "Init sys from wrapper, using " << initbuffer << std::endl; @@ -153,14 +153,14 @@ int main(int argc, char *argv[]) { assert(8 == sizeof(double)); // get box - readbuffer_(&socket, (char *)(cell_h), 9 * sizeof(double)); - readbuffer_(&socket, (char *)(cell_ih), 9 * sizeof(double)); + readbuffer_(&socket, (char*)(cell_h), 9 * sizeof(double)); + readbuffer_(&socket, (char*)(cell_ih), 9 * sizeof(double)); for (int dd = 0; dd < 9; ++dd) { dbox[dd] = cell_h[(dd % 3) * 3 + (dd / 3)] * cvt_len; } // get number of atoms - readbuffer_(&socket, (char *)(&cbuf), sizeof(int32_t)); + readbuffer_(&socket, (char*)(&cbuf), sizeof(int32_t)); if (natoms < 0) { natoms = cbuf; if (b_verb) { @@ -176,7 +176,7 @@ int main(int argc, char *argv[]) { } // get coord - readbuffer_(&socket, (char *)(msg_buff), natoms * 3 * sizeof(double)); + readbuffer_(&socket, (char*)(msg_buff), natoms * 3 * sizeof(double)); for (int ii = 0; ii < natoms * 3; ++ii) { dcoord_tmp[ii] = msg_buff[ii] * cvt_len; } @@ -199,12 +199,12 @@ int main(int argc, char *argv[]) { << std::setprecision(10) << dener << std::endl; } writebuffer_(&socket, msg_forceready, MSGLEN); - writebuffer_(&socket, (char *)(&ener), sizeof(double)); - writebuffer_(&socket, (char *)(&natoms), sizeof(int32_t)); - writebuffer_(&socket, (char *)(msg_buff), 3 * natoms * sizeof(double)); - writebuffer_(&socket, (char *)(virial), 9 * sizeof(double)); + writebuffer_(&socket, (char*)(&ener), sizeof(double)); + writebuffer_(&socket, (char*)(&natoms), sizeof(int32_t)); + writebuffer_(&socket, (char*)(msg_buff), 3 * natoms * sizeof(double)); + writebuffer_(&socket, (char*)(virial), 9 * sizeof(double)); cbuf = 7; - writebuffer_(&socket, (char *)(&cbuf), sizeof(int32_t)); + writebuffer_(&socket, (char*)(&cbuf), sizeof(int32_t)); writebuffer_(&socket, msg_nothing, 7); hasdata = false; } else { diff --git a/source/ipi/include/sockets.h b/source/ipi/include/sockets.h index 08f24c68ed..150b7c1a69 100644 --- a/source/ipi/include/sockets.h +++ b/source/ipi/include/sockets.h @@ -15,7 +15,7 @@ extern "C" { #endif -void error(const char *msg); +void error(const char* msg); /* Opens a socket. Note that fortran passes an extra argument for the string length, but this is @@ -29,7 +29,7 @@ void error(const char *msg); recommended. host: The name of the host server. */ -void open_socket_(int *psockfd, int *inet, int *port, const char *host); +void open_socket_(int* psockfd, int* inet, int* port, const char* host); /* Writes to a socket. Args: @@ -37,7 +37,7 @@ void open_socket_(int *psockfd, int *inet, int *port, const char *host); data: The data to be written to the socket. plen: The length of the data in bytes. */ -void writebuffer_(int *psockfd, char *data, int len); +void writebuffer_(int* psockfd, char* data, int len); /* Reads from a socket. Args: @@ -45,7 +45,7 @@ void writebuffer_(int *psockfd, char *data, int len); data: The storage array for data read from the socket. plen: The length of the data in bytes. */ -void readbuffer_(int *psockfd, char *data, int len); +void readbuffer_(int* psockfd, char* data, int len); #ifdef __cplusplus } diff --git a/source/ipi/src/sockets.c b/source/ipi/src/sockets.c index d9a2b8a865..1d45849f1a 100644 --- a/source/ipi/src/sockets.c +++ b/source/ipi/src/sockets.c @@ -45,14 +45,14 @@ Can be linked to a FORTRAN code that does not support sockets natively. #include #include -void error(const char *msg) +void error(const char* msg) // Prints an error message and then exits. { perror(msg); exit(-1); } -void open_socket_(int *psockfd, int *inet, int *port, const char *host) +void open_socket_(int* psockfd, int* inet, int* port, const char* host) /* Opens a socket. Note that fortran passes an extra argument for the string length, but this is @@ -70,14 +70,14 @@ ignored here for C compatibility. { int sockfd, portno, n; - struct hostent *server; + struct hostent* server; - struct sockaddr *psock; + struct sockaddr* psock; int ssock; if (*inet > 0) { // creates an internet socket struct sockaddr_in serv_addr; - psock = (struct sockaddr *)&serv_addr; + psock = (struct sockaddr*)&serv_addr; ssock = sizeof(serv_addr); sockfd = socket(AF_INET, SOCK_STREAM, 0); if (sockfd < 0) { @@ -90,9 +90,9 @@ ignored here for C compatibility. exit(-1); } - bzero((char *)&serv_addr, sizeof(serv_addr)); + bzero((char*)&serv_addr, sizeof(serv_addr)); serv_addr.sin_family = AF_INET; - bcopy((char *)server->h_addr, (char *)&serv_addr.sin_addr.s_addr, + bcopy((char*)server->h_addr, (char*)&serv_addr.sin_addr.s_addr, server->h_length); serv_addr.sin_port = htons(*port); if (connect(sockfd, psock, ssock) < 0) { @@ -100,10 +100,10 @@ ignored here for C compatibility. } } else { // creates a unix socket struct sockaddr_un serv_addr; - psock = (struct sockaddr *)&serv_addr; + psock = (struct sockaddr*)&serv_addr; ssock = sizeof(serv_addr); sockfd = socket(AF_UNIX, SOCK_STREAM, 0); - bzero((char *)&serv_addr, sizeof(serv_addr)); + bzero((char*)&serv_addr, sizeof(serv_addr)); serv_addr.sun_family = AF_UNIX; strcpy(serv_addr.sun_path, "/tmp/ipi_"); strcpy(serv_addr.sun_path + 9, host); @@ -115,7 +115,7 @@ ignored here for C compatibility. *psockfd = sockfd; } -void writebuffer_(int *psockfd, char *data, int len) +void writebuffer_(int* psockfd, char* data, int len) /* Writes to a socket. Args: @@ -134,7 +134,7 @@ void writebuffer_(int *psockfd, char *data, int len) } } -void readbuffer_(int *psockfd, char *data, int len) +void readbuffer_(int* psockfd, char* data, int len) /* Reads from a socket. Args: diff --git a/source/ipi/tests/test_driver.py b/source/ipi/tests/test_driver.py index 54c632a2b6..9a30c2c8ce 100644 --- a/source/ipi/tests/test_driver.py +++ b/source/ipi/tests/test_driver.py @@ -53,6 +53,10 @@ def write_input(self, atoms, **kwargs) -> None: atoms.write(self.xyz_file, format="xyz") +@unittest.skipIf( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) class TestDPIPI(unittest.TestCase): # copy from test_deeppot_a.py @classmethod @@ -253,6 +257,10 @@ def test_normalize_coords(self) -> None: np.testing.assert_almost_equal(ee.ravel(), expected_se.ravel(), default_places) +@unittest.skipIf( + os.environ.get("ENABLE_PYTORCH", "1") != "1", + reason="Skip test because PyTorch support is not enabled.", +) class TestDPIPIPt(TestDPIPI): @classmethod def setUpClass(cls) -> None: diff --git a/source/lib/include/ComputeDescriptor.h b/source/lib/include/ComputeDescriptor.h index 733cb1ee0c..edede310b6 100644 --- a/source/lib/include/ComputeDescriptor.h +++ b/source/lib/include/ComputeDescriptor.h @@ -9,100 +9,100 @@ #include "switcher.h" #include "utilities.h" -inline void compute_descriptor(std::vector &descrpt_a, - std::vector &descrpt_r, - std::vector &rot_mat, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &fmt_nlist_r, - const std::vector &sec_a, - const std::vector &sec_r, +inline void compute_descriptor(std::vector& descrpt_a, + std::vector& descrpt_r, + std::vector& rot_mat, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& fmt_nlist_r, + const std::vector& sec_a, + const std::vector& sec_r, const int axis0_type, const int axis0_idx, const int axis1_type, const int axis1_idx); -inline void compute_descriptor(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &descrpt_r, - std::vector &descrpt_r_deriv, - std::vector &rij_a, - std::vector &rij_r, - std::vector &rot_mat, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &fmt_nlist_r, - const std::vector &sec_a, - const std::vector &sec_r, +inline void compute_descriptor(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& descrpt_r, + std::vector& descrpt_r_deriv, + std::vector& rij_a, + std::vector& rij_r, + std::vector& rot_mat, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& fmt_nlist_r, + const std::vector& sec_a, + const std::vector& sec_r, const int axis0_type, const int axis0_idx, const int axis1_type, const int axis1_idx); -inline void compute_descriptor_se_a_extf(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const std::vector &efield, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &sec_a, - const double &rmin, - const double &rmax); +inline void compute_descriptor_se_a_extf(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const std::vector& efield, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& sec_a, + const double& rmin, + const double& rmax); inline void compute_descriptor_se_a_ef_para( - std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const std::vector &efield, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &sec_a, - const double &rmin, - const double &rmax); + std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const std::vector& efield, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& sec_a, + const double& rmin, + const double& rmax); inline void compute_descriptor_se_a_ef_vert( - std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const std::vector &efield, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &sec_a, - const double &rmin, - const double &rmax); + std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const std::vector& efield, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& sec_a, + const double& rmin, + const double& rmax); static void compute_dRdT(double (*dRdT)[9], - const double *r1, - const double *r2, - const double *rot) { - double *dRdT0 = dRdT[0]; - double *dRdT1 = dRdT[1]; - double *dRdT2 = dRdT[2]; - const double *xx = rot; - const double *yy = rot + 3; + const double* r1, + const double* r2, + const double* rot) { + double* dRdT0 = dRdT[0]; + double* dRdT1 = dRdT[1]; + double* dRdT2 = dRdT[2]; + const double* xx = rot; + const double* yy = rot + 3; double nr1 = sqrt(deepmd::dot3(r1, r1)); double nr12 = nr1 * nr1; @@ -160,14 +160,14 @@ static void compute_dRdT(double (*dRdT)[9], } static void compute_dRdT_1(double (*dRdT)[9], - const double *r1, - const double *r2, - const double *rot) { - double *dRdT0 = dRdT[0]; - double *dRdT1 = dRdT[1]; - double *dRdT2 = dRdT[2]; - const double *xx = rot; - const double *yy = rot + 3; + const double* r1, + const double* r2, + const double* rot) { + double* dRdT0 = dRdT[0]; + double* dRdT1 = dRdT[1]; + double* dRdT2 = dRdT[2]; + const double* xx = rot; + const double* yy = rot + 3; double nr1 = sqrt(deepmd::dot3(r1, r1)); double nr12 = nr1 * nr1; @@ -225,14 +225,14 @@ static void compute_dRdT_1(double (*dRdT)[9], } static void compute_dRdT_2(double (*dRdT)[9], - const double *r1, - const double *r2, - const double *rot) { - double *dRdT0 = dRdT[0]; - double *dRdT1 = dRdT[1]; - double *dRdT2 = dRdT[2]; - const double *xx = rot; - const double *yy = rot + 3; + const double* r1, + const double* r2, + const double* rot) { + double* dRdT0 = dRdT[0]; + double* dRdT1 = dRdT[1]; + double* dRdT2 = dRdT[2]; + const double* xx = rot; + const double* yy = rot + 3; double nr1 = sqrt(deepmd::dot3(r1, r1)); double nr12 = nr1 * nr1; @@ -287,23 +287,23 @@ static void compute_dRdT_2(double (*dRdT)[9], // n_sel_r_nei x 12 // (1./rr, cos_theta, cos_phi, sin_phi) x 4 x (x, y, z) + //(1./rr) x 4 x (x, y, z) -void compute_descriptor(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &descrpt_r, - std::vector &descrpt_r_deriv, - std::vector &rij_a, - std::vector &rij_r, - std::vector &rot_mat, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &fmt_nlist_r, - const std::vector &sec_a, - const std::vector &sec_r, +void compute_descriptor(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& descrpt_r, + std::vector& descrpt_r_deriv, + std::vector& rij_a, + std::vector& rij_r, + std::vector& rot_mat, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& fmt_nlist_r, + const std::vector& sec_a, + const std::vector& sec_r, const int axis0_type, const int axis0_idx, const int axis1_type, @@ -318,7 +318,7 @@ void compute_descriptor(std::vector &descrpt_a, break; } sel_a_diff[jj].resize(3); - const int &j_idx = fmt_nlist_a[jj]; + const int& j_idx = fmt_nlist_a[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -344,7 +344,7 @@ void compute_descriptor(std::vector &descrpt_a, break; } sel_r_diff[jj].resize(3); - const int &j_idx = fmt_nlist_r[jj]; + const int& j_idx = fmt_nlist_r[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -411,9 +411,9 @@ void compute_descriptor(std::vector &descrpt_a, // rotation matrix double rot[9]; - double *xx = rot; - double *yy = rot + 3; - double *zz = rot + 6; + double* xx = rot; + double* yy = rot + 3; + double* zz = rot + 6; for (unsigned dd = 0; dd < 3; ++dd) { xx[dd] = r1[dd]; yy[dd] = r2[dd]; @@ -472,7 +472,7 @@ void compute_descriptor(std::vector &descrpt_a, if (fmt_nlist_r[jj] < 0) { break; } - const double *rdiff = &sel_r_diff[jj][0]; + const double* rdiff = &sel_r_diff[jj][0]; double rr = sqrt(deepmd::dot3(rdiff, rdiff)); descrpt_r[jj] = 1. / rr; } @@ -503,7 +503,7 @@ void compute_descriptor(std::vector &descrpt_a, } // drdS, stored in transposed form double dtrdST[4][3]; - double *rr = &sel_a_diff[nei_iter][0]; + double* rr = &sel_a_diff[nei_iter][0]; double tr[3]; deepmd::dotmv3(tr, rot, rr); double nr2 = deepmd::dot3(tr, tr); @@ -638,7 +638,7 @@ void compute_descriptor(std::vector &descrpt_a, break; } - const double *rr = &sel_r_diff[nei_iter][0]; + const double* rr = &sel_r_diff[nei_iter][0]; double nr = sqrt(deepmd::dot3(rr, rr)); double nr3 = nr * nr * nr; int idx = nei_iter * 12; @@ -658,19 +658,19 @@ void compute_descriptor(std::vector &descrpt_a, } } -void compute_descriptor(std::vector &descrpt_a, - std::vector &descrpt_r, - std::vector &rot_mat, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &fmt_nlist_r, - const std::vector &sec_a, - const std::vector &sec_r, +void compute_descriptor(std::vector& descrpt_a, + std::vector& descrpt_r, + std::vector& rot_mat, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& fmt_nlist_r, + const std::vector& sec_a, + const std::vector& sec_r, const int axis0_type, const int axis0_idx, const int axis1_type, @@ -683,7 +683,7 @@ void compute_descriptor(std::vector &descrpt_a, break; } sel_a_diff[jj].resize(3); - const int &j_idx = fmt_nlist_a[jj]; + const int& j_idx = fmt_nlist_a[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -703,7 +703,7 @@ void compute_descriptor(std::vector &descrpt_a, break; } sel_r_diff[jj].resize(3); - const int &j_idx = fmt_nlist_r[jj]; + const int& j_idx = fmt_nlist_r[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -734,9 +734,9 @@ void compute_descriptor(std::vector &descrpt_a, // rotation matrix double rot[9]; - double *xx = rot; - double *yy = rot + 3; - double *zz = rot + 6; + double* xx = rot; + double* yy = rot + 3; + double* zz = rot + 6; for (unsigned dd = 0; dd < 3; ++dd) { xx[dd] = r1[dd]; yy[dd] = r2[dd]; @@ -805,21 +805,21 @@ void compute_descriptor(std::vector &descrpt_a, // output deriv size: n_sel_a_nei x 4 x 12 // (1./rr, cos_theta, cos_phi, sin_phi) x 4 x (x, y, z) -void compute_descriptor_se_a_extf(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const std::vector &efield, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &sec_a, - const double &rmin, - const double &rmax) { - const double *ef_ = &efield[i_idx * 3 + 0]; +void compute_descriptor_se_a_extf(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const std::vector& efield, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& sec_a, + const double& rmin, + const double& rmax) { + const double* ef_ = &efield[i_idx * 3 + 0]; double ef[3] = {0.}; if (std::isnan(ef_[0]) || std::isnan(ef_[1]) || std::isnan(ef_[2])) { ef[0] = 1.; @@ -842,7 +842,7 @@ void compute_descriptor_se_a_extf(std::vector &descrpt_a, break; } sel_a_diff[jj].resize(3); - const int &j_idx = fmt_nlist_a[jj]; + const int& j_idx = fmt_nlist_a[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -872,7 +872,7 @@ void compute_descriptor_se_a_extf(std::vector &descrpt_a, if (fmt_nlist_a[nei_iter] < 0) { break; } - const double *rr = &sel_a_diff[nei_iter][0]; + const double* rr = &sel_a_diff[nei_iter][0]; // check validity of ef double nr2 = deepmd::dot3(rr, rr); double inr = 1. / sqrt(nr2); @@ -946,21 +946,21 @@ void compute_descriptor_se_a_extf(std::vector &descrpt_a, // output deriv size: n_sel_a_nei x 4 x 12 // (1./rr, cos_theta, cos_phi, sin_phi) x 4 x (x, y, z) -void compute_descriptor_se_a_ef_para(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const std::vector &efield, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &sec_a, - const double &rmin, - const double &rmax) { - const double *ef_ = &efield[i_idx * 3 + 0]; +void compute_descriptor_se_a_ef_para(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const std::vector& efield, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& sec_a, + const double& rmin, + const double& rmax) { + const double* ef_ = &efield[i_idx * 3 + 0]; double ef[3] = {0.}; if (std::isnan(ef_[0]) || std::isnan(ef_[1]) || std::isnan(ef_[2])) { ef[0] = 1.; @@ -983,7 +983,7 @@ void compute_descriptor_se_a_ef_para(std::vector &descrpt_a, break; } sel_a_diff[jj].resize(3); - const int &j_idx = fmt_nlist_a[jj]; + const int& j_idx = fmt_nlist_a[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -1013,7 +1013,7 @@ void compute_descriptor_se_a_ef_para(std::vector &descrpt_a, if (fmt_nlist_a[nei_iter] < 0) { break; } - const double *rr = &sel_a_diff[nei_iter][0]; + const double* rr = &sel_a_diff[nei_iter][0]; // check validity of ef double nr2 = deepmd::dot3(rr, rr); double inr = 1. / sqrt(nr2); @@ -1083,21 +1083,21 @@ void compute_descriptor_se_a_ef_para(std::vector &descrpt_a, // output deriv size: n_sel_a_nei x 4 x 12 // (1./rr, cos_theta, cos_phi, sin_phi) x 4 x (x, y, z) -void compute_descriptor_se_a_ef_vert(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const std::vector &efield, - const int &i_idx, - const std::vector &fmt_nlist_a, - const std::vector &sec_a, - const double &rmin, - const double &rmax) { - const double *ef_ = &efield[i_idx * 3 + 0]; +void compute_descriptor_se_a_ef_vert(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const std::vector& efield, + const int& i_idx, + const std::vector& fmt_nlist_a, + const std::vector& sec_a, + const double& rmin, + const double& rmax) { + const double* ef_ = &efield[i_idx * 3 + 0]; double ef[3] = {0.}; if (std::isnan(ef_[0]) || std::isnan(ef_[1]) || std::isnan(ef_[2])) { ef[0] = 1.; @@ -1120,7 +1120,7 @@ void compute_descriptor_se_a_ef_vert(std::vector &descrpt_a, break; } sel_a_diff[jj].resize(3); - const int &j_idx = fmt_nlist_a[jj]; + const int& j_idx = fmt_nlist_a[jj]; if (b_pbc) { region.diffNearestNeighbor( posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], @@ -1150,7 +1150,7 @@ void compute_descriptor_se_a_ef_vert(std::vector &descrpt_a, if (fmt_nlist_a[nei_iter] < 0) { break; } - const double *rr = &sel_a_diff[nei_iter][0]; + const double* rr = &sel_a_diff[nei_iter][0]; // check validity of ef double nr2 = deepmd::dot3(rr, rr); double inr = 1. / sqrt(nr2); diff --git a/source/lib/include/SimulationRegion.h b/source/lib/include/SimulationRegion.h index 7cc853d25b..377a115dc0 100644 --- a/source/lib/include/SimulationRegion.h +++ b/source/lib/include/SimulationRegion.h @@ -13,82 +13,82 @@ class SimulationRegion { const static int SPACENDIM = MOASPNDIM; public: - void reinitBox(const double *boxv); - void affineTransform(const double *affine_map); - void reinitOrigin(const double *orig); - void reinitOrigin(const std::vector &orig); + void reinitBox(const double* boxv); + void affineTransform(const double* affine_map); + void reinitOrigin(const double* orig); + void reinitOrigin(const std::vector& orig); void backup(); void recover(); public: SimulationRegion(); ~SimulationRegion(); - double *getBoxTensor() { return boxt; }; - const double *getBoxTensor() const { return boxt; }; - double *getRecBoxTensor() { return rec_boxt; } - const double *getRecBoxTensor() const { return rec_boxt; } - double *getBoxOrigin() { return origin; } - const double *getBoxOrigin() const { return origin; } + double* getBoxTensor() { return boxt; }; + const double* getBoxTensor() const { return boxt; }; + double* getRecBoxTensor() { return rec_boxt; } + const double* getRecBoxTensor() const { return rec_boxt; } + double* getBoxOrigin() { return origin; } + const double* getBoxOrigin() const { return origin; } double getVolume() const { return volume; } public: - void toFaceDistance(double *dd) const; + void toFaceDistance(double* dd) const; public: - void phys2Inter(double *i_v, const VALUETYPE *p_v) const; - void inter2Phys(VALUETYPE *p_v, const double *i_v) const; + void phys2Inter(double* i_v, const VALUETYPE* p_v) const; + void inter2Phys(VALUETYPE* p_v, const double* i_v) const; public: bool isPeriodic(const int dim) const { return is_periodic[dim]; } - static int compactIndex(const int *idx); - double *getShiftVec(const int index = 0); - const double *getShiftVec(const int index = 0) const; - int getShiftIndex(const int *idx) const; + static int compactIndex(const int* idx); + double* getShiftVec(const int index = 0); + const double* getShiftVec(const int index = 0) const; + int getShiftIndex(const int* idx) const; int getNullShiftIndex() const; - void shiftCoord(const int *idx, - VALUETYPE &x, - VALUETYPE &y, - VALUETYPE &z) const; + void shiftCoord(const int* idx, + VALUETYPE& x, + VALUETYPE& y, + VALUETYPE& z) const; static int getNumbShiftVec() { return shift_info_size; } static int getShiftVecTotalSize() { return shift_vec_size; } public: - void diffNearestNeighbor(const VALUETYPE *r0, - const VALUETYPE *r1, - VALUETYPE *phys) const; + void diffNearestNeighbor(const VALUETYPE* r0, + const VALUETYPE* r1, + VALUETYPE* phys) const; virtual void diffNearestNeighbor(const VALUETYPE x0, const VALUETYPE y0, const VALUETYPE z0, const VALUETYPE x1, const VALUETYPE y1, const VALUETYPE z1, - VALUETYPE &dx, - VALUETYPE &dy, - VALUETYPE &dz) const; + VALUETYPE& dx, + VALUETYPE& dy, + VALUETYPE& dz) const; virtual void diffNearestNeighbor(const VALUETYPE x0, const VALUETYPE y0, const VALUETYPE z0, const VALUETYPE x1, const VALUETYPE y1, const VALUETYPE z1, - VALUETYPE &dx, - VALUETYPE &dy, - VALUETYPE &dz, - int &shift_x, - int &shift_y, - int &shift_z) const; + VALUETYPE& dx, + VALUETYPE& dy, + VALUETYPE& dz, + int& shift_x, + int& shift_y, + int& shift_z) const; virtual void diffNearestNeighbor(const VALUETYPE x0, const VALUETYPE y0, const VALUETYPE z0, const VALUETYPE x1, const VALUETYPE y1, const VALUETYPE z1, - VALUETYPE &dx, - VALUETYPE &dy, - VALUETYPE &dz, - VALUETYPE &shift_x, - VALUETYPE &shift_y, - VALUETYPE &shift_z) const; + VALUETYPE& dx, + VALUETYPE& dy, + VALUETYPE& dz, + VALUETYPE& shift_x, + VALUETYPE& shift_y, + VALUETYPE& shift_z) const; private: void computeVolume(); @@ -118,25 +118,25 @@ class SimulationRegion { static int index3to1(const int tx, const int ty, const int tz) { return (NBOX_ZZ * (NBOX_YY * (tx + DBOX_XX) + ty + DBOX_YY) + tz + DBOX_ZZ); } - double *getInterShiftVec(const int index = 0); - const double *getInterShiftVec(const int index = 0) const; + double* getInterShiftVec(const int index = 0); + const double* getInterShiftVec(const int index = 0) const; private: - void copy(double *o_v, const double *i_v) const; - void naiveTensorDotVector(double *out, - const double *i_t, - const double *i_v) const; - void naiveTensorTransDotVector(double *out, - const double *i_t, - const double *i_v) const; - void tensorDotVector(double *out, const double *i_t, const double *i_v) const; - void tensorTransDotVector(double *out, - const double *i_t, - const double *i_v) const; - void getFromRestart(double *my_boxv, double *my_orig, bool *period) const; - void defaultInitBox(double *my_boxv, double *my_orig, bool *period) const; - void apply_periodic(int dim, double *dd) const; - void apply_periodic(int dim, double *dd, int &shift) const; + void copy(double* o_v, const double* i_v) const; + void naiveTensorDotVector(double* out, + const double* i_t, + const double* i_v) const; + void naiveTensorTransDotVector(double* out, + const double* i_t, + const double* i_v) const; + void tensorDotVector(double* out, const double* i_t, const double* i_v) const; + void tensorTransDotVector(double* out, + const double* i_t, + const double* i_v) const; + void getFromRestart(double* my_boxv, double* my_orig, bool* period) const; + void defaultInitBox(double* my_boxv, double* my_orig, bool* period) const; + void apply_periodic(int dim, double* dd) const; + void apply_periodic(int dim, double* dd, int& shift) const; private: std::fstream fp; diff --git a/source/lib/include/SimulationRegion_Impl.h b/source/lib/include/SimulationRegion_Impl.h index cab06087e3..7b4c3dbb4d 100644 --- a/source/lib/include/SimulationRegion_Impl.h +++ b/source/lib/include/SimulationRegion_Impl.h @@ -23,9 +23,9 @@ SimulationRegion::SimulationRegion() { } template -void SimulationRegion::defaultInitBox(double *my_boxv, - double *my_orig, - bool *period) const { +void SimulationRegion::defaultInitBox(double* my_boxv, + double* my_orig, + bool* period) const { // by default is a 1,1,1 logical box for (int ii = 0; ii < SPACENDIM; ++ii) { for (int jj = 0; jj < SPACENDIM; ++jj) { @@ -55,7 +55,7 @@ void SimulationRegion::recover() { } template -inline void SimulationRegion::reinitBox(const double *boxv_) { +inline void SimulationRegion::reinitBox(const double* boxv_) { for (int ii = 0; ii < SPACENDIM * SPACENDIM; ++ii) { boxt[ii] = boxv_[ii]; } @@ -66,7 +66,7 @@ inline void SimulationRegion::reinitBox(const double *boxv_) { template inline void SimulationRegion::affineTransform( - const double *affine_map) { + const double* affine_map) { tensorDotVector(boxt + SPACENDIM * 0, affine_map, boxt + SPACENDIM * 0); tensorDotVector(boxt + SPACENDIM * 1, affine_map, boxt + SPACENDIM * 1); tensorDotVector(boxt + SPACENDIM * 2, affine_map, boxt + SPACENDIM * 2); @@ -76,7 +76,7 @@ inline void SimulationRegion::affineTransform( } template -inline void SimulationRegion::reinitOrigin(const double *orig) { +inline void SimulationRegion::reinitOrigin(const double* orig) { for (int ii = 0; ii < SPACENDIM; ++ii) { origin[ii] = orig[ii]; } @@ -84,7 +84,7 @@ inline void SimulationRegion::reinitOrigin(const double *orig) { template inline void SimulationRegion::reinitOrigin( - const std::vector &orig) { + const std::vector& orig) { for (int ii = 0; ii < SPACENDIM; ++ii) { origin[ii] = orig[ii]; } @@ -93,14 +93,14 @@ inline void SimulationRegion::reinitOrigin( template void SimulationRegion::computeShiftVec() { int tmp_idx[3]; - int &ii(tmp_idx[0]); - int &jj(tmp_idx[1]); - int &kk(tmp_idx[2]); + int& ii(tmp_idx[0]); + int& jj(tmp_idx[1]); + int& kk(tmp_idx[2]); for (ii = -DBOX_XX; ii <= DBOX_XX; ++ii) { for (jj = -DBOX_YY; jj <= DBOX_YY; ++jj) { for (kk = -DBOX_ZZ; kk <= DBOX_ZZ; ++kk) { - double *posi = getShiftVec(getShiftIndex(tmp_idx)); - double *inter_posi = getInterShiftVec(getShiftIndex(tmp_idx)); + double* posi = getShiftVec(getShiftIndex(tmp_idx)); + double* inter_posi = getInterShiftVec(getShiftIndex(tmp_idx)); inter_posi[0] = ii; inter_posi[1] = jj; inter_posi[2] = kk; @@ -112,29 +112,29 @@ void SimulationRegion::computeShiftVec() { } template -inline double *SimulationRegion::getShiftVec(const int index) { +inline double* SimulationRegion::getShiftVec(const int index) { return shift_vec + SPACENDIM * index; } template -inline const double *SimulationRegion::getShiftVec( +inline const double* SimulationRegion::getShiftVec( const int index) const { return shift_vec + SPACENDIM * index; } template -inline double *SimulationRegion::getInterShiftVec(const int index) { +inline double* SimulationRegion::getInterShiftVec(const int index) { return inter_shift_vec + SPACENDIM * index; } template -inline const double *SimulationRegion::getInterShiftVec( +inline const double* SimulationRegion::getInterShiftVec( const int index) const { return inter_shift_vec + SPACENDIM * index; } template -inline int SimulationRegion::getShiftIndex(const int *idx) const { +inline int SimulationRegion::getShiftIndex(const int* idx) const { return index3to1(idx[0], idx[1], idx[2]); } @@ -144,16 +144,16 @@ inline int SimulationRegion::getNullShiftIndex() const { } template -inline int SimulationRegion::compactIndex(const int *idx) { +inline int SimulationRegion::compactIndex(const int* idx) { return index3to1(idx[0], idx[1], idx[2]); } template -inline void SimulationRegion::shiftCoord(const int *idx, - VALUETYPE &x, - VALUETYPE &y, - VALUETYPE &z) const { - const double *shift = getShiftVec(getShiftIndex(idx)); +inline void SimulationRegion::shiftCoord(const int* idx, + VALUETYPE& x, + VALUETYPE& y, + VALUETYPE& z) const { + const double* shift = getShiftVec(getShiftIndex(idx)); x += shift[0]; y += shift[1]; z += shift[2]; @@ -199,7 +199,7 @@ inline void SimulationRegion::shiftCoord(const int *idx, template inline void SimulationRegion::apply_periodic(int dim, - double *dd) const { + double* dd) const { if (!is_periodic[dim]) { return; } @@ -212,8 +212,8 @@ inline void SimulationRegion::apply_periodic(int dim, template inline void SimulationRegion::apply_periodic(int dim, - double *dd, - int &shift) const { + double* dd, + int& shift) const { shift = 0; if (!is_periodic[dim]) { return; @@ -229,7 +229,7 @@ inline void SimulationRegion::apply_periodic(int dim, template inline void SimulationRegion::diffNearestNeighbor( - const VALUETYPE *r0, const VALUETYPE *r1, VALUETYPE *phys) const { + const VALUETYPE* r0, const VALUETYPE* r1, VALUETYPE* phys) const { double inter[3]; for (int dd = 0; dd < 3; ++dd) { phys[dd] = r0[dd] - r1[dd]; @@ -249,9 +249,9 @@ inline void SimulationRegion::diffNearestNeighbor( const VALUETYPE x1, const VALUETYPE y1, const VALUETYPE z1, - VALUETYPE &dx, - VALUETYPE &dy, - VALUETYPE &dz) const { + VALUETYPE& dx, + VALUETYPE& dy, + VALUETYPE& dz) const { // diffNearestNeighbor (0, x0, x1, dx); // diffNearestNeighbor (1, y0, y1, dy); // diffNearestNeighbor (2, z0, z1, dz); @@ -278,12 +278,12 @@ inline void SimulationRegion::diffNearestNeighbor( const VALUETYPE x1, const VALUETYPE y1, const VALUETYPE z1, - VALUETYPE &dx, - VALUETYPE &dy, - VALUETYPE &dz, - int &shift_x, - int &shift_y, - int &shift_z) const { + VALUETYPE& dx, + VALUETYPE& dy, + VALUETYPE& dz, + int& shift_x, + int& shift_y, + int& shift_z) const { // diffNearestNeighbor (0, x0, x1, dx, shift_x); // diffNearestNeighbor (1, y0, y1, dy, shift_y); // diffNearestNeighbor (2, z0, z1, dz, shift_z); @@ -310,12 +310,12 @@ inline void SimulationRegion::diffNearestNeighbor( const VALUETYPE x1, const VALUETYPE y1, const VALUETYPE z1, - VALUETYPE &dx, - VALUETYPE &dy, - VALUETYPE &dz, - VALUETYPE &shift_x, - VALUETYPE &shift_y, - VALUETYPE &shift_z) const { + VALUETYPE& dx, + VALUETYPE& dy, + VALUETYPE& dz, + VALUETYPE& shift_x, + VALUETYPE& shift_y, + VALUETYPE& shift_z) const { // diffNearestNeighbor (0, x0, x1, dx, shift_x); // diffNearestNeighbor (1, y0, y1, dy, shift_y); // diffNearestNeighbor (2, z0, z1, dz, shift_z); @@ -333,7 +333,7 @@ inline void SimulationRegion::diffNearestNeighbor( dx = phys[0]; dy = phys[1]; dz = phys[2]; - const double *tmp_shift( + const double* tmp_shift( getShiftVec(index3to1(i_shift_x, i_shift_y, i_shift_z))); shift_x = tmp_shift[0]; shift_y = tmp_shift[1]; @@ -342,7 +342,7 @@ inline void SimulationRegion::diffNearestNeighbor( template inline void SimulationRegion::phys2Inter( - double *i_v, const VALUETYPE *p_v_) const { + double* i_v, const VALUETYPE* p_v_) const { double p_v[3]; for (int dd = 0; dd < 3; ++dd) { p_v[dd] = p_v_[dd]; @@ -351,8 +351,8 @@ inline void SimulationRegion::phys2Inter( } template -inline void SimulationRegion::inter2Phys(VALUETYPE *p_v_, - const double *i_v) const { +inline void SimulationRegion::inter2Phys(VALUETYPE* p_v_, + const double* i_v) const { double p_v[3]; tensorTransDotVector(p_v, boxt, i_v); for (int dd = 0; dd < 3; ++dd) { @@ -361,7 +361,7 @@ inline void SimulationRegion::inter2Phys(VALUETYPE *p_v_, } template -inline void SimulationRegion::toFaceDistance(double *dd) const { +inline void SimulationRegion::toFaceDistance(double* dd) const { double tmp[3]; deepmd::cprod(boxt + 3, boxt + 6, tmp); dd[0] = volume * deepmd::invsqrt(deepmd::dot3(tmp, tmp)); @@ -374,8 +374,8 @@ inline void SimulationRegion::toFaceDistance(double *dd) const { // static int tmp_count = 0; template -inline void SimulationRegion::copy(double *o_v, - const double *i_v) const { +inline void SimulationRegion::copy(double* o_v, + const double* i_v) const { #ifdef DEBUG_CHECK_ASSERTIONS assert(o_v != i_v); #endif @@ -386,7 +386,7 @@ inline void SimulationRegion::copy(double *o_v, template inline void SimulationRegion::naiveTensorDotVector( - double *o_v, const double *i_t, const double *i_v) const { + double* o_v, const double* i_t, const double* i_v) const { o_v[0] = i_v[0] * i_t[0 * 3 + 0] + i_v[1] * i_t[0 * 3 + 1] + i_v[2] * i_t[0 * 3 + 2]; o_v[1] = i_v[0] * i_t[1 * 3 + 0] + i_v[1] * i_t[1 * 3 + 1] + @@ -397,7 +397,7 @@ inline void SimulationRegion::naiveTensorDotVector( template inline void SimulationRegion::naiveTensorTransDotVector( - double *o_v, const double *i_t, const double *i_v) const { + double* o_v, const double* i_t, const double* i_v) const { o_v[0] = i_v[0] * i_t[0 * 3 + 0] + i_v[1] * i_t[1 * 3 + 0] + i_v[2] * i_t[2 * 3 + 0]; o_v[1] = i_v[0] * i_t[0 * 3 + 1] + i_v[1] * i_t[1 * 3 + 1] + @@ -408,7 +408,7 @@ inline void SimulationRegion::naiveTensorTransDotVector( template inline void SimulationRegion::tensorDotVector( - double *o_v, const double *i_t, const double *i_v) const { + double* o_v, const double* i_t, const double* i_v) const { // the compiler will auto-matically optimize the following code away... // const double * tmp_v (i_v); // if (o_v == i_v){ @@ -421,7 +421,7 @@ inline void SimulationRegion::tensorDotVector( template inline void SimulationRegion::tensorTransDotVector( - double *o_v, const double *i_t, const double *i_v) const { + double* o_v, const double* i_t, const double* i_v) const { naiveTensorTransDotVector(o_v, i_t, i_v); } diff --git a/source/lib/include/env_mat_nvnmd.h b/source/lib/include/env_mat_nvnmd.h index d3c18270cf..ce391a9563 100644 --- a/source/lib/include/env_mat_nvnmd.h +++ b/source/lib/include/env_mat_nvnmd.h @@ -28,16 +28,16 @@ date: 2021-12-6 namespace deepmd { template -void env_mat_a_nvnmd_quantize_cpu(std::vector &descrpt_a, - std::vector &descrpt_a_deriv, - std::vector &rij_a, - const std::vector &posi, - const std::vector &type, - const int &i_idx, - const std::vector &fmt_nlist, - const std::vector &sec, - const float &rmin, - const float &rmax); +void env_mat_a_nvnmd_quantize_cpu(std::vector& descrpt_a, + std::vector& descrpt_a_deriv, + std::vector& rij_a, + const std::vector& posi, + const std::vector& type, + const int& i_idx, + const std::vector& fmt_nlist, + const std::vector& sec, + const float& rmin, + const float& rmax); } union U_Flt64_Int64 { @@ -59,7 +59,7 @@ union U_Flt64_Int64 { split double into sign, expo, and frac */ template // float and double -void split_flt(T x, int64_t &sign, int64_t &expo, int64_t &mant) { +void split_flt(T x, int64_t& sign, int64_t& expo, int64_t& mant) { U_Flt64_Int64 ufi; ufi.nflt = x; sign = (ufi.nint >> 63) & 0x01; @@ -71,7 +71,7 @@ void split_flt(T x, int64_t &sign, int64_t &expo, int64_t &mant) { find the max exponent for float array x */ template // float and double -void find_max_expo(int64_t &max_expo, T *x, int64_t M) { +void find_max_expo(int64_t& max_expo, T* x, int64_t M) { int ii, jj, kk; U_Flt64_Int64 ufi; int64_t expo; @@ -87,7 +87,7 @@ void find_max_expo(int64_t &max_expo, T *x, int64_t M) { find the max exponent for float array x */ template // float and double -void find_max_expo(int64_t &max_expo, T *x, int64_t N, int64_t M) { +void find_max_expo(int64_t& max_expo, T* x, int64_t N, int64_t M) { int ii, jj, kk; U_Flt64_Int64 ufi; int64_t expo; @@ -103,7 +103,7 @@ void find_max_expo(int64_t &max_expo, T *x, int64_t N, int64_t M) { dot multiply */ template // float and double -void dotmul_flt_nvnmd(T &y, T *x1, T *x2, int64_t M) { +void dotmul_flt_nvnmd(T& y, T* x1, T* x2, int64_t M) { int ii, jj, kk; U_Flt64_Int64 ufi; // @@ -146,7 +146,7 @@ void dotmul_flt_nvnmd(T &y, T *x1, T *x2, int64_t M) { multiply */ template // float and double -void mul_flt_nvnmd(T &y, T x1, T x2) { +void mul_flt_nvnmd(T& y, T x1, T x2) { U_Flt64_Int64 ufi1, ufi2, ufi3; ufi1.nflt = x1; ufi1.nint &= FLT_MASK; @@ -161,7 +161,7 @@ void mul_flt_nvnmd(T &y, T x1, T x2) { add */ template // float and double -void add_flt_nvnmd(T &y, T x1, T x2) { +void add_flt_nvnmd(T& y, T x1, T x2) { U_Flt64_Int64 ufi1, ufi2, ufi3; int64_t sign1, sign2, sign3; int64_t expo1, expo2, expo3; diff --git a/source/lib/include/gpu_cuda.h b/source/lib/include/gpu_cuda.h index 9504a95b7a..8fc7781f4c 100644 --- a/source/lib/include/gpu_cuda.h +++ b/source/lib/include/gpu_cuda.h @@ -23,7 +23,7 @@ DPAssert((res), __FILE__, __LINE__); \ } inline void DPAssert(cudaError_t code, - const char *file, + const char* file, int line, bool abort = true) { if (code != cudaSuccess) { @@ -61,21 +61,21 @@ inline void DPAssert(cudaError_t code, nborAssert((res), __FILE__, __LINE__); \ } inline void nborAssert(cudaError_t code, - const char *file, + const char* file, int line, bool abort = true) { if (code != cudaSuccess) { std::string error_msg = "DeePMD-kit: Illegal nbor list sorting: "; try { DPAssert(code, file, line, true); - } catch (deepmd::deepmd_exception_oom &e) { + } catch (deepmd::deepmd_exception_oom& e) { error_msg += e.what(); if (abort) { throw deepmd::deepmd_exception_oom(error_msg); } else { fprintf(stderr, "%s\n", error_msg.c_str()); } - } catch (deepmd::deepmd_exception &e) { + } catch (deepmd::deepmd_exception& e) { error_msg += e.what(); if (abort) { throw deepmd::deepmd_exception(error_msg); @@ -87,8 +87,8 @@ inline void nborAssert(cudaError_t code, } #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 600 -static __inline__ __device__ double atomicAdd(double *address, double val) { - unsigned long long int *address_as_ull = (unsigned long long int *)address; +static __inline__ __device__ double atomicAdd(double* address, double val) { + unsigned long long int* address_as_ull = (unsigned long long int*)address; unsigned long long int old = *address_as_ull, assumed; do { assumed = old; @@ -103,68 +103,68 @@ static __inline__ __device__ double atomicAdd(double *address, double val) { namespace deepmd { -inline void DPGetDeviceCount(int &gpu_num) { cudaGetDeviceCount(&gpu_num); } +inline void DPGetDeviceCount(int& gpu_num) { cudaGetDeviceCount(&gpu_num); } inline cudaError_t DPSetDevice(int rank) { return cudaSetDevice(rank); } template -void memcpy_host_to_device(FPTYPE *device, const std::vector &host) { +void memcpy_host_to_device(FPTYPE* device, const std::vector& host) { DPErrcheck(cudaMemcpy(device, &host[0], sizeof(FPTYPE) * host.size(), cudaMemcpyHostToDevice)); } template -void memcpy_host_to_device(FPTYPE *device, const FPTYPE *host, const int size) { +void memcpy_host_to_device(FPTYPE* device, const FPTYPE* host, const int size) { DPErrcheck( cudaMemcpy(device, host, sizeof(FPTYPE) * size, cudaMemcpyHostToDevice)); } template -void memcpy_device_to_host(const FPTYPE *device, std::vector &host) { +void memcpy_device_to_host(const FPTYPE* device, std::vector& host) { DPErrcheck(cudaMemcpy(&host[0], device, sizeof(FPTYPE) * host.size(), cudaMemcpyDeviceToHost)); } template -void memcpy_device_to_host(const FPTYPE *device, FPTYPE *host, const int size) { +void memcpy_device_to_host(const FPTYPE* device, FPTYPE* host, const int size) { DPErrcheck( cudaMemcpy(host, device, sizeof(FPTYPE) * size, cudaMemcpyDeviceToHost)); } template -void malloc_device_memory(FPTYPE *&device, const std::vector &host) { - DPErrcheck(cudaMalloc((void **)&device, sizeof(FPTYPE) * host.size())); +void malloc_device_memory(FPTYPE*& device, const std::vector& host) { + DPErrcheck(cudaMalloc((void**)&device, sizeof(FPTYPE) * host.size())); } template -void malloc_device_memory(FPTYPE *&device, const int size) { - DPErrcheck(cudaMalloc((void **)&device, sizeof(FPTYPE) * size)); +void malloc_device_memory(FPTYPE*& device, const int size) { + DPErrcheck(cudaMalloc((void**)&device, sizeof(FPTYPE) * size)); } template -void malloc_device_memory_sync(FPTYPE *&device, - const std::vector &host) { - DPErrcheck(cudaMalloc((void **)&device, sizeof(FPTYPE) * host.size())); +void malloc_device_memory_sync(FPTYPE*& device, + const std::vector& host) { + DPErrcheck(cudaMalloc((void**)&device, sizeof(FPTYPE) * host.size())); memcpy_host_to_device(device, host); } template -void malloc_device_memory_sync(FPTYPE *&device, - const FPTYPE *host, +void malloc_device_memory_sync(FPTYPE*& device, + const FPTYPE* host, const int size) { - DPErrcheck(cudaMalloc((void **)&device, sizeof(FPTYPE) * size)); + DPErrcheck(cudaMalloc((void**)&device, sizeof(FPTYPE) * size)); memcpy_host_to_device(device, host, size); } template -void delete_device_memory(FPTYPE *&device) { +void delete_device_memory(FPTYPE*& device) { if (device != NULL) { DPErrcheck(cudaFree(device)); } } template -void memset_device_memory(FPTYPE *device, const int var, const int size) { +void memset_device_memory(FPTYPE* device, const int var, const int size) { DPErrcheck(cudaMemset(device, var, sizeof(FPTYPE) * size)); } } // end of namespace deepmd diff --git a/source/lib/include/gpu_rocm.h b/source/lib/include/gpu_rocm.h index abb7ddfa62..c522c6aed4 100644 --- a/source/lib/include/gpu_rocm.h +++ b/source/lib/include/gpu_rocm.h @@ -25,7 +25,7 @@ DPAssert((res), __FILE__, __LINE__); \ } inline void DPAssert(hipError_t code, - const char *file, + const char* file, int line, bool abort = true) { if (code != hipSuccess) { @@ -46,14 +46,14 @@ inline void DPAssert(hipError_t code, nborAssert((res), __FILE__, __LINE__); \ } inline void nborAssert(hipError_t code, - const char *file, + const char* file, int line, bool abort = true) { if (code != hipSuccess) { std::string error_msg = "DeePMD-kit: Illegal nbor list sorting: "; try { DPAssert(code, file, line, true); - } catch (deepmd::deepmd_exception &e) { + } catch (deepmd::deepmd_exception& e) { error_msg += e.what(); if (abort) { throw deepmd::deepmd_exception(error_msg); @@ -65,65 +65,65 @@ inline void nborAssert(hipError_t code, } namespace deepmd { -inline void DPGetDeviceCount(int &gpu_num) { hipGetDeviceCount(&gpu_num); } +inline void DPGetDeviceCount(int& gpu_num) { hipGetDeviceCount(&gpu_num); } inline hipError_t DPSetDevice(int rank) { return hipSetDevice(rank); } template -void memcpy_host_to_device(FPTYPE *device, std::vector &host) { +void memcpy_host_to_device(FPTYPE* device, std::vector& host) { DPErrcheck(hipMemcpy(device, &host[0], sizeof(FPTYPE) * host.size(), hipMemcpyHostToDevice)); } template -void memcpy_host_to_device(FPTYPE *device, const FPTYPE *host, const int size) { +void memcpy_host_to_device(FPTYPE* device, const FPTYPE* host, const int size) { DPErrcheck( hipMemcpy(device, host, sizeof(FPTYPE) * size, hipMemcpyHostToDevice)); } template -void memcpy_device_to_host(const FPTYPE *device, std::vector &host) { +void memcpy_device_to_host(const FPTYPE* device, std::vector& host) { DPErrcheck(hipMemcpy(&host[0], device, sizeof(FPTYPE) * host.size(), hipMemcpyDeviceToHost)); } template -void memcpy_device_to_host(const FPTYPE *device, FPTYPE *host, const int size) { +void memcpy_device_to_host(const FPTYPE* device, FPTYPE* host, const int size) { DPErrcheck( hipMemcpy(host, device, sizeof(FPTYPE) * size, hipMemcpyDeviceToHost)); } template -void malloc_device_memory(FPTYPE *&device, std::vector &host) { - DPErrcheck(hipMalloc((void **)&device, sizeof(FPTYPE) * host.size())); +void malloc_device_memory(FPTYPE*& device, std::vector& host) { + DPErrcheck(hipMalloc((void**)&device, sizeof(FPTYPE) * host.size())); } template -void malloc_device_memory(FPTYPE *&device, const int size) { - DPErrcheck(hipMalloc((void **)&device, sizeof(FPTYPE) * size)); +void malloc_device_memory(FPTYPE*& device, const int size) { + DPErrcheck(hipMalloc((void**)&device, sizeof(FPTYPE) * size)); } template -void malloc_device_memory_sync(FPTYPE *&device, std::vector &host) { - DPErrcheck(hipMalloc((void **)&device, sizeof(FPTYPE) * host.size())); +void malloc_device_memory_sync(FPTYPE*& device, std::vector& host) { + DPErrcheck(hipMalloc((void**)&device, sizeof(FPTYPE) * host.size())); memcpy_host_to_device(device, host); } template -void malloc_device_memory_sync(FPTYPE *&device, - const FPTYPE *host, +void malloc_device_memory_sync(FPTYPE*& device, + const FPTYPE* host, const int size) { - DPErrcheck(hipMalloc((void **)&device, sizeof(FPTYPE) * size)); + DPErrcheck(hipMalloc((void**)&device, sizeof(FPTYPE) * size)); memcpy_host_to_device(device, host, size); } template -void delete_device_memory(FPTYPE *&device) { +void delete_device_memory(FPTYPE*& device) { if (device != NULL) { DPErrcheck(hipFree(device)); } } template -void memset_device_memory(FPTYPE *device, const int var, const int size) { +void memset_device_memory(FPTYPE* device, const int var, const int size) { DPErrcheck(hipMemset(device, var, sizeof(FPTYPE) * size)); } } // namespace deepmd diff --git a/source/lib/include/pairwise.h b/source/lib/include/pairwise.h index bbb4119e59..f711bd6f88 100644 --- a/source/lib/include/pairwise.h +++ b/source/lib/include/pairwise.h @@ -10,8 +10,8 @@ namespace deepmd { * @param[in] idxs The indexes of the fragment that each atom belongs to. -1 * will be ignored. */ -void group_atoms_cpu(std::vector> &fragments, - const std::vector &idxs); +void group_atoms_cpu(std::vector>& fragments, + const std::vector& idxs); /** * DPRc pairwise map. * @@ -30,15 +30,15 @@ void group_atoms_cpu(std::vector> &fragments, * @param[in] nloc The number of local atoms. * @param[in] nall The number of all atoms, including local and ghost atoms. */ -void dprc_pairwise_map_cpu(std::vector &forward_qm_map, - std::vector &backward_qm_map, - std::vector &forward_qmmm_map, - std::vector &backward_qmmm_map, - int &nloc_qm, - int &nloc_qmmm, - int &nall_qm, - int &nall_qmmm, - const std::vector> &fragments, +void dprc_pairwise_map_cpu(std::vector& forward_qm_map, + std::vector& backward_qm_map, + std::vector& forward_qmmm_map, + std::vector& backward_qmmm_map, + int& nloc_qm, + int& nloc_qmmm, + int& nall_qm, + int& nall_qmmm, + const std::vector>& fragments, const int nloc, const int nall); } // namespace deepmd diff --git a/source/lib/include/prod_env_mat.h b/source/lib/include/prod_env_mat.h index 60da638d68..d8ca4d1861 100644 --- a/source/lib/include/prod_env_mat.h +++ b/source/lib/include/prod_env_mat.h @@ -8,34 +8,34 @@ namespace deepmd { template -void prod_env_mat_a_cpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &inlist, +void prod_env_mat_a_cpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type = NULL); + const int* f_type = NULL); template -void prod_env_mat_r_cpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &inlist, +void prod_env_mat_r_cpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, @@ -44,49 +44,49 @@ void prod_env_mat_r_cpu(FPTYPE *em, #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM template -void prod_env_mat_a_gpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &gpu_inlist, - int *array_int, - unsigned long long *array_longlong, +void prod_env_mat_a_gpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& gpu_inlist, + int* array_int, + unsigned long long* array_longlong, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type = NULL); + const int* f_type = NULL); template -void prod_env_mat_r_gpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &gpu_inlist, - int *array_int, - unsigned long long *array_longlong, +void prod_env_mat_r_gpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& gpu_inlist, + int* array_int, + unsigned long long* array_longlong, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec); -void env_mat_nbor_update(InputNlist &inlist, - InputNlist &gpu_inlist, - int &max_nbor_size, - int *&nbor_list_dev, - const int *mesh, +void env_mat_nbor_update(InputNlist& inlist, + InputNlist& gpu_inlist, + int& max_nbor_size, + int*& nbor_list_dev, + const int* mesh, const int size); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/source/lib/include/region.cuh b/source/lib/include/region.cuh index 0feafad49e..6dc71861f1 100644 --- a/source/lib/include/region.cuh +++ b/source/lib/include/region.cuh @@ -1,9 +1,9 @@ #pragma once template -__device__ inline void tensorDotVector(FPTYPE *o_v, - const FPTYPE *i_v, - const FPTYPE *i_t) { +__device__ inline void tensorDotVector(FPTYPE* o_v, + const FPTYPE* i_v, + const FPTYPE* i_t) { o_v[0] = i_v[0] * i_t[0 * 3 + 0] + i_v[1] * i_t[0 * 3 + 1] + i_v[2] * i_t[0 * 3 + 2]; o_v[1] = i_v[0] * i_t[1 * 3 + 0] + i_v[1] * i_t[1 * 3 + 1] + @@ -12,9 +12,9 @@ __device__ inline void tensorDotVector(FPTYPE *o_v, i_v[2] * i_t[2 * 3 + 2]; } template -__device__ inline void tensorTransDotVector(FPTYPE *o_v, - const FPTYPE *i_v, - const FPTYPE *i_t) { +__device__ inline void tensorTransDotVector(FPTYPE* o_v, + const FPTYPE* i_v, + const FPTYPE* i_t) { o_v[0] = i_v[0] * i_t[0 * 3 + 0] + i_v[1] * i_t[1 * 3 + 0] + i_v[2] * i_t[2 * 3 + 0]; o_v[1] = i_v[0] * i_t[0 * 3 + 1] + i_v[1] * i_t[1 * 3 + 1] + @@ -23,19 +23,19 @@ __device__ inline void tensorTransDotVector(FPTYPE *o_v, i_v[2] * i_t[2 * 3 + 2]; } template -__device__ inline void phys2Inter(FPTYPE *inter, - const FPTYPE *phys, - const FPTYPE *rec_boxt) { +__device__ inline void phys2Inter(FPTYPE* inter, + const FPTYPE* phys, + const FPTYPE* rec_boxt) { tensorDotVector(inter, phys, rec_boxt); } template -__device__ inline void inter2Phys(FPTYPE *phys, - const FPTYPE *inter, - const FPTYPE *boxt) { +__device__ inline void inter2Phys(FPTYPE* phys, + const FPTYPE* inter, + const FPTYPE* boxt) { tensorTransDotVector(phys, inter, boxt); } template -__device__ inline FPTYPE compute_volume(const FPTYPE *boxt) { +__device__ inline FPTYPE compute_volume(const FPTYPE* boxt) { FPTYPE volume = boxt[0 * 3 + 0] * (boxt[1 * 3 + 1] * boxt[2 * 3 + 2] - boxt[2 * 3 + 1] * boxt[1 * 3 + 2]) - boxt[0 * 3 + 1] * (boxt[1 * 3 + 0] * boxt[2 * 3 + 2] - diff --git a/source/lib/include/tabulate.h b/source/lib/include/tabulate.h index 47c3062449..ab57ddf6e9 100644 --- a/source/lib/include/tabulate.h +++ b/source/lib/include/tabulate.h @@ -111,6 +111,41 @@ void tabulate_fusion_se_r_grad_grad_cpu(FPTYPE* dz_dy, const int nnei, const int last_layer_size); +template +void tabulate_fusion_se_t_tebd_cpu(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template +void tabulate_fusion_se_t_tebd_grad_cpu(FPTYPE* dy_dem_x, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template +void tabulate_fusion_se_t_tebd_grad_grad_cpu(FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM template void tabulate_fusion_se_a_gpu(FPTYPE* out, @@ -219,5 +254,40 @@ void tabulate_fusion_se_r_grad_grad_gpu(FPTYPE* dz_dy, const int nloc, const int nnei, const int last_layer_size); + +template +void tabulate_fusion_se_t_tebd_gpu(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template +void tabulate_fusion_se_t_tebd_grad_gpu(FPTYPE* dy_dem_x, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template +void tabulate_fusion_se_t_tebd_grad_grad_gpu(FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); #endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM } // namespace deepmd diff --git a/source/lib/src/fmt_nlist.cc b/source/lib/src/fmt_nlist.cc index 2bf3e78e99..3965585cf8 100644 --- a/source/lib/src/fmt_nlist.cc +++ b/source/lib/src/fmt_nlist.cc @@ -18,26 +18,26 @@ struct NeighborInfo { int index; NeighborInfo() : type(0), dist(0), index(0) {} NeighborInfo(int tt, FPTYPE dd, int ii) : type(tt), dist(dd), index(ii) {} - bool operator<(const NeighborInfo &b) const { + bool operator<(const NeighborInfo& b) const { return (type < b.type || (type == b.type && (dist < b.dist || (dist == b.dist && index < b.index)))); } }; -int format_nlist_i_fill_a(std::vector &fmt_nei_idx_a, - std::vector &fmt_nei_idx_r, - const std::vector &posi, - const int &ntypes, - const std::vector &type, - const SimulationRegion ®ion, - const bool &b_pbc, - const int &i_idx, - const std::vector &nei_idx_a, - const std::vector &nei_idx_r, - const double &rcut, - const std::vector &sec_a, - const std::vector &sec_r) { +int format_nlist_i_fill_a(std::vector& fmt_nei_idx_a, + std::vector& fmt_nei_idx_r, + const std::vector& posi, + const int& ntypes, + const std::vector& type, + const SimulationRegion& region, + const bool& b_pbc, + const int& i_idx, + const std::vector& nei_idx_a, + const std::vector& nei_idx_r, + const double& rcut, + const std::vector& sec_a, + const std::vector& sec_r) { #ifdef DEBUG assert(sec_a.size() == ntypes + 1); assert(sec_r.size() == ntypes + 1); @@ -57,7 +57,7 @@ int format_nlist_i_fill_a(std::vector &fmt_nei_idx_a, sel_nei.reserve(nei_idx_a.size() + nei_idx_r.size()); for (unsigned kk = 0; kk < nei_idx.size(); ++kk) { double diff[3]; - const int &j_idx = nei_idx[kk]; + const int& j_idx = nei_idx[kk]; if (b_pbc) { region.diffNearestNeighbor(posi[j_idx * 3 + 0], posi[j_idx * 3 + 1], posi[j_idx * 3 + 2], posi[i_idx * 3 + 0], @@ -78,7 +78,7 @@ int format_nlist_i_fill_a(std::vector &fmt_nei_idx_a, std::vector nei_iter = sec_a; int overflowed = -1; for (unsigned kk = 0; kk < sel_nei.size(); ++kk) { - const int &nei_type = sel_nei[kk].type; + const int& nei_type = sel_nei[kk].type; if (nei_iter[nei_type] >= sec_a[nei_type + 1]) { int r_idx_iter = (nei_iter[nei_type]++) - sec_a[nei_type + 1] + sec_r[nei_type]; @@ -96,13 +96,13 @@ int format_nlist_i_fill_a(std::vector &fmt_nei_idx_a, } template -int format_nlist_i_cpu(std::vector &fmt_nei_idx_a, - const std::vector &posi, - const std::vector &type, - const int &i_idx, - const std::vector &nei_idx_a, - const float &rcut, - const std::vector &sec_a) { +int format_nlist_i_cpu(std::vector& fmt_nei_idx_a, + const std::vector& posi, + const std::vector& type, + const int& i_idx, + const std::vector& nei_idx_a, + const float& rcut, + const std::vector& sec_a) { fmt_nei_idx_a.resize(sec_a.back()); fill(fmt_nei_idx_a.begin(), fmt_nei_idx_a.end(), -1); @@ -115,7 +115,7 @@ int format_nlist_i_cpu(std::vector &fmt_nei_idx_a, for (unsigned kk = 0; kk < nei_idx.size(); ++kk) { // rcut is float in this function, so float rr is enough float diff[3]; - const int &j_idx = nei_idx[kk]; + const int& j_idx = nei_idx[kk]; if (type[j_idx] < 0) { continue; } @@ -132,7 +132,7 @@ int format_nlist_i_cpu(std::vector &fmt_nei_idx_a, std::vector nei_iter = sec_a; int overflowed = -1; for (unsigned kk = 0; kk < sel_nei.size(); ++kk) { - const int &nei_type = sel_nei[kk].type; + const int& nei_type = sel_nei[kk].type; if (nei_iter[nei_type] < sec_a[nei_type + 1]) { fmt_nei_idx_a[nei_iter[nei_type]++] = sel_nei[kk].index; } else { @@ -143,10 +143,10 @@ int format_nlist_i_cpu(std::vector &fmt_nei_idx_a, } template -void deepmd::format_nlist_cpu(int *nlist, - const InputNlist &in_nlist, - const FPTYPE *coord, - const int *type, +void deepmd::format_nlist_cpu(int* nlist, + const InputNlist& in_nlist, + const FPTYPE* coord, + const int* type, const int nloc, const int nall, const float rcut, @@ -165,7 +165,7 @@ void deepmd::format_nlist_cpu(int *nlist, std::copy(in_nlist.firstneigh[ii], in_nlist.firstneigh[ii] + i_num, ilist.begin()); format_nlist_i_cpu(fmt_ilist, posi_, type_, i_idx, ilist, rcut, sec); - int *cur_nlist = nlist + i_idx * nnei; + int* cur_nlist = nlist + i_idx * nnei; if (fmt_ilist.size() != nnei) { std::cerr << "FATAL: formatted nlist of i have length " << fmt_ilist.size() << " which does not match " << nnei @@ -176,37 +176,37 @@ void deepmd::format_nlist_cpu(int *nlist, } } -template int format_nlist_i_cpu(std::vector &fmt_nei_idx_a, - const std::vector &posi, - const std::vector &type, - const int &i_idx, - const std::vector &nei_idx_a, - const float &rcut, - const std::vector &sec_a); - -template int format_nlist_i_cpu(std::vector &fmt_nei_idx_a, - const std::vector &posi, - const std::vector &type, - const int &i_idx, - const std::vector &nei_idx_a, - const float &rcut, - const std::vector &sec_a); +template int format_nlist_i_cpu(std::vector& fmt_nei_idx_a, + const std::vector& posi, + const std::vector& type, + const int& i_idx, + const std::vector& nei_idx_a, + const float& rcut, + const std::vector& sec_a); + +template int format_nlist_i_cpu(std::vector& fmt_nei_idx_a, + const std::vector& posi, + const std::vector& type, + const int& i_idx, + const std::vector& nei_idx_a, + const float& rcut, + const std::vector& sec_a); template void deepmd::format_nlist_cpu( - int *nlist, - const deepmd::InputNlist &in_nlist, - const double *coord, - const int *type, + int* nlist, + const deepmd::InputNlist& in_nlist, + const double* coord, + const int* type, const int nloc, const int nall, const float rcut, const std::vector sec); template void deepmd::format_nlist_cpu( - int *nlist, - const deepmd::InputNlist &in_nlist, - const float *coord, - const int *type, + int* nlist, + const deepmd::InputNlist& in_nlist, + const float* coord, + const int* type, const int nloc, const int nall, const float rcut, diff --git a/source/lib/src/gpu/CMakeLists.txt b/source/lib/src/gpu/CMakeLists.txt index 0d176dc320..e51181315c 100644 --- a/source/lib/src/gpu/CMakeLists.txt +++ b/source/lib/src/gpu/CMakeLists.txt @@ -1,6 +1,7 @@ if(USE_CUDA_TOOLKIT) - # required cmake version 3.23: CMAKE_CUDA_ARCHITECTURES all - cmake_minimum_required(VERSION 3.23) + # required cmake version 3.25: CMAKE_CUDA_ARCHITECTURES all, + # RUN_OUTPUT_STDOUT_VARIABLE + cmake_minimum_required(VERSION 3.25.2) # project name project(deepmd_op_cuda) set(GPU_LIB_NAME deepmd_op_cuda) @@ -25,7 +26,7 @@ if(USE_CUDA_TOOLKIT) # cub has been included in CUDA Toolkit 11, we do not need to include it any # more see https://github.com/NVIDIA/cub - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_LESS "11") + if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS "11") include(FetchContent) FetchContent_Declare( cub_download @@ -38,10 +39,25 @@ if(USE_CUDA_TOOLKIT) endif() include_directories(${CUB_SOURCE_ROOT}) endif() - if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_LESS "9") + if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS "9") message(FATAL_ERROR "CUDA version must be >= 9.0") endif() + # NVCC compilation errors with gcc-14 and c++11 Cases in other repos: + # https://gitlab.archlinux.org/archlinux/packaging/packages/cuda/-/issues/12 + if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL "14") + set_if_higher(CMAKE_CUDA_STANDARD 14) + endif() + + # CUDA 13.0+ requires C++17 + if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL "13.0") + set_if_higher(CMAKE_CUDA_STANDARD 17) + message( + STATUS + "CUDA ${CMAKE_CUDA_COMPILER_VERSION} detected, setting C++ standard to 17" + ) + endif() + message(STATUS "NVCC version is " ${CMAKE_CUDA_COMPILER_VERSION}) # arch will be configured by CMAKE_CUDA_ARCHITECTURES @@ -54,7 +70,7 @@ if(USE_CUDA_TOOLKIT) elseif(USE_ROCM_TOOLKIT) # required cmake version - cmake_minimum_required(VERSION 3.21) + cmake_minimum_required(VERSION 3.25.2) # project name project(deepmd_op_rocm) enable_language(HIP) diff --git a/source/lib/src/gpu/coord.cu b/source/lib/src/gpu/coord.cu index 52ec9ff09d..5030f67caf 100644 --- a/source/lib/src/gpu/coord.cu +++ b/source/lib/src/gpu/coord.cu @@ -2,22 +2,22 @@ #include "device.h" #include "region.cuh" -__device__ inline int collapse_index(const int *idx, const int *size) { +__device__ inline int collapse_index(const int* idx, const int* size) { return (idx[0] * size[1] + idx[1]) * size[2] + idx[2]; } __device__ inline void index_recover(const int in_idx, - const int *size, - int *idx) { + const int* size, + int* idx) { idx[2] = in_idx % size[2]; idx[1] = int(in_idx / size[2]) % size[1]; idx[0] = int(int(in_idx / size[2]) / size[1]); } -__device__ inline void idx_addshift(int *idx, const int *shift) { +__device__ inline void idx_addshift(int* idx, const int* shift) { for (int dd = 0; dd < 3; dd++) { idx[dd] += shift[dd]; } } -__device__ inline void idx_unshift(int *idx, const int *shift) { +__device__ inline void idx_unshift(int* idx, const int* shift) { for (int dd = 0; dd < 3; dd++) { idx[dd] -= shift[dd]; } @@ -42,9 +42,9 @@ __device__ inline double _fmod(double x, double y) { return fmod(x, y); } __device__ inline float _fmod(float x, float y) { return fmodf(x, y); } template -__global__ void normalize_one(FPTYPE *out_c, - const FPTYPE *boxt, - const FPTYPE *rec_boxt, +__global__ void normalize_one(FPTYPE* out_c, + const FPTYPE* boxt, + const FPTYPE* rec_boxt, const int nall) { // <<>> int idy = blockIdx.x * blockDim.x + threadIdx.x; @@ -63,14 +63,14 @@ __global__ void normalize_one(FPTYPE *out_c, } template -__global__ void _fill_idx_cellmap(int *idx_cellmap, - int *idx_cellmap_noshift, - const FPTYPE *in_c, - const FPTYPE *rec_boxt, - const int *nat_stt, - const int *nat_end, - const int *ext_stt, - const int *ext_end, +__global__ void _fill_idx_cellmap(int* idx_cellmap, + int* idx_cellmap_noshift, + const FPTYPE* in_c, + const FPTYPE* rec_boxt, + const int* nat_stt, + const int* nat_end, + const int* ext_stt, + const int* ext_end, const int nloc) { int idy = blockIdx.x * blockDim.x + threadIdx.x; int ext_ncell[3]; @@ -107,9 +107,9 @@ __global__ void _fill_idx_cellmap(int *idx_cellmap, } } -__global__ void _fill_loc_cellnum_map(int *temp_idx_order, - int *loc_cellnum_map, - const int *idx_cellmap_noshift, +__global__ void _fill_loc_cellnum_map(int* temp_idx_order, + int* loc_cellnum_map, + const int* idx_cellmap_noshift, const int nloc, const int loc_cellnum) { int idy = blockIdx.x * blockDim.x + threadIdx.x; @@ -125,15 +125,15 @@ __global__ void _fill_loc_cellnum_map(int *temp_idx_order, } } -__global__ void _fill_total_cellnum_map(int *total_cellnum_map, - int *mask_cellnum_map, - int *cell_map, - int *cell_shift_map, - const int *nat_stt, - const int *nat_end, - const int *ext_stt, - const int *ext_end, - const int *loc_cellnum_map, +__global__ void _fill_total_cellnum_map(int* total_cellnum_map, + int* mask_cellnum_map, + int* cell_map, + int* cell_shift_map, + const int* nat_stt, + const int* nat_end, + const int* ext_stt, + const int* ext_end, + const int* loc_cellnum_map, const int total_cellnum) { int idy = blockIdx.x * blockDim.x + threadIdx.x; int ext_ncell[3]; @@ -145,7 +145,7 @@ __global__ void _fill_total_cellnum_map(int *total_cellnum_map, idx_orig_shift[dd] = nat_stt[dd] - ext_stt[dd]; } if (idy < total_cellnum) { - int *shift = cell_shift_map + idy * 3; + int* shift = cell_shift_map + idy * 3; int idx[3]; index_recover(idy, ext_ncell, idx); idx_unshift(idx, idx_orig_shift); @@ -169,36 +169,36 @@ __global__ void _fill_total_cellnum_map(int *total_cellnum_map, } } -__global__ void _build_loc_clist(int *clist, - const int *idx_cellmap, - const int *idx_order, - const int *sec_num_map, +__global__ void _build_loc_clist(int* clist, + const int* idx_cellmap, + const int* idx_order, + const int* sec_num_map, const int nloc) { int idy = blockIdx.x * blockDim.x + threadIdx.x; if (idy >= nloc) { return; } int cell_idx = idx_cellmap[idy]; - int *clist_row = clist + sec_num_map[cell_idx]; + int* clist_row = clist + sec_num_map[cell_idx]; clist_row[idx_order[idy]] = idy; } template -__global__ void _copy_coord(FPTYPE *out_c, - int *out_t, - int *mapping, - const FPTYPE *in_c, - const int *in_t, - const int *cell_map, - const int *cell_shift_map, - const int *sec_loc_cellnum_map, - const int *sec_total_cellnum_map, - const int *loc_clist, +__global__ void _copy_coord(FPTYPE* out_c, + int* out_t, + int* mapping, + const FPTYPE* in_c, + const int* in_t, + const int* cell_map, + const int* cell_shift_map, + const int* sec_loc_cellnum_map, + const int* sec_total_cellnum_map, + const int* loc_clist, const int nloc, const int nall, const int total_cellnum, - const FPTYPE *boxt, - const FPTYPE *rec_boxt) { + const FPTYPE* boxt, + const FPTYPE* rec_boxt) { int idy = blockIdx.x * blockDim.x + threadIdx.x; if (idy >= nall) { return; @@ -241,26 +241,26 @@ __global__ void _copy_coord(FPTYPE *out_c, } template -void compute_int_data(int *int_data, - const FPTYPE *in_c, - const int *cell_info, - const deepmd::Region ®ion, +void compute_int_data(int* int_data, + const FPTYPE* in_c, + const int* cell_info, + const deepmd::Region& region, const int nloc, const int loc_cellnum, const int total_cellnum) { - int *idx_cellmap = int_data; - int *idx_cellmap_noshift = idx_cellmap + nloc; - int *temp_idx_order = idx_cellmap_noshift + nloc; - int *loc_cellnum_map = temp_idx_order + nloc; - int *total_cellnum_map = loc_cellnum_map + loc_cellnum; - int *mask_cellnum_map = total_cellnum_map + total_cellnum; - int *cell_map = mask_cellnum_map + total_cellnum; - int *cell_shift_map = cell_map + total_cellnum; - const int *nat_stt = cell_info; - const int *nat_end = cell_info + 3; - const int *ext_stt = cell_info + 6; - const int *ext_end = cell_info + 9; - const FPTYPE *rec_boxt = region.rec_boxt; + int* idx_cellmap = int_data; + int* idx_cellmap_noshift = idx_cellmap + nloc; + int* temp_idx_order = idx_cellmap_noshift + nloc; + int* loc_cellnum_map = temp_idx_order + nloc; + int* total_cellnum_map = loc_cellnum_map + loc_cellnum; + int* mask_cellnum_map = total_cellnum_map + total_cellnum; + int* cell_map = mask_cellnum_map + total_cellnum; + int* cell_shift_map = cell_map + total_cellnum; + const int* nat_stt = cell_info; + const int* nat_end = cell_info + 3; + const int* ext_stt = cell_info + 6; + const int* ext_end = cell_info + 9; + const FPTYPE* rec_boxt = region.rec_boxt; const int nblock_loc = (nloc + TPB - 1) / TPB; _fill_idx_cellmap<<>>(idx_cellmap, idx_cellmap_noshift, in_c, @@ -283,17 +283,17 @@ void compute_int_data(int *int_data, DPErrcheck(gpuDeviceSynchronize()); } -void build_loc_clist(int *int_data, +void build_loc_clist(int* int_data, const int nloc, const int loc_cellnum, const int total_cellnum) { const int nblock = (nloc + TPB - 1) / TPB; - const int *idx_cellmap_noshift = int_data + nloc; - const int *temp_idx_order = idx_cellmap_noshift + nloc; - const int *sec_loc_cellnum_map = temp_idx_order + nloc + loc_cellnum + + const int* idx_cellmap_noshift = int_data + nloc; + const int* temp_idx_order = idx_cellmap_noshift + nloc; + const int* sec_loc_cellnum_map = temp_idx_order + nloc + loc_cellnum + 2 * total_cellnum + total_cellnum + 3 * total_cellnum; - int *loc_clist = int_data + nloc * 3 + loc_cellnum + total_cellnum * 3 + + int* loc_clist = int_data + nloc * 3 + loc_cellnum + total_cellnum * 3 + total_cellnum * 3 + loc_cellnum + 1 + total_cellnum + 1; _build_loc_clist<<>>(loc_clist, idx_cellmap_noshift, temp_idx_order, sec_loc_cellnum_map, nloc); @@ -302,26 +302,26 @@ void build_loc_clist(int *int_data, } template -void copy_coord(FPTYPE *out_c, - int *out_t, - int *mapping, - const int *int_data, - const FPTYPE *in_c, - const int *in_t, +void copy_coord(FPTYPE* out_c, + int* out_t, + int* mapping, + const int* int_data, + const FPTYPE* in_c, + const int* in_t, const int nloc, const int nall, const int loc_cellnum, const int total_cellnum, - const deepmd::Region ®ion) { + const deepmd::Region& region) { const int nblock = (nall + TPB - 1) / TPB; - const int *cell_map = int_data + 3 * nloc + loc_cellnum + 2 * total_cellnum; - const int *cell_shift_map = cell_map + total_cellnum; - const int *sec_loc_cellnum_map = cell_shift_map + 3 * total_cellnum; - const int *sec_total_cellnum_map = sec_loc_cellnum_map + loc_cellnum + 1; - const int *loc_clist = sec_total_cellnum_map + total_cellnum + 1; + const int* cell_map = int_data + 3 * nloc + loc_cellnum + 2 * total_cellnum; + const int* cell_shift_map = cell_map + total_cellnum; + const int* sec_loc_cellnum_map = cell_shift_map + 3 * total_cellnum; + const int* sec_total_cellnum_map = sec_loc_cellnum_map + loc_cellnum + 1; + const int* loc_clist = sec_total_cellnum_map + total_cellnum + 1; - const FPTYPE *boxt = region.boxt; - const FPTYPE *rec_boxt = region.rec_boxt; + const FPTYPE* boxt = region.boxt; + const FPTYPE* rec_boxt = region.rec_boxt; _copy_coord<<>>(out_c, out_t, mapping, in_c, in_t, cell_map, cell_shift_map, sec_loc_cellnum_map, sec_total_cellnum_map, loc_clist, nloc, nall, @@ -332,13 +332,13 @@ void copy_coord(FPTYPE *out_c, namespace deepmd { template -void normalize_coord_gpu(FPTYPE *coord, +void normalize_coord_gpu(FPTYPE* coord, const int natom, - const Region ®ion) { + const Region& region) { DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); - const FPTYPE *boxt = region.boxt; - const FPTYPE *rec_boxt = region.rec_boxt; + const FPTYPE* boxt = region.boxt; + const FPTYPE* rec_boxt = region.rec_boxt; const int nblock = (natom + TPB - 1) / TPB; normalize_one<<>>(coord, boxt, rec_boxt, natom); DPErrcheck(gpuGetLastError()); @@ -349,35 +349,35 @@ void normalize_coord_gpu(FPTYPE *coord, // memory):idx_map,idx_map_noshift,temp_idx_order,loc_cellnum_map,total_cellnum_map,mask_cellnum_map, // cell_map,cell_shift_map,sec_loc_cellnum_map,sec_total_cellnum_map,loc_clist template -int copy_coord_gpu(FPTYPE *out_c, - int *out_t, - int *mapping, - int *nall, - int *int_data, - const FPTYPE *in_c, - const int *in_t, - const int &nloc, - const int &mem_nall, - const int &loc_cellnum, - const int &total_cellnum, - const int *cell_info, - const Region ®ion) { +int copy_coord_gpu(FPTYPE* out_c, + int* out_t, + int* mapping, + int* nall, + int* int_data, + const FPTYPE* in_c, + const int* in_t, + const int& nloc, + const int& mem_nall, + const int& loc_cellnum, + const int& total_cellnum, + const int* cell_info, + const Region& region) { DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); compute_int_data(int_data, in_c, cell_info, region, nloc, loc_cellnum, total_cellnum); - int *int_data_cpu = new int + int* int_data_cpu = new int [loc_cellnum + 2 * total_cellnum + loc_cellnum + 1 + total_cellnum + 1]; // loc_cellnum_map,total_cellnum_map,mask_cellnum_map,sec_loc_cellnum_map,sec_total_cellnum_map DPErrcheck(gpuMemcpy(int_data_cpu, int_data + 3 * nloc, sizeof(int) * (loc_cellnum + 2 * total_cellnum), gpuMemcpyDeviceToHost)); DPErrcheck(gpuGetLastError()); - int *loc_cellnum_map = int_data_cpu; - int *total_cellnum_map = loc_cellnum_map + loc_cellnum; - int *mask_cellnum_map = total_cellnum_map + total_cellnum; - int *sec_loc_cellnum_map = mask_cellnum_map + total_cellnum; - int *sec_total_cellnum_map = sec_loc_cellnum_map + loc_cellnum + 1; + int* loc_cellnum_map = int_data_cpu; + int* total_cellnum_map = loc_cellnum_map + loc_cellnum; + int* mask_cellnum_map = total_cellnum_map + total_cellnum; + int* sec_loc_cellnum_map = mask_cellnum_map + total_cellnum; + int* sec_total_cellnum_map = sec_loc_cellnum_map + loc_cellnum + 1; sec_loc_cellnum_map[0] = 0; sec_total_cellnum_map[0] = nloc; int max_cell = 0; @@ -412,36 +412,36 @@ int copy_coord_gpu(FPTYPE *out_c, return 0; } -template void normalize_coord_gpu(float *coord, +template void normalize_coord_gpu(float* coord, const int natom, - const Region ®ion); -template void normalize_coord_gpu(double *coord, + const Region& region); +template void normalize_coord_gpu(double* coord, const int natom, - const Region ®ion); -template int copy_coord_gpu(float *out_c, - int *out_t, - int *mapping, - int *nall, - int *int_data, - const float *in_c, - const int *in_t, - const int &nloc, - const int &mem_nall, - const int &loc_cellnum, - const int &total_cellnum, - const int *cell_info, - const Region ®ion); -template int copy_coord_gpu(double *out_c, - int *out_t, - int *mapping, - int *nall, - int *int_data, - const double *in_c, - const int *in_t, - const int &nloc, - const int &mem_nall, - const int &loc_cellnum, - const int &total_cellnum, - const int *cell_info, - const Region ®ion); + const Region& region); +template int copy_coord_gpu(float* out_c, + int* out_t, + int* mapping, + int* nall, + int* int_data, + const float* in_c, + const int* in_t, + const int& nloc, + const int& mem_nall, + const int& loc_cellnum, + const int& total_cellnum, + const int* cell_info, + const Region& region); +template int copy_coord_gpu(double* out_c, + int* out_t, + int* mapping, + int* nall, + int* int_data, + const double* in_c, + const int* in_t, + const int& nloc, + const int& mem_nall, + const int& loc_cellnum, + const int& total_cellnum, + const int* cell_info, + const Region& region); } // namespace deepmd diff --git a/source/lib/src/gpu/cudart/cudart_stub.cc b/source/lib/src/gpu/cudart/cudart_stub.cc index cfbabd6f5e..222cdeb942 100644 --- a/source/lib/src/gpu/cudart/cudart_stub.cc +++ b/source/lib/src/gpu/cudart/cudart_stub.cc @@ -16,12 +16,12 @@ static cudaError_t DP_CudartGetSymbolNotFoundError() { return cudaErrorSharedObjectSymbolNotFound; } -void *DP_cudart_dlopen(char *libname) { - static auto handle = [](std::string libname) -> void * { +void* DP_cudart_dlopen(char* libname) { + static auto handle = [](std::string libname) -> void* { #if defined(_WIN32) - void *dso_handle = LoadLibrary(libname.c_str()); + void* dso_handle = LoadLibrary(libname.c_str()); #else - void *dso_handle = dlopen(libname.c_str(), RTLD_NOW | RTLD_LOCAL); + void* dso_handle = dlopen(libname.c_str(), RTLD_NOW | RTLD_LOCAL); #endif if (!dso_handle) { std::cerr << "DeePMD-kit: Cannot find " << libname << std::endl; @@ -37,15 +37,15 @@ void *DP_cudart_dlopen(char *libname) { return handle; } -void *DP_cudart_dlsym(void *handle, const char *sym_name) { +void* DP_cudart_dlsym(void* handle, const char* sym_name) { // check if the handle is nullptr, if so, return a function that // returns cudaErrorSharedObjectSymbolNotFound if (!handle) { - return reinterpret_cast(&DP_CudartGetSymbolNotFoundError); + return reinterpret_cast(&DP_CudartGetSymbolNotFoundError); } - void *symbol = dlsym(handle, sym_name); + void* symbol = dlsym(handle, sym_name); if (!symbol) { - return reinterpret_cast(&DP_CudartGetSymbolNotFoundError); + return reinterpret_cast(&DP_CudartGetSymbolNotFoundError); } return symbol; }; diff --git a/source/lib/src/gpu/neighbor_list.cu b/source/lib/src/gpu/neighbor_list.cu index fc4e784915..70bc406f5a 100644 --- a/source/lib/src/gpu/neighbor_list.cu +++ b/source/lib/src/gpu/neighbor_list.cu @@ -28,9 +28,9 @@ struct parallel_prefix_scan_op { }; template -__global__ void parallel_prefix_scan(int *numneigh, - int *nei_order, - const int *temp_nlist, +__global__ void parallel_prefix_scan(int* numneigh, + int* nei_order, + const int* temp_nlist, const int mem_size, const int nloc, const int nall) { @@ -67,14 +67,14 @@ __global__ void parallel_prefix_scan(int *numneigh, } template -__device__ inline FPTYPE dev_dot(FPTYPE *arr1, FPTYPE *arr2) { +__device__ inline FPTYPE dev_dot(FPTYPE* arr1, FPTYPE* arr2) { return arr1[0] * arr2[0] + arr1[1] * arr2[1] + arr1[2] * arr2[2]; } template -__global__ void build_nlist(int *ilist, - int *temp_nlist, - const FPTYPE *c_cpy, +__global__ void build_nlist(int* ilist, + int* temp_nlist, + const FPTYPE* c_cpy, const FPTYPE rcut2, const int nloc, const int nall, @@ -82,12 +82,12 @@ __global__ void build_nlist(int *ilist, const unsigned int atom_idx = blockIdx.x; const unsigned int neighbor_idx = blockIdx.y * blockDim.y + threadIdx.y; if (neighbor_idx < nall) { - int *neighbor_row = temp_nlist + atom_idx * mem_size; + int* neighbor_row = temp_nlist + atom_idx * mem_size; if (neighbor_idx == atom_idx) { ilist[atom_idx] = atom_idx; } else { - const FPTYPE *ccoord = c_cpy + atom_idx * 3; - const FPTYPE *ncoord = c_cpy + neighbor_idx * 3; + const FPTYPE* ccoord = c_cpy + atom_idx * 3; + const FPTYPE* ncoord = c_cpy + neighbor_idx * 3; FPTYPE diff[3]; for (int kk = 0; kk < 3; kk++) { diff[kk] = ccoord[kk] - ncoord[kk]; @@ -100,16 +100,16 @@ __global__ void build_nlist(int *ilist, } } -__global__ void fill_nlist(int **firstneigh, - const int *temp_nlist, - const int *nei_order, +__global__ void fill_nlist(int** firstneigh, + const int* temp_nlist, + const int* nei_order, const int mem_size, const int nall) { const unsigned int atom_idx = blockIdx.x; const unsigned int neighbor_idx = blockIdx.y * blockDim.y + threadIdx.y; if (neighbor_idx < nall) { - const int *in_row = temp_nlist + atom_idx * mem_size; - int *out_row = firstneigh[atom_idx]; + const int* in_row = temp_nlist + atom_idx * mem_size; + int* out_row = firstneigh[atom_idx]; int nei = in_row[neighbor_idx]; if (nei != -1) { out_row[nei_order[atom_idx * mem_size + neighbor_idx]] = nei; @@ -117,8 +117,8 @@ __global__ void fill_nlist(int **firstneigh, } } -__global__ void map_nlist(int *nlist, - const int *nlist_map, +__global__ void map_nlist(int* nlist, + const int* nlist_map, const int nloc, const int nnei) { int atom_idx = blockIdx.x; @@ -133,11 +133,11 @@ __global__ void map_nlist(int *nlist, } } -__global__ void map_nei_info(int *nlist, - int *ntype, - bool *nmask, - const int *type, - const int *nlist_map, +__global__ void map_nei_info(int* nlist, + int* ntype, + bool* nmask, + const int* type, + const int* nlist_map, const int nloc, const int nnei, const int ntypes) { @@ -159,10 +159,10 @@ __global__ void map_nei_info(int *nlist, } } -__global__ void map_nei_info_noconvert(int *nlist, - int *ntype, - bool *nmask, - const int *type, +__global__ void map_nei_info_noconvert(int* nlist, + int* ntype, + bool* nmask, + const int* type, const int nloc, const int nnei, const int ntypes) { @@ -183,26 +183,26 @@ __global__ void map_nei_info_noconvert(int *nlist, namespace deepmd { template -int build_nlist_gpu(InputNlist &nlist, - int *max_list_size, - int *nlist_data, - const FPTYPE *c_cpy, - const int &nloc, - const int &nall, - const int &mem_size, - const float &rcut) { +int build_nlist_gpu(InputNlist& nlist, + int* max_list_size, + int* nlist_data, + const FPTYPE* c_cpy, + const int& nloc, + const int& nall, + const int& mem_size, + const float& rcut) { if (mem_size < nall) { return 1; } DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); const int nblock = (nall + TPB - 1) / TPB; - int *ilist = nlist.ilist; - int *numneigh = nlist.numneigh; - int **firstneigh = nlist.firstneigh; + int* ilist = nlist.ilist; + int* numneigh = nlist.numneigh; + int** firstneigh = nlist.firstneigh; DPErrcheck(gpuMemset(nlist_data, -1, sizeof(int) * 2 * nloc * mem_size)); - int *temp_nlist = nlist_data; // nloc*mem_size - int *nei_order = temp_nlist + nloc * mem_size; + int* temp_nlist = nlist_data; // nloc*mem_size + int* nei_order = temp_nlist + nloc * mem_size; nlist.inum = nloc; FPTYPE rcut2 = rcut * rcut; @@ -220,7 +220,7 @@ int build_nlist_gpu(InputNlist &nlist, mem_size, nall); DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); - int *numneigh_host = new int[nloc]; + int* numneigh_host = new int[nloc]; DPErrcheck(gpuMemcpy(numneigh_host, numneigh, sizeof(int) * nloc, gpuMemcpyDeviceToHost)); int max_nei = 0; @@ -234,8 +234,8 @@ int build_nlist_gpu(InputNlist &nlist, return 0; } -void use_nlist_map(int *nlist, - const int *nlist_map, +void use_nlist_map(int* nlist, + const int* nlist_map, const int nloc, const int nnei) { DPErrcheck(gpuGetLastError()); @@ -248,11 +248,11 @@ void use_nlist_map(int *nlist, DPErrcheck(gpuDeviceSynchronize()); } -void use_nei_info_gpu(int *nlist, - int *ntype, - bool *nmask, - const int *type, - const int *nlist_map, +void use_nei_info_gpu(int* nlist, + int* ntype, + bool* nmask, + const int* type, + const int* nlist_map, const int nloc, const int nnei, const int ntypes, @@ -275,25 +275,25 @@ void use_nei_info_gpu(int *nlist, DPErrcheck(gpuDeviceSynchronize()); } -template int build_nlist_gpu(InputNlist &nlist, - int *max_list_size, - int *nlist_data, - const float *c_cpy, - const int &nloc, - const int &nall, - const int &mem_size, - const float &rcut); -template int build_nlist_gpu(InputNlist &nlist, - int *max_list_size, - int *nlist_data, - const double *c_cpy, - const int &nloc, - const int &nall, - const int &mem_size, - const float &rcut); +template int build_nlist_gpu(InputNlist& nlist, + int* max_list_size, + int* nlist_data, + const float* c_cpy, + const int& nloc, + const int& nall, + const int& mem_size, + const float& rcut); +template int build_nlist_gpu(InputNlist& nlist, + int* max_list_size, + int* nlist_data, + const double* c_cpy, + const int& nloc, + const int& nall, + const int& mem_size, + const float& rcut); -__global__ void map_filter_ftype(int *ftype_out, - const int *ftype_in, +__global__ void map_filter_ftype(int* ftype_out, + const int* ftype_in, const int nloc) { int ii = blockIdx.x * blockDim.x + threadIdx.x; if (ii < nloc) { @@ -301,7 +301,7 @@ __global__ void map_filter_ftype(int *ftype_out, } } -void filter_ftype_gpu(int *ftype_out, const int *ftype_in, const int nloc) { +void filter_ftype_gpu(int* ftype_out, const int* ftype_in, const int nloc) { DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); int nblock = (nloc + TPB - 1) / TPB; diff --git a/source/lib/src/gpu/region.cu b/source/lib/src/gpu/region.cu index 849eecfc3e..45fb8a2802 100644 --- a/source/lib/src/gpu/region.cu +++ b/source/lib/src/gpu/region.cu @@ -3,30 +3,30 @@ #include "region.h" template -__global__ void _phys2Inter(FPTYPE *inter, - const FPTYPE *phys, - const FPTYPE *rec_boxt) { +__global__ void _phys2Inter(FPTYPE* inter, + const FPTYPE* phys, + const FPTYPE* rec_boxt) { phys2Inter(inter, phys, rec_boxt); } template -__global__ void _inter2Phys(FPTYPE *phys, - const FPTYPE *inter, - const FPTYPE *boxt) { +__global__ void _inter2Phys(FPTYPE* phys, + const FPTYPE* inter, + const FPTYPE* boxt) { inter2Phys(phys, inter, boxt); } template -__global__ void _compute_volume(FPTYPE *volume, const FPTYPE *boxt) { +__global__ void _compute_volume(FPTYPE* volume, const FPTYPE* boxt) { volume[0] = compute_volume(boxt); } namespace deepmd { // only for unittest template -void convert_to_inter_gpu(FPTYPE *ri, - const Region ®ion, - const FPTYPE *rp) { +void convert_to_inter_gpu(FPTYPE* ri, + const Region& region, + const FPTYPE* rp) { DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); _phys2Inter<<<1, 1>>>(ri, rp, region.rec_boxt); @@ -35,9 +35,9 @@ void convert_to_inter_gpu(FPTYPE *ri, } template -void convert_to_phys_gpu(FPTYPE *rp, - const Region ®ion, - const FPTYPE *ri) { +void convert_to_phys_gpu(FPTYPE* rp, + const Region& region, + const FPTYPE* ri) { DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); _inter2Phys<<<1, 1>>>(rp, ri, region.boxt); @@ -46,7 +46,7 @@ void convert_to_phys_gpu(FPTYPE *rp, } template -void volume_gpu(FPTYPE *volume, const Region ®ion) { +void volume_gpu(FPTYPE* volume, const Region& region) { DPErrcheck(gpuGetLastError()); DPErrcheck(gpuDeviceSynchronize()); _compute_volume<<<1, 1>>>(volume, region.boxt); @@ -54,18 +54,18 @@ void volume_gpu(FPTYPE *volume, const Region ®ion) { DPErrcheck(gpuDeviceSynchronize()); } -template void convert_to_inter_gpu(float *ri, - const Region ®ion, - const float *rp); -template void convert_to_inter_gpu(double *ri, - const Region ®ion, - const double *rp); -template void convert_to_phys_gpu(float *rp, - const Region ®ion, - const float *ri); -template void convert_to_phys_gpu(double *rp, - const Region ®ion, - const double *ri); -template void volume_gpu(float *volume, const Region ®ion); -template void volume_gpu(double *volume, const Region ®ion); +template void convert_to_inter_gpu(float* ri, + const Region& region, + const float* rp); +template void convert_to_inter_gpu(double* ri, + const Region& region, + const double* rp); +template void convert_to_phys_gpu(float* rp, + const Region& region, + const float* ri); +template void convert_to_phys_gpu(double* rp, + const Region& region, + const double* ri); +template void volume_gpu(float* volume, const Region& region); +template void volume_gpu(double* volume, const Region& region); } // namespace deepmd diff --git a/source/lib/src/gpu/tabulate.cu b/source/lib/src/gpu/tabulate.cu index e0723b81af..72c4c7a4e1 100644 --- a/source/lib/src/gpu/tabulate.cu +++ b/source/lib/src/gpu/tabulate.cu @@ -103,6 +103,38 @@ __forceinline__ __device__ void locate_xx_se_t(FPTYPE& xx, } } +// same with locate_xx_se_t +template +__forceinline__ __device__ void locate_xx_se_t_tebd(FPTYPE& xx, + int& table_idx, + const FPTYPE& lower, + const FPTYPE& upper, + const FPTYPE& min, + const FPTYPE& max, + const FPTYPE& stride0, + const FPTYPE& stride1) { + if (xx < min) { + table_idx = 0; + xx = (FPTYPE)0.; + } else if (xx < lower) { + table_idx = (int)((xx - min) / stride1); + xx -= (table_idx * stride1 + min); + } else if (xx < upper) { + int first_stride = int((lower - min) / stride1); + table_idx = first_stride + (int)((xx - lower) / stride0); + xx -= ((table_idx - first_stride) * stride0 + lower); + } else if (xx < max) { + int first_stride = + int((lower - min) / stride1) + int((upper - lower) / stride0); + table_idx = first_stride + (int)((xx - upper) / stride1); + xx -= ((table_idx - first_stride) * stride1 + upper); + } else { + table_idx = int((lower - min) / stride1) + int((upper - lower) / stride0) + + (int)((max - upper) / stride1) - 1; + xx = (FPTYPE)0.; + } +} + template __forceinline__ __device__ void locate_xx_se_r(FPTYPE& xx, int& table_idx, @@ -599,7 +631,7 @@ __global__ void tabulate_fusion_se_t_grad_grad_fifth_order_polynomial( FPTYPE sum = (FPTYPE)0.; for (int ii = 0; ii < nnei_i; ii++) { int mark_table_idx = -1; - for (int jj = 0; ii < nnei_j; jj++) { + for (int jj = 0; jj < nnei_j; jj++) { FPTYPE xx = em_x[block_idx * nnei_i * nnei_j + ii * nnei_j + jj]; FPTYPE tmp = xx; FPTYPE dz_xx = @@ -630,6 +662,194 @@ __global__ void tabulate_fusion_se_t_grad_grad_fifth_order_polynomial( dz_dy[block_idx * last_layer_size + thread_idx] = sum; } +// Apply Grid-Stride Loop +template +__global__ void tabulate_fusion_se_t_tebd_fifth_order_polynomial( + FPTYPE* out, + const FPTYPE* table, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE lower, + const FPTYPE upper, + const FPTYPE max, + const FPTYPE stride0, + const FPTYPE stride1, + const int nnei_i, + const int nnei_j, + const int last_layer_size, + const int_64 total_work) { + // NOT USED: em: (nfnl, nnei_i, nnei_j) + // em_x: (nfnl * nnei_i * nnei_j, 1) flat version of em + // total_work = nloc * nnei_i * nnei_j + // Grid-Stride Loop + for (int_64 i = (int_64)blockIdx.x * blockDim.x + threadIdx.x; i < total_work; + i += (int_64)gridDim.x * blockDim.x) { + // Decompose the 1D index 'i' to get atom and neighbor indices + const int_64 block_idx = i / (nnei_i * nnei_j); + const int_64 local_idx = i % (nnei_i * nnei_j); + const int_64 ii = local_idx / nnei_j; + const int_64 jj = local_idx % nnei_j; + + // Read the input value xx for this specific neighbor pair. + FPTYPE xx = em_x[i]; + + // Determine the table index based on the value of xx. + int table_idx = 0; + locate_xx_se_t_tebd(xx, table_idx, lower, upper, -max, max, stride0, + stride1); + + // Serially loop through the 'last_layer_size' dimension to calculate all + // features. + for (int idx = 0; idx < last_layer_size; idx++) { + FPTYPE var[6]; + load_polynomial_params(var, table, table_idx, idx, last_layer_size); + FPTYPE res = + var[0] + + (var[1] + + (var[2] + (var[3] + (var[4] + var[5] * xx) * xx) * xx) * xx) * + xx; + // Calculate the unique 1D output index for the 4D tensor (block_idx, ii, + // jj, idx). + const int_64 out_idx = + (int_64)block_idx * nnei_i * nnei_j * last_layer_size + + (int_64)ii * nnei_j * last_layer_size + (int_64)jj * last_layer_size + + idx; + // Write the result to the global output memory. + out[out_idx] = res; + } + } +} + +// Apply Grid-Stride Loop +template +__global__ void tabulate_fusion_se_t_tebd_grad_fifth_order_polynomial( + FPTYPE* dy_dem_x, + const FPTYPE* table, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dy, + const FPTYPE lower, + const FPTYPE upper, + const FPTYPE max, + const FPTYPE stride0, + const FPTYPE stride1, + const int nnei_i, + const int nnei_j, + const int last_layer_size, + const int_64 total_work) { + // total_work = nloc * nnei_i * nnei_j + // Grid-Stride Loop + for (int_64 i = (int_64)blockIdx.x * blockDim.x + threadIdx.x; i < total_work; + i += (int_64)gridDim.x * blockDim.x) { + // Decompose the 1D index 'i' to get atom and neighbor indices + const int_64 block_idx = i / (nnei_i * nnei_j); + const int_64 local_idx = i % (nnei_i * nnei_j); + const int ii = local_idx / nnei_j; + const int jj = local_idx % nnei_j; + + // Determine the table index based on the value of xx. + FPTYPE xx = em_x[i]; + int table_idx = 0; + locate_xx_se_t_tebd(xx, table_idx, lower, upper, -max, max, stride0, + stride1); + + // Accumulate the gradient contributions from all features. + FPTYPE grad_sum = 0.0; + for (int idx = 0; idx < last_layer_size; idx++) { + FPTYPE var[6]; + load_polynomial_params(var, table, table_idx, idx, last_layer_size); + + // Calculate the derivative of the polynomial with respect to xx. + FPTYPE dres_dxx = + var[1] + ((FPTYPE)2. * var[2] + + ((FPTYPE)3. * var[3] + + ((FPTYPE)4. * var[4] + (FPTYPE)5. * var[5] * xx) * xx) * + xx) * + xx; + + // Read the incoming gradient from the previous layer. + const int_64 dy_idx = + (int_64)block_idx * nnei_i * nnei_j * last_layer_size + + (int_64)ii * nnei_j * last_layer_size + (int_64)jj * last_layer_size + + idx; + FPTYPE dy_val = dy[dy_idx]; + + // Apply the chain rule: dL/dxx = sum over idx [ (dL/d_res_mm) * + // (d_res_mm/dxx) ] + grad_sum += dy_val * dres_dxx; + } + + // Write the final summed gradient to the output array. + dy_dem_x[i] = grad_sum; + } +} + +// Apply Grid-Stride Loop +template +__global__ void tabulate_fusion_se_t_tebd_grad_grad_fifth_order_polynomial( + FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dz_dy_dem_x, + const FPTYPE lower, + const FPTYPE upper, + const FPTYPE max, + const FPTYPE stride0, + const FPTYPE stride1, + const int nnei_i, + const int nnei_j, + const int last_layer_size, + const int_64 total_work) { + // total_work = nloc * nnei_i * nnei_j + // Grid-Stride Loop + for (int_64 i = (int_64)blockIdx.x * blockDim.x + threadIdx.x; i < total_work; + i += (int_64)gridDim.x * blockDim.x) { + // Decompose the 1D index 'i' to get atom and neighbor indices + const int_64 block_idx = i / (nnei_i * nnei_j); + const int_64 local_idx = i % (nnei_i * nnei_j); + const int ii = local_idx / nnei_j; + const int jj = local_idx % nnei_j; + + FPTYPE xx = em_x[i]; + + // Read the incoming gradient for xx. This value is the same for all 'idx' + // features. + const FPTYPE dz_dy_dem_x_val = dz_dy_dem_x[i]; + + // Determine the table index based on the value of xx. + int table_idx = 0; + locate_xx_se_t_tebd(xx, table_idx, lower, upper, -max, max, stride0, + stride1); + + // Serially loop through the 'last_layer_size' dimension. + for (int idx = 0; idx < last_layer_size; idx++) { + FPTYPE var[6]; + load_polynomial_params(var, table, table_idx, idx, last_layer_size); + + // Calculate the derivative of the polynomial with respect to xx. + FPTYPE dres_dxx = + var[1] + ((FPTYPE)2. * var[2] + + ((FPTYPE)3. * var[3] + + ((FPTYPE)4. * var[4] + (FPTYPE)5. * var[5] * xx) * xx) * + xx) * + xx; + + // Apply the chain rule: dz/dy_idx = (dz/dxx) * (dxx/dy_idx) + // which simplifies to dz_dy_dem_x_val * dres_dxx + FPTYPE out_grad = dz_dy_dem_x_val * dres_dxx; + + // Calculate the unique 1D output index for the 4D tensor (block_idx, ii, + // jj, idx). + const int_64 out_idx = + (int_64)block_idx * nnei_i * nnei_j * last_layer_size + + (int_64)ii * nnei_j * last_layer_size + (int_64)jj * last_layer_size + + idx; + dz_dy[out_idx] = out_grad; + } + } +} + template __global__ void tabulate_fusion_se_r_fifth_order_polynomial( FPTYPE* out, @@ -923,6 +1143,99 @@ void tabulate_fusion_se_t_grad_grad_gpu(FPTYPE* dz_dy, DPErrcheck(gpuDeviceSynchronize()); } +// SE_T_TEBD kernels - preserve full nt_i x nt_j structure unlike SE_T +template +void tabulate_fusion_se_t_tebd_gpu(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size) { + if (nloc <= 0 || nnei_i <= 0 || nnei_j <= 0) { + return; + } + const int_64 total_work = (int_64)nloc * nnei_i * nnei_j; + // Use fixed number of threads per block + const int num_threads = TPB; + // Calculate number of blocks needed + const int num_blocks = (total_work + num_threads - 1) / num_threads; + + DPErrcheck(gpuGetLastError()); + DPErrcheck(gpuDeviceSynchronize()); + tabulate_fusion_se_t_tebd_fifth_order_polynomial + <<>>(out, table, em_x, em, table_info[0], + table_info[1], table_info[2], table_info[3], + table_info[4], nnei_i, nnei_j, + last_layer_size, total_work); + DPErrcheck(gpuGetLastError()); + DPErrcheck(gpuDeviceSynchronize()); +} + +template +void tabulate_fusion_se_t_tebd_grad_gpu(FPTYPE* dy_dem_x, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size) { + if (nloc <= 0 || nnei_i <= 0 || nnei_j <= 0) { + return; + } + const int_64 total_work = (int_64)nloc * nnei_i * nnei_j; + const int num_threads = TPB; + const int num_blocks = (total_work + num_threads - 1) / num_threads; + + DPErrcheck(gpuGetLastError()); + DPErrcheck(gpuDeviceSynchronize()); + DPErrcheck(gpuMemset(dy_dem_x, 0, sizeof(FPTYPE) * total_work)); + tabulate_fusion_se_t_tebd_grad_fifth_order_polynomial + <<>>(dy_dem_x, table, em_x, em, dy, + table_info[0], table_info[1], table_info[2], + table_info[3], table_info[4], nnei_i, + nnei_j, last_layer_size, total_work); + DPErrcheck(gpuGetLastError()); + DPErrcheck(gpuDeviceSynchronize()); +} + +template +void tabulate_fusion_se_t_tebd_grad_grad_gpu(FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size) { + if (nloc <= 0 || nnei_i <= 0 || nnei_j <= 0) { + return; + } + const int_64 total_work = (int_64)nloc * nnei_i * nnei_j; + const int num_threads = TPB; + const int num_blocks = (total_work + num_threads - 1) / num_threads; + + DPErrcheck(gpuGetLastError()); + DPErrcheck(gpuDeviceSynchronize()); + DPErrcheck( + gpuMemset(dz_dy, 0, sizeof(FPTYPE) * total_work * last_layer_size)); + + tabulate_fusion_se_t_tebd_grad_grad_fifth_order_polynomial + <<>>(dz_dy, table, em_x, em, dz_dy_dem_x, + table_info[0], table_info[1], table_info[2], + table_info[3], table_info[4], nnei_i, + nnei_j, last_layer_size, total_work); + DPErrcheck(gpuGetLastError()); + DPErrcheck(gpuDeviceSynchronize()); +} + template void tabulate_fusion_se_r_gpu(FPTYPE* out, const FPTYPE* table, @@ -1181,4 +1494,73 @@ template void tabulate_fusion_se_r_grad_grad_gpu( const int nnei, const int last_layer_size); +// Template instantiations for SE_T_TEBD GPU functions +template void tabulate_fusion_se_t_tebd_gpu(float* out, + const float* table, + const float* table_info, + const float* em_x, + const float* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template void tabulate_fusion_se_t_tebd_gpu(double* out, + const double* table, + const double* table_info, + const double* em_x, + const double* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template void tabulate_fusion_se_t_tebd_grad_gpu( + float* dy_dem_x, + const float* table, + const float* table_info, + const float* em_x, + const float* em, + const float* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template void tabulate_fusion_se_t_tebd_grad_gpu( + double* dy_dem_x, + const double* table, + const double* table_info, + const double* em_x, + const double* em, + const double* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template void tabulate_fusion_se_t_tebd_grad_grad_gpu( + float* dz_dy, + const float* table, + const float* table_info, + const float* em_x, + const float* em, + const float* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + +template void tabulate_fusion_se_t_tebd_grad_grad_gpu( + double* dz_dy, + const double* table, + const double* table_info, + const double* em_x, + const double* em, + const double* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); + } // namespace deepmd diff --git a/source/lib/src/pairwise.cc b/source/lib/src/pairwise.cc index f5b21d9856..b4a68b00b7 100644 --- a/source/lib/src/pairwise.cc +++ b/source/lib/src/pairwise.cc @@ -8,7 +8,7 @@ #include "errors.h" template -std::vector sort_indexes(const std::vector &v) { +std::vector sort_indexes(const std::vector& v) { // https://stackoverflow.com/a/12399290/9567349 // by Lukasz Wiklendt under CC BY-SA 4.0 std::vector idx(v.size()); @@ -18,8 +18,8 @@ std::vector sort_indexes(const std::vector &v) { return idx; } -void deepmd::group_atoms_cpu(std::vector> &fragments, - const std::vector &idxs) { +void deepmd::group_atoms_cpu(std::vector>& fragments, + const std::vector& idxs) { int natoms = idxs.size(); // sort idxs std::vector idxs_idx = sort_indexes(idxs); @@ -41,15 +41,15 @@ void deepmd::group_atoms_cpu(std::vector> &fragments, } void deepmd::dprc_pairwise_map_cpu( - std::vector &forward_qm_map, - std::vector &backward_qm_map, - std::vector &forward_qmmm_map, - std::vector &backward_qmmm_map, - int &nloc_qm, - int &nloc_qmmm, - int &nall_qm, - int &nall_qmmm, - const std::vector> &fragments, + std::vector& forward_qm_map, + std::vector& backward_qm_map, + std::vector& forward_qmmm_map, + std::vector& backward_qmmm_map, + int& nloc_qm, + int& nloc_qmmm, + int& nall_qm, + int& nall_qmmm, + const std::vector>& fragments, const int nloc, const int nall) { int nfragments = fragments.size(); diff --git a/source/lib/src/prod_env_mat.cc b/source/lib/src/prod_env_mat.cc index 81984c78e4..302fac4bc9 100644 --- a/source/lib/src/prod_env_mat.cc +++ b/source/lib/src/prod_env_mat.cc @@ -12,22 +12,22 @@ using namespace deepmd; template -void deepmd::prod_env_mat_a_cpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &inlist, +void deepmd::prod_env_mat_a_cpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type) { + const int* f_type) { if (f_type == NULL) { f_type = type; } @@ -108,16 +108,16 @@ void deepmd::prod_env_mat_a_cpu(FPTYPE *em, } template -void deepmd::prod_env_mat_r_cpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &inlist, +void deepmd::prod_env_mat_r_cpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, @@ -191,66 +191,66 @@ void deepmd::prod_env_mat_r_cpu(FPTYPE *em, } } -template void deepmd::prod_env_mat_a_cpu(double *em, - double *em_deriv, - double *rij, - int *nlist, - const double *coord, - const int *type, - const InputNlist &inlist, +template void deepmd::prod_env_mat_a_cpu(double* em, + double* em_deriv, + double* rij, + int* nlist, + const double* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const double *avg, - const double *std, + const double* avg, + const double* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type); + const int* f_type); -template void deepmd::prod_env_mat_a_cpu(float *em, - float *em_deriv, - float *rij, - int *nlist, - const float *coord, - const int *type, - const InputNlist &inlist, +template void deepmd::prod_env_mat_a_cpu(float* em, + float* em_deriv, + float* rij, + int* nlist, + const float* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const float *avg, - const float *std, + const float* avg, + const float* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type); + const int* f_type); -template void deepmd::prod_env_mat_r_cpu(double *em, - double *em_deriv, - double *rij, - int *nlist, - const double *coord, - const int *type, - const InputNlist &inlist, +template void deepmd::prod_env_mat_r_cpu(double* em, + double* em_deriv, + double* rij, + int* nlist, + const double* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const double *avg, - const double *std, + const double* avg, + const double* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec); -template void deepmd::prod_env_mat_r_cpu(float *em, - float *em_deriv, - float *rij, - int *nlist, - const float *coord, - const int *type, - const InputNlist &inlist, +template void deepmd::prod_env_mat_r_cpu(float* em, + float* em_deriv, + float* rij, + int* nlist, + const float* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const float *avg, - const float *std, + const float* avg, + const float* std, const int nloc, const int nall, const float rcut, @@ -258,17 +258,17 @@ template void deepmd::prod_env_mat_r_cpu(float *em, const std::vector sec); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM -void deepmd::env_mat_nbor_update(InputNlist &inlist, - InputNlist &gpu_inlist, - int &max_nbor_size, - int *&nbor_list_dev, - const int *mesh, +void deepmd::env_mat_nbor_update(InputNlist& inlist, + InputNlist& gpu_inlist, + int& max_nbor_size, + int*& nbor_list_dev, + const int* mesh, const int size) { - int *mesh_host = new int[size]; + int* mesh_host = new int[size]; memcpy_device_to_host(mesh, mesh_host, size); - memcpy(&inlist.ilist, 4 + mesh_host, sizeof(int *)); - memcpy(&inlist.numneigh, 8 + mesh_host, sizeof(int *)); - memcpy(&inlist.firstneigh, 12 + mesh_host, sizeof(int **)); + memcpy(&inlist.ilist, 4 + mesh_host, sizeof(int*)); + memcpy(&inlist.numneigh, 8 + mesh_host, sizeof(int*)); + memcpy(&inlist.firstneigh, 12 + mesh_host, sizeof(int**)); const int ago = mesh_host[0]; if (ago == 0 || gpu_inlist.inum < inlist.inum) { const int inum = inlist.inum; @@ -306,7 +306,7 @@ void deepmd::env_mat_nbor_update(InputNlist &inlist, // copy nbor list from host to the device std::vector nbor_list_host(static_cast(inum) * max_nbor_size, 0); - int **_firstneigh = (int **)malloc(sizeof(int *) * inum); + int** _firstneigh = (int**)malloc(sizeof(int*) * inum); for (int ii = 0; ii < inum; ii++) { _firstneigh[ii] = nbor_list_dev + ii * max_nbor_size; for (int jj = 0; jj < inlist.numneigh[ii]; jj++) { diff --git a/source/lib/src/prod_env_mat_nvnmd.cc b/source/lib/src/prod_env_mat_nvnmd.cc index d7d98b71d5..a8bf5ce29e 100644 --- a/source/lib/src/prod_env_mat_nvnmd.cc +++ b/source/lib/src/prod_env_mat_nvnmd.cc @@ -43,22 +43,22 @@ using namespace deepmd; */ template -void deepmd::prod_env_mat_a_nvnmd_quantize_cpu(FPTYPE *em, - FPTYPE *em_deriv, - FPTYPE *rij, - int *nlist, - const FPTYPE *coord, - const int *type, - const InputNlist &inlist, +void deepmd::prod_env_mat_a_nvnmd_quantize_cpu(FPTYPE* em, + FPTYPE* em_deriv, + FPTYPE* rij, + int* nlist, + const FPTYPE* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const FPTYPE *avg, - const FPTYPE *std, + const FPTYPE* avg, + const FPTYPE* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type) { + const int* f_type) { if (f_type == NULL) { f_type = type; } @@ -143,40 +143,40 @@ void deepmd::prod_env_mat_a_nvnmd_quantize_cpu(FPTYPE *em, } template void deepmd::prod_env_mat_a_nvnmd_quantize_cpu( - double *em, - double *em_deriv, - double *rij, - int *nlist, - const double *coord, - const int *type, - const InputNlist &inlist, + double* em, + double* em_deriv, + double* rij, + int* nlist, + const double* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const double *avg, - const double *std, + const double* avg, + const double* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type); + const int* f_type); template void deepmd::prod_env_mat_a_nvnmd_quantize_cpu( - float *em, - float *em_deriv, - float *rij, - int *nlist, - const float *coord, - const int *type, - const InputNlist &inlist, + float* em, + float* em_deriv, + float* rij, + int* nlist, + const float* coord, + const int* type, + const InputNlist& inlist, const int max_nbor_size, - const float *avg, - const float *std, + const float* avg, + const float* std, const int nloc, const int nall, const float rcut, const float rcut_smth, const std::vector sec, - const int *f_type); + const int* f_type); #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM // UNDEFINE diff --git a/source/lib/src/tabulate.cc b/source/lib/src/tabulate.cc index 9352980351..e3b1b770ca 100644 --- a/source/lib/src/tabulate.cc +++ b/source/lib/src/tabulate.cc @@ -541,6 +541,156 @@ void deepmd::tabulate_fusion_se_t_grad_grad_cpu(FPTYPE* dz_dy, } } +template +void deepmd::tabulate_fusion_se_t_tebd_cpu(FPTYPE* out, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size) { + memset(out, 0, sizeof(FPTYPE) * nloc * nnei_i * nnei_j * last_layer_size); + const FPTYPE lower = table_info[0]; + const FPTYPE upper = table_info[1]; + const FPTYPE _max = table_info[2]; + const FPTYPE stride0 = table_info[3]; + const FPTYPE stride1 = table_info[4]; + +#pragma omp parallel for + for (int ii = 0; ii < nloc; ii++) { + for (int jj = 0; jj < nnei_i; jj++) { + for (int kk = 0; kk < nnei_j; kk++) { + FPTYPE xx = em_x[ii * nnei_i * nnei_j + jj * nnei_j + kk]; + int table_idx = 0; + locate_xx_se_t(lower, upper, -_max, _max, stride0, stride1, xx, + table_idx); + + // For SE_TEBD, we preserve the full nt_i x nt_j x ng structure + // instead of reducing it like SE_T does + for (int mm = 0; mm < last_layer_size; mm++) { + FPTYPE a0 = table[table_idx * last_layer_size * 6 + 6 * mm + 0]; + FPTYPE a1 = table[table_idx * last_layer_size * 6 + 6 * mm + 1]; + FPTYPE a2 = table[table_idx * last_layer_size * 6 + 6 * mm + 2]; + FPTYPE a3 = table[table_idx * last_layer_size * 6 + 6 * mm + 3]; + FPTYPE a4 = table[table_idx * last_layer_size * 6 + 6 * mm + 4]; + FPTYPE a5 = table[table_idx * last_layer_size * 6 + 6 * mm + 5]; + + FPTYPE res = a0 + a1 * xx + a2 * xx * xx + a3 * xx * xx * xx + + a4 * xx * xx * xx * xx + a5 * xx * xx * xx * xx * xx; + + // Store result preserving the nt_i x nt_j structure + out[ii * nnei_i * nnei_j * last_layer_size + + jj * nnei_j * last_layer_size + kk * last_layer_size + mm] = res; + } + } + } + } +} + +template +void deepmd::tabulate_fusion_se_t_tebd_grad_cpu(FPTYPE* dy_dem_x, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size) { + memset(dy_dem_x, 0, sizeof(FPTYPE) * nloc * nnei_i * nnei_j); + const FPTYPE lower = table_info[0]; + const FPTYPE upper = table_info[1]; + const FPTYPE _max = table_info[2]; + const FPTYPE stride0 = table_info[3]; + const FPTYPE stride1 = table_info[4]; + +#pragma omp parallel for + for (int ii = 0; ii < nloc; ii++) { + for (int jj = 0; jj < nnei_i; jj++) { + for (int kk = 0; kk < nnei_j; kk++) { + FPTYPE xx = em_x[ii * nnei_i * nnei_j + jj * nnei_j + kk]; + int table_idx = 0; + locate_xx_se_t(lower, upper, -_max, _max, stride0, stride1, xx, + table_idx); + + FPTYPE grad_sum = 0.0; + for (int mm = 0; mm < last_layer_size; mm++) { + FPTYPE a1 = table[table_idx * last_layer_size * 6 + 6 * mm + 1]; + FPTYPE a2 = table[table_idx * last_layer_size * 6 + 6 * mm + 2]; + FPTYPE a3 = table[table_idx * last_layer_size * 6 + 6 * mm + 3]; + FPTYPE a4 = table[table_idx * last_layer_size * 6 + 6 * mm + 4]; + FPTYPE a5 = table[table_idx * last_layer_size * 6 + 6 * mm + 5]; + + FPTYPE dres_dxx = a1 + 2.0 * a2 * xx + 3.0 * a3 * xx * xx + + 4.0 * a4 * xx * xx * xx + + 5.0 * a5 * xx * xx * xx * xx; + + FPTYPE dy_val = + dy[ii * nnei_i * nnei_j * last_layer_size + + jj * nnei_j * last_layer_size + kk * last_layer_size + mm]; + grad_sum += dy_val * dres_dxx; + } + + dy_dem_x[ii * nnei_i * nnei_j + jj * nnei_j + kk] = grad_sum; + } + } + } +} + +template +void deepmd::tabulate_fusion_se_t_tebd_grad_grad_cpu( + FPTYPE* dz_dy, + const FPTYPE* table, + const FPTYPE* table_info, + const FPTYPE* em_x, + const FPTYPE* em, + const FPTYPE* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size) { + memset(dz_dy, 0, sizeof(FPTYPE) * nloc * nnei_i * nnei_j * last_layer_size); + const FPTYPE lower = table_info[0]; + const FPTYPE upper = table_info[1]; + const FPTYPE _max = table_info[2]; + const FPTYPE stride0 = table_info[3]; + const FPTYPE stride1 = table_info[4]; + +#pragma omp parallel for + for (int ii = 0; ii < nloc; ii++) { + for (int jj = 0; jj < nnei_i; jj++) { + for (int kk = 0; kk < nnei_j; kk++) { + FPTYPE xx = em_x[ii * nnei_i * nnei_j + jj * nnei_j + kk]; + int table_idx = 0; + locate_xx_se_t(lower, upper, -_max, _max, stride0, stride1, xx, + table_idx); + + FPTYPE dz_dy_dem_x_val = + dz_dy_dem_x[ii * nnei_i * nnei_j + jj * nnei_j + kk]; + + for (int mm = 0; mm < last_layer_size; mm++) { + FPTYPE a1 = table[table_idx * last_layer_size * 6 + 6 * mm + 1]; + FPTYPE a2 = table[table_idx * last_layer_size * 6 + 6 * mm + 2]; + FPTYPE a3 = table[table_idx * last_layer_size * 6 + 6 * mm + 3]; + FPTYPE a4 = table[table_idx * last_layer_size * 6 + 6 * mm + 4]; + FPTYPE a5 = table[table_idx * last_layer_size * 6 + 6 * mm + 5]; + + FPTYPE dres_dxx = a1 + 2.0 * a2 * xx + 3.0 * a3 * xx * xx + + 4.0 * a4 * xx * xx * xx + + 5.0 * a5 * xx * xx * xx * xx; + + dz_dy[ii * nnei_i * nnei_j * last_layer_size + + jj * nnei_j * last_layer_size + kk * last_layer_size + mm] = + dz_dy_dem_x_val * dres_dxx; + } + } + } + } +} + template void deepmd::tabulate_fusion_se_r_cpu(FPTYPE* out, const FPTYPE* table, @@ -810,7 +960,70 @@ template void deepmd::tabulate_fusion_se_t_grad_grad_cpu( const int nnei_i, const int nnei_j, const int last_layer_size); - +template void deepmd::tabulate_fusion_se_t_tebd_cpu( + float* out, + const float* table, + const float* table_info, + const float* em_x, + const float* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); +template void deepmd::tabulate_fusion_se_t_tebd_cpu( + double* out, + const double* table, + const double* table_info, + const double* em_x, + const double* em, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); +template void deepmd::tabulate_fusion_se_t_tebd_grad_cpu( + float* dy_dem_x, + const float* table, + const float* table_info, + const float* em_x, + const float* em, + const float* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); +template void deepmd::tabulate_fusion_se_t_tebd_grad_cpu( + double* dy_dem_x, + const double* table, + const double* table_info, + const double* em_x, + const double* em, + const double* dy, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); +template void deepmd::tabulate_fusion_se_t_tebd_grad_grad_cpu( + float* dz_dy, + const float* table, + const float* table_info, + const float* em_x, + const float* em, + const float* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); +template void deepmd::tabulate_fusion_se_t_tebd_grad_grad_cpu( + double* dz_dy, + const double* table, + const double* table_info, + const double* em_x, + const double* em, + const double* dz_dy_dem_x, + const int nloc, + const int nnei_i, + const int nnei_j, + const int last_layer_size); template void deepmd::tabulate_fusion_se_r_cpu( float* out, const float* table, diff --git a/source/lib/tests/CMakeLists.txt b/source/lib/tests/CMakeLists.txt index f2fd8969d9..b8fbdcc00e 100644 --- a/source/lib/tests/CMakeLists.txt +++ b/source/lib/tests/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 3.9) +cmake_minimum_required(VERSION 3.25.2) project(libdeepmd_test) file(GLOB TEST_SRC test_*.cc) diff --git a/source/lib/tests/test_env_mat_a.cc b/source/lib/tests/test_env_mat_a.cc index d041d1a0a1..3c309ca9ae 100644 --- a/source/lib/tests/test_env_mat_a.cc +++ b/source/lib/tests/test_env_mat_a.cc @@ -500,7 +500,7 @@ TEST_F(TestEnvMatA, prod_cpu) { } } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); deepmd::convert_nlist(inlist, nlist_a_cpy); @@ -536,7 +536,7 @@ TEST_F(TestEnvMatA, prod_cpu_equal_cpu) { } } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); convert_nlist(inlist, nlist_a_cpy); std::vector em(static_cast(nloc) * ndescrpt), @@ -612,7 +612,7 @@ TEST_F(TestEnvMatA, prod_gpu) { max_nbor_size = 4096; } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]), gpu_inlist; convert_nlist(inlist, nlist_a_cpy); @@ -626,7 +626,7 @@ TEST_F(TestEnvMatA, prod_gpu) { double *posi_cpy_dev = NULL, *avg_dev = NULL, *std_dev = NULL; int *atype_cpy_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL, *memory_dev = NULL; - uint_64 *array_longlong_dev = NULL; + uint_64* array_longlong_dev = NULL; deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(em_deriv_dev, em_deriv); deepmd::malloc_device_memory_sync(rij_dev, rij); @@ -690,7 +690,7 @@ TEST_F(TestEnvMatA, prod_gpu_equal_cpu) { max_nbor_size = 4096; } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]), gpu_inlist; convert_nlist(inlist, nlist_a_cpy); @@ -704,7 +704,7 @@ TEST_F(TestEnvMatA, prod_gpu_equal_cpu) { double *posi_cpy_dev = NULL, *avg_dev = NULL, *std_dev = NULL; int *atype_cpy_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL, *memory_dev = NULL; - uint_64 *array_longlong_dev = NULL; + uint_64* array_longlong_dev = NULL; deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(em_deriv_dev, em_deriv); deepmd::malloc_device_memory_sync(rij_dev, rij); diff --git a/source/lib/tests/test_env_mat_a_mix.cc b/source/lib/tests/test_env_mat_a_mix.cc index d7e6cc88eb..e96311dafd 100644 --- a/source/lib/tests/test_env_mat_a_mix.cc +++ b/source/lib/tests/test_env_mat_a_mix.cc @@ -528,7 +528,7 @@ TEST_F(TestEnvMatAMix, prod_cpu) { } } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); deepmd::convert_nlist(inlist, nlist_a_cpy); @@ -537,7 +537,7 @@ TEST_F(TestEnvMatAMix, prod_cpu) { rij(static_cast(nloc) * nnei * 3); std::vector nlist(static_cast(nloc) * nnei); std::vector ntype(static_cast(nloc) * nnei); - bool *nmask = new bool[static_cast(nloc) * nnei]; + bool* nmask = new bool[static_cast(nloc) * nnei]; memset(nmask, 0, sizeof(bool) * nloc * nnei); std::vector avg(ntypes * ndescrpt, 0); std::vector std(ntypes * ndescrpt, 1); @@ -573,7 +573,7 @@ TEST_F(TestEnvMatAMix, prod_cpu_equal_cpu) { } } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); convert_nlist(inlist, nlist_a_cpy); std::vector em(static_cast(nloc) * ndescrpt), @@ -650,7 +650,7 @@ TEST_F(TestEnvMatAMix, prod_gpu) { max_nbor_size = 4096; } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]), gpu_inlist; convert_nlist(inlist, nlist_a_cpy); @@ -659,18 +659,18 @@ TEST_F(TestEnvMatAMix, prod_gpu) { rij(static_cast(nloc) * nnei * 3, 0.0); std::vector nlist(static_cast(nloc) * nnei, 0); std::vector ntype(static_cast(nloc) * nnei, 0); - bool *nmask = new bool[static_cast(nloc) * nnei]; + bool* nmask = new bool[static_cast(nloc) * nnei]; memset(nmask, 0, sizeof(bool) * nloc * nnei); std::vector avg(ntypes * ndescrpt, 0); std::vector std(ntypes * ndescrpt, 1); double *em_dev = NULL, *em_deriv_dev = NULL, *rij_dev = NULL; - bool *nmask_dev = NULL; + bool* nmask_dev = NULL; double *posi_cpy_dev = NULL, *avg_dev = NULL, *std_dev = NULL; int *f_atype_cpy_dev = NULL, *atype_dev = NULL, *nlist_dev = NULL, *ntype_dev = NULL, *mapping_dev = NULL, *array_int_dev = NULL, *memory_dev = NULL; - uint_64 *array_longlong_dev = NULL; + uint_64* array_longlong_dev = NULL; deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(em_deriv_dev, em_deriv); deepmd::malloc_device_memory_sync(rij_dev, rij); @@ -751,7 +751,7 @@ TEST_F(TestEnvMatAMix, prod_gpu_equal_cpu) { max_nbor_size = 4096; } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]), gpu_inlist; convert_nlist(inlist, nlist_a_cpy); @@ -765,7 +765,7 @@ TEST_F(TestEnvMatAMix, prod_gpu_equal_cpu) { double *posi_cpy_dev = NULL, *avg_dev = NULL, *std_dev = NULL; int *f_atype_cpy_dev = NULL, *atype_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL, *memory_dev = NULL; - uint_64 *array_longlong_dev = NULL; + uint_64* array_longlong_dev = NULL; deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(em_deriv_dev, em_deriv); deepmd::malloc_device_memory_sync(rij_dev, rij); diff --git a/source/lib/tests/test_env_mat_r.cc b/source/lib/tests/test_env_mat_r.cc index 3024e651d9..96da7e6963 100644 --- a/source/lib/tests/test_env_mat_r.cc +++ b/source/lib/tests/test_env_mat_r.cc @@ -278,7 +278,7 @@ TEST_F(TestEnvMatR, prod_cpu) { } } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); convert_nlist(inlist, nlist_a_cpy); @@ -313,7 +313,7 @@ TEST_F(TestEnvMatR, prod_cpu_equal_cpu) { } } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]); convert_nlist(inlist, nlist_a_cpy); std::vector em(nloc * ndescrpt), em_deriv(nloc * ndescrpt * 3), @@ -378,7 +378,7 @@ TEST_F(TestEnvMatR, prod_gpu) { max_nbor_size = 4096; } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]), gpu_inlist; convert_nlist(inlist, nlist_a_cpy); @@ -392,7 +392,7 @@ TEST_F(TestEnvMatR, prod_gpu) { double *posi_cpy_dev = NULL, *avg_dev = NULL, *std_dev = NULL; int *atype_cpy_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL, *memory_dev = NULL; - uint_64 *array_longlong_dev = NULL; + uint_64* array_longlong_dev = NULL; deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(em_deriv_dev, em_deriv); deepmd::malloc_device_memory_sync(rij_dev, rij); @@ -457,7 +457,7 @@ TEST_F(TestEnvMatR, prod_gpu_equal_cpu) { max_nbor_size = 4096; } std::vector ilist(nloc), numneigh(nloc); - std::vector firstneigh(nloc); + std::vector firstneigh(nloc); deepmd::InputNlist inlist(nloc, &ilist[0], &numneigh[0], &firstneigh[0]), gpu_inlist; convert_nlist(inlist, nlist_a_cpy); @@ -471,7 +471,7 @@ TEST_F(TestEnvMatR, prod_gpu_equal_cpu) { double *posi_cpy_dev = NULL, *avg_dev = NULL, *std_dev = NULL; int *atype_cpy_dev = NULL, *nlist_dev = NULL, *array_int_dev = NULL, *memory_dev = NULL; - uint_64 *array_longlong_dev = NULL; + uint_64* array_longlong_dev = NULL; deepmd::malloc_device_memory_sync(em_dev, em); deepmd::malloc_device_memory_sync(em_deriv_dev, em_deriv); deepmd::malloc_device_memory_sync(rij_dev, rij); diff --git a/source/lib/tests/test_main.cc b/source/lib/tests/test_main.cc index df7815b694..2ce083b175 100644 --- a/source/lib/tests/test_main.cc +++ b/source/lib/tests/test_main.cc @@ -1,7 +1,7 @@ // SPDX-License-Identifier: LGPL-3.0-or-later #include -int main(int argc, char **argv) { +int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); return RUN_ALL_TESTS(); } diff --git a/source/lib/tests/test_tabulate_se_a.cc b/source/lib/tests/test_tabulate_se_a.cc index ce2defb22c..66a77f41fd 100644 --- a/source/lib/tests/test_tabulate_se_a.cc +++ b/source/lib/tests/test_tabulate_se_a.cc @@ -777,7 +777,7 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_gpu) { EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter[jj]), 1e-5); } - double *two_embed_dev = nullptr; + double* two_embed_dev = nullptr; deepmd::malloc_device_memory_sync(two_embed_dev, two_embed); deepmd::malloc_device_memory_sync(xyz_scatter_dev, xyz_scatter); deepmd::tabulate_fusion_se_a_gpu(xyz_scatter_dev, table_dev, &info[0], @@ -831,7 +831,7 @@ TEST_F(TestTabulateSeA, tabulate_fusion_se_a_grad_gpu) { EXPECT_LT(fabs(dy_dem[jj] - expected_dy_dem[jj]), 1e-5); } - double *two_embed_dev = nullptr; + double* two_embed_dev = nullptr; deepmd::malloc_device_memory_sync(two_embed_dev, two_embed); deepmd::malloc_device_memory_sync(dy_dem_x_dev, dy_dem_x); deepmd::malloc_device_memory_sync(dy_dem_dev, dy_dem); diff --git a/source/lib/tests/test_tabulate_se_t_tebd.cc b/source/lib/tests/test_tabulate_se_t_tebd.cc new file mode 100644 index 0000000000..fe7f314119 --- /dev/null +++ b/source/lib/tests/test_tabulate_se_t_tebd.cc @@ -0,0 +1,767 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later +#include + +#include +#include + +#include "device.h" +#include "tabulate.h" +#include "utilities.h" + +class TestTabulateSeTTebd : public ::testing::Test { + protected: + std::vector table = { + -1.0600000163027882e02, 7.7059358807135015e02, -5.6954714749735385e03, + 1.2167808756610991e03, -7.6199102434332218e01, 1.0706136029373441e00, + -1.0600000164528124e02, 7.7059358630452323e02, -5.6954715659539552e03, + 1.2167808757436076e03, -7.6199099707724926e01, 1.0706134206080884e00, + -1.0600000163027882e02, 7.7059358807135015e02, -5.6954714749735385e03, + 1.2167808756610991e03, -7.6199102434332218e01, 1.0706136029373441e00, + -1.0600000164528124e02, 7.7059358630452323e02, -5.6954715659539552e03, + 1.2167808757436076e03, -7.6199099707724926e01, 1.0706134206080884e00, + -9.6000006759336443e01, 6.2969719646863621e02, -4.2053706363664551e03, + 9.0372155784831205e02, -5.7600014239472898e01, 8.6528676197113796e-01, + -9.6000006828502180e01, 6.2969718981238339e02, -4.2053709121998018e03, + 9.0372156236848912e02, -5.7600006817493266e01, 8.6528625106787871e-01, + -9.6000006759336443e01, 6.2969719646863621e02, -4.2053706363664551e03, + 9.0372155784831205e02, -5.7600014239472898e01, 8.6528676197113796e-01, + -9.6000006828502180e01, 6.2969718981238339e02, -4.2053709121998018e03, + 9.0372156236848912e02, -5.7600006817493266e01, 8.6528625106787871e-01, + -8.6000028021606425e01, 5.0303296429845562e02, -3.0008648248894533e03, + 6.4939597734382562e02, -4.2250984019314707e01, 6.8180015607155764e-01, + -8.6000028340480625e01, 5.0303293978396903e02, -3.0008656209622986e03, + 6.4939600529391078e02, -4.2250965541906716e01, 6.8179882734268982e-01, + -8.6000028021606425e01, 5.0303296429845562e02, -3.0008648248894533e03, + 6.4939597734382562e02, -4.2250984019314707e01, 6.8180015607155764e-01, + -8.6000028340480625e01, 5.0303293978396903e02, -3.0008656209622986e03, + 6.4939600529353049e02, -4.2250965541830588e01, 6.8179882733888086e-01, + -7.6000116148038558e01, 3.9060139597613619e02, -2.0515743554479322e03, + 4.4772754091167945e02, -2.9848087537832814e01, 5.2014755686537917e-01, + -7.6000117618125429e01, 3.9060130821883052e02, -2.0515765138621105e03, + 4.4772766653712006e02, -2.9848047259266409e01, 5.2014443989116910e-01, + -7.6000116148038558e01, 3.9060139597613619e02, -2.0515743554479322e03, + 4.4772754091167945e02, -2.9848087537832814e01, 5.2014755686537917e-01, + -7.6000117618125742e01, 3.9060130821877993e02, -2.0515765138659344e03, + 4.4772766652483722e02, -2.9848047256692499e01, 5.2014443976043645e-01, + -6.6000481290731443e01, 2.9240425245900917e02, -1.3271250821434478e03, + 2.9263955624337893e02, -2.0087224005740719e01, 3.8031147992206349e-01, + -6.6000488067863742e01, 2.9240394960550276e02, -1.3271304743966571e03, + 2.9264002765325057e02, -2.0087154325946980e01, 3.8030522013794582e-01, + -6.6000481290731443e01, 2.9240425245900917e02, -1.3271250821434478e03, + 2.9263955624337893e02, -2.0087224005740719e01, 3.8031147992206349e-01, + -6.6000488067883694e01, 2.9240394960308691e02, -1.3271304745319526e03, + 2.9264002727267626e02, -2.0087154245656002e01, 3.8030521605011575e-01, + -5.6001992867343972e01, 2.0844745574402617e02, -7.9715799906587699e02, + 1.7805563184427194e02, -1.2663929104029080e01, 2.6224978307822894e-01, + -5.6002024103130161e01, 2.0844646075692629e02, -7.9717003898786652e02, + 1.7805715054974732e02, -1.2663864677938077e01, 2.6224029170957303e-01, + -5.6001992867343972e01, 2.0844745574402617e02, -7.9715799906587699e02, + 1.7805563184427194e02, -1.2663929104029080e01, 2.6224978307822894e-01, + -5.6002024104383771e01, 2.0844646064871867e02, -7.9717004324410516e02, + 1.7805714044473001e02, -1.2663862524337585e01, 2.6224018166598279e-01, + -4.6008230210744550e01, 1.3874976550319553e02, -4.3134867537287749e02, + 9.7902623595157010e01, -7.2734403121911884e00, 1.6589123996688057e-01, + -4.6008373996710617e01, 1.3874671965012058e02, -4.3137141216256458e02, + 9.7906861443792735e01, -7.2735856084076280e00, 1.6588642735924275e-01, + -4.6008230210744550e01, 1.3874976550319553e02, -4.3134867537287749e02, + 9.7902623595157010e01, -7.2734403121911884e00, 1.6589123996688057e-01, + -4.6008374075307870e01, 1.3874671513440606e02, -4.3137152784492957e02, + 9.7906652364871050e01, -7.2735401377994249e00, 1.6588408717348646e-01, + -3.6033642533368131e01, 8.3364086172019398e01, -1.9942175516407502e02, + 4.6124022747838069e01, -3.6130563858549958e00, 9.1249773312287188e-02, + -3.6034298111245583e01, 8.3355843868269616e01, -1.9945266030093268e02, + 4.6135000705962462e01, -3.6142786797647353e00, 9.1293932043118198e-02, + -3.6033642533368131e01, 8.3364086172019398e01, -1.9942175516407502e02, + 4.6124022747838069e01, -3.6130563858549958e00, 9.1249773312287188e-02, + -3.6034302998781108e01, 8.3355675173745269e01, -1.9945516784358935e02, + 4.6132303200740992e01, -3.6136582565667807e00, 9.1261386291659793e-02, + -2.6132076703837274e01, 4.2398929436319683e01, -7.1037171119057973e01, + 1.3425662262407457e01, -7.5172495708992593e-01, 7.7522572203268742e-03, + -2.6134776894873077e01, 4.2384732735328775e01, -7.1030526549717337e01, + 1.3431455085299461e01, -7.5302028721199155e-01, 7.8186246126207160e-03, + -2.6132076703837274e01, 4.2398929436319683e01, -7.1037171119057973e01, + 1.3425662262405055e01, -7.5172495708944420e-01, 7.7522572203027138e-03, + -2.6135071381093578e01, 4.2379566840123424e01, -7.1067162844830236e01, + 1.3434603316099608e01, -7.5251233833488806e-01, 7.7734884077347950e-03, + -2.2221480705551805e01, 3.0067218434037404e01, -4.1779705297521097e01, + -1.9077757705724110e02, 3.6413466026808294e02, -1.6067397401486718e02, + -2.2225430071703467e01, 3.0060809113889512e01, -4.1712800191721314e01, + -1.9084786311022177e02, 3.6410062714257685e02, -1.6063028238785057e02, + -2.2221480705551830e01, 3.0067218434036263e01, -4.1779705297545611e01, + -1.9077757705723738e02, 3.6413466026815809e02, -1.6067397401492047e02, + -2.2226913938674084e01, 3.0042371820589185e01, -4.1801582285426832e01, + -1.9048619249019526e02, 3.6373874557858261e02, -1.6052358406417352e02, + -2.1250858373060836e01, 2.7343847665267702e01, -3.6044215009418814e01, + -1.7618484800469861e02, 3.3120085405644409e02, -1.4534825256321494e02, + -2.1254939505030809e01, 2.7342716030835884e01, -3.5955450545431681e01, + -1.7635550119316844e02, 3.3127447930769307e02, -1.4533876561022046e02, + -2.1250858373060954e01, 2.7343847665262818e01, -3.6044215009514119e01, + -1.7618484800464822e02, 3.3120085405666612e02, -1.4534825256338749e02, + -2.1257155379297881e01, 2.7317691772612619e01, -3.6063526926252166e01, + -1.7588696592837897e02, 3.3079005662384850e02, -1.4519086534447842e02, + -2.0283472228681301e01, 2.4763027042036295e01, -3.0876160316998963e01, + -1.6184864900381874e02, 2.9976970905591691e02, -1.3084395423768876e02, + -2.0287461515322455e01, 2.4769400540137131e01, -3.0762734380983186e01, + -1.6214886052089241e02, 2.9998995088792128e02, -1.3088331758129965e02, + -2.0283472228681809e01, 2.4763027042017129e01, -3.0876160317336627e01, + -1.6184864900359682e02, 2.9976970905662938e02, -1.3084395423826805e02, + -2.0290765181946348e01, 2.4735639907973120e01, -3.0892738413082597e01, + -1.6154574482310053e02, 2.9934595420013272e02, -1.3068028494926122e02, + -1.9319499689234629e01, 2.2323824431805683e01, -2.6243395369841849e01, + -1.4782286378121026e02, 2.6985759662396487e02, -1.1715474197881395e02, + -1.9323022570439292e01, 2.2340565860680357e01, -2.6102786429129356e01, + -1.4828764857305418e02, 2.7027298759214750e02, -1.1726163007473576e02, + -1.9319499689236839e01, 2.2323824431730525e01, -2.6243395371031539e01, + -1.4782286378021576e02, 2.6985759662609979e02, -1.1715474198068593e02, + -1.9327939259284843e01, 2.2295320666731183e01, -2.6257097174199931e01, + -1.4751677383623073e02, 2.6942341041084092e02, -1.1698575776762208e02, + -1.8359079763330211e01, 2.0025118950280675e01, -2.2113826757823226e01, + -1.3415932552431914e02, 2.4147795894487624e02, -1.0427314537549884e02, + -1.8361534194530734e01, 2.0055847278170305e01, -2.1944107342764479e01, + -1.3482982214648752e02, 2.4214772485703989e02, -1.0447085300268679e02, + -1.8359079763339750e01, 2.0025118949989704e01, -2.2113826761939308e01, + -1.3415932552009582e02, 2.4147795895089951e02, -1.0427314538136979e02, + -1.8368836959765495e01, 1.9995657614892380e01, -2.2124533894067383e01, + -1.3385233293246981e02, 2.4103659293914149e02, -1.0410011400771683e02, + -1.7402299525814517e01, 1.7865597763687486e01, -1.8455503416511757e01, + -1.2090765118569301e02, 2.1464125749038132e02, -9.2190581022134992e01, + -1.7402744551259310e01, 1.7914800567904472e01, -1.8255754666855470e01, + -1.2183089355280822e02, 2.1563582256173194e02, -9.2507405324257306e01, + -1.7402299525855486e01, 1.7865597762572605e01, -1.8455503430527756e01, + -1.2090765116826699e02, 2.1464125750558804e02, -9.2190581039770791e01, + -1.7413567239985614e01, 1.7835392747330133e01, -1.8463115133795956e01, + -1.2060260469703572e02, 2.1419685510959093e02, -9.2015134441585104e01, + -1.6449179896085464e01, 1.5843762224435309e01, -1.5236722252652665e01, + -1.0811515163854509e02, 1.8935506712501905e02, -8.0897437157402223e01, + -1.6446174965543889e01, 1.5916874201410112e01, -1.5007553197461570e01, + -1.0934291295595986e02, 1.9075532567542470e02, -8.1366596347119696e01, + -1.6449179896260411e01, 1.5843762220214204e01, -1.5236722299508587e01, + -1.0811515156878269e02, 1.8935506715588940e02, -8.0897437207525684e01, + -1.6462173655481337e01, 1.5813096619069219e01, -1.5241142983208677e01, + -1.0781563484017332e02, 1.8891289499393798e02, -8.0721658713418606e01, + -1.5499661595231082e01, 1.3957945516559789e01, -1.2426145992195885e01, + -9.5826844741964834e01, 1.6562434781973772e02, -7.0383233416004117e01, + -1.5491037589250178e01, 1.4061349904707843e01, -1.2170301483989650e01, + -9.7412966929875139e01, 1.6751874597575440e02, -7.1041920384880939e01, + -1.5499661595973759e01, 1.3957945500778198e01, -1.2426146145776961e01, + -9.5826844470313858e01, 1.6562434784656404e02, -7.0383233547510557e01, + -1.5514618579274794e01, 1.3927192540790591e01, -1.2427264674287118e01, + -9.5537423121432880e01, 1.6519113036542510e02, -7.0209783384625098e01, + -1.4553592409098401e01, 1.2206343505203831e01, -9.9929274597052196e00, + -8.4085595900823435e01, 1.4345191724964303e02, -6.0636862050381758e01, + -1.4536130507533649e01, 1.2347228125716077e01, -9.7159302678980044e00, + -8.6081002959763751e01, 1.4592996741513730e02, -6.1523840242331410e01, + -1.4553592412232879e01, 1.2206343446986155e01, -9.9929279524397305e00, + -8.4085594870780753e01, 1.4345191706222485e02, -6.0636862352071532e01, + -1.4570766853404239e01, 1.2175998366492486e01, -9.9905856922863112e00, + -8.3812185051328299e01, 1.4303633648493073e02, -6.0469165577726159e01, + -1.3610717065161962e01, 1.0587059629986399e01, -7.9068321681349163e00, + -7.2932404423885004e01, 1.2283913327111270e02, -5.1646910322317169e01, + -1.3579708436673444e01, 1.0773027159520954e01, -7.6175370796795425e00, + -7.5376833196183071e01, 1.2597958225245242e02, -5.2797863799745748e01, + -1.3610717078313911e01, 1.0587059418306087e01, -7.9068337121483454e00, + -7.2932400620636059e01, 1.2283913169238102e02, -5.1646910832841897e01, + -1.3630368323321786e01, 1.0557789879027116e01, -7.9007777139483810e00, + -7.2682825476758552e01, 1.2245259140017740e02, -5.1489446559796768e01, + -1.2670671078399982e01, 9.0981634949263963e00, -6.1383490362855788e00, + -6.2406844162279825e01, 1.0378677653422224e02, -4.3402055519687693e01, + -1.2619333100308433e01, 9.3364634226935799e00, -5.8491811509717584e00, + -6.5316414528433455e01, 1.0763857666200300e02, -4.4841832720191050e01, + -1.2670671133253135e01, 9.0981627374157021e00, -6.1383537481895356e00, + -6.2406830503476570e01, 1.0378676818216074e02, -4.3402055529436716e01, + -1.2693036794620980e01, 9.0708908225804148e00, -6.1281713411274001e00, + -6.2191660620037396e01, 1.0344456594081470e02, -4.3260806640248063e01, + -1.1732979767504439e01, 7.7377614739662697e00, -4.6587775146685351e00, + -5.2547655563671029e01, 8.6296103981829802e01, -3.5891515805495345e01, + -1.1651721415208119e01, 8.0340005825064456e00, -4.3852919661646119e00, + -5.5898160750405737e01, 9.0851291378134590e01, -3.7622755083739385e01, + -1.1732979994779518e01, 7.7377588120662892e00, -4.6587914600219875e00, + -5.2547607987974565e01, 8.6296066930227624e01, -3.5891510429190419e01, + -1.1758218632638741e01, 7.7137968422318544e00, -4.6438239588320966e00, + -5.2381405657406454e01, 8.6019170302439520e01, -3.5774653697918737e01, + -1.0797063195543267e01, 6.5040766534586290e00, -3.4402783696562169e00, + -4.3393478931462226e01, 7.0370032342568010e01, -2.9105535302381853e01, + -1.0672637254876815e01, 6.8603244928014488e00, -3.1995767859681346e00, + -4.7101348454718874e01, 7.5530774605740319e01, -3.1094453979913311e01, + -1.0797064129672576e01, 6.5040675030570139e00, -3.4403181344841500e00, + -4.3393319126804485e01, 7.0369884883020177e01, -2.9105501594155889e01, + -1.0825134802124644e01, 6.4853446725127366e00, -3.4195560956016346e00, + -4.3296381389022351e01, 7.0187483762520671e01, -2.9024415860031247e01, + -9.8622468030169337e00, 5.3955359781222549e00, -2.4558741324534137e00, + -3.4983728078555984e01, 5.6014425934291204e01, -2.3035887876475471e01, + -9.6769173769353625e00, 5.8079540801032961e00, -2.2635143148159220e00, + -3.8890523502249145e01, 6.1563046720547966e01, -2.5198820521877391e01, + -9.8622505990399034e00, 5.3955054149765509e00, -2.4559821583353774e00, + -3.4983216045684472e01, 5.6013889382190079e01, -2.3035736114340502e01, + -9.8926597117464805e00, 5.3849440641688187e00, -2.4279562878572039e00, + -3.4983707025980287e01, 5.5966629574570753e01, -2.3006306589550750e01, + -8.9277749780883457e00, 4.4108678323349286e00, -1.6793815271288624e00, + -2.7359655656676122e01, 4.3239544183593061e01, -1.7676416286664047e01, + -8.6587749152265552e00, 4.8674392165289442e00, -1.5450097170494306e00, + -3.1230915545542118e01, 4.8829474992442343e01, -1.9874755288141955e01, + -8.9277901202336185e00, 4.4107699183102085e00, -1.6796551456533098e00, + -2.7358123514289456e01, 4.3237769027728554e01, -1.7675844947587926e01, + -8.9590559763951383e00, 4.4128957610428623e00, -1.6423658138809611e00, + -2.7493743583145054e01, 4.3380518846300511e01, -1.7719639183506050e01, + -7.9928164326293913e00, 3.5492331091008302e00, -1.0852462622393610e00, + -2.0565792757352423e01, 3.2061909496398073e01, -1.3023704651715642e01, + -7.6125412569887647e00, 4.0287966748633526e00, -1.0084592804412351e00, + -2.4116992333062022e01, 3.7252797603904497e01, -1.5077495076198684e01, + -7.9928747817255603e00, 3.5489404571097585e00, -1.0858609980296849e00, + -2.0561701094768868e01, 3.2056747083970720e01, -1.3021877019728107e01, + -8.0213899495838241e00, 3.5708128515175943e00, -1.0368753205735253e00, + -2.0877831538201836e01, 3.2456559535389509e01, -1.3165540198118645e01, + -7.0564174984379102e00, 2.8104770395789380e00, -6.4821407306458223e-01, + -1.4652118176169953e01, 2.2507145963021038e01, -9.0780963613608154e00, + -6.5338936679228468e00, 3.2846161494194233e00, -6.1760141818709846e-01, + -1.7606122820367215e01, 2.6855555289500277e01, -1.0803821410528570e01, + -7.0566263531717324e00, 2.8097184139861691e00, -6.4925197579297411e-01, + -1.4643483271177150e01, 2.2495243692983838e01, -9.0734373052814821e00, + -7.0742646195707266e00, 2.8621047467298468e00, -5.8641470402843421e-01, + -1.5178915176777426e01, 2.3211717123277591e01, -9.3414295847965061e00, + -6.1172231064332783e00, 2.1957964102200167e00, -3.4265643705632465e-01, + -9.6769153352706798e00, 1.4613873405033004e01, -5.8450824172251430e00, + -5.4212678780860326e00, 2.6341589573018260e00, -3.4085224757280796e-01, + -1.1835854891340576e01, 1.7794701474942944e01, -7.1075278532253687e00, + -6.1178367984533244e00, 2.1945528943967396e00, -3.4261268423617658e-01, + -9.6695829134679272e00, 1.4600877298870854e01, -5.8381668136523013e00, + -6.1072022151656586e00, 2.2922503774685161e00, -2.6715334266026142e-01, + -1.0408120531614587e01, 1.5617405440391840e01, -6.2270636615178061e00, + -5.1722074807324017e00, 1.7098190643016411e00, -1.4098618492175408e-01, + -5.7061337346696464e00, 8.4331806866534098e00, -3.3349192888568142e00, + -4.2766424379800121e00, 2.0860564217794284e00, -1.5548660419053545e-01, + -7.0034949575065015e00, 1.0332245608764421e01, -4.0873492185766374e00, + -5.1727690165421372e00, 1.7132539127425084e00, -1.2776576793785877e-01, + -5.7565343018918274e00, 8.4941254548170697e00, -3.3479852132230872e00, + -5.0998839330979591e00, 1.8678855512825561e00, -5.7718910331047868e-02, + -6.5095346397755423e00, 9.5462002113817768e00, -3.7632628689263172e00, + -4.2112469382255613e00, 1.3675717927787789e00, -9.4961575783498800e-03, + -2.7877417589321136e00, 3.9953503912711956e00, -1.5499906707437840e00, + -3.1046711877098376e00, 1.6568346830533449e00, -4.5990009889900242e-02, + -3.3140676307068091e00, 4.7472200808709299e00, -1.8492173878772247e00, + -4.1976749320353317e00, 1.4246952243441517e00, 8.7531923058200650e-02, + -3.0996975434049761e00, 4.4668738099197531e00, -1.7103055321708385e00, + -4.0163145894665320e00, 1.5923303121893606e00, 5.8249749369824022e-02, + -3.3748048713195491e00, 4.7925769874900315e00, -1.8598420111853879e00, + -3.1955533414298376e00, 1.2168024121915868e00, 9.9474205814620603e-02, + -8.6811124876189694e-01, 1.1994338853723501e00, -4.4837238870567747e-01, + -1.9098914522594992e00, 1.3654451552507061e00, 2.9537044429980407e-03, + -9.3701125207094127e-01, 1.2575365835116745e00, -4.7248060681970733e-01, + -3.0285770502890443e00, 1.6166340190704305e00, 4.8662683065338386e-01, + -1.2308607057515726e00, 1.6114560066217587e00, -6.5896729332189652e-01, + -2.8078044229222514e00, 1.4555130910035559e00, 9.0876948497501955e-02, + -1.0566809618626720e00, 1.3938154223720176e00, -5.2279617091852160e-01, + -1.9963264755188566e00, 1.3672906754961440e00, 2.0801988470625002e-01, + 2.0083818728351077e-02, -1.5135587406137185e-02, -1.4175240342178652e-02, + -6.9344786794476854e-01, 1.2280621078720415e00, 1.2333381103148277e-02, + -1.0895386066093759e-02, 2.1764282171790141e-02, -1.0106900291744604e-02, + -1.2036881930169383e00, 2.0482931230000392e00, -1.2689218008973949e-01, + -5.0580690719339239e-01, 3.4047786101030464e-01, -7.0959386937004015e-02, + -1.4470760938303664e00, 1.4285049373060201e00, 5.5764887956399375e-02, + -2.9461990750009881e-02, 2.3005167601875431e-02, -1.0760396189439407e-02, + -4.3024292433642597e-01, 1.7121633497582587e00, 3.5705413032693957e-02, + -9.9216800479772127e-01, 1.5115432403429119e00, -6.3985596276149748e-01, + 5.4770961684437192e-01, 1.2565653391084903e00, 9.1639130181564755e-03, + -6.8547618650262643e-01, 1.2037212931265591e00, -5.1526772142324506e-01, + 4.8142431677326969e-01, 1.2842025505965851e00, -3.1103960497811806e-01, + -3.8667287940463613e-01, 9.2663039525338942e-01, -4.1330437951972537e-01, + 1.9976512094478704e-02, 1.4898674304290889e00, -2.1940405767858565e-03, + -8.0791207141984167e-01, 1.3979310081478775e00, -5.9845265079421794e-01, + 1.1971451112382212e00, 1.6539633089946477e00, -2.7009878691796618e-01, + -2.8868139196850624e00, 4.7294193613612734e00, -1.9578020397520424e00, + 1.8164162541717044e00, 1.4570111710269262e00, 2.2385898037164991e-02, + -3.1195681762439769e00, 4.9723722392038878e00, -2.0423972644796100e00, + 1.5812403987207633e00, 1.1421043858413655e00, -4.4319666868952730e-02, + -2.3144705949527720e00, 3.7448930479898297e00, -1.5426803544433196e00, + 1.4992161878806018e00, 1.6612039136364238e00, -2.2870713891204597e-02, + -3.4442115437939465e00, 5.5057190995408973e00, -2.2657208348376137e00, + 2.4658130352390710e00, 1.5819912227884063e00, -1.3204477532594588e-01, + -5.7752803465671017e00, 9.0677018990478242e00, -3.6843468204828174e00, + 3.1062201217160963e00, 1.8205810727868250e00, 7.3942159732456811e-02, + -7.3418038323250947e00, 1.1309154676354810e01, -4.5733470083866452e00, + 2.5667672162869133e00, 1.3762236869878626e00, 5.4823291778512563e-02, + -5.5558964069977943e00, 8.5620133672289516e00, -3.4575259608624478e00, + 2.9333361085351610e00, 1.9771000784477066e00, 2.1600903596218385e-02, + -7.7786452012965430e00, 1.2026327126407146e01, -4.8722408979121159e00, + 3.5238342146994350e00, 1.8411341262124141e00, 1.0485737443151430e-01, + -1.0316470080846322e01, 1.5628354265192609e01, -6.2547428286449396e00, + 4.3947471898784478e00, 2.3129375587624681e00, 1.6998863701958250e-01, + -1.3069120913924280e01, 1.9764673064124775e01, -7.9234176878170990e00, + 3.5464051944219954e00, 1.7786047141550632e00, 1.8395466553434961e-01, + -1.0256713338978345e01, 1.5450540198835597e01, -6.1709943751208902e00, + 4.3074781177775723e00, 2.4284702978185178e00, 1.2121907902830774e-01, + -1.3510697720561426e01, 2.0490823414440431e01, -8.2265504110307699e00, + 4.5269670710447079e00, 2.3411415500822019e00, 3.7814443659878427e-01, + -1.6533454371385766e01, 2.4532574055181296e01, -9.7222898630871342e00, + 5.6498078480438974e00, 2.8871559084424092e00, 3.1648740182441881e-01, + -1.9832336139347099e01, 2.9630584562783888e01, -1.1804975183138390e01, + 4.5317970588477650e00, 2.3235629480266455e00, 4.0711209040396701e-01, + -1.6523611973754900e01, 2.4482080409856291e01, -9.6968326211377835e00, + 5.6107427774726322e00, 2.9693568967987254e00, 2.6856229367890733e-01, + -2.0186235796983127e01, 3.0228033555488111e01, -1.2057362656117963e01, + 5.5230828784340904e00, 3.0159142144119913e00, 7.5032702265793638e-01, + -2.4452361306480910e01, 3.5745746299744695e01, -1.4059387633540990e01, + 6.8467243986091164e00, 3.5205846294935204e00, 5.5323452910250115e-01, + -2.7424447720726722e01, 4.0542113968978946e01, -1.6058340606199877e01, + 5.5241079122419858e00, 3.0111097413061287e00, 7.6043241689918206e-01, + -2.4453330947201032e01, 3.5733842835424838e01, -1.4052622761934279e01, + 6.8330970703372866e00, 3.5730950345697865e00, 5.0442967447855436e-01, + -2.7630302835415993e01, 4.0921397061842079e01, -1.6223699529825666e01, + 6.5233214752268127e00, 3.8455313715589599e00, 1.2738445662734672e00, + -3.4142511056048967e01, 4.9288751118195229e01, -1.9258816488331760e01, + 7.9798691992574877e00, 4.2304633704347614e00, 9.4916911879724064e-01, + -3.6082800915305256e01, 5.2740474636382487e01, -2.0757970588732530e01, + 6.5235391967368317e00, 3.8442392655293900e00, 1.2772689685023881e00, + -3.4144245582802192e01, 4.9286600694030149e01, -1.9257235266278844e01, + 7.9780164759860508e00, 4.2581364755189171e00, 9.0490824102641643e-01, + -3.6146890048111374e01, 5.2902251888236343e01, -2.0834714063750525e01, + 7.5301209868737518e00, 4.8266093670811516e00, 1.9906532239804082e00, + -4.5696171225139402e01, 6.5222794336738914e01, -2.5330008845677121e01, + 9.0592048208341964e00, 5.0524444639807982e00, 1.5639083038511417e00, + -4.6227354827270197e01, 6.6742768625790532e01, -2.6090733281390481e01, + 7.5301672757177256e00, 4.8262668988539703e00, 1.9917837214882572e00, + -4.5697152262800707e01, 6.5222641787790508e01, -2.5329699752317662e01, + 9.0617089689058279e00, 5.0627200474303731e00, 1.5306087886050987e00, + -4.6201245261995687e01, 6.6753711704174307e01, -2.6103836713323240e01, + 8.5439978438576958e00, 5.9605352581937785e00, 2.9388171122244109e00, + -5.9213652478598007e01, 8.3623964589400401e01, -3.2288651007290504e01, + 1.0100238105795977e01, 6.0156046860821641e00, 2.4311227628788585e00, + -5.8189717323516248e01, 8.2972590004142106e01, -3.2212869674305303e01, + 8.5440076687321067e00, 5.9604459430021439e00, 2.9391801366526531e00, + -5.9214078468041464e01, 8.3624068891376510e01, -3.2288610777657510e01, + 1.0103667533796683e01, 6.0158650887345448e00, 2.4107760944314816e00, + -5.8125625048064265e01, 8.2906979417176174e01, -3.2191629006406409e01, + 9.5650113177877785e00, 7.2498153679976820e00, 4.1551371399277919e00, + -7.4795843598083408e01, 1.0457037732454131e02, -4.0151433068943419e01, + 1.1116968561077568e01, 7.1347098863330896e00, 3.5688140741297674e00, + -7.2151486218593305e01, 1.0165680693075836e02, -3.9206269356622016e01, + 9.5650133940644455e00, 7.2497924894015711e00, 4.1552503042122613e00, + -7.4796005009548836e01, 1.0457044971811401e02, -4.0151435976986221e01, + 1.1120034079668221e01, 7.1303147700774092e00, 3.5594873892317103e00, + -7.2082067018068685e01, 1.0156598726189708e02, -3.9171834664292227e01, + 1.0593064483227742e01, 8.6969028070512202e00, 5.6755396034912966e00, + -9.2539537763180832e01, 1.2813560149579646e02, -4.8933613418447223e01, + 1.2119543877083460e01, 8.4137603187360543e00, 4.9925034366798311e00, + -8.8194505075704640e01, 1.2287993196505218e02, -4.7096724506223822e01, + 1.0593064919257221e01, 8.6968970567044934e00, 5.6755738143875760e00, + -9.2539593640863643e01, 1.2813563331215474e02, -4.8933618162805772e01, + 1.2121921818513506e01, 8.4078642204619420e00, 4.9908632634858190e00, + -8.8134432374832016e01, 1.2279086550380391e02, -4.7060844505587738e01, + 1.1627957207938659e01, 1.0303707615441018e01, 7.5344011042552923e00, + -1.1253294830348190e02, 1.5438372244089408e02, -5.8647453529357783e01, + 1.3114510015623049e01, 9.8513572940713416e00, 6.7213349376406626e00, + -1.0635738219113546e02, 1.4665751311861146e02, -5.5881528760137869e01, + 1.1627957298834614e01, 1.0303706197478814e01, 7.5344111366673712e00, + -1.1253296638384563e02, 1.5438373415898508e02, -5.8647455853629580e01, + 1.3116237925845430e01, 9.8455331102145145e00, 6.7243141059359051e00, + -1.0631074264006560e02, 1.4658112805680690e02, -5.5849452095162235e01, + 1.2669386535689361e01, 1.2071287030293307e01, 9.7633555455962835e00, + -1.3485075345900265e02, 1.8336444946299886e02, -6.9300787627414508e01, + 1.4105804414673191e01, 1.1444289269702800e01, 8.7789794745243590e00, + -1.2666835962860844e02, 1.7298274034188972e02, -6.5547771558832267e01, + 1.2669386554490638e01, 1.2071286687068984e01, 9.7633584027450482e00, + -1.3485075900242089e02, 1.8336445335820781e02, -6.9300788508071975e01, + 1.4107018463574896e01, 1.1439185153305873e01, 8.7843335749580440e00, + -1.2663444344319166e02, 1.7292158897636148e02, -6.5521162694327174e01, + 1.3716937488160630e01, 1.3999597459400730e01, 1.2389915672436279e01, + -1.5954894249539399e02, 2.1510813446746886e02, -8.0895567204040049e01, + 1.5095682313349364e01, 1.3189272906323732e01, 1.1192627051714643e01, + -1.4915916817312757e02, 2.0184825850919157e02, -7.6081293415969839e01, + 1.3716937492019641e01, 1.3999597377767842e01, 1.2389916464009524e01, + -1.5954894412085929e02, 2.1510813567394996e02, -8.0895567498068928e01, + 1.5096520030681436e01, 1.3185064407456906e01, 1.1198910160279951e01, + -1.4913565617175487e02, 2.0180124290250004e02, -7.6060129778156622e01, + 1.4770075388032444e01, 1.6087303167766446e01, 1.5436222950666867e01, + -1.8666021493779203e02, 2.4962122089688103e02, -9.3426463524457304e01, + 1.6085379191481852e01, 1.5083589447287226e01, 1.3991739427782750e01, + -1.7386892459375579e02, 2.3325385095807121e02, -8.7470099643500802e01, + 1.4770075388818769e01, 1.6087303148664304e01, 1.5436223164442264e01, + -1.8666021539675981e02, 2.4962122125116741e02, -9.3426463615076329e01, + 1.6085951551006787e01, 1.5080238931969067e01, 1.3998101278449143e01, + -1.7385331837944693e02, 2.3321864790104019e02, -8.7453697552144448e01, + 1.5828143941097450e01, 1.8331670220961666e01, 1.8918268274003861e01, + -2.1619095210442941e02, 2.8688297635978756e02, -1.0687973526499771e02, + 1.7075534787366465e01, 1.7125200136366264e01, 1.7207074959934751e01, + -2.0084388544719391e02, 2.6720765911058965e02, -9.9705133726570395e01, + 1.5828143941256627e01, 1.8331670216557445e01, 1.8918268330404022e01, + -2.1619095222989833e02, 2.8688297645950814e02, -1.0687973529137253e02, + 1.7075923730873765e01, 1.7122590193964911e01, 1.7213058024904747e01, + -2.0083402645820061e02, 2.6718180837697332e02, -9.9692640534772679e01, + 1.6890371426423382e01, 2.0728579569842751e01, 2.2845917469463828e01, + -2.4812083435502871e02, 3.2684448823688496e02, -1.2123263616047282e02, + 1.8066449820492846e01, 1.9312661524160735e01, 2.0870036016187061e01, + -2.3013589616073858e02, 3.0372498377642154e02, -1.1277999824352135e02, + 1.6890371426455424e01, 2.0728579568840633e01, 2.2845917484032956e01, + -2.4812083438838550e02, 3.2684448826399682e02, -1.2123263616782057e02, + 1.8066713333743454e01, 1.9310657703202459e01, 2.0875423564416035e01, + -2.3013008228413184e02, 3.0370630494679148e02, -1.1277060230387309e02, + 1.7955886187113396e01, 2.3272683588860026e01, 2.7223982220959247e01, + -2.8240595076334000e02, 3.6943078590316281e02, -1.3645364576977221e02, + 1.9058236733002300e01, 2.1644988962398710e01, 2.5012267757287322e01, + -2.6180071928343307e02, 3.4282650121799617e02, -1.2669036882336400e02, + 1.7955886187119816e01, 2.3272683588634656e01, 2.7223982224651898e01, + -2.8240595077199526e02, 3.6943078591032139e02, -1.3645364577174797e02, + 1.9058414960148450e01, 2.1643466247439289e01, 2.5016983354038196e01, + -2.6179767020610126e02, 3.4281320617581565e02, -1.2668337355331974e02, + 1.9023741366983238e01, 2.5957710504548576e01, 3.2054387652193789e01, + -3.1898571318422574e02, 4.1454655650462962e02, -1.5250373535684176e02, + 2.0050906563887416e01, 2.4121527381838824e01, 2.9665428981325245e01, + -2.9589665055055406e02, 3.8453661583827250e02, -1.4143340987287985e02, + 1.9023741366984520e01, 2.5957710504498362e01, 3.2054387653114766e01, + -3.1898571318642672e02, 4.1454655650647550e02, -1.5250373535735841e02, + 2.0051026978020587e01, 2.4120379273875816e01, 2.9669474257430963e01, + -2.9589543070583102e02, 3.8452729731205977e02, -1.4142824748467820e02, + 2.0092947487287756e01, 2.8776895490568755e01, 3.7339233558876920e01, + -9.8781982607414882e00, 7.0916635282296292e-01, -1.2340880155534291e-02, + 2.1044418341890132e01, 2.6741847681518077e01, 3.4861073630499796e01, + -9.1700568642165461e00, 6.5220324713443967e-01, -1.1045071585279443e-02, + 2.0092947487288011e01, 2.8776895490557653e01, 3.7339233559103448e01, + -9.8781982608033179e00, 7.0916635282857932e-01, -1.2340880155703077e-02, + 2.1044499630877905e01, 2.6740987496092696e01, 3.4864491165514394e01, + -9.1707199731434574e00, 6.5223741134844682e-01, -1.1045188698410773e-02, + 2.1162510215379026e01, 3.1723491960797684e01, 4.3084295875067085e01, + -4.1033675985379521e00, -6.6095139594000130e-01, 6.0977735530407223e-02, + 2.2038706806958309e01, 2.9505670300337073e01, 4.0630600131872811e01, + -2.7905442844326718e00, -8.3885972791335117e-01, 6.8309956404426039e-02, + 2.1162510215379076e01, 3.1723491960795304e01, 4.3084295875120795e01, + -4.1033675985539224e00, -6.6095139593840913e-01, 6.0977735530354210e-02, + 2.2038761643178379e01, 2.9505029336592230e01, 4.0633451796171073e01, + -2.7913314472201640e00, -8.3878528163749511e-01, 6.8307595298566767e-02, + 3.1719012432820758e01, 6.7480322661109355e01, 1.3318978565899991e02, + -1.6791944323404795e01, -1.0181217992701848e00, 1.2989592638281225e-01, + 3.2009499874031789e01, 6.5013296175889408e01, 1.3669799889514238e02, + -1.7009031615065428e01, -1.0689880784706638e00, 1.3388972346122466e-01, + 3.1719012432820758e01, 6.7480322661109355e01, 1.3318978565899991e02, + -1.6791944323404795e01, -1.0181217992701848e00, 1.2989592638281225e-01, + 3.2009500887769519e01, 6.5013269472322307e01, 1.3669829238273672e02, + -1.7009116366540379e01, -1.0689798256828462e00, 1.3388945486998777e-01, + 4.1931127118492086e01, 1.1600186087954401e02, 3.1751764022286790e02, + -4.6438894455748802e01, -8.7599401950869438e-01, 2.2297105562740663e-01, + 4.2002297497564768e01, 1.1479764873768737e02, 3.2393143797302810e02, + -4.7847299173836262e01, -7.8150712905299369e-01, 2.2131248436241077e-01, + 4.1931127118492086e01, 1.1600186087954401e02, 3.1751764022286790e02, + -4.6438894455748802e01, -8.7599401950869438e-01, 2.2297105562740663e-01, + 4.2002297514594851e01, 1.1479764793294436e02, 3.2393145467669495e02, + -4.7847304068128608e01, -7.8150664807362491e-01, 2.2131246858403722e-01, + 5.1984670105634827e01, 1.7926303194781252e02, 6.2846495111925287e02, + -1.0034649475039414e02, 2.4606292097951082e-01, 3.3256752105517051e-01, + 5.2000554052128159e01, 1.7883235795593501e02, 6.3273302895025176e02, + -1.0138733878813618e02, 3.2804187851642969e-01, 3.3055293107858102e-01, + 5.1984670105634827e01, 1.7926303194781252e02, 6.2846495111925287e02, + -1.0034649475039414e02, 2.4606292097951082e-01, 3.3256752105517051e-01, + 5.2000554052402805e01, 1.7883235793562420e02, 6.3273302962903426e02, + -1.0138733898825184e02, 3.2804189825766372e-01, 3.3055293042886030e-01, + 6.1996666427075382e01, 2.5724136589119979e02, 1.0913830717468406e03, + -1.8317243758181812e02, 2.5193786568880601e00, 4.6277932792022042e-01, + 6.2000133522892554e01, 2.5710536851489377e02, 1.0934673032018356e03, + -1.8370056934287794e02, 2.5630609198690104e00, 4.6162176037505448e-01, + 6.1996666427075382e01, 2.5724136589119979e02, 1.0913830717468406e03, + -1.8317243758181812e02, 2.5193786568880601e00, 4.6277932792022042e-01, + 6.2000133522896938e01, 2.5710536851442714e02, 1.0934673032246803e03, + -1.8370056934963364e02, 2.5630609205366826e00, 4.6162176035304603e-01, + 7.1999279107664492e01, 3.4965254984584158e02, 1.7356304176273381e03, + -3.0063395678020430e02, 6.2079056750108883e00, 6.1505333334154833e-01, + 7.2000032172982571e01, 3.4961232791697932e02, 1.7365043785874466e03, + -3.0086002522613632e02, 6.2270725229979789e00, 6.1452738833821030e-01, + 7.1999279107664492e01, 3.4965254984584158e02, 1.7356304176273381e03, + -3.0063395678020430e02, 6.2079056750108883e00, 6.1505333334154833e-01, + 7.2000032172982642e01, 3.4961232791696904e02, 1.7365043785881401e03, + -3.0086002522634379e02, 6.2270725230187063e00, 6.1452738833751985e-01, + 8.1999844359310714e01, 4.5636323545227941e02, 2.5918884526432239e03, + -4.5885344883307727e02, 1.1616256691917803e01, 7.8948404417119522e-01, + 8.2000007751936337e01, 4.5635184072744744e02, 2.5922210189842476e03, + -4.5894061525528980e02, 1.1623761628208563e01, 7.8927378661620728e-01, + 8.1999844359310714e01, 4.5636323545227941e02, 2.5918884526432239e03, + -4.5885344883307727e02, 1.1616256691917803e01, 7.8948404417119522e-01, + 8.2000007751936337e01, 4.5635184072744744e02, 2.5922210189842476e03, + -4.5894061525528980e02, 1.1623761628208563e01, 7.8927378661620728e-01}; + + // Table info parameters + std::vector table_info = { + -2.1000000000000000e01, 2.1000000000000000e01, 1.0500000000000000e02, + 1.0000000000000000e00, 1.0000000000000000e01, -1.0000000000000000e00}; + + // Environment matrix data (em_x) + std::vector em_x = { + 9.3816147034272368e-01, -1.6703373029862567e-01, -4.4294526064601734e-02, + -2.8798505489184573e-01, -1.6703373029862567e-01, 9.2489218226366088e-01, + -2.8928196536572048e-01, -4.7833509099876154e-01, -4.4294526064601734e-02, + -2.8928196536572048e-01, 5.7034320185695120e-01, 1.8771147911830000e-01, + -2.8798505489184573e-01, -4.7833509099876154e-01, 1.8771147911830000e-01, + 4.0174654365823070e-01, 8.4370316144902313e-01, -3.7813146789689916e-02, + -3.6989397568296523e-01, -4.0554075086539937e-01, -3.7813146789689916e-02, + 6.5766402633747112e-01, -4.2312966361682885e-01, 1.2685067374257861e-01, + -3.6989397568296523e-01, -4.2312966361682885e-01, 6.0171576901660107e-01, + 9.8283160997298613e-02, -4.0554075086539937e-01, 1.2685067374257861e-01, + 9.8283160997298613e-02, 2.1324148100625978e-01, 9.7843596341516559e-01, + -1.0492833888237871e-01, -1.0538688914576379e-01, -2.0453551592353389e-01, + -1.0492833888237871e-01, 7.7943976693565231e-01, -1.5898500035781410e-01, + 9.4834209331437741e-02, -1.0538688914576379e-01, -1.5898500035781410e-01, + 7.4778071691708869e-01, -6.1895255142095873e-01, -2.0453551592353389e-01, + 9.4834209331437741e-02, -6.1895255142095873e-01, 6.0844713798743799e-01, + 1.0079020879244640e00, -2.3855984150631487e-01, -3.4608276043004524e-02, + -4.7448768267289088e-01, -2.3855984150631487e-01, 4.9732018171028253e-01, + -3.1320787082485729e-01, -1.4528004145602180e-01, -3.4608276043004524e-02, + -3.1320787082485729e-01, 4.7696729363954582e-01, 1.1723268074231248e-01, + -4.7448768267289088e-01, -1.4528004145602180e-01, 1.1723268074231248e-01, + 4.0511515406019899e-01}; + + // Environment matrix data (em) - same as em_x reshaped to 4x4x4 + std::vector em = em_x; + + // Expected outputs + std::vector expected_xyz_scatter = { + -2.079966061311932624e00, -7.692086176232872941e-01, + -1.330535805782471437e00, -1.535049115492904015e00, + -3.387998402723848113e00, -2.128016321720361592e00, + -3.277106691660685023e00, -3.038592571524457941e00, + -3.249069620732517372e00, -1.970139560184089200e00, + -3.099042807698376478e00, -2.871870824895611740e00, + -3.505396485462592526e00, -2.262865452602483174e00, + -3.422716794166150045e00, -3.179802052814081126e00, + -3.387998402723848113e00, -2.128016321720361592e00, + -3.277106691660685023e00, -3.038592571524457941e00, + -2.097644546971316970e00, -7.853817517282856153e-01, + -1.357716376463428354e00, -1.553790419377535814e00, + -3.506551159314416299e00, -2.264198868197136782e00, + -3.424124719375644776e00, -3.181191625971385051e00, + -3.660169689351369549e00, -2.443224371314374199e00, + -3.605953674751581683e00, -3.365824124673929774e00, + -3.249069620732517372e00, -1.970139560184089200e00, + -3.099042807698376478e00, -2.871870824895611740e00, + -3.506551159314416299e00, -2.264198868197136782e00, + -3.424124719375644776e00, -3.181191625971385051e00, + -2.530401019792652750e00, -1.199448737241302965e00, + -2.045857780281324256e00, -2.028209748272746360e00, + -2.967997665475758939e00, -1.658223979429474859e00, + -2.714263642699213275e00, -2.536766176574209553e00, + -3.505396485462592526e00, -2.262865452602483174e00, + -3.422716794166150045e00, -3.179802052814081126e00, + -3.660169689351369549e00, -2.443224371314374199e00, + -3.605953674751581683e00, -3.365824124673929774e00, + -2.967997665475758939e00, -1.658223979429474859e00, + -2.714263642699213275e00, -2.536766176574209553e00, + -2.720389222561622855e00, -1.393795402757894042e00, + -2.345287611062552902e00, -2.246069186600956336e00, + -2.203412666452015500e00, -8.832922757908097111e-01, + -1.522624038797550927e00, -1.666955513084583895e00, + -3.241305103472522209e00, -1.961375736752476673e00, + -3.088893233422514317e00, -2.862569539842237276e00, + -3.574866904199356821e00, -2.343386670122550530e00, + -3.506415378108253655e00, -3.263388258370226502e00, + -3.603390794940432240e00, -2.376642828344087466e00, + -3.540123958621633538e00, -3.297678764698442677e00, + -3.241305103472522209e00, -1.961375736752476673e00, + -3.088893233422514317e00, -2.862569539842237276e00, + -2.429999277303452665e00, -1.100023053295984932e00, + -1.884632599239140616e00, -1.915409451705063049e00, + -3.617225590985869044e00, -2.392818241566808002e00, + -3.556322178869617989e00, -3.314299042117806238e00, + -3.041076611338295876e00, -1.738238794499435524e00, + -2.817792369079202164e00, -2.623522438271889357e00, + -3.574866904199356821e00, -2.343386670122550530e00, + -3.506415378108253655e00, -3.263388258370226502e00, + -3.617225590985869044e00, -2.392818241566808002e00, + -3.556322178869617989e00, -3.314299042117806238e00, + -2.494628058196212361e00, -1.163767319481099349e00, + -1.988521437257453117e00, -1.987826554800365741e00, + -3.075717632972006399e00, -1.776439224176796783e00, + -2.866012767430887731e00, -2.664752088723873147e00, + -3.603390794940432240e00, -2.376642828344087466e00, + -3.540123958621633538e00, -3.297678764698442677e00, + -3.041076611338295876e00, -1.738238794499435524e00, + -2.817792369079202164e00, -2.623522438271889357e00, + -3.075717632972006399e00, -1.776439224176796783e00, + -2.866012767430887731e00, -2.664752088723873147e00, + -2.937692581645545253e00, -1.625281447224701248e00, + -2.670609348993408627e00, -2.500890937682363369e00, + -2.025708717938600856e00, -7.199180117105589138e-01, + -1.247903242801984502e00, -1.477847643583491077e00, + -3.319906298979989234e00, -2.050399877432793616e00, + -3.190625595262678527e00, -2.956812988881384996e00, + -3.320427161433143759e00, -2.050991936909778435e00, + -3.191292681270106346e00, -2.957438103019588649e00, + -3.426553578385282073e00, -2.172154331535164129e00, + -3.325432621990271009e00, -3.084944527573738160e00, + -3.319906298979989234e00, -2.050399877432793616e00, + -3.190625595262678527e00, -2.956812988881384996e00, + -2.284064566837665566e00, -9.593000707401009297e-01, + -1.650522869993226838e00, -1.754438009603198623e00, + -3.379465923571113084e00, -2.118266849034519783e00, + -3.266349102610927257e00, -3.028338829309447533e00, + -3.079911054493713163e00, -1.781075057846572385e00, + -2.871812857785094764e00, -2.669747195759931468e00, + -3.320427161433143759e00, -2.050991936909778435e00, + -3.191292681270106346e00, -2.957438103019588649e00, + -3.379465923571113084e00, -2.118266849034519783e00, + -3.266349102610927257e00, -3.028338829309447533e00, + -2.322820890996963428e00, -9.962581166315247483e-01, + -1.712462166775041483e00, -1.796847406967182925e00, + -3.773974160346314832e00, -2.578147069133936320e00, + -3.733153897787112818e00, -3.501723916740695497e00, + -3.426553578385282073e00, -2.172154331535164129e00, + -3.325432621990271009e00, -3.084944527573738160e00, + -3.079911054493713163e00, -1.781075057846572385e00, + -2.871812857785094764e00, -2.669747195759931468e00, + -3.773974160346314832e00, -2.578147069133936320e00, + -3.733153897787112818e00, -3.501723916740695497e00, + -2.486913710376590814e00, -1.156110267723822860e00, + -1.976136357370776642e00, -1.979145885578800446e00, + -1.985509025147957018e00, -6.837428483499845555e-01, + -1.187510572446436763e00, -1.435784454536765731e00, + -3.459824255787086411e00, -2.210355755872039030e00, + -3.366749809979232122e00, -3.124964584278840274e00, + -3.237455343632644311e00, -1.957033195714938278e00, + -3.083852092710132720e00, -2.857958563155091891e00, + -3.657182553291094873e00, -2.439708494346079171e00, + -3.602533119650989324e00, -3.362243802860200859e00, + -3.459824255787086411e00, -2.210355755872039030e00, + -3.366749809979232122e00, -3.124964584278840274e00, + -2.612858501561634128e00, -1.282799717023056063e00, + -2.177105827922439119e00, -2.122088127377702360e00, + -3.527506999611477845e00, -2.288426034556684474e00, + -3.449583412378998126e00, -3.206410310763728511e00, + -3.364731487733799042e00, -2.101446556578417724e00, + -3.247719091239923639e00, -3.010635929634002395e00, + -3.237455343632644311e00, -1.957033195714938278e00, + -3.083852092710132720e00, -2.857958563155091891e00, + -3.527506999611477845e00, -2.288426034556684474e00, + -3.449583412378998126e00, -3.206410310763728511e00, + -2.635737441843280848e00, -1.306200303853801703e00, + -2.213214111467705525e00, -2.148325634854432398e00, + -3.052719273633249131e00, -1.751058696872038523e00, + -2.834060115075742292e00, -2.637372572375544344e00, + -3.657182553291094873e00, -2.439708494346079171e00, + -3.602533119650989324e00, -3.362243802860200859e00, + -3.364731487733799042e00, -2.101446556578417724e00, + -3.247719091239923639e00, -3.010635929634002395e00, + -3.052719273633249131e00, -1.751058696872038523e00, + -2.834060115075742292e00, -2.637372572375544344e00, + -2.716586630714696859e00, -1.389827110741064020e00, + -2.339411933160714607e00, -2.241657223573843893e00}; + + // Expected gradients + std::vector expected_dy_dem_x = { + 6.020806901503564568e00, 4.855212102289546472e00, 5.539480671183889093e00, + 4.079184573583559903e00, 4.855212102289546472e00, 6.002622647924519939e00, + 4.071866106011439612e00, 3.510130130958049044e00, 5.539480671183889093e00, + 4.071866106011439612e00, 5.377211408293748107e00, 5.608493562913938746e00, + 4.079184573583559903e00, 3.510130130958049044e00, 5.608493562913938746e00, + 5.361908159053506928e00, 5.856733593151207451e00, 5.562527732431956196e00, + 3.696970497751218332e00, 3.591262868531118535e00, 5.562527732431956196e00, + 5.486518130714558517e00, 3.554623757729209022e00, 5.671203227676853587e00, + 3.696970497751218332e00, 3.554623757729209022e00, 5.409085428112309302e00, + 5.689207994509548172e00, 3.591262868531118535e00, 5.671203227676853587e00, + 5.689207994509548172e00, 5.576190638038706737e00, 6.061353798127412595e00, + 5.246377090750031336e00, 5.243735237316045961e00, 4.603911062145699340e00, + 5.246377090750031336e00, 5.721012903047327569e00, 4.908561477952471108e00, + 5.690718304708298980e00, 5.243735237316045961e00, 4.908561477952471108e00, + 5.654457931163397788e00, 3.881456101028720962e00, 4.603911062145699340e00, + 5.690718304708298980e00, 3.881456101028720962e00, 5.417053162665033561e00, + 6.074411415516496326e00, 4.379884232572607949e00, 5.573202499423647538e00, + 3.509692162941947302e00, 4.379884232572607949e00, 5.338347197725587634e00, + 3.943361243591621879e00, 4.998151625670068654e00, 5.573202499423647538e00, + 3.943361243591621879e00, 5.336491405093880758e00, 5.678289531805588375e00, + 3.509692162941947302e00, 4.998151625670068654e00, 5.678289531805588375e00, + 5.359752988507883842e00}; + + // Test dimensions + const int nloc = 4; + const int nnei_i = 4; + const int nnei_j = 4; + const int last_layer_size = 4; + + void SetUp() override {} + void TearDown() override {} +}; + +TEST_F(TestTabulateSeTTebd, tabulate_fusion_se_t_tebd_cpu) { + std::vector xyz_scatter(nloc * nnei_i * nnei_j * last_layer_size, 0); + deepmd::tabulate_fusion_se_t_tebd_cpu( + &xyz_scatter[0], &table[0], &table_info[0], &em_x[0], &em[0], nloc, + nnei_i, nnei_j, last_layer_size); + EXPECT_EQ(xyz_scatter.size(), nloc * nnei_i * nnei_j * last_layer_size); + EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); + for (int jj = 0; jj < xyz_scatter.size(); ++jj) { + EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter[jj]), 1e-5); + } +} + +TEST_F(TestTabulateSeTTebd, tabulate_fusion_se_t_tebd_grad_cpu) { + std::vector dy_dem_x(em_x.size()); + std::vector dy(nloc * nnei_i * nnei_j * last_layer_size, 1.0); + + deepmd::tabulate_fusion_se_t_tebd_grad_cpu( + &dy_dem_x[0], &table[0], &table_info[0], &em_x[0], &em[0], &dy[0], nloc, + nnei_i, nnei_j, last_layer_size); + + EXPECT_EQ(dy_dem_x.size(), nloc * nnei_i * nnei_j); + EXPECT_EQ(dy_dem_x.size(), expected_dy_dem_x.size()); + + for (int jj = 0; jj < dy_dem_x.size(); ++jj) { + EXPECT_LT(fabs(dy_dem_x[jj] - expected_dy_dem_x[jj]), 1e-5); + } +} + +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM +TEST_F(TestTabulateSeTTebd, tabulate_fusion_se_t_tebd_gpu) { + std::vector xyz_scatter(nloc * nnei_i * nnei_j * last_layer_size, 0); + double *xyz_scatter_dev = NULL, *table_dev = NULL, *em_x_dev = NULL, + *em_dev = NULL; + + deepmd::malloc_device_memory_sync(xyz_scatter_dev, xyz_scatter); + deepmd::malloc_device_memory_sync(table_dev, table); + deepmd::malloc_device_memory_sync(em_x_dev, em_x); + deepmd::malloc_device_memory_sync(em_dev, em); + + deepmd::tabulate_fusion_se_t_tebd_gpu( + xyz_scatter_dev, table_dev, &table_info[0], em_x_dev, em_dev, nloc, + nnei_i, nnei_j, last_layer_size); + + deepmd::memcpy_device_to_host(xyz_scatter_dev, xyz_scatter); + + deepmd::delete_device_memory(xyz_scatter_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + + EXPECT_EQ(xyz_scatter.size(), nloc * nnei_i * nnei_j * last_layer_size); + EXPECT_EQ(xyz_scatter.size(), expected_xyz_scatter.size()); + for (int jj = 0; jj < xyz_scatter.size(); ++jj) { + EXPECT_LT(fabs(xyz_scatter[jj] - expected_xyz_scatter[jj]), 1e-5); + } +} + +TEST_F(TestTabulateSeTTebd, tabulate_fusion_se_t_tebd_grad_gpu) { + std::vector dy_dem_x(em_x.size(), 0.0); + std::vector dy(nloc * nnei_i * nnei_j * last_layer_size, 1.0); + + double *dy_dem_x_dev = NULL, *table_dev = NULL, *em_x_dev = NULL, + *em_dev = NULL, *dy_dev = NULL; + + deepmd::malloc_device_memory_sync(dy_dem_x_dev, dy_dem_x); + deepmd::malloc_device_memory_sync(table_dev, table); + deepmd::malloc_device_memory_sync(em_x_dev, em_x); + deepmd::malloc_device_memory_sync(em_dev, em); + deepmd::malloc_device_memory_sync(dy_dev, dy); + + deepmd::tabulate_fusion_se_t_tebd_grad_gpu( + dy_dem_x_dev, table_dev, &table_info[0], em_x_dev, em_dev, dy_dev, nloc, + nnei_i, nnei_j, last_layer_size); + + deepmd::memcpy_device_to_host(dy_dem_x_dev, dy_dem_x); + + deepmd::delete_device_memory(dy_dem_x_dev); + deepmd::delete_device_memory(table_dev); + deepmd::delete_device_memory(em_x_dev); + deepmd::delete_device_memory(em_dev); + deepmd::delete_device_memory(dy_dev); + + EXPECT_EQ(dy_dem_x.size(), nloc * nnei_i * nnei_j); + EXPECT_EQ(dy_dem_x.size(), expected_dy_dem_x.size()); + + for (int jj = 0; jj < dy_dem_x.size(); ++jj) { + EXPECT_LT(fabs(dy_dem_x[jj] - expected_dy_dem_x[jj]), 1e-5); + } +} +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM diff --git a/source/lmp/builtin.cmake b/source/lmp/builtin.cmake index f29e9d3319..e051e5c24a 100644 --- a/source/lmp/builtin.cmake +++ b/source/lmp/builtin.cmake @@ -5,7 +5,52 @@ # assume LAMMPS CMake file has been executed, so these target/variables exist: # lammps LAMMPS_SOURCE_DIR get_lammps_version -get_lammps_version(${LAMMPS_SOURCE_DIR}/version.h LAMMPS_VERSION_NUMBER) +# Since May 15, 2025, the output of get_lammps_version is changed. We vendor the +# old get_lammps_version +# https://github.com/lammps/lammps/commit/b3e7121535863df3db487cd3e6a68c080bf2a6b4#diff-1214db0d1c015a50103f61f8ff7896053dec7ebc1edb930d6ef8bb07282f52abR75 +function(_get_lammps_version version_header variable) + file(STRINGS ${version_header} line REGEX LAMMPS_VERSION) + set(MONTHS + x + Jan + Feb + Mar + Apr + May + Jun + Jul + Aug + Sep + Oct + Nov + Dec) + string(REGEX + REPLACE "#define LAMMPS_VERSION \"([0-9]+) ([A-Za-z]+) ([0-9]+)\"" + "\\1" day "${line}") + string(REGEX + REPLACE "#define LAMMPS_VERSION \"([0-9]+) ([A-Za-z]+) ([0-9]+)\"" + "\\2" month "${line}") + string(REGEX + REPLACE "#define LAMMPS_VERSION \"([0-9]+) ([A-Za-z]+) ([0-9]+)\"" + "\\3" year "${line}") + string(STRIP ${day} day) + string(STRIP ${month} month) + string(STRIP ${year} year) + list(FIND MONTHS "${month}" month) + string(LENGTH ${day} day_length) + string(LENGTH ${month} month_length) + if(day_length EQUAL 1) + set(day "0${day}") + endif() + if(month_length EQUAL 1) + set(month "0${month}") + endif() + set(${variable} + "${year}${month}${day}" + PARENT_SCOPE) +endfunction() + +_get_lammps_version(${LAMMPS_SOURCE_DIR}/version.h LAMMPS_VERSION_NUMBER) configure_file("${CMAKE_CURRENT_LIST_DIR}/deepmd_version.h.in" "${CMAKE_CURRENT_BINARY_DIR}/deepmd_version.h" @ONLY) diff --git a/source/lmp/compute_deeptensor_atom.cpp b/source/lmp/compute_deeptensor_atom.cpp index 68c97a629e..f38279d936 100644 --- a/source/lmp/compute_deeptensor_atom.cpp +++ b/source/lmp/compute_deeptensor_atom.cpp @@ -24,7 +24,7 @@ using namespace LAMMPS_NS; /* ---------------------------------------------------------------------- */ -ComputeDeeptensorAtom::ComputeDeeptensorAtom(LAMMPS *lmp, int narg, char **arg) +ComputeDeeptensorAtom::ComputeDeeptensorAtom(LAMMPS* lmp, int narg, char** arg) : Compute(lmp, narg, arg), dp(lmp), tensor(nullptr) { if (strcmp(update->unit_style, "lj") == 0) { error->all(FLERR, @@ -45,7 +45,7 @@ ComputeDeeptensorAtom::ComputeDeeptensorAtom(LAMMPS *lmp, int narg, char **arg) int gpu_rank = dp.get_node_rank(); try { dt.init(model_file, gpu_rank); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } sel_types = dt.sel_types(); @@ -83,7 +83,7 @@ void ComputeDeeptensorAtom::init() { #endif } -void ComputeDeeptensorAtom::init_list(int /*id*/, NeighList *ptr) { +void ComputeDeeptensorAtom::init_list(int /*id*/, NeighList* ptr) { list = ptr; } @@ -101,10 +101,10 @@ void ComputeDeeptensorAtom::compute_peratom() { array_atom = tensor; } - double **x = atom->x; - double **f = atom->f; - int *type = atom->type; - int *mask = atom->mask; + double** x = atom->x; + double** f = atom->f; + int* type = atom->type; + int* mask = atom->mask; int nlocal = atom->nlocal; int nghost = atom->nghost; int nall = nlocal + nghost; @@ -145,7 +145,7 @@ void ComputeDeeptensorAtom::compute_peratom() { try { dt.compute(gtensor, force, virial, atensor, avirial, dcoord, dtype, dbox, nghost, lmp_list); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } diff --git a/source/lmp/compute_deeptensor_atom.h b/source/lmp/compute_deeptensor_atom.h index a90283aa9e..aeba8c11f4 100644 --- a/source/lmp/compute_deeptensor_atom.h +++ b/source/lmp/compute_deeptensor_atom.h @@ -30,19 +30,19 @@ namespace LAMMPS_NS { class ComputeDeeptensorAtom : public Compute { public: - ComputeDeeptensorAtom(class LAMMPS *, int, char **); + ComputeDeeptensorAtom(class LAMMPS*, int, char**); ~ComputeDeeptensorAtom() override; void init() override; void compute_peratom() override; double memory_usage() override; - void init_list(int, class NeighList *) override; + void init_list(int, class NeighList*) override; double dist_unit_cvt_factor; private: int nmax; - double **tensor; + double** tensor; PairDeepMD dp; - class NeighList *list; + class NeighList* list; deepmd_compat::DeepTensor dt; std::vector sel_types; }; diff --git a/source/lmp/fix_dplr.cpp b/source/lmp/fix_dplr.cpp index ac161730db..90cb4f4bba 100644 --- a/source/lmp/fix_dplr.cpp +++ b/source/lmp/fix_dplr.cpp @@ -24,7 +24,7 @@ using namespace LAMMPS_NS; using namespace FixConst; using namespace std; -static bool is_key(const string &input) { +static bool is_key(const string& input) { vector keys; keys.push_back("model"); keys.push_back("type_associate"); @@ -39,7 +39,7 @@ static bool is_key(const string &input) { return false; } -FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg) +FixDPLR::FixDPLR(LAMMPS* lmp, int narg, char** arg) : Fix(lmp, narg, arg), xstr(nullptr), ystr(nullptr), @@ -145,11 +145,11 @@ FixDPLR::FixDPLR(LAMMPS *lmp, int narg, char **arg) try { dpt.init(model, 0, "dipole_charge"); dtm.init(model, 0, "dipole_charge"); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } - pair_deepmd = (PairDeepMD *)force->pair_match("deepmd", 1, pair_deepmd_index); + pair_deepmd = (PairDeepMD*)force->pair_match("deepmd", 1, pair_deepmd_index); if (!pair_deepmd) { error->all(FLERR, "pair_style deepmd should be set before this fix\n"); } @@ -305,7 +305,7 @@ void FixDPLR::init() { /* ---------------------------------------------------------------------- */ void FixDPLR::setup_post_neighbor() { - double **x = atom->x; + double** x = atom->x; vector > valid_pairs; get_valid_pairs(valid_pairs, true); @@ -358,7 +358,7 @@ void FixDPLR::min_setup(int vflag) { setup(vflag); } /* ---------------------------------------------------------------------- */ -void FixDPLR::get_valid_pairs(vector > &pairs, bool is_setup) { +void FixDPLR::get_valid_pairs(vector >& pairs, bool is_setup) { pairs.clear(); int nlocal = atom->nlocal; @@ -366,12 +366,12 @@ void FixDPLR::get_valid_pairs(vector > &pairs, bool is_setup) { int nall = nlocal + nghost; vector dtype(nall); // get type - int *type = atom->type; + int* type = atom->type; for (int ii = 0; ii < nall; ++ii) { dtype[ii] = type_idx_map[type[ii] - 1]; } - int **bondlist = neighbor->bondlist; + int** bondlist = neighbor->bondlist; int nbondlist = neighbor->nbondlist; for (int ii = 0; ii < nbondlist; ++ii) { int idx0 = -1, idx1 = -1; @@ -437,9 +437,9 @@ void FixDPLR::get_valid_pairs(vector > &pairs, bool is_setup) { /* ---------------------------------------------------------------------- */ void FixDPLR::pre_exchange() { - double **x = atom->x; - double **v = atom->v; - int *type = atom->type; + double** x = atom->x; + double** v = atom->v; + int* type = atom->type; int nlocal = atom->nlocal; int nghost = atom->nghost; int nall = nlocal + nghost; @@ -461,8 +461,8 @@ void FixDPLR::pre_exchange() { /* ---------------------------------------------------------------------- */ void FixDPLR::pre_force(int vflag) { - double **x = atom->x; - int *type = atom->type; + double** x = atom->x; + int* type = atom->type; int nlocal = atom->nlocal; int nghost = atom->nghost; int nall = nlocal + nghost; @@ -503,7 +503,7 @@ void FixDPLR::pre_force(int vflag) { } } // get lammps nlist - NeighList *list = pair_deepmd->list; + NeighList* list = pair_deepmd->list; deepmd_compat::InputNlist lmp_list(list->inum, list->ilist, list->numneigh, list->firstneigh); lmp_list.set_mask(NEIGHMASK); @@ -515,7 +515,7 @@ void FixDPLR::pre_force(int vflag) { // compute try { dpt.compute(tensor, dcoord, dtype, dbox, nghost, lmp_list); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } // cout << "tensor of size " << tensor.size() << endl; @@ -607,7 +607,7 @@ void FixDPLR::post_force(int vflag) { update_efield_variables(); } - PPPMDPLR *pppm_dplr = (PPPMDPLR *)force->kspace_match("pppm/dplr", 1); + PPPMDPLR* pppm_dplr = (PPPMDPLR*)force->kspace_match("pppm/dplr", 1); int nlocal = atom->nlocal; int nghost = atom->nghost; int nall = nlocal + nghost; @@ -616,7 +616,7 @@ void FixDPLR::post_force(int vflag) { vector dtype(nall, 0); // set values for dcoord, dbox, dfele { - int *type = atom->type; + int* type = atom->type; for (int ii = 0; ii < nall; ++ii) { dtype[ii] = type_idx_map[type[ii] - 1]; } @@ -627,7 +627,7 @@ void FixDPLR::post_force(int vflag) { dbox[6] = domain->h[4] / dist_unit_cvt_factor; // zx dbox[3] = domain->h[5] / dist_unit_cvt_factor; // yx // get coord - double **x = atom->x; + double** x = atom->x; for (int ii = 0; ii < nall; ++ii) { for (int dd = 0; dd < 3; ++dd) { dcoord[ii * 3 + dd] = @@ -636,15 +636,15 @@ void FixDPLR::post_force(int vflag) { } // revise force according to efield if (pppm_dplr) { - const vector &dfele_(pppm_dplr->get_fele()); + const vector& dfele_(pppm_dplr->get_fele()); assert(dfele_.size() == nlocal * 3); for (int ii = 0; ii < nlocal * 3; ++ii) { dfele[ii] += dfele_[ii]; } } // revise force and virial according to efield - double *q = atom->q; - imageint *image = atom->image; + double* q = atom->q; + imageint* image = atom->image; double unwrap[3]; double v[6]; efield_fsum[0] = efield_fsum[1] = efield_fsum[2] = efield_fsum[3] = 0.0; @@ -675,7 +675,7 @@ void FixDPLR::post_force(int vflag) { } } // lmp nlist - NeighList *list = pair_deepmd->list; + NeighList* list = pair_deepmd->list; deepmd_compat::InputNlist lmp_list(list->inum, list->ilist, list->numneigh, list->firstneigh); // bonded pairs @@ -696,7 +696,7 @@ void FixDPLR::post_force(int vflag) { for (int ii = 0; ii < 9; ++ii) { dvcorr[ii] *= ener_unit_cvt_factor; } - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } assert(dfcorr.size() == dcoord.size()); @@ -726,7 +726,7 @@ void FixDPLR::post_force(int vflag) { // cout << endl; // } // apply the force correction - double **f = atom->f; + double** f = atom->f; for (int ii = 0; ii < nlocal; ++ii) { for (int dd = 0; dd < 3; ++dd) { f[ii][dd] += dfcorr[ii * 3 + dd]; @@ -778,7 +778,7 @@ void FixDPLR::min_post_force(int vflag) { post_force(vflag); } /* ---------------------------------------------------------------------- */ -int FixDPLR::pack_reverse_comm(int n, int first, double *buf) { +int FixDPLR::pack_reverse_comm(int n, int first, double* buf) { int m = 0; int last = first + n; for (int i = first; i < last; i++) { @@ -791,7 +791,7 @@ int FixDPLR::pack_reverse_comm(int n, int first, double *buf) { /* ---------------------------------------------------------------------- */ -void FixDPLR::unpack_reverse_comm(int n, int *list, double *buf) { +void FixDPLR::unpack_reverse_comm(int n, int* list, double* buf) { int m = 0; for (int i = 0; i < n; i++) { int j = list[i]; diff --git a/source/lmp/fix_dplr.h b/source/lmp/fix_dplr.h index 5f1161fda6..cd2c54f9d9 100644 --- a/source/lmp/fix_dplr.h +++ b/source/lmp/fix_dplr.h @@ -37,7 +37,7 @@ namespace deepmd_compat = deepmd::hpp; namespace LAMMPS_NS { class FixDPLR : public Fix { public: - FixDPLR(class LAMMPS *, int, char **); + FixDPLR(class LAMMPS*, int, char**); ~FixDPLR() override; int setmask() override; void init() override; @@ -52,14 +52,14 @@ class FixDPLR : public Fix { void min_pre_exchange() override; void min_pre_force(int) override; void min_post_force(int) override; - int pack_reverse_comm(int, int, double *) override; - void unpack_reverse_comm(int, int *, double *) override; + int pack_reverse_comm(int, int, double*) override; + void unpack_reverse_comm(int, int*, double*) override; double compute_scalar(void) override; double compute_vector(int) override; double ener_unit_cvt_factor, dist_unit_cvt_factor, force_unit_cvt_factor; private: - PairDeepMD *pair_deepmd; + PairDeepMD* pair_deepmd; deepmd_compat::DeepTensor dpt; deepmd_compat::DipoleChargeModifier dtm; std::string model; @@ -74,7 +74,7 @@ class FixDPLR : public Fix { std::vector efield; std::vector efield_fsum, efield_fsum_all; int efield_force_flag; - void get_valid_pairs(std::vector > &pairs, bool is_setup); + void get_valid_pairs(std::vector >& pairs, bool is_setup); int varflag; char *xstr, *ystr, *zstr; int xvar, yvar, zvar, xstyle, ystyle, zstyle; diff --git a/source/lmp/fix_ttm_dp.h b/source/lmp/fix_ttm_dp.h index 168f880226..3eb4ccd533 100644 --- a/source/lmp/fix_ttm_dp.h +++ b/source/lmp/fix_ttm_dp.h @@ -13,6 +13,6 @@ class FixTTMDP : public FixTTM { tmp[2] = nzgrid; return tmp; }; - double ***const get_T_electron() const { return T_electron; }; + double*** const get_T_electron() const { return T_electron; }; }; } // namespace LAMMPS_NS diff --git a/source/lmp/pair_base.cpp b/source/lmp/pair_base.cpp index a62956bbe4..ab60ccc780 100644 --- a/source/lmp/pair_base.cpp +++ b/source/lmp/pair_base.cpp @@ -35,9 +35,9 @@ using namespace LAMMPS_NS; using namespace std; -static int stringCmp(const void *a, const void *b) { - char *m = (char *)a; - char *n = (char *)b; +static int stringCmp(const void* a, const void* b) { + char* m = (char*)a; + char* n = (char*)b; int i, sum = 0; for (i = 0; i < MPI_MAX_PROCESSOR_NAME; i++) { @@ -98,7 +98,7 @@ int PairDeepBaseModel::get_node_rank() { return looprank; } -std::string PairDeepBaseModel::get_file_content(const std::string &model) { +std::string PairDeepBaseModel::get_file_content(const std::string& model) { int myrank = 0, root = 0; MPI_Comm_rank(MPI_COMM_WORLD, &myrank); int nchar = 0; @@ -108,7 +108,7 @@ std::string PairDeepBaseModel::get_file_content(const std::string &model) { nchar = file_content.size(); } MPI_Bcast(&nchar, 1, MPI_INT, root, MPI_COMM_WORLD); - char *buff = (char *)malloc(sizeof(char) * nchar); + char* buff = (char*)malloc(sizeof(char) * nchar); if (myrank == root) { memcpy(buff, file_content.c_str(), sizeof(char) * nchar); } @@ -122,7 +122,7 @@ std::string PairDeepBaseModel::get_file_content(const std::string &model) { } std::vector PairDeepBaseModel::get_file_content( - const std::vector &models) { + const std::vector& models) { std::vector file_contents(models.size()); for (unsigned ii = 0; ii < models.size(); ++ii) { file_contents[ii] = get_file_content(models[ii]); @@ -130,11 +130,11 @@ std::vector PairDeepBaseModel::get_file_content( return file_contents; } -void PairDeepBaseModel::make_fparam_from_compute(vector &fparam) { +void PairDeepBaseModel::make_fparam_from_compute(vector& fparam) { assert(do_compute_fparam); int icompute = modify->find_compute(compute_fparam_id); - Compute *compute = modify->compute[icompute]; + Compute* compute = modify->compute[icompute]; if (!compute) { error->all(FLERR, "compute id is not found: " + compute_fparam_id); @@ -152,18 +152,18 @@ void PairDeepBaseModel::make_fparam_from_compute(vector &fparam) { compute->compute_vector(); compute->invoked_flag |= Compute::INVOKED_VECTOR; } - double *cvector = compute->vector; + double* cvector = compute->vector; for (int jj = 0; jj < dim_fparam; ++jj) { fparam[jj] = cvector[jj]; } } } -void PairDeepBaseModel::make_aparam_from_compute(vector &aparam) { +void PairDeepBaseModel::make_aparam_from_compute(vector& aparam) { assert(do_compute_aparam); int icompute = modify->find_compute(compute_aparam_id); - Compute *compute = modify->compute[icompute]; + Compute* compute = modify->compute[icompute]; if (!compute) { error->all(FLERR, "compute id is not found: " + compute_aparam_id); @@ -176,10 +176,10 @@ void PairDeepBaseModel::make_aparam_from_compute(vector &aparam) { compute->invoked_flag |= Compute::INVOKED_PERATOM; } if (dim_aparam == 1) { - double *cvector = compute->vector_atom; + double* cvector = compute->vector_atom; aparam.assign(cvector, cvector + nlocal); } else if (dim_aparam > 1) { - double **carray = compute->array_atom; + double** carray = compute->array_atom; for (int ii = 0; ii < nlocal; ++ii) { for (int jj = 0; jj < dim_aparam; ++jj) { aparam[ii * dim_aparam + jj] = carray[ii][jj]; @@ -189,13 +189,13 @@ void PairDeepBaseModel::make_aparam_from_compute(vector &aparam) { } #ifdef USE_TTM -void PairDeepBaseModel::make_ttm_fparam(vector &fparam) { +void PairDeepBaseModel::make_ttm_fparam(vector& fparam) { assert(do_ttm); // get ttm_fix - const FixTTMDP *ttm_fix = NULL; + const FixTTMDP* ttm_fix = NULL; for (int ii = 0; ii < modify->nfix; ii++) { if (string(modify->fix[ii]->id) == ttm_fix_id) { - ttm_fix = dynamic_cast(modify->fix[ii]); + ttm_fix = dynamic_cast(modify->fix[ii]); } } if (!ttm_fix) { @@ -208,7 +208,7 @@ void PairDeepBaseModel::make_ttm_fparam(vector &fparam) { int nxnodes = nnodes[0]; int nynodes = nnodes[1]; int nznodes = nnodes[2]; - double ***const T_electron = ttm_fix->get_T_electron(); + double*** const T_electron = ttm_fix->get_T_electron(); int numb_effective_nodes = 0; double total_Te = 0; @@ -230,27 +230,27 @@ void PairDeepBaseModel::make_ttm_fparam(vector &fparam) { #endif #ifdef USE_TTM -void PairDeepBaseModel::make_ttm_aparam(vector &daparam) { +void PairDeepBaseModel::make_ttm_aparam(vector& daparam) { assert(do_ttm); // get ttm_fix - const FixTTMDP *ttm_fix = NULL; + const FixTTMDP* ttm_fix = NULL; for (int ii = 0; ii < modify->nfix; ii++) { if (string(modify->fix[ii]->id) == ttm_fix_id) { - ttm_fix = dynamic_cast(modify->fix[ii]); + ttm_fix = dynamic_cast(modify->fix[ii]); } } if (!ttm_fix) { error->all(FLERR, "fix ttm id is not found: " + ttm_fix_id); } // modify - double **x = atom->x; - int *mask = atom->mask; + double** x = atom->x; + int* mask = atom->mask; int nlocal = atom->nlocal; vector nnodes = ttm_fix->get_nodes(); int nxnodes = nnodes[0]; int nynodes = nnodes[1]; int nznodes = nnodes[2]; - double ***const T_electron = ttm_fix->get_T_electron(); + double*** const T_electron = ttm_fix->get_T_electron(); double dx = domain->xprd / nxnodes; double dy = domain->yprd / nynodes; double dz = domain->zprd / nynodes; @@ -275,8 +275,8 @@ void PairDeepBaseModel::make_ttm_aparam(vector &daparam) { } #endif -void PairDeepBaseModel::cum_sum(std::map &sum, - std::map &vec) { +void PairDeepBaseModel::cum_sum(std::map& sum, + std::map& vec) { sum[0] = 0; for (int ii = 1; ii < vec.size(); ++ii) { sum[ii] = sum[ii - 1] + vec[ii - 1]; @@ -284,10 +284,10 @@ void PairDeepBaseModel::cum_sum(std::map &sum, } PairDeepBaseModel::PairDeepBaseModel( - LAMMPS *lmp, - const char *cite_user_package, - deepmd_compat::DeepBaseModel &deep_model, - deepmd_compat::DeepBaseModelDevi &deep_model_devi) + LAMMPS* lmp, + const char* cite_user_package, + deepmd_compat::DeepBaseModel& deep_model, + deepmd_compat::DeepBaseModelDevi& deep_model_devi) : Pair(lmp), deep_base(deep_model), deep_base_model_devi(deep_model_devi) @@ -349,7 +349,7 @@ void PairDeepBaseModel::print_summary(const string pre) const { // capture cout to a string, then call LAMMPS's utils::logmesg // https://stackoverflow.com/a/4043813/9567349 std::stringstream buffer; - std::streambuf *sbuf = std::cout.rdbuf(); + std::streambuf* sbuf = std::cout.rdbuf(); std::cout.rdbuf(buffer.rdbuf()); cout << "Summary of lammps deepmd module ..." << endl; @@ -405,9 +405,9 @@ void PairDeepBaseModel::allocate() { } } -void PairDeepBaseModel::read_restart(FILE *) { is_restart = true; } +void PairDeepBaseModel::read_restart(FILE*) { is_restart = true; } -void PairDeepBaseModel::write_restart(FILE *) { +void PairDeepBaseModel::write_restart(FILE*) { // pass } @@ -454,23 +454,23 @@ double PairDeepBaseModel::init_one(int i, int j) { return cutoff; } -void *PairDeepBaseModel::extract(const char *str, int &dim) { +void* PairDeepBaseModel::extract(const char* str, int& dim) { if (strcmp(str, "cut_coul") == 0) { dim = 0; - return (void *)&cutoff; + return (void*)&cutoff; } if (strcmp(str, "scale") == 0) { dim = 2; - return (void *)scale; + return (void*)scale; } return NULL; } -void ana_st(double &max, - double &min, - double &sum, - const vector &vec, - const int &nloc) { +void ana_st(double& max, + double& min, + double& sum, + const vector& vec, + const int& nloc) { if (nloc == 0) { return; } @@ -488,9 +488,9 @@ void ana_st(double &max, } } -void make_uniform_aparam(vector &daparam, - const vector &aparam, - const int &nlocal) { +void make_uniform_aparam(vector& daparam, + const vector& aparam, + const int& nlocal) { unsigned dim_aparam = aparam.size(); daparam.resize(static_cast(dim_aparam) * nlocal); for (int ii = 0; ii < nlocal; ++ii) { diff --git a/source/lmp/pair_base.h b/source/lmp/pair_base.h index 055b45d20e..1dd4b84041 100644 --- a/source/lmp/pair_base.h +++ b/source/lmp/pair_base.h @@ -30,23 +30,23 @@ namespace deepmd_compat = deepmd::hpp; namespace LAMMPS_NS { class PairDeepBaseModel : public Pair { public: - PairDeepBaseModel(class LAMMPS *, - const char *, - deepmd_compat::DeepBaseModel &, - deepmd_compat::DeepBaseModelDevi &); + PairDeepBaseModel(class LAMMPS*, + const char*, + deepmd_compat::DeepBaseModel&, + deepmd_compat::DeepBaseModelDevi&); virtual ~PairDeepBaseModel() override; - void *extract(const char *, int &) override; + void* extract(const char*, int&) override; void init_style() override; - void write_restart(FILE *) override; - void read_restart(FILE *) override; + void write_restart(FILE*) override; + void read_restart(FILE*) override; double init_one(int i, int j) override; void print_summary(const std::string pre) const; int get_node_rank(); - void cum_sum(std::map &, std::map &); + void cum_sum(std::map&, std::map&); - std::string get_file_content(const std::string &model); + std::string get_file_content(const std::string& model); std::vector get_file_content( - const std::vector &models); + const std::vector& models); std::vector type_names; double ener_unit_cvt_factor, dist_unit_cvt_factor, force_unit_cvt_factor; @@ -54,7 +54,7 @@ class PairDeepBaseModel : public Pair { deepmd_compat::DeepBaseModel deep_base; deepmd_compat::DeepBaseModelDevi deep_base_model_devi; virtual void allocate(); - double **scale; + double** scale; unsigned numb_models; double cutoff; int numb_types; @@ -83,16 +83,16 @@ class PairDeepBaseModel : public Pair { double eps; double eps_v; - void make_fparam_from_compute(std::vector &fparam); + void make_fparam_from_compute(std::vector& fparam); bool do_compute_fparam; std::string compute_fparam_id; - void make_aparam_from_compute(std::vector &aparam); + void make_aparam_from_compute(std::vector& aparam); bool do_compute_aparam; std::string compute_aparam_id; - void make_ttm_fparam(std::vector &fparam); + void make_ttm_fparam(std::vector& fparam); - void make_ttm_aparam(std::vector &dparam); + void make_ttm_aparam(std::vector& dparam); bool do_ttm; std::string ttm_fix_id; int *counts, *displacements; @@ -103,13 +103,13 @@ class PairDeepBaseModel : public Pair { } // namespace LAMMPS_NS -void make_uniform_aparam(std::vector &daparam, - const std::vector &aparam, - const int &nlocal); -void ana_st(double &max, - double &min, - double &sum, - const std::vector &vec, - const int &nloc); +void make_uniform_aparam(std::vector& daparam, + const std::vector& aparam, + const int& nlocal); +void ana_st(double& max, + double& min, + double& sum, + const std::vector& vec, + const int& nloc); #endif diff --git a/source/lmp/pair_deepmd.cpp b/source/lmp/pair_deepmd.cpp index a11ad7f99c..3684c38dd9 100644 --- a/source/lmp/pair_deepmd.cpp +++ b/source/lmp/pair_deepmd.cpp @@ -117,7 +117,7 @@ static const char cite_user_deepmd_package[] = " doi = {10.1021/acs.jctc.5c00340},\n" "}\n\n"; -PairDeepMD::PairDeepMD(LAMMPS *lmp) +PairDeepMD::PairDeepMD(LAMMPS* lmp) : PairDeepBaseModel( lmp, cite_user_deepmd_package, deep_pot, deep_pot_model_devi) { // Constructor body can be empty @@ -141,10 +141,10 @@ void PairDeepMD::compute(int eflag, int vflag) { } bool do_ghost = true; // dpa2 communication - commdata_ = (CommBrickDeepMD *)comm; - double **x = atom->x; - double **f = atom->f; - int *type = atom->type; + commdata_ = (CommBrickDeepMD*)comm; + double** x = atom->x; + double** f = atom->f; + int* type = atom->type; int nlocal = atom->nlocal; int nghost = 0; if (do_ghost) { @@ -249,7 +249,7 @@ void PairDeepMD::compute(int eflag, int vflag) { try { deep_pot.compute(dener, dforce, dvirial, dcoord, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } @@ -260,7 +260,7 @@ void PairDeepMD::compute(int eflag, int vflag) { try { deep_pot.compute(dener, dforce, dvirial, deatom, dvatom, dcoord, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } if (eflag_atom) { @@ -312,7 +312,7 @@ void PairDeepMD::compute(int eflag, int vflag) { deep_pot_model_devi.compute(all_energy, all_force, all_virial, dcoord, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } else { @@ -321,7 +321,7 @@ void PairDeepMD::compute(int eflag, int vflag) { all_atom_energy, all_atom_virial, dcoord, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } @@ -449,7 +449,7 @@ void PairDeepMD::compute(int eflag, int vflag) { if (out_each == 1) { vector std_f_all(atom->natoms); // Gather std_f and tags - tagint *tag = atom->tag; + tagint* tag = atom->tag; int nprocs = comm->nprocs; // Grow arrays if necessary if (atom->natoms > stdf_comm_buff_size) { @@ -496,7 +496,7 @@ void PairDeepMD::compute(int eflag, int vflag) { if (numb_models == 1) { try { deep_pot.compute(dener, dforce, dvirial, dcoord, dtype, dbox); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } else { @@ -525,7 +525,7 @@ void PairDeepMD::compute(int eflag, int vflag) { } } -static bool is_key(const string &input) { +static bool is_key(const string& input) { vector keys; keys.push_back("out_freq"); keys.push_back("out_file"); @@ -548,7 +548,7 @@ static bool is_key(const string &input) { return false; } -void PairDeepMD::settings(int narg, char **arg) { +void PairDeepMD::settings(int narg, char** arg) { if (narg <= 0) { error->all(FLERR, "Illegal pair_style command"); } @@ -568,7 +568,7 @@ void PairDeepMD::settings(int narg, char **arg) { if (numb_models == 1) { try { deep_pot.init(arg[0], get_node_rank(), get_file_content(arg[0])); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } cutoff = deep_pot.cutoff() * dist_unit_cvt_factor; @@ -581,7 +581,7 @@ void PairDeepMD::settings(int narg, char **arg) { deep_pot.init(arg[0], get_node_rank(), get_file_content(arg[0])); deep_pot_model_devi.init(models, get_node_rank(), get_file_content(models)); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } cutoff = deep_pot_model_devi.cutoff() * dist_unit_cvt_factor; @@ -798,7 +798,7 @@ void PairDeepMD::settings(int narg, char **arg) { set coeffs for one or more type pairs ------------------------------------------------------------------------- */ -void PairDeepMD::coeff(int narg, char **arg) { +void PairDeepMD::coeff(int narg, char** arg) { if (!allocated) { allocate(); } @@ -889,7 +889,7 @@ void PairDeepMD::coeff(int narg, char **arg) { /* ---------------------------------------------------------------------- */ -int PairDeepMD::pack_reverse_comm(int n, int first, double *buf) { +int PairDeepMD::pack_reverse_comm(int n, int first, double* buf) { int i, m, last; m = 0; @@ -913,7 +913,7 @@ int PairDeepMD::pack_reverse_comm(int n, int first, double *buf) { /* ---------------------------------------------------------------------- */ -void PairDeepMD::unpack_reverse_comm(int n, int *list, double *buf) { +void PairDeepMD::unpack_reverse_comm(int n, int* list, double* buf) { int i, j, m; m = 0; diff --git a/source/lmp/pair_deepmd.h b/source/lmp/pair_deepmd.h index a8b3c13f4c..6d54a69fe6 100644 --- a/source/lmp/pair_deepmd.h +++ b/source/lmp/pair_deepmd.h @@ -42,20 +42,20 @@ class CommBrickDeepMD : public CommBrick { }; class PairDeepMD : public PairDeepBaseModel { public: - PairDeepMD(class LAMMPS *); + PairDeepMD(class LAMMPS*); ~PairDeepMD() override; - void settings(int, char **) override; - void coeff(int, char **) override; + void settings(int, char**) override; + void coeff(int, char**) override; void compute(int, int) override; - int pack_reverse_comm(int, int, double *) override; - void unpack_reverse_comm(int, int *, double *) override; + int pack_reverse_comm(int, int, double*) override; + void unpack_reverse_comm(int, int*, double*) override; protected: deepmd_compat::DeepPot deep_pot; deepmd_compat::DeepPotModelDevi deep_pot_model_devi; private: - CommBrickDeepMD *commdata_; + CommBrickDeepMD* commdata_; }; } // namespace LAMMPS_NS diff --git a/source/lmp/pair_deepspin.cpp b/source/lmp/pair_deepspin.cpp index accdce4c79..494ddcfb68 100644 --- a/source/lmp/pair_deepspin.cpp +++ b/source/lmp/pair_deepspin.cpp @@ -117,7 +117,7 @@ static const char cite_user_deepmd_package[] = " doi = {10.1021/acs.jctc.5c00340},\n" "}\n\n"; -PairDeepSpin::PairDeepSpin(LAMMPS *lmp) +PairDeepSpin::PairDeepSpin(LAMMPS* lmp) : PairDeepBaseModel( lmp, cite_user_deepmd_package, deep_spin, deep_spin_model_devi) { // Constructor body can be empty @@ -141,10 +141,10 @@ void PairDeepSpin::compute(int eflag, int vflag) { } bool do_ghost = true; // dpa2 communication - commdata_ = (CommBrickDeepSpin *)comm; - double **x = atom->x; - double **f = atom->f; - int *type = atom->type; + commdata_ = (CommBrickDeepSpin*)comm; + double** x = atom->x; + double** f = atom->f; + int* type = atom->type; int nlocal = atom->nlocal; int nghost = 0; if (do_ghost) { @@ -155,8 +155,8 @@ void PairDeepSpin::compute(int eflag, int vflag) { vector dspin(nall * 3, 0.); vector dfm(nall * 3, 0.); - double **sp = atom->sp; - double **fm = atom->fm; + double** sp = atom->sp; + double** fm = atom->fm; // spin initialize if (atom->sp_flag) { // get spin @@ -251,7 +251,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { deep_spin.compute(dener, dforce, dforce_mag, dvirial, dcoord, dspin, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } @@ -263,7 +263,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { deep_spin.compute(dener, dforce, dforce_mag, dvirial, deatom, dvatom, dcoord, dspin, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } if (eflag_atom) { @@ -315,7 +315,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { deep_spin_model_devi.compute(all_energy, all_force, all_force_mag, all_virial, dcoord, dspin, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } else { @@ -324,7 +324,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { all_energy, all_force, all_force_mag, all_virial, all_atom_energy, all_atom_virial, dcoord, dspin, dtype, dbox, nghost, lmp_list, ago, fparam, daparam); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } @@ -473,7 +473,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { // need support for spin atomic force. vector std_f_all(atom->natoms); // Gather std_f and tags - tagint *tag = atom->tag; + tagint* tag = atom->tag; int nprocs = comm->nprocs; // Grow arrays if necessary if (atom->natoms > stdf_comm_buff_size) { @@ -521,7 +521,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { try { deep_spin.compute(dener, dforce, dforce_mag, dvirial, dcoord, dspin, dtype, dbox); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } } else { @@ -558,7 +558,7 @@ void PairDeepSpin::compute(int eflag, int vflag) { } } -static bool is_key(const string &input) { +static bool is_key(const string& input) { vector keys; keys.push_back("out_freq"); keys.push_back("out_file"); @@ -581,7 +581,7 @@ static bool is_key(const string &input) { return false; } -void PairDeepSpin::settings(int narg, char **arg) { +void PairDeepSpin::settings(int narg, char** arg) { if (narg <= 0) { error->all(FLERR, "Illegal pair_style command"); } @@ -601,7 +601,7 @@ void PairDeepSpin::settings(int narg, char **arg) { if (numb_models == 1) { try { deep_spin.init(arg[0], get_node_rank(), get_file_content(arg[0])); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } cutoff = deep_spin.cutoff() * dist_unit_cvt_factor; @@ -614,7 +614,7 @@ void PairDeepSpin::settings(int narg, char **arg) { deep_spin.init(arg[0], get_node_rank(), get_file_content(arg[0])); deep_spin_model_devi.init(models, get_node_rank(), get_file_content(models)); - } catch (deepmd_compat::deepmd_exception &e) { + } catch (deepmd_compat::deepmd_exception& e) { error->one(FLERR, e.what()); } cutoff = deep_spin_model_devi.cutoff() * dist_unit_cvt_factor; @@ -828,7 +828,7 @@ void PairDeepSpin::settings(int narg, char **arg) { set coeffs for one or more type pairs ------------------------------------------------------------------------- */ -void PairDeepSpin::coeff(int narg, char **arg) { +void PairDeepSpin::coeff(int narg, char** arg) { if (!allocated) { allocate(); } @@ -919,7 +919,7 @@ void PairDeepSpin::coeff(int narg, char **arg) { /* ---------------------------------------------------------------------- */ -int PairDeepSpin::pack_reverse_comm(int n, int first, double *buf) { +int PairDeepSpin::pack_reverse_comm(int n, int first, double* buf) { int i, m, last; m = 0; @@ -946,7 +946,7 @@ int PairDeepSpin::pack_reverse_comm(int n, int first, double *buf) { /* ---------------------------------------------------------------------- */ -void PairDeepSpin::unpack_reverse_comm(int n, int *list, double *buf) { +void PairDeepSpin::unpack_reverse_comm(int n, int* list, double* buf) { int i, j, m; m = 0; diff --git a/source/lmp/pair_deepspin.h b/source/lmp/pair_deepspin.h index 47d6678441..cc31db8bf5 100644 --- a/source/lmp/pair_deepspin.h +++ b/source/lmp/pair_deepspin.h @@ -42,13 +42,13 @@ class CommBrickDeepSpin : public CommBrick { }; class PairDeepSpin : public PairDeepBaseModel { public: - PairDeepSpin(class LAMMPS *); + PairDeepSpin(class LAMMPS*); ~PairDeepSpin() override; - void settings(int, char **) override; - void coeff(int, char **) override; + void settings(int, char**) override; + void coeff(int, char**) override; void compute(int, int) override; - int pack_reverse_comm(int, int, double *) override; - void unpack_reverse_comm(int, int *, double *) override; + int pack_reverse_comm(int, int, double*) override; + void unpack_reverse_comm(int, int*, double*) override; protected: deepmd_compat::DeepSpin deep_spin; @@ -56,7 +56,7 @@ class PairDeepSpin : public PairDeepBaseModel { std::vector > all_force_mag; private: - CommBrickDeepSpin *commdata_; + CommBrickDeepSpin* commdata_; }; } // namespace LAMMPS_NS diff --git a/source/lmp/plugin/CMakeLists.txt b/source/lmp/plugin/CMakeLists.txt index a0998b3ce9..8f32af3e3e 100644 --- a/source/lmp/plugin/CMakeLists.txt +++ b/source/lmp/plugin/CMakeLists.txt @@ -2,7 +2,7 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION) message(STATUS "enable LAMMPS plugin mode") add_library(lammps_interface INTERFACE) if(DEFINED LAMMPS_VERSION) - cmake_minimum_required(VERSION 3.14) + cmake_minimum_required(VERSION 3.25.2) include(FetchContent) FetchContent_Declare( lammps_download @@ -38,8 +38,54 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION) # get_lammps_version # https://github.com/lammps/lammps/blob/c2a12f97c5f665852fb38fdd4922f7dd2e77a0a1/cmake/Modules/LAMMPSUtils.cmake#L27-L46 - include(${LAMMPS_SOURCE_ROOT}/cmake/Modules/LAMMPSUtils.cmake) - get_lammps_version(${LAMMPS_HEADER_DIR}/version.h LAMMPS_VERSION_NUMBER) + # include(${LAMMPS_SOURCE_ROOT}/cmake/Modules/LAMMPSUtils.cmake) Since May 15, + # 2025, the output of get_lammps_version is changed. We vendor the old + # get_lammps_version + # https://github.com/lammps/lammps/commit/b3e7121535863df3db487cd3e6a68c080bf2a6b4#diff-1214db0d1c015a50103f61f8ff7896053dec7ebc1edb930d6ef8bb07282f52abR75 + + function(_get_lammps_version version_header variable) + file(STRINGS ${version_header} line REGEX LAMMPS_VERSION) + set(MONTHS + x + Jan + Feb + Mar + Apr + May + Jun + Jul + Aug + Sep + Oct + Nov + Dec) + string(REGEX + REPLACE "#define LAMMPS_VERSION \"([0-9]+) ([A-Za-z]+) ([0-9]+)\"" + "\\1" day "${line}") + string(REGEX + REPLACE "#define LAMMPS_VERSION \"([0-9]+) ([A-Za-z]+) ([0-9]+)\"" + "\\2" month "${line}") + string(REGEX + REPLACE "#define LAMMPS_VERSION \"([0-9]+) ([A-Za-z]+) ([0-9]+)\"" + "\\3" year "${line}") + string(STRIP ${day} day) + string(STRIP ${month} month) + string(STRIP ${year} year) + list(FIND MONTHS "${month}" month) + string(LENGTH ${day} day_length) + string(LENGTH ${month} month_length) + if(day_length EQUAL 1) + set(day "0${day}") + endif() + if(month_length EQUAL 1) + set(month "0${month}") + endif() + set(${variable} + "${year}${month}${day}" + PARENT_SCOPE) + endfunction() + + _get_lammps_version(${LAMMPS_HEADER_DIR}/version.h LAMMPS_VERSION_NUMBER) set(LAMMPS_VERSION_NUMBER ${LAMMPS_VERSION_NUMBER} PARENT_SCOPE) @@ -75,8 +121,7 @@ if(DEFINED LAMMPS_SOURCE_ROOT OR DEFINED LAMMPS_VERSION) target_link_libraries(${libname} PUBLIC ${LIB_DEEPMD_C}) target_precompile_headers(${libname} PUBLIC [["deepmd.hpp"]]) remove_definitions(-D_GLIBCXX_USE_CXX11_ABI=${OP_CXX_ABI}) - if("$ENV{CIBUILDWHEEL}" STREQUAL "1" OR "$ENV{LMP_CXX11_ABI_0}" STREQUAL - "1") + if("$ENV{LMP_CXX11_ABI_0}" STREQUAL "1") add_definitions(-D_GLIBCXX_USE_CXX11_ABI=0) endif() else() diff --git a/source/lmp/plugin/deepmdplugin.cpp b/source/lmp/plugin/deepmdplugin.cpp index 4f62cb3944..d3b54f8e41 100644 --- a/source/lmp/plugin/deepmdplugin.cpp +++ b/source/lmp/plugin/deepmdplugin.cpp @@ -15,22 +15,22 @@ using namespace LAMMPS_NS; -static Pair *pairdeepmd(LAMMPS *lmp) { return new PairDeepMD(lmp); } -static Pair *pairdeepspin(LAMMPS *lmp) { return new PairDeepSpin(lmp); } +static Pair* pairdeepmd(LAMMPS* lmp) { return new PairDeepMD(lmp); } +static Pair* pairdeepspin(LAMMPS* lmp) { return new PairDeepSpin(lmp); } -static Compute *computedeepmdtensoratom(LAMMPS *lmp, int narg, char **arg) { +static Compute* computedeepmdtensoratom(LAMMPS* lmp, int narg, char** arg) { return new ComputeDeeptensorAtom(lmp, narg, arg); } -static Fix *fixdplr(LAMMPS *lmp, int narg, char **arg) { +static Fix* fixdplr(LAMMPS* lmp, int narg, char** arg) { return new FixDPLR(lmp, narg, arg); } #if LAMMPS_VERSION_NUMBER >= 20220328 -static KSpace *pppmdplr(LAMMPS *lmp) { return new PPPMDPLR(lmp); } +static KSpace* pppmdplr(LAMMPS* lmp) { return new PPPMDPLR(lmp); } #endif -extern "C" void lammpsplugin_init(void *lmp, void *handle, void *regfunc) { +extern "C" void lammpsplugin_init(void* lmp, void* handle, void* regfunc) { lammpsplugin_t plugin; lammpsplugin_regfunc register_plugin = (lammpsplugin_regfunc)regfunc; @@ -39,7 +39,7 @@ extern "C" void lammpsplugin_init(void *lmp, void *handle, void *regfunc) { plugin.name = "deepmd"; plugin.info = "deepmd pair style " STR_GIT_SUMM; plugin.author = "Han Wang"; - plugin.creator.v1 = (lammpsplugin_factory1 *)&pairdeepmd; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdeepmd; plugin.handle = handle; (*register_plugin)(&plugin, lmp); @@ -48,7 +48,7 @@ extern "C" void lammpsplugin_init(void *lmp, void *handle, void *regfunc) { plugin.name = "deepspin"; plugin.info = "deepspin pair style " STR_GIT_SUMM; plugin.author = "Duo Zhang"; - plugin.creator.v1 = (lammpsplugin_factory1 *)&pairdeepspin; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdeepspin; plugin.handle = handle; (*register_plugin)(&plugin, lmp); @@ -56,14 +56,14 @@ extern "C" void lammpsplugin_init(void *lmp, void *handle, void *regfunc) { plugin.name = "deeptensor/atom"; plugin.info = "compute deeptensor/atom " STR_GIT_SUMM; plugin.author = "Han Wang"; - plugin.creator.v2 = (lammpsplugin_factory2 *)&computedeepmdtensoratom; + plugin.creator.v2 = (lammpsplugin_factory2*)&computedeepmdtensoratom; (*register_plugin)(&plugin, lmp); plugin.style = "fix"; plugin.name = "dplr"; plugin.info = "fix dplr " STR_GIT_SUMM; plugin.author = "Han Wang"; - plugin.creator.v2 = (lammpsplugin_factory2 *)&fixdplr; + plugin.creator.v2 = (lammpsplugin_factory2*)&fixdplr; (*register_plugin)(&plugin, lmp); #if LAMMPS_VERSION_NUMBER >= 20220328 @@ -72,7 +72,7 @@ extern "C" void lammpsplugin_init(void *lmp, void *handle, void *regfunc) { plugin.name = "pppm/dplr"; plugin.info = "kspace pppm/dplr " STR_GIT_SUMM; plugin.author = "Han Wang"; - plugin.creator.v1 = (lammpsplugin_factory1 *)&pppmdplr; + plugin.creator.v1 = (lammpsplugin_factory1*)&pppmdplr; (*register_plugin)(&plugin, lmp); #endif } diff --git a/source/lmp/pppm_dplr.cpp b/source/lmp/pppm_dplr.cpp index e1bdb828af..3597a31548 100644 --- a/source/lmp/pppm_dplr.cpp +++ b/source/lmp/pppm_dplr.cpp @@ -36,10 +36,10 @@ enum { FORWARD_IK, FORWARD_AD, FORWARD_IK_PERATOM, FORWARD_AD_PERATOM }; #if LAMMPS_VERSION_NUMBER < 20181109 // See lammps/lammps#1165 -PPPMDPLR::PPPMDPLR(LAMMPS *lmp, int narg, char **arg) +PPPMDPLR::PPPMDPLR(LAMMPS* lmp, int narg, char** arg) : PPPM(lmp, narg, arg) #else -PPPMDPLR::PPPMDPLR(LAMMPS *lmp) +PPPMDPLR::PPPMDPLR(LAMMPS* lmp) : PPPM(lmp) #endif { @@ -232,7 +232,7 @@ void PPPMDPLR::compute(int eflag, int vflag) { // ntotal accounts for TIP4P tallying eatom/vatom for ghost atoms if (evflag_atom) { - double *q = atom->q; + double* q = atom->q; int nlocal = atom->nlocal; int ntotal = nlocal; if (tip4pflag) { @@ -288,8 +288,8 @@ void PPPMDPLR::fieldforce_ik() { // (mx,my,mz) = global coords of moving stencil pt // ek = 3 components of E-field on particle - double *q = atom->q; - double **x = atom->x; + double* q = atom->q; + double** x = atom->x; // double **f = atom->f; int nlocal = atom->nlocal; @@ -347,7 +347,7 @@ void PPPMDPLR::fieldforce_ad() { FFT_SCALAR ekx, eky, ekz; double s1, s2, s3; double sf = 0.0; - double *prd; + double* prd; prd = domain->prd; double xprd = prd[0]; @@ -364,8 +364,8 @@ void PPPMDPLR::fieldforce_ad() { // (mx,my,mz) = global coords of moving stencil pt // ek = 3 components of E-field on particle - double *q = atom->q; - double **x = atom->x; + double* q = atom->q; + double** x = atom->x; // double **f = atom->f; int nlocal = atom->nlocal; diff --git a/source/lmp/pppm_dplr.h b/source/lmp/pppm_dplr.h index b7e221c686..79a9a9ce37 100644 --- a/source/lmp/pppm_dplr.h +++ b/source/lmp/pppm_dplr.h @@ -21,14 +21,14 @@ class PPPMDPLR : public PPPM { public: #if LAMMPS_VERSION_NUMBER < 20181109 // See lammps/lammps#1165 - PPPMDPLR(class LAMMPS *, int, char **); + PPPMDPLR(class LAMMPS*, int, char**); #else - PPPMDPLR(class LAMMPS *); + PPPMDPLR(class LAMMPS*); #endif ~PPPMDPLR() override {}; void init() override; - const std::vector &get_fele() const { return fele; }; - std::vector &get_fele() { return fele; } + const std::vector& get_fele() const { return fele; }; + std::vector& get_fele() { return fele; } protected: void compute(int, int) override; diff --git a/source/lmp/tests/test_deeptensor.py b/source/lmp/tests/test_deeptensor.py index 6fb7cde746..20be3033b8 100644 --- a/source/lmp/tests/test_deeptensor.py +++ b/source/lmp/tests/test_deeptensor.py @@ -66,6 +66,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) # TODO # write_lmp_data(box, coord, type_HO, data_type_map_file) @@ -142,7 +146,7 @@ def test_compute_deeptensor_atom(lammps) -> None: lammps.variable("tensor atom c_tensor[1]") lammps.dump("1 all custom 1 dump id c_tensor[1]") lammps.run(0) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 assert np.array(lammps.variables["tensor"].value) == pytest.approx( expected_d[idx_map] ) @@ -155,7 +159,7 @@ def test_compute_deeptensor_atom_si(lammps_si) -> None: lammps_si.variable("tensor atom c_tensor[1]") lammps_si.dump("1 all custom 1 dump id c_tensor[1]") lammps_si.run(0) - idx_map = lammps_si.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_si.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 assert np.array(lammps_si.variables["tensor"].value) == pytest.approx( expected_d[idx_map] * constants.dist_metal2si ) diff --git a/source/lmp/tests/test_dplr.py b/source/lmp/tests/test_dplr.py index 21d1f18658..dd0c03aabe 100644 --- a/source/lmp/tests/test_dplr.py +++ b/source/lmp/tests/test_dplr.py @@ -271,6 +271,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data_full( box, coord, mol_list, type_OH, charge, data_file, bond_list, mass_list ) @@ -357,7 +361,7 @@ def test_pair_deepmd_sr(lammps) -> None: lammps.pair_coeff("* *") lammps.run(0) assert lammps.eval("pe") == pytest.approx(expected_e_sr) - id_list = lammps.lmp.numpy.extract_atom("id") + id_list = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] for ii in range(6): assert lammps.atoms[np.where(id_list == (ii + 1))[0][0]].force == pytest.approx( expected_f_sr[ii] @@ -378,7 +382,7 @@ def test_pair_deepmd_sr_virial(lammps) -> None: ) lammps.dump_modify("1 sort id") lammps.run(0) - id_list = lammps.lmp.numpy.extract_atom("id") + id_list = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] idx_list = [np.where(id_list == i)[0][0] for i in range(1, 7)] assert lammps.eval("pe") == pytest.approx(expected_e_sr) for ii in range(6): @@ -445,7 +449,7 @@ def test_pair_deepmd_lr_efield_constant(lammps) -> None: ) lammps.fix_modify("0 energy yes virial yes") lammps.run(0) - id_list = lammps.lmp.numpy.extract_atom("id") + id_list = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] assert lammps.eval("evdwl") == pytest.approx(expected_evdwl_lr_efield_constant) assert lammps.eval("f_0") == pytest.approx(expected_e_efield_constant) assert lammps.eval("pe") == pytest.approx(expected_e_lr_efield_constant) @@ -481,7 +485,7 @@ def test_pair_deepmd_lr_efield_variable(lammps) -> None: ) lammps.fix_modify("0 energy yes virial yes") lammps.run(0) - id_list = lammps.lmp.numpy.extract_atom("id") + id_list = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] assert lammps.eval("evdwl") == pytest.approx(expected_evdwl_lr_efield_variable) assert lammps.eval("f_0") == pytest.approx(expected_e_efield_variable) assert lammps.eval("pe") == pytest.approx(expected_e_lr_efield_variable) diff --git a/source/lmp/tests/test_lammps.py b/source/lmp/tests/test_lammps.py index ad8f8cdaac..05de1e69fa 100644 --- a/source/lmp/tests/test_lammps.py +++ b/source/lmp/tests/test_lammps.py @@ -230,6 +230,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -340,7 +344,7 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 assert np.array(lammps.variables["eatom"].value) == pytest.approx( expected_ae[idx_map] ) @@ -408,7 +412,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 assert np.array(lammps.variables["eatom"].value) == pytest.approx( expected_ae[idx_map] ) @@ -545,7 +549,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -604,7 +608,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value diff --git a/source/lmp/tests/test_lammps_3types.py b/source/lmp/tests/test_lammps_3types.py index f0cbe19ddf..9156914dbc 100644 --- a/source/lmp/tests/test_lammps_3types.py +++ b/source/lmp/tests/test_lammps_3types.py @@ -253,6 +253,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) @@ -320,7 +324,7 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -371,7 +375,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value diff --git a/source/lmp/tests/test_lammps_dpa_jax.py b/source/lmp/tests/test_lammps_dpa_jax.py index 4867b5f84e..51b2d56742 100644 --- a/source/lmp/tests/test_lammps_dpa_jax.py +++ b/source/lmp/tests/test_lammps_dpa_jax.py @@ -228,6 +228,10 @@ def setup_module(): + if os.environ.get("ENABLE_JAX", "1") != "1": + pytest.skip( + "Skip test because JAX support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -334,13 +338,17 @@ def test_pair_deepmd_virial(lammps): assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps): lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -367,6 +375,10 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps): lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -385,7 +397,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -405,6 +417,10 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps): relative = 1.0 lammps.pair_style( @@ -434,6 +450,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps): relative = 1.0 lammps.pair_style( @@ -511,7 +531,7 @@ def test_pair_deepmd_virial_real(lammps_real): assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -520,6 +540,10 @@ def test_pair_deepmd_virial_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real): lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -550,6 +574,10 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real): lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -570,7 +598,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -594,6 +622,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): relative = 1.0 lammps_real.pair_style( @@ -627,6 +659,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): relative = 1.0 lammps_real.pair_style( diff --git a/source/lmp/tests/test_lammps_dpa_pt.py b/source/lmp/tests/test_lammps_dpa_pt.py index e66b93e09e..6ba71fe52b 100644 --- a/source/lmp/tests/test_lammps_dpa_pt.py +++ b/source/lmp/tests/test_lammps_dpa_pt.py @@ -226,6 +226,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -330,13 +334,17 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -363,6 +371,10 @@ def test_pair_deepmd_model_devi(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -381,7 +393,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -401,6 +413,10 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -430,6 +446,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -507,7 +527,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -516,6 +536,10 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -546,6 +570,10 @@ def test_pair_deepmd_model_devi_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -566,7 +594,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -590,6 +618,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -623,6 +655,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -682,6 +718,10 @@ def test_pair_deepmd_si(lammps_si) -> None: ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/lmp/tests/test_lammps_dpa_pt_nopbc.py b/source/lmp/tests/test_lammps_dpa_pt_nopbc.py index 563650c714..989a782b5f 100644 --- a/source/lmp/tests/test_lammps_dpa_pt_nopbc.py +++ b/source/lmp/tests/test_lammps_dpa_pt_nopbc.py @@ -224,6 +224,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -328,13 +332,17 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -361,6 +369,10 @@ def test_pair_deepmd_model_devi(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -379,7 +391,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -399,6 +411,10 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -428,6 +444,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -505,7 +525,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -514,6 +534,10 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -544,6 +568,10 @@ def test_pair_deepmd_model_devi_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -564,7 +592,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -588,6 +616,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -621,6 +653,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -680,6 +716,10 @@ def test_pair_deepmd_si(lammps_si) -> None: ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/lmp/tests/test_lammps_dpa_sel_pt.py b/source/lmp/tests/test_lammps_dpa_sel_pt.py index 9ff2883fc1..f65c710409 100644 --- a/source/lmp/tests/test_lammps_dpa_sel_pt.py +++ b/source/lmp/tests/test_lammps_dpa_sel_pt.py @@ -229,6 +229,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -333,13 +337,17 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -366,6 +374,10 @@ def test_pair_deepmd_model_devi(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -384,7 +396,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -404,6 +416,10 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -433,6 +449,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -510,7 +530,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -519,6 +539,10 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -549,6 +573,10 @@ def test_pair_deepmd_model_devi_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -569,7 +597,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -593,6 +621,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -626,6 +658,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( diff --git a/source/lmp/tests/test_lammps_faparam.py b/source/lmp/tests/test_lammps_faparam.py index 4206aa68fb..1a614c3d24 100644 --- a/source/lmp/tests/test_lammps_faparam.py +++ b/source/lmp/tests/test_lammps_faparam.py @@ -140,6 +140,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) @@ -213,7 +217,7 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value diff --git a/source/lmp/tests/test_lammps_jax.py b/source/lmp/tests/test_lammps_jax.py index 5d88cfca12..e3d0e5ce74 100644 --- a/source/lmp/tests/test_lammps_jax.py +++ b/source/lmp/tests/test_lammps_jax.py @@ -228,6 +228,10 @@ def setup_module(): + if os.environ.get("ENABLE_JAX", "1") != "1": + pytest.skip( + "Skip test because JAX support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -332,13 +336,17 @@ def test_pair_deepmd_virial(lammps): assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps): lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -365,6 +373,10 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps): lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -383,7 +395,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -403,6 +415,10 @@ def test_pair_deepmd_model_devi_virial(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps): relative = 1.0 lammps.pair_style( @@ -432,6 +448,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps): relative = 1.0 lammps.pair_style( @@ -509,7 +529,7 @@ def test_pair_deepmd_virial_real(lammps_real): assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -518,6 +538,10 @@ def test_pair_deepmd_virial_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real): lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -548,6 +572,10 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real): lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -568,7 +596,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -592,6 +620,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): relative = 1.0 lammps_real.pair_style( @@ -625,6 +657,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): relative = 1.0 lammps_real.pair_style( @@ -684,6 +720,10 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list): with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/lmp/tests/test_lammps_pd.py b/source/lmp/tests/test_lammps_pd.py index 31ee2e482a..85275c4027 100644 --- a/source/lmp/tests/test_lammps_pd.py +++ b/source/lmp/tests/test_lammps_pd.py @@ -229,6 +229,10 @@ def setup_module(): + if os.environ.get("ENABLE_PADDLE", "1") != "1": + pytest.skip( + "Skip test because Paddle support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -333,13 +337,17 @@ def test_pair_deepmd_virial(lammps): assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii], RTOL, ATOL) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps): lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -366,6 +374,10 @@ def test_pair_deepmd_model_devi(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps): lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -384,7 +396,7 @@ def test_pair_deepmd_model_devi_virial(lammps): assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1], RTOL, ATOL ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -406,6 +418,10 @@ def test_pair_deepmd_model_devi_virial(lammps): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps): relative = 1.0 lammps.pair_style( @@ -435,6 +451,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps): assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps): relative = 1.0 lammps.pair_style( @@ -516,7 +536,7 @@ def test_pair_deepmd_virial_real(lammps_real): RTOL, ATOL, ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -527,6 +547,10 @@ def test_pair_deepmd_virial_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real): lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -557,6 +581,10 @@ def test_pair_deepmd_model_devi_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real): lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -581,7 +609,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): RTOL, ATOL, ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -621,6 +649,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): relative = 1.0 lammps_real.pair_style( @@ -654,6 +686,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real): ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real): relative = 1.0 lammps_real.pair_style( @@ -714,6 +750,10 @@ def test_pair_deepmd_si(lammps_si): ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list): with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/lmp/tests/test_lammps_pt.py b/source/lmp/tests/test_lammps_pt.py index 9aed014b62..f6fb8f949b 100644 --- a/source/lmp/tests/test_lammps_pt.py +++ b/source/lmp/tests/test_lammps_pt.py @@ -226,6 +226,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) write_lmp_data(box, coord, type_OH, data_file) write_lmp_data(box, coord, type_HO, data_type_map_file) write_lmp_data( @@ -330,13 +334,17 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -363,6 +371,10 @@ def test_pair_deepmd_model_devi(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial(lammps) -> None: lammps.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -381,7 +393,7 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps.variables[f"virial{ii}"].value @@ -401,6 +413,10 @@ def test_pair_deepmd_model_devi_virial(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -430,6 +446,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: assert md[3] == pytest.approx(np.sqrt(np.mean(np.square(expected_md_v)))) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -507,7 +527,7 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -516,6 +536,10 @@ def test_pair_deepmd_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -546,6 +570,10 @@ def test_pair_deepmd_model_devi_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: lammps_real.pair_style( f"deepmd {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1 atomic" @@ -566,7 +594,7 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: assert lammps_real.atoms[ii].force == pytest.approx( expected_f[lammps_real.atoms[ii].id - 1] * constants.force_metal2real ) - idx_map = lammps_real.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps_real.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 for ii in range(9): assert np.array( lammps_real.variables[f"virial{ii}"].value @@ -590,6 +618,10 @@ def test_pair_deepmd_model_devi_virial_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -623,6 +655,10 @@ def test_pair_deepmd_model_devi_atomic_relative_real(lammps_real) -> None: ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative_v_real(lammps_real) -> None: relative = 1.0 lammps_real.pair_style( @@ -682,6 +718,10 @@ def test_pair_deepmd_si(lammps_si) -> None: ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/lmp/tests/test_lammps_spin.py b/source/lmp/tests/test_lammps_spin.py index 39e12b03fc..4bced5e3f0 100644 --- a/source/lmp/tests/test_lammps_spin.py +++ b/source/lmp/tests/test_lammps_spin.py @@ -100,6 +100,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data_spin(box, coord, spin, type_NiO, data_file) @@ -172,7 +176,7 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 assert np.array(lammps.variables["eatom"].value) == pytest.approx( expected_ae[idx_map] ) diff --git a/source/lmp/tests/test_lammps_spin_nopbc.py b/source/lmp/tests/test_lammps_spin_nopbc.py index 3ef9a7a355..0147b3da23 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc.py +++ b/source/lmp/tests/test_lammps_spin_nopbc.py @@ -99,6 +99,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_TENSORFLOW", "1") != "1": + pytest.skip( + "Skip test because TensorFlow support is not enabled.", + ) write_lmp_data_spin(box, coord, spin, type_NiO, data_file) diff --git a/source/lmp/tests/test_lammps_spin_nopbc_pt.py b/source/lmp/tests/test_lammps_spin_nopbc_pt.py index 8a4dc7a7f1..236e94b191 100644 --- a/source/lmp/tests/test_lammps_spin_nopbc_pt.py +++ b/source/lmp/tests/test_lammps_spin_nopbc_pt.py @@ -96,6 +96,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) write_lmp_data_spin(box, coord, spin, type_NiO, data_file) @@ -146,6 +150,10 @@ def test_pair_deepmd(lammps) -> None: lammps.run(1) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1" @@ -169,6 +177,10 @@ def test_pair_deepmd_model_devi(lammps) -> None: assert md[9] == pytest.approx(np.mean(expected_md_fm)) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -207,6 +219,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/lmp/tests/test_lammps_spin_pt.py b/source/lmp/tests/test_lammps_spin_pt.py index e15a10ee72..fd46cb6514 100644 --- a/source/lmp/tests/test_lammps_spin_pt.py +++ b/source/lmp/tests/test_lammps_spin_pt.py @@ -96,6 +96,10 @@ def setup_module() -> None: + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip( + "Skip test because PyTorch support is not enabled.", + ) write_lmp_data_spin(box, coord, spin, type_NiO, data_file) @@ -168,7 +172,7 @@ def test_pair_deepmd_virial(lammps) -> None: assert lammps.atoms[ii].force == pytest.approx( expected_f[lammps.atoms[ii].id - 1] ) - idx_map = lammps.lmp.numpy.extract_atom("id") - 1 + idx_map = lammps.lmp.numpy.extract_atom("id")[: coord.shape[0]] - 1 assert np.array(lammps.variables["eatom"].value) == pytest.approx( expected_ae[idx_map] ) @@ -186,6 +190,10 @@ def test_pair_deepmd_virial(lammps) -> None: # ) / constants.nktv2p == pytest.approx(expected_v[idx_map, ii]) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi(lammps) -> None: lammps.pair_style( f"deepspin {pb_file.resolve()} {pb_file2.resolve()} out_file {md_file.resolve()} out_freq 1" @@ -209,6 +217,10 @@ def test_pair_deepmd_model_devi(lammps) -> None: assert md[9] == pytest.approx(np.mean(expected_md_fm)) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: relative = 1.0 lammps.pair_style( @@ -247,6 +259,10 @@ def test_pair_deepmd_model_devi_atomic_relative(lammps) -> None: ("balance_args",), [(["--balance"],), ([],)], ) +@pytest.mark.skipif( + os.environ.get("ENABLE_TENSORFLOW", "1") != "1", + reason="Skip test because TensorFlow support is not enabled.", +) def test_pair_deepmd_mpi(balance_args: list) -> None: with tempfile.NamedTemporaryFile() as f: sp.check_call( diff --git a/source/op/pd/CMakeLists.txt b/source/op/pd/CMakeLists.txt new file mode 100644 index 0000000000..194ceb4061 --- /dev/null +++ b/source/op/pd/CMakeLists.txt @@ -0,0 +1,72 @@ +file(GLOB OP_SRC comm.cc) + +add_library(deepmd_op_pd SHARED ${OP_SRC}) + +if(NOT DEFINED PADDLE_INFERENCE_DIR) + message( + FATAL_ERROR + "please set PADDLE_INFERENCE_DIR with -DPADDLE_INFERENCE_DIR=/path/to/paddle_inference" + ) +endif() + +set(PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH + ${PADDLE_INFERENCE_DIR}/third_party/install) +include_directories(${PADDLE_INFERENCE_DIR}) +include_directories(${PADDLE_INFERENCE_DIR}/paddle/include) +include_directories(${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/protobuf/include) +include_directories(${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/glog/include) +include_directories(${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/gflags/include) +include_directories(${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/xxhash/include) + +set(PADDLE_INFER_LIB ${PADDLE_INFERENCE_DIR}/paddle/lib/libpaddle_inference.so) +set(MKLML_LIB + ${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/mklml/lib/libmklml_intel.so + ${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/mklml/lib/libiomp5.so) + +target_link_libraries(deepmd_op_pd PRIVATE ${LIB_DEEPMD}) + +if(APPLE) + set_target_properties(deepmd_op_pd PROPERTIES INSTALL_RPATH "@loader_path") +else() + set_target_properties(deepmd_op_pd PROPERTIES INSTALL_RPATH "$ORIGIN") +endif() + +find_package(MPI) +if(MPI_FOUND) + include(CheckCXXSymbolExists) + set(CMAKE_REQUIRED_INCLUDES ${MPI_CXX_INCLUDE_DIRS}) + set(CMAKE_REQUIRED_LIBRARIES ${MPI_CXX_LIBRARIES}) + check_cxx_symbol_exists(MPIX_Query_cuda_support "mpi.h" CUDA_AWARE) + if(NOT CUDA_AWARE) + check_cxx_symbol_exists(MPIX_Query_cuda_support "mpi.h;mpi-ext.h" OMP_CUDA) + if(NOT OMP_CUDA) + target_compile_definitions(deepmd_op_pd PRIVATE NO_CUDA_AWARE) + endif() + endif() + target_link_libraries(deepmd_op_pd PRIVATE MPI::MPI_CXX) + target_compile_definitions(deepmd_op_pd PRIVATE USE_MPI) +endif() + +if(CMAKE_TESTING_ENABLED) + target_link_libraries(deepmd_op_pd PRIVATE coverage_config) +endif() + +target_link_libraries( + deepmd_op_pd + PRIVATE ${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/glog/lib/libglog.a + ${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/gflags/lib/libgflags.a + ${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/xxhash/lib/libxxhash.a + ${PADDLE_INFERENCE_DIR_THIRD_PARTY_PATH}/protobuf/lib/libprotobuf.a + ${PADDLE_INFER_LIB} + ${MKLML_LIB} + dl + pthread) + +if(BUILD_PY_IF) + install(TARGETS deepmd_op_pd DESTINATION deepmd/lib/) +else(BUILD_PY_IF) + install( + TARGETS deepmd_op_pd + EXPORT DeePMDTargets + DESTINATION lib/) +endif(BUILD_PY_IF) diff --git a/source/op/pd/comm.cc b/source/op/pd/comm.cc new file mode 100644 index 0000000000..548e5db83a --- /dev/null +++ b/source/op/pd/comm.cc @@ -0,0 +1,479 @@ +// SPDX-License-Identifier: LGPL-3.0-or-later + +#ifdef USE_MPI +#include +#ifdef OMPI_MPI_H +#include +#endif +#endif +#include + +#include "paddle/extension.h" + +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) +#include "device.h" +#endif + +#ifdef USE_MPI +template +static MPI_Datatype get_mpi_type(); + +template <> +MPI_Datatype get_mpi_type() { + return MPI_FLOAT; +} + +template <> +MPI_Datatype get_mpi_type() { + return MPI_DOUBLE; +} +#endif + +#ifdef USE_MPI +static void unpack_communicator(const paddle::Tensor& communicator_tensor, + MPI_Comm& mpi_comm) { +#ifdef OMPI_MPI_H + const int64_t* communicator = communicator_tensor.data(); +#else + const int64_t* ptr = communicator_tensor.data(); + const int* communicator = reinterpret_cast(ptr); +#endif + mpi_comm = reinterpret_cast(*communicator); +} +#endif + +template +void Border_forward_t(const paddle::Tensor& sendlist_tensor, + const paddle::Tensor& sendproc_tensor, + const paddle::Tensor& recvproc_tensor, + const paddle::Tensor& sendnum_tensor, + const paddle::Tensor& recvnum_tensor, + paddle::Tensor& g1, + const paddle::Tensor& communicator_tensor, + const paddle::Tensor& nlocal_tensor, + const paddle::Tensor& nghost_tensor) { + int64_t send_list_len = sendlist_tensor.numel(); + + paddle::Tensor cpu_sendlist = paddle::empty( + {send_list_len}, paddle::DataType::INT64, paddle::CPUPlace()); + cpu_sendlist.copy_(sendlist_tensor, paddle::CPUPlace(), true); + int64_t* sendlist = cpu_sendlist.data(); + + int nswap = sendproc_tensor.dims()[0]; + + paddle::Tensor cpu_sendproc = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_sendproc.copy_(sendproc_tensor, paddle::CPUPlace(), true); + int* sendproc = cpu_sendproc.data(); + + paddle::Tensor cpu_recvproc = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_recvproc.copy_(recvproc_tensor, paddle::CPUPlace(), true); + int* recvproc = cpu_recvproc.data(); + + paddle::Tensor cpu_sendnum = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_sendnum.copy_(sendnum_tensor, paddle::CPUPlace(), true); + int* sendnum = cpu_sendnum.data(); + + paddle::Tensor cpu_recvnum = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_recvnum.copy_(recvnum_tensor, paddle::CPUPlace(), true); + int* recvnum = cpu_recvnum.data(); + + int tensor_size = g1.dims()[1]; + + paddle::Tensor cpu_nlocal = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_nlocal.copy_(nlocal_tensor, paddle::CPUPlace(), true); + int nlocal = *(cpu_nlocal.data()); + + paddle::Tensor cpu_nghost = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_nghost.copy_(nghost_tensor, paddle::CPUPlace(), true); + int nghost = *(cpu_nghost.data()); + + int ntotal = nlocal + nghost; + + paddle::Tensor recv_g1_tensor = g1; + +#ifdef USE_MPI + // MPI initialization check + int mpi_init = 0; + MPI_Initialized(&mpi_init); + int cuda_aware = 1; + int me = 0; + MPI_Comm world; + int world_size = 0; + + if (mpi_init) { + unpack_communicator(communicator_tensor, world); + MPI_Comm_rank(world, &me); + MPI_Comm_size(world, &world_size); + } + + MPI_Datatype mpi_type = get_mpi_type(); + MPI_Request request; + +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + if (world_size >= 1) { + int version, subversion; + MPI_Get_version(&version, &subversion); + if (version >= 4) { +#ifdef NO_CUDA_AWARE + cuda_aware = 0; +#else + cuda_aware = MPIX_Query_cuda_support(); +#endif + } else { + cuda_aware = 0; + } + + if (cuda_aware == 0) { + recv_g1_tensor = paddle::empty_like(g1, g1.dtype(), paddle::CPUPlace()); + recv_g1_tensor.copy_(g1, recv_g1_tensor.place(), true); + } + } +#endif + +#endif // USE_MPI + FPTYPE* recv_g1 = recv_g1_tensor.data() + nlocal * tensor_size; + + for (int iswap = 0; iswap < nswap; ++iswap) { + int nrecv = recvnum[iswap]; + int nsend = sendnum[iswap]; + paddle::Tensor isendlist; + paddle::Tensor send_g1_tensor; + FPTYPE* send_g1 = nullptr; + + if (nsend != 0) { + std::intptr_t addr = static_cast(sendlist[iswap]); + int* isendlist_ptr = reinterpret_cast(addr); + isendlist = + paddle::from_blob(isendlist_ptr, {nsend}, paddle::DataType::INT32, + phi::DataLayout::NCHW, paddle::CPUPlace()) + .copy_to(recv_g1_tensor.place(), true); + send_g1_tensor = + paddle::experimental::index_select(recv_g1_tensor, isendlist, 0); + send_g1 = send_g1_tensor.data(); + } + +#ifdef USE_MPI + if (sendproc[iswap] != me) { + if (nrecv) { + MPI_Irecv(recv_g1, nrecv * tensor_size, mpi_type, recvproc[iswap], 0, + world, &request); + } + if (nsend) { + MPI_Send(send_g1, nsend * tensor_size, mpi_type, sendproc[iswap], 0, + world); + } + if (nrecv) { + MPI_Wait(&request, MPI_STATUS_IGNORE); + } + } else { +#endif + +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) +#ifdef USE_MPI + if (cuda_aware == 0) { + memcpy(recv_g1, send_g1, + (unsigned long)nsend * tensor_size * sizeof(FPTYPE)); + } else { + gpuMemcpy(recv_g1, send_g1, + (unsigned long)nsend * tensor_size * sizeof(FPTYPE), + gpuMemcpyDeviceToDevice); + } +#else + gpuMemcpy(recv_g1, send_g1, + (unsigned long)nsend * tensor_size * sizeof(FPTYPE), + gpuMemcpyDeviceToDevice); +#endif + +#else + memcpy(recv_g1, send_g1, + (unsigned long)nsend * tensor_size * sizeof(FPTYPE)); +#endif + +#ifdef USE_MPI + } +#endif + recv_g1 += nrecv * tensor_size; + } + +#ifdef USE_MPI +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + if (cuda_aware == 0) { + g1.copy_(recv_g1_tensor, g1.place(), true); + } +#endif +#endif +} + +void Border_forward(const paddle::Tensor& sendlist_tensor, + const paddle::Tensor& sendproc_tensor, + const paddle::Tensor& recvproc_tensor, + const paddle::Tensor& sendnum_tensor, + const paddle::Tensor& recvnum_tensor, + paddle::Tensor& g1_tensor, + const paddle::Tensor& communicator_tensor, + const paddle::Tensor& nlocal_tensor, + const paddle::Tensor& nghost_tensor) { + bool type_flag = (g1_tensor.dtype() == phi::DataType::FLOAT64) ? true : false; + if (type_flag) { + Border_forward_t(sendlist_tensor, sendproc_tensor, recvproc_tensor, + sendnum_tensor, recvnum_tensor, g1_tensor, + communicator_tensor, nlocal_tensor, nghost_tensor); + } else { + Border_forward_t(sendlist_tensor, sendproc_tensor, recvproc_tensor, + sendnum_tensor, recvnum_tensor, g1_tensor, + communicator_tensor, nlocal_tensor, nghost_tensor); + } +} + +template +void Border_backward_t(const paddle::Tensor& sendlist_tensor, + const paddle::Tensor& sendproc_tensor, + const paddle::Tensor& recvproc_tensor, + const paddle::Tensor& sendnum_tensor, + const paddle::Tensor& recvnum_tensor, + const paddle::Tensor& g1_tensor, + const paddle::Tensor& communicator_tensor, + const paddle::Tensor& nlocal_tensor, + const paddle::Tensor& nghost_tensor, + paddle::Tensor& recv_g1_tensor_grad) { +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + gpuDeviceSynchronize(); +#endif + paddle::Tensor d_local_g1_tensor = + paddle::empty(recv_g1_tensor_grad.shape(), recv_g1_tensor_grad.dtype(), + recv_g1_tensor_grad.place()); + d_local_g1_tensor.copy_(recv_g1_tensor_grad.contiguous(), + d_local_g1_tensor.place(), true); + +#ifdef USE_MPI + int mpi_init = 0, world_size = 0, me = 0, cuda_aware = 1; + MPI_Initialized(&mpi_init); + + MPI_Comm world; + if (mpi_init) { + unpack_communicator(communicator_tensor, world); + MPI_Comm_rank(world, &me); + MPI_Comm_size(world, &world_size); + } + + auto mpi_type = get_mpi_type(); + MPI_Request request; + +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + if (world_size >= 1) { + int version, subversion; + MPI_Get_version(&version, &subversion); + + if (version >= 4) { +#ifdef NO_CUDA_AWARE + cuda_aware = 0; +#else + cuda_aware = MPIX_Query_cuda_support(); +#endif + } else { + cuda_aware = 0; + } + + if (cuda_aware == 0) { + d_local_g1_tensor = paddle::empty_like( + recv_g1_tensor_grad, recv_g1_tensor_grad.dtype(), paddle::CPUPlace()); + d_local_g1_tensor.copy_(recv_g1_tensor_grad, d_local_g1_tensor.place(), + true); + } + } +#endif +#endif // USE_MPI + int64_t send_list_len = sendlist_tensor.numel(); + paddle::Tensor cpu_sendlist = paddle::empty( + {send_list_len}, paddle::DataType::INT64, paddle::CPUPlace()); + cpu_sendlist.copy_(sendlist_tensor, paddle::CPUPlace(), true); + int64_t* recvlist = cpu_sendlist.data(); + + int nswap = sendproc_tensor.dims()[0]; + // swap send and recv here + paddle::Tensor cpu_recvproc = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_recvproc.copy_(recvproc_tensor, paddle::CPUPlace(), true); + int* recvproc = cpu_recvproc.data(); + + paddle::Tensor cpu_sendproc = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_sendproc.copy_(sendproc_tensor, paddle::CPUPlace(), true); + int* sendproc = cpu_sendproc.data(); + + paddle::Tensor cpu_sendnum = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_sendnum.copy_(sendnum_tensor, paddle::CPUPlace(), true); + int* recvnum = cpu_sendnum.data(); + + paddle::Tensor cpu_recvnum = + paddle::empty({nswap}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_recvnum.copy_(recvnum_tensor, paddle::CPUPlace(), true); + int* sendnum = cpu_recvnum.data(); + + FPTYPE* local_g1 = d_local_g1_tensor.data(); + int tensor_size = d_local_g1_tensor.dims()[1]; + + paddle::Tensor cpu_nlocal = + paddle::empty({1}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_nlocal.copy_(nlocal_tensor, paddle::CPUPlace(), true); + int nlocal = *cpu_nlocal.data(); + + paddle::Tensor cpu_nghost = + paddle::empty({1}, paddle::DataType::INT32, paddle::CPUPlace()); + cpu_nghost.copy_(nghost_tensor, paddle::CPUPlace(), true); + int nghost = *cpu_nghost.data(); + int ntotal = nlocal + nghost; + + paddle::Tensor send_g1_tensor, recv_g1_tensor; + FPTYPE *recv_g1 = nullptr, *send_g1 = nullptr; + + if (nswap != 0) { + send_g1_tensor = d_local_g1_tensor; + + int max_recvnum = + *(paddle::experimental::max(cpu_sendnum, {}, false).data()); + recv_g1_tensor = + paddle::empty({max_recvnum, tensor_size}, d_local_g1_tensor.dtype(), + d_local_g1_tensor.place()); + recv_g1 = recv_g1_tensor.data(); + send_g1 = send_g1_tensor.data() + ntotal * tensor_size; + } + + for (int iswap = nswap - 1; iswap >= 0; --iswap) { + int nrecv = recvnum[iswap]; + int nsend = sendnum[iswap]; + + paddle::Tensor irecvlist; + if (nrecv) { + std::intptr_t addr = static_cast(recvlist[iswap]); + int* irecvlist_ptr = reinterpret_cast(addr); + irecvlist = + paddle::from_blob(irecvlist_ptr, {nrecv}, paddle::DataType::INT32, + paddle::DataLayout::NCHW, paddle::CPUPlace()) + .copy_to(d_local_g1_tensor.place(), true); + } + + if (nsend) { + send_g1 -= nsend * tensor_size; + } + +#ifdef USE_MPI + if (sendproc[iswap] != me) { + if (nrecv) { + MPI_Irecv(recv_g1, nrecv * tensor_size, mpi_type, recvproc[iswap], 0, + world, &request); + } + if (nsend) { + MPI_Send(send_g1, nsend * tensor_size, mpi_type, sendproc[iswap], 0, + world); + } + if (nrecv) { + MPI_Wait(&request, MPI_STATUS_IGNORE); + } + } else { +#endif + if (nrecv) { +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) +#ifdef USE_MPI + if (cuda_aware == 0) { + memcpy(recv_g1, send_g1, + (unsigned long)nrecv * tensor_size * sizeof(FPTYPE)); + } else { + gpuMemcpy(recv_g1, send_g1, + (unsigned long)nrecv * tensor_size * sizeof(FPTYPE), + gpuMemcpyDeviceToDevice); + } +#else + gpuMemcpy(recv_g1, send_g1, + (unsigned long)nrecv * tensor_size * sizeof(FPTYPE), + gpuMemcpyDeviceToDevice); +#endif +#else + memcpy(recv_g1, send_g1, + (unsigned long)nrecv * tensor_size * sizeof(FPTYPE)); +#endif + } +#ifdef USE_MPI + } +#endif + if (nrecv) { + d_local_g1_tensor = paddle::experimental::index_add_( + d_local_g1_tensor, irecvlist, recv_g1_tensor.slice(0, nrecv), 0); + } + } +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + gpuDeviceSynchronize(); +#endif + +#ifdef USE_MPI +#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) + if (cuda_aware == 0) { + recv_g1_tensor_grad.copy_(d_local_g1_tensor, recv_g1_tensor_grad.place(), + true); + } +#endif +#endif +} + +void Border_backward(const paddle::Tensor& sendlist_tensor, + const paddle::Tensor& sendproc_tensor, + const paddle::Tensor& recvproc_tensor, + const paddle::Tensor& sendnum_tensor, + const paddle::Tensor& recvnum_tensor, + const paddle::Tensor& g1_tensor, + const paddle::Tensor& communicator_tensor, + const paddle::Tensor& nlocal_tensor, + const paddle::Tensor& nghost_tensor, + paddle::Tensor& recv_g1_tensor_grad) { + bool type_flag = + (recv_g1_tensor_grad.dtype() == paddle::DataType::FLOAT64) ? true : false; + if (type_flag) { + Border_backward_t(sendlist_tensor, sendproc_tensor, recvproc_tensor, + sendnum_tensor, recvnum_tensor, g1_tensor, + communicator_tensor, nlocal_tensor, nghost_tensor, + recv_g1_tensor_grad); + } else { + Border_backward_t(sendlist_tensor, sendproc_tensor, recvproc_tensor, + sendnum_tensor, recvnum_tensor, g1_tensor, + communicator_tensor, nlocal_tensor, nghost_tensor, + recv_g1_tensor_grad); + } +} + +/** + * @brief communicate the latest g1_tensor info to other lmp proc + * @param[in] sendlist_tensor list of atoms to send in each swap + * @param[in] sendproc_tensor proc to send to at each swap + * @param[in] recvproc_tensor proc to recv from at each swap + * @param[in] sendnum_tensor # of atoms to send in each swap + * @param[in] recvnum_tensor # of atoms to recv in each swap + * @param[in] g1_tensor tensor to store g1_tensor info + * @param[in] communicator_tensor MPI_comm data in lmp + * @param[in] nlocal_tensor # of local atoms + * @param[in] nghost_tensor # of nghost atoms + * @param[out] recv_g1_tensor g1_tensor after communication + **/ +PD_BUILD_OP(border_op) + .Inputs({"sendlist_tensor", "sendproc_tensor", "recvproc_tensor", + "sendnum_tensor", "recvnum_tensor", "g1_tensor", + "communicator_tensor", "nlocal_tensor", "nghost_tensor"}) + .Outputs({"recv_g1_tensor"}) + .SetKernelFn(PD_KERNEL(Border_forward)) + .SetInplaceMap({{"g1_tensor", "recv_g1_tensor"}}); + +PD_BUILD_GRAD_OP(border_op) + .Inputs({"sendlist_tensor", "sendproc_tensor", "recvproc_tensor", + "sendnum_tensor", "recvnum_tensor", "g1_tensor", + "communicator_tensor", "nlocal_tensor", "nghost_tensor", + paddle::Grad("recv_g1_tensor")}) + .Outputs({paddle::Grad("g1_tensor")}) + .SetInplaceMap({{paddle::Grad("recv_g1_tensor"), + paddle::Grad("g1_tensor")}}) + .SetKernelFn(PD_KERNEL(Border_backward)); diff --git a/source/op/pd/setup.py b/source/op/pd/setup.py new file mode 100644 index 0000000000..951b50de9c --- /dev/null +++ b/source/op/pd/setup.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os + + +def main(): + current_dir = os.path.abspath(os.getcwd()) + script_dir = os.path.abspath(os.path.dirname(__file__)) + + if current_dir != script_dir: + raise RuntimeError( + f"[ERROR] Please run this script under directory: `{script_dir}`" + ) + + from paddle.utils.cpp_extension import ( + CppExtension, + setup, + ) + + setup(name="deepmd_op_pd", ext_modules=CppExtension(sources=["comm.cc"])) + + +if __name__ == "__main__": + main() diff --git a/source/op/pt/comm.cc b/source/op/pt/comm.cc index 71a2b0e118..97466a4833 100644 --- a/source/op/pt/comm.cc +++ b/source/op/pt/comm.cc @@ -86,7 +86,7 @@ class Border : public torch::autograd::Function { #ifdef USE_MPI int mpi_init = 0; MPI_Initialized(&mpi_init); - int cuda_aware = 1; + int cuda_aware = 0; int me = 0; MPI_Comm world; int world_size = 0; @@ -99,17 +99,9 @@ class Border : public torch::autograd::Function { MPI_Request request; #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) if (world_size >= 1) { - int version, subversion; - MPI_Get_version(&version, &subversion); - if (version >= 4) { -#ifdef NO_CUDA_AWARE - cuda_aware = 0; -#else - cuda_aware = MPIX_Query_cuda_support(); +#ifndef NO_CUDA_AWARE + cuda_aware = MPIX_Query_cuda_support(); #endif - } else { - cuda_aware = 0; - } if (cuda_aware == 0) { recv_g1_tensor = torch::empty_like(g1).to(torch::kCPU); recv_g1_tensor.copy_(g1); @@ -193,10 +185,6 @@ class Border : public torch::autograd::Function { static torch::autograd::variable_list backward_t( torch::autograd::AutogradContext* ctx, torch::autograd::variable_list grad_output) { -#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) - gpuDeviceSynchronize(); -#endif - torch::autograd::variable_list saved_variables = ctx->get_saved_variables(); torch::Tensor sendlist_tensor = saved_variables[0]; torch::Tensor sendproc_tensor = saved_variables[1]; @@ -212,7 +200,7 @@ class Border : public torch::autograd::Function { int mpi_init = 0; MPI_Initialized(&mpi_init); int world_size = 0; - int cuda_aware = 1; + int cuda_aware = 0; int me = 0; MPI_Comm world; if (mpi_init) { @@ -224,17 +212,9 @@ class Border : public torch::autograd::Function { MPI_Request request; #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) if (world_size >= 1) { - int version, subversion; - MPI_Get_version(&version, &subversion); - if (version >= 4) { -#ifdef NO_CUDA_AWARE - cuda_aware = 0; -#else - cuda_aware = MPIX_Query_cuda_support(); +#ifndef NO_CUDA_AWARE + cuda_aware = MPIX_Query_cuda_support(); #endif - } else { - cuda_aware = 0; - } if (cuda_aware == 0) { d_local_g1_tensor = torch::empty_like(grad_output[0]).to(torch::kCPU); d_local_g1_tensor.copy_(grad_output[0]); @@ -329,9 +309,6 @@ class Border : public torch::autograd::Function { recv_g1_tensor.slice(0, 0, nrecv)); } } -#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) - gpuDeviceSynchronize(); -#endif #ifdef USE_MPI #if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM) if (cuda_aware == 0) { diff --git a/source/op/pt/tabulate_multi_device.cc b/source/op/pt/tabulate_multi_device.cc index feae37af81..530e9ddf4d 100644 --- a/source/op/pt/tabulate_multi_device.cc +++ b/source/op/pt/tabulate_multi_device.cc @@ -335,6 +335,150 @@ void TabulateFusionSeTGradGradForward(const torch::Tensor& table_tensor, } } +template +void TabulateFusionSeTTebdForward(const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + int64_t last_layer_size, + torch::Tensor& descriptor_tensor) { + // check input shape + if (table_tensor.dim() != 2) { + throw std::invalid_argument("Dim of table should be 2"); + } + if (em_x_tensor.dim() != 2) { + throw std::invalid_argument("Dim of em_x should be 2"); + } + if (em_tensor.dim() != 3) { + throw std::invalid_argument("Dim of em should be 3"); + } + // get the device + std::string device; + GetTensorDevice(table_tensor, device); + // flat the tensors + FPTYPE* descriptor = descriptor_tensor.view({-1}).data_ptr(); + + const FPTYPE* table = table_tensor.view({-1}).data_ptr(); + const FPTYPE* table_info = table_info_tensor.view({-1}).data_ptr(); + const FPTYPE* em_x = em_x_tensor.view({-1}).data_ptr(); + const FPTYPE* em = em_tensor.view({-1}).data_ptr(); + + const int64_t nloc = em_tensor.size(0); + const int64_t nnei_i = em_tensor.size(1); + const int64_t nnei_j = em_tensor.size(2); + // compute + if (device == "GPU") { +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + deepmd::tabulate_fusion_se_t_tebd_gpu(descriptor, table, table_info, em_x, + em, nloc, nnei_i, nnei_j, + last_layer_size); +#else + throw std::runtime_error( + "The input tensor is on the GPU, but the GPU support for the " + "customized OP library is not enabled."); +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + } else if (device == "CPU") { + deepmd::tabulate_fusion_se_t_tebd_cpu(descriptor, table, table_info, em_x, + em, nloc, nnei_i, nnei_j, + last_layer_size); + } +} + +template +void TabulateFusionSeTTebdGradForward(const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + const torch::Tensor& dy_tensor, + const torch::Tensor& descriptor_tensor, + torch::Tensor& dy_dem_x_tensor) { + // check input shape + if (dy_tensor.dim() != 4) { + throw std::invalid_argument("Dim of dy_tensor should be 4"); + } + std::string device; + GetTensorDevice(table_tensor, device); + // flat the tensors + FPTYPE* dy_dem_x = dy_dem_x_tensor.view({-1}).data_ptr(); + + const FPTYPE* table = table_tensor.view({-1}).data_ptr(); + const FPTYPE* table_info = table_info_tensor.view({-1}).data_ptr(); + const FPTYPE* em_x = em_x_tensor.view({-1}).data_ptr(); + const FPTYPE* em = em_tensor.view({-1}).data_ptr(); + const FPTYPE* dy = dy_tensor.view({-1}).data_ptr(); + + const int64_t nloc = em_tensor.size(0); + const int64_t nnei_i = em_tensor.size(1); + const int64_t nnei_j = em_tensor.size(2); + const int64_t last_layer_size = descriptor_tensor.size(3); + + // compute + if (device == "GPU") { +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + deepmd::tabulate_fusion_se_t_tebd_grad_gpu(dy_dem_x, table, table_info, + em_x, em, dy, nloc, nnei_i, + nnei_j, last_layer_size); +#else + throw std::runtime_error( + "The input tensor is on the GPU, but the GPU support for the " + "customized OP library is not enabled."); +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + } else if (device == "CPU") { + deepmd::tabulate_fusion_se_t_tebd_grad_cpu(dy_dem_x, table, table_info, + em_x, em, dy, nloc, nnei_i, + nnei_j, last_layer_size); + } +} + +template +void TabulateFusionSeTTebdGradGradForward( + const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + const torch::Tensor& dz_dy_dem_x_tensor, + const torch::Tensor& descriptor_tensor, + torch::Tensor& dz_dy_tensor) { + // Check input shape + if (dz_dy_dem_x_tensor.dim() != 3) { + throw std::invalid_argument("Dim of dz_dy_dem_x should be 3"); + } + // get the device + std::string device; + GetTensorDevice(table_tensor, device); + // flat the tensors + FPTYPE* dz_dy = dz_dy_tensor.view({-1}).data_ptr(); + + const FPTYPE* table = table_tensor.view({-1}).data_ptr(); + const FPTYPE* table_info = table_info_tensor.view({-1}).data_ptr(); + const FPTYPE* em_x = em_x_tensor.view({-1}).data_ptr(); + const FPTYPE* em = em_tensor.view({-1}).data_ptr(); + const FPTYPE* dz_dy_dem_x = dz_dy_dem_x_tensor.view({-1}).data_ptr(); + const int64_t nloc = em_tensor.size(0); + const int64_t nnei_i = em_tensor.size(1); + const int64_t nnei_j = em_tensor.size(2); + const int64_t last_layer_size = descriptor_tensor.size(3); + // compute + if (device == "GPU") { +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM + deepmd::tabulate_fusion_se_t_tebd_grad_grad_gpu( + dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j, + last_layer_size); +#else + throw std::runtime_error( + "The input tensor is on the GPU, but the GPU support for the " + "customized OP library is not enabled."); +#endif // GOOGLE_CUDA || TENSORFLOW_USE_ROCM + TORCH_CHECK(last_layer_size <= 1024, + "In the process of model compression, the size of the " + "last layer of embedding net must be less than 1024!"); + } else if (device == "CPU") { + deepmd::tabulate_fusion_se_t_tebd_grad_grad_cpu( + dz_dy, table, table_info, em_x, em, dz_dy_dem_x, nloc, nnei_i, nnei_j, + last_layer_size); + } +} + template void TabulateFusionSeRForward(const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, @@ -939,6 +1083,90 @@ class TabulateFusionSeROp } }; +class TabulateFusionSeTTebdOp + : public torch::autograd::Function { + public: + static std::vector forward( + torch::autograd::AutogradContext* ctx, + const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + int64_t last_layer_size) { + bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false; + if (type_flag) { + return forward_t(ctx, table_tensor, table_info_tensor, + em_x_tensor, em_tensor, last_layer_size); + } else { + return forward_t(ctx, table_tensor, table_info_tensor, em_x_tensor, + em_tensor, last_layer_size); + } + } + + template + static torch::autograd::variable_list forward_t( + torch::autograd::AutogradContext* ctx, + const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + int64_t last_layer_size) { + // allocate output tensors + auto options = torch::TensorOptions() + .dtype(table_tensor.dtype()) + .device(table_tensor.device()); + torch::Tensor descriptor_tensor = + torch::empty({em_tensor.size(0), em_tensor.size(1), em_tensor.size(2), + last_layer_size}, + options); + // compute + TabulateFusionSeTTebdForward(table_tensor, table_info_tensor, + em_x_tensor, em_tensor, + last_layer_size, descriptor_tensor); + // save data + ctx->save_for_backward({table_tensor, table_info_tensor, em_x_tensor, + em_tensor, descriptor_tensor}); + return {descriptor_tensor}; + } + + static torch::autograd::variable_list backward( + torch::autograd::AutogradContext* ctx, + torch::autograd::variable_list grad_output) { + torch::autograd::variable_list saved_variables = ctx->get_saved_variables(); + torch::Tensor table_tensor = saved_variables[0]; + bool type_flag = (table_tensor.dtype() == torch::kDouble) ? true : false; + if (type_flag) { + return backward_t(ctx, grad_output); + } else { + return backward_t(ctx, grad_output); + } + } + + template + static torch::autograd::variable_list backward_t( + torch::autograd::AutogradContext* ctx, + torch::autograd::variable_list grad_output) { + // load data + torch::autograd::variable_list saved_variables = ctx->get_saved_variables(); + torch::Tensor table_tensor = saved_variables[0]; + torch::Tensor table_info_tensor = saved_variables[1]; + torch::Tensor em_x_tensor = saved_variables[2]; + torch::Tensor em_tensor = saved_variables[3]; + torch::Tensor descriptor_tensor = saved_variables[4]; + + torch::Tensor dy_tensor = grad_output[0].contiguous(); + // allocate output tensors + torch::Tensor dy_dem_x_tensor = torch::zeros_like(em_x_tensor); + // compute + TabulateFusionSeTTebdGradForward( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, dy_tensor, + descriptor_tensor, dy_dem_x_tensor); + + return {at::Tensor(), at::Tensor(), dy_dem_x_tensor, at::Tensor(), + at::Tensor()}; + } +}; + std::vector tabulate_fusion_se_a( const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, // only cpu @@ -972,6 +1200,16 @@ std::vector tabulate_fusion_se_t( em_x_tensor, em_tensor, last_layer_size); } +std::vector tabulate_fusion_se_t_tebd( + const torch::Tensor& table_tensor, + const torch::Tensor& table_info_tensor, // only cpu + const torch::Tensor& em_x_tensor, + const torch::Tensor& em_tensor, + int64_t last_layer_size) { + return TabulateFusionSeTTebdOp::apply( + table_tensor, table_info_tensor, em_x_tensor, em_tensor, last_layer_size); +} + std::vector tabulate_fusion_se_r( const torch::Tensor& table_tensor, const torch::Tensor& table_info_tensor, // only cpu @@ -990,6 +1228,9 @@ TORCH_LIBRARY_FRAGMENT(deepmd, m) { TORCH_LIBRARY_FRAGMENT(deepmd, m) { m.def("tabulate_fusion_se_t", tabulate_fusion_se_t); } +TORCH_LIBRARY_FRAGMENT(deepmd, m) { + m.def("tabulate_fusion_se_t_tebd", tabulate_fusion_se_t_tebd); +} TORCH_LIBRARY_FRAGMENT(deepmd, m) { m.def("tabulate_fusion_se_r", tabulate_fusion_se_r); } diff --git a/source/op/tf/descrpt_se_a_mask.cc b/source/op/tf/descrpt_se_a_mask.cc index 28e4a575db..7f8bcd9411 100644 --- a/source/op/tf/descrpt_se_a_mask.cc +++ b/source/op/tf/descrpt_se_a_mask.cc @@ -32,7 +32,7 @@ struct NeighborInfo { int index; NeighborInfo() : type(0), dist(0), index(0) {} NeighborInfo(int tt, FPTYPE dd, int ii) : type(tt), dist(dd), index(ii) {} - bool operator<(const NeighborInfo &b) const { + bool operator<(const NeighborInfo& b) const { return (type < b.type || (type == b.type && (dist < b.dist || (dist == b.dist && index < b.index)))); @@ -42,24 +42,24 @@ struct NeighborInfo { template class DescrptSeAMaskOp : public OpKernel { public: - explicit DescrptSeAMaskOp(OpKernelConstruction *context) : OpKernel(context) { + explicit DescrptSeAMaskOp(OpKernelConstruction* context) : OpKernel(context) { // OP_REQUIRES_OK(context); } - void Compute(OpKernelContext *context) override { + void Compute(OpKernelContext* context) override { deepmd::safe_compute( - context, [this](OpKernelContext *context) { this->_Compute(context); }); + context, [this](OpKernelContext* context) { this->_Compute(context); }); } - void _Compute(OpKernelContext *context) { + void _Compute(OpKernelContext* context) { // Grab the input tensor int context_input_index = 0; - const Tensor &coord_tensor = context->input(context_input_index++); - const Tensor &type_tensor = context->input(context_input_index++); - const Tensor &mask_matrix_tensor = context->input(context_input_index++); - const Tensor &box_tensor = context->input(context_input_index++); - const Tensor &natoms_tensor = context->input(context_input_index++); - const Tensor &mesh_tensor = context->input(context_input_index++); + const Tensor& coord_tensor = context->input(context_input_index++); + const Tensor& type_tensor = context->input(context_input_index++); + const Tensor& mask_matrix_tensor = context->input(context_input_index++); + const Tensor& box_tensor = context->input(context_input_index++); + const Tensor& natoms_tensor = context->input(context_input_index++); + const Tensor& mesh_tensor = context->input(context_input_index++); // set size of the sample OP_REQUIRES(context, (coord_tensor.shape().dims() == 2), @@ -109,18 +109,18 @@ class DescrptSeAMaskOp : public OpKernel { nlist_shape.AddDim(static_cast(total_atom_num) * total_atom_num); int context_output_index = 0; - Tensor *descrpt_tensor = NULL; + Tensor* descrpt_tensor = NULL; OP_REQUIRES_OK( context, context->allocate_output(context_output_index++, descrpt_shape, &descrpt_tensor)); - Tensor *descrpt_deriv_tensor = NULL; + Tensor* descrpt_deriv_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, descrpt_deriv_shape, &descrpt_deriv_tensor)); - Tensor *rij_tensor = NULL; + Tensor* rij_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, rij_shape, &rij_tensor)); - Tensor *nlist_tensor = NULL; + Tensor* nlist_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(context_output_index++, nlist_shape, &nlist_tensor)); @@ -317,9 +317,9 @@ class DescrptSeAMaskOp : public OpKernel { compute_t max_distance = 10000.0; void buildAndSortNeighborList(int i_idx, const std::vector d_coord3, - std::vector &d_type, - std::vector &d_mask, - std::vector &sorted_nlist, + std::vector& d_type, + std::vector& d_mask, + std::vector& sorted_nlist, int total_atom_num) { // sorted_nlist.resize(total_atom_num); std::vector> sel_nei; diff --git a/source/op/tf/dotmul_flt_nvnmd.cc b/source/op/tf/dotmul_flt_nvnmd.cc index 1aca3e8bf8..ecfac60a0a 100644 --- a/source/op/tf/dotmul_flt_nvnmd.cc +++ b/source/op/tf/dotmul_flt_nvnmd.cc @@ -37,15 +37,15 @@ modw = 1: normalize w[hh, : , kk] using namespace tensorflow; template -void split_flt(T x, int64_t &sign, int64_t &expo, int64_t &mant); +void split_flt(T x, int64_t& sign, int64_t& expo, int64_t& mant); // read matmul_flt_nvnmd.cc template // float and double -void find_max_expo(int64_t &max_expo, T *x, int64_t M); +void find_max_expo(int64_t& max_expo, T* x, int64_t M); // read matmul_flt_nvnmd.cc template // float and double -void find_max_expo(int64_t &max_expo, T *x, int64_t N, int64_t M); +void find_max_expo(int64_t& max_expo, T* x, int64_t N, int64_t M); //- register the operator REGISTER_OP("DotmulFltNvnmd") @@ -60,19 +60,19 @@ template class DotmulFltNvnmdOp : public OpKernel { public: /// Constructor. - explicit DotmulFltNvnmdOp(OpKernelConstruction *context) + explicit DotmulFltNvnmdOp(OpKernelConstruction* context) : OpKernel(context) {}; /// Compute the descriptor /// param: context - void Compute(OpKernelContext *context) override { + void Compute(OpKernelContext* context) override { // check DCHECK_EQ(2, context->num_inputs()); - const Tensor &X = context->input(0); - const Tensor &W = context->input(1); + const Tensor& X = context->input(0); + const Tensor& W = context->input(1); - const TensorShape &shX = X.shape(); - const TensorShape &shW = W.shape(); + const TensorShape& shX = X.shape(); + const TensorShape& shW = W.shape(); TensorShape shY; DCHECK_EQ(shW.dims(), shX.dims()); @@ -104,7 +104,7 @@ class DotmulFltNvnmdOp : public OpKernel { } // create output - Tensor *Y = NULL; + Tensor* Y = NULL; OP_REQUIRES_OK(context, context->allocate_output(0, shY, &Y)); // compute @@ -131,8 +131,8 @@ class DotmulFltNvnmdOp : public OpKernel { for (ii = 0; ii < H * N; ii++) { // find x max exponnet - find_max_expo(expo_max1, (FPTYPE *)&x[ii * M], M); - find_max_expo(expo_max2, (FPTYPE *)&w[ii * M], M); + find_max_expo(expo_max1, (FPTYPE*)&x[ii * M], M); + find_max_expo(expo_max2, (FPTYPE*)&w[ii * M], M); // s = 0; for (jj = 0; jj < M; jj++) { diff --git a/source/op/tf/matmul_flt_nvnmd.cc b/source/op/tf/matmul_flt_nvnmd.cc index 22ed23c0a3..c2821096c1 100644 --- a/source/op/tf/matmul_flt_nvnmd.cc +++ b/source/op/tf/matmul_flt_nvnmd.cc @@ -37,15 +37,15 @@ modw = 1: normalize w[hh, : , kk] using namespace tensorflow; template -void split_flt(T x, int64_t &sign, int64_t &expo, int64_t &mant); +void split_flt(T x, int64_t& sign, int64_t& expo, int64_t& mant); // read matmul_flt_nvnmd.cc template // float and double -void find_max_expo(int64_t &max_expo, T *x, int64_t M); +void find_max_expo(int64_t& max_expo, T* x, int64_t M); // read matmul_flt_nvnmd.cc template // float and double -void find_max_expo(int64_t &max_expo, T *x, int64_t N, int64_t M); +void find_max_expo(int64_t& max_expo, T* x, int64_t N, int64_t M); //- register the operator REGISTER_OP("MatmulFltNvnmd") @@ -62,21 +62,21 @@ template class MatmulFltNvnmdOp : public OpKernel { public: /// Constructor. - explicit MatmulFltNvnmdOp(OpKernelConstruction *context) : OpKernel(context) { + explicit MatmulFltNvnmdOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("normx", &normx)); OP_REQUIRES_OK(context, context->GetAttr("normw", &normw)); }; /// Compute the descriptor /// param: context - void Compute(OpKernelContext *context) override { + void Compute(OpKernelContext* context) override { // check DCHECK_EQ(2, context->num_inputs()); - const Tensor &X = context->input(0); - const Tensor &W = context->input(1); + const Tensor& X = context->input(0); + const Tensor& W = context->input(1); - const TensorShape &shX = X.shape(); - const TensorShape &shW = W.shape(); + const TensorShape& shX = X.shape(); + const TensorShape& shW = W.shape(); TensorShape shY; DCHECK_EQ(shW.dims(), shX.dims()); @@ -103,7 +103,7 @@ class MatmulFltNvnmdOp : public OpKernel { } // create output - Tensor *Y = NULL; + Tensor* Y = NULL; OP_REQUIRES_OK(context, context->allocate_output(0, shY, &Y)); // compute @@ -130,7 +130,7 @@ class MatmulFltNvnmdOp : public OpKernel { for (hh = 0; hh < H; hh++) { // find x max exponnet if ((normx & 0x0f) == 0) { // normalize x[:,:] - find_max_expo(expo_max1, (FPTYPE *)&x[hh * N * M], + find_max_expo(expo_max1, (FPTYPE*)&x[hh * N * M], static_cast(N) * M); for (ii = 0; ii < N; ii++) { expo_max1s[ii] = expo_max1; @@ -138,14 +138,14 @@ class MatmulFltNvnmdOp : public OpKernel { } else { // normalize x[ii,:] for (ii = 0; ii < N; ii++) { - find_max_expo(expo_max1, (FPTYPE *)&x[hh * N * M + ii * M], M); + find_max_expo(expo_max1, (FPTYPE*)&x[hh * N * M + ii * M], M); expo_max1s[ii] = expo_max1; } } // find w max exponnet if ((normw & 0x0f) == 0) { // normalize w[:,:] - find_max_expo(expo_max2, (FPTYPE *)&w[hh * M * K], + find_max_expo(expo_max2, (FPTYPE*)&w[hh * M * K], static_cast(M) * K); for (kk = 0; kk < K; kk++) { expo_max2s[kk] = expo_max2; @@ -153,7 +153,7 @@ class MatmulFltNvnmdOp : public OpKernel { } else { // normalize w[:,kk] for (kk = 0; kk < K; kk++) { - find_max_expo(expo_max2, (FPTYPE *)&w[hh * M * K + kk], M, K); + find_max_expo(expo_max2, (FPTYPE*)&w[hh * M * K + kk], M, K); expo_max2s[kk] = expo_max2; } } diff --git a/source/op/tf/optimizer/parallel.cc b/source/op/tf/optimizer/parallel.cc index f5b7c62b6a..87a53b18ae 100644 --- a/source/op/tf/optimizer/parallel.cc +++ b/source/op/tf/optimizer/parallel.cc @@ -27,7 +27,7 @@ // based on tensorflow/core/grappler/optimizers/remapper.cc struct RemapperContext { - explicit RemapperContext(GrapplerItem *item, Status *status) + explicit RemapperContext(GrapplerItem* item, Status* status) : nodes_to_preserve(item->NodesToPreserve()), graph_view(&item->graph, status) {} @@ -35,11 +35,11 @@ struct RemapperContext { utils::MutableGraphView graph_view; }; -bool IsProdForce(const NodeDef &node) { return node.op() == "ProdForceSeA"; } +bool IsProdForce(const NodeDef& node) { return node.op() == "ProdForceSeA"; } -bool FindProdForce(RemapperContext *ctx, int node_index) { - const auto *node_view = ctx->graph_view.GetNode(node_index); - const auto *node_def = node_view->node(); +bool FindProdForce(RemapperContext* ctx, int node_index) { + const auto* node_view = ctx->graph_view.GetNode(node_index); + const auto* node_def = node_view->node(); return IsProdForce(*node_def); } @@ -55,17 +55,17 @@ TF_INT64 GetNThreads() { return tot; } -Status ParallelProdForce(RemapperContext *ctx, +Status ParallelProdForce(RemapperContext* ctx, int node_index, - std::vector *invalidated_nodes, - std::vector *nodes_to_delete) { + std::vector* invalidated_nodes, + std::vector* nodes_to_delete) { // skip on GPUs if (GetNumAvailableGPUs() > 0) { return Status(); } - const NodeDef *ori_node = ctx->graph_view.GetNode(node_index)->node(); - auto &src_attr = ori_node->attr(); + const NodeDef* ori_node = ctx->graph_view.GetNode(node_index)->node(); + auto& src_attr = ori_node->attr(); TF_INT64 tot = GetNThreads(); if (tot <= 1) { return Status(); @@ -75,11 +75,11 @@ Status ParallelProdForce(RemapperContext *ctx, sum_node.set_name(ori_node->name()); sum_node.set_op("AddN"); sum_node.set_device(ori_node->device()); - auto *sum_attr = sum_node.mutable_attr(); + auto* sum_attr = sum_node.mutable_attr(); (*sum_attr)["N"].set_i(tot); (*sum_attr)["T"] = src_attr.at("T"); - utils::Mutation *mutation = ctx->graph_view.GetMutationBuilder(); + utils::Mutation* mutation = ctx->graph_view.GetMutationBuilder(); Status status; for (int ii = 0; ii < tot; ++ii) { @@ -92,7 +92,7 @@ Status ParallelProdForce(RemapperContext *ctx, sub_node.add_input(ori_node->input(jj)); } // set frac - auto *sub_attr = sub_node.mutable_attr(); + auto* sub_attr = sub_node.mutable_attr(); (*sub_attr)["T"] = src_attr.at("T"); (*sub_attr)["n_a_sel"] = src_attr.at("n_a_sel"); (*sub_attr)["n_r_sel"] = src_attr.at("n_r_sel"); @@ -111,9 +111,9 @@ Status ParallelProdForce(RemapperContext *ctx, return Status(); } -Status DPParallel::Optimize(Cluster *cluster, - const GrapplerItem &item, - GraphDef *optimized_graph) { +Status DPParallel::Optimize(Cluster* cluster, + const GrapplerItem& item, + GraphDef* optimized_graph) { GrapplerItem mutable_item = item; Status status; RemapperContext ctx(&mutable_item, &status); @@ -147,7 +147,7 @@ Status DPParallel::Optimize(Cluster *cluster, } // Remove invalidated nodes. - utils::Mutation *mutation = ctx.graph_view.GetMutationBuilder(); + utils::Mutation* mutation = ctx.graph_view.GetMutationBuilder(); for (int i = 0; i < num_nodes; ++i) { if (nodes_to_delete[i]) { mutation->RemoveNode(ctx.graph_view.GetNode(i)); diff --git a/source/op/tf/prod_force_se_a_mask.cc b/source/op/tf/prod_force_se_a_mask.cc index a7b08ae664..6c938f88e0 100644 --- a/source/op/tf/prod_force_se_a_mask.cc +++ b/source/op/tf/prod_force_se_a_mask.cc @@ -17,23 +17,23 @@ using CPUDevice = Eigen::ThreadPoolDevice; template class ProdForceSeAMaskOp : public OpKernel { public: - explicit ProdForceSeAMaskOp(OpKernelConstruction *context) + explicit ProdForceSeAMaskOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("total_atom_num", &total_atom_num)); } - void Compute(OpKernelContext *context) override { + void Compute(OpKernelContext* context) override { deepmd::safe_compute( - context, [this](OpKernelContext *context) { this->_Compute(context); }); + context, [this](OpKernelContext* context) { this->_Compute(context); }); } - void _Compute(OpKernelContext *context) { + void _Compute(OpKernelContext* context) { // Grab the input tensor - const Tensor &net_deriv_tensor = context->input(0); - const Tensor &in_deriv_tensor = context->input(1); - const Tensor &mask_tensor = context->input(2); - const Tensor &nlist_tensor = context->input(3); + const Tensor& net_deriv_tensor = context->input(0); + const Tensor& in_deriv_tensor = context->input(1); + const Tensor& mask_tensor = context->input(2); + const Tensor& nlist_tensor = context->input(3); // set size of the sample OP_REQUIRES(context, (net_deriv_tensor.shape().dims() == 2), @@ -67,7 +67,7 @@ class ProdForceSeAMaskOp : public OpKernel { force_shape.AddDim(3 * static_cast(nall)); // std::cout << "forcesahpe " << force_shape.dim_size(0) << " " << // force_shape.dim_size(1) << std::endl; - Tensor *force_tensor = NULL; + Tensor* force_tensor = NULL; OP_REQUIRES_OK(context, context->allocate_output(0, force_shape, &force_tensor)); diff --git a/source/op/tf/prod_force_se_a_mask_grad.cc b/source/op/tf/prod_force_se_a_mask_grad.cc index a01919199f..c7ff091857 100644 --- a/source/op/tf/prod_force_se_a_mask_grad.cc +++ b/source/op/tf/prod_force_se_a_mask_grad.cc @@ -16,24 +16,24 @@ using CPUDevice = Eigen::ThreadPoolDevice; template class ProdForceSeAMaskGradOp : public OpKernel { public: - explicit ProdForceSeAMaskGradOp(OpKernelConstruction *context) + explicit ProdForceSeAMaskGradOp(OpKernelConstruction* context) : OpKernel(context) { OP_REQUIRES_OK(context, context->GetAttr("total_atom_num", &total_atom_num)); } - void Compute(OpKernelContext *context) override { + void Compute(OpKernelContext* context) override { deepmd::safe_compute( - context, [this](OpKernelContext *context) { this->_Compute(context); }); + context, [this](OpKernelContext* context) { this->_Compute(context); }); } - void _Compute(OpKernelContext *context) { + void _Compute(OpKernelContext* context) { // Grab the input tensor - const Tensor &grad_tensor = context->input(0); - const Tensor &net_deriv_tensor = context->input(1); - const Tensor &in_deriv_tensor = context->input(2); - const Tensor &mask_tensor = context->input(3); - const Tensor &nlist_tensor = context->input(4); + const Tensor& grad_tensor = context->input(0); + const Tensor& net_deriv_tensor = context->input(1); + const Tensor& in_deriv_tensor = context->input(2); + const Tensor& mask_tensor = context->input(3); + const Tensor& nlist_tensor = context->input(4); // set size of the sample TensorShape grad_shape = grad_tensor.shape(); @@ -82,7 +82,7 @@ class ProdForceSeAMaskGradOp : public OpKernel { grad_net_shape.AddDim(static_cast(nloc) * ndescrpt); // allocate the output tensor - Tensor *grad_net_tensor = NULL; + Tensor* grad_net_tensor = NULL; OP_REQUIRES_OK( context, context->allocate_output(0, grad_net_shape, &grad_net_tensor)); diff --git a/source/op/tf/unaggregated_grad.cc b/source/op/tf/unaggregated_grad.cc index cf645f6c21..329e25b2d2 100644 --- a/source/op/tf/unaggregated_grad.cc +++ b/source/op/tf/unaggregated_grad.cc @@ -75,6 +75,10 @@ FPTYPE grad(const FPTYPE xbar, case 6: { return y * (1 - y); } + case 7: { + const FPTYPE sig = 1.0 / (1.0 + exp(-xbar)); + return sig + xbar * sig * (1 - sig); + } default: return -1; } @@ -105,6 +109,11 @@ FPTYPE grad_grad(const FPTYPE xbar, const FPTYPE y, const int functype) { case 6: { return y * (1 - y) * (1 - 2 * y); } + case 7: { + const FPTYPE sig = 1.0 / (1.0 + exp(-xbar)); + const FPTYPE d_sig = sig * (1 - sig); + return 2 * d_sig + xbar * d_sig * (1 - 2 * sig); + } default: return -1; } diff --git a/source/tests/__init__.py b/source/tests/__init__.py index 6ceb116d85..16149c2cd0 100644 --- a/source/tests/__init__.py +++ b/source/tests/__init__.py @@ -1 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later + +import os + +if os.environ.get("DP_CI_IMPORT_PADDLE_BEFORE_TF", "0") == "1": + import paddle # noqa: F401 + import tensorflow # noqa: F401 diff --git a/source/tests/array_api_strict/common.py b/source/tests/array_api_strict/common.py index 28f67a97f6..50109ded86 100644 --- a/source/tests/array_api_strict/common.py +++ b/source/tests/array_api_strict/common.py @@ -1,13 +1,10 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, -) import array_api_strict import numpy as np -def to_array_api_strict_array(array: Optional[np.ndarray]): +def to_array_api_strict_array(array: np.ndarray | None): """Convert a numpy array to a JAX array. Parameters diff --git a/source/tests/array_api_strict/fitting/fitting.py b/source/tests/array_api_strict/fitting/fitting.py index 323a49cfe8..af0e57375b 100644 --- a/source/tests/array_api_strict/fitting/fitting.py +++ b/source/tests/array_api_strict/fitting/fitting.py @@ -31,6 +31,8 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any: "fparam_inv_std", "aparam_avg", "aparam_inv_std", + "case_embd", + "default_fparam_tensor", }: value = to_array_api_strict_array(value) elif name == "emask": diff --git a/source/tests/common/dpmodel/test_fitting_stat.py b/source/tests/common/dpmodel/test_fitting_stat.py new file mode 100644 index 0000000000..101d2a9ad7 --- /dev/null +++ b/source/tests/common/dpmodel/test_fitting_stat.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import numpy as np + +from deepmd.dpmodel.descriptor import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + EnergyFittingNet, +) + + +def _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds): + merged_output_stat = [] + nsys = len(sys_natoms) + ndof = len(avgs) + for ii in range(nsys): + sys_dict = {} + tmp_data_f = [] + tmp_data_a = [] + for jj in range(ndof): + rng = np.random.default_rng(2025 * ii + 220 * jj) + tmp_data_f.append( + rng.normal(loc=avgs[jj], scale=stds[jj], size=(sys_nframes[ii], 1)) + ) + rng = np.random.default_rng(220 * ii + 1636 * jj) + tmp_data_a.append( + rng.normal( + loc=avgs[jj], scale=stds[jj], size=(sys_nframes[ii], sys_natoms[ii]) + ) + ) + tmp_data_f = np.transpose(tmp_data_f, (1, 2, 0)) + tmp_data_a = np.transpose(tmp_data_a, (1, 2, 0)) + sys_dict["fparam"] = tmp_data_f + sys_dict["aparam"] = tmp_data_a + merged_output_stat.append(sys_dict) + return merged_output_stat + + +def _brute_fparam_pt(data, ndim): + adata = [ii["fparam"] for ii in data] + all_data = [] + for ii in adata: + tmp = np.reshape(ii, [-1, ndim]) + if len(all_data) == 0: + all_data = np.array(tmp) + else: + all_data = np.concatenate((all_data, tmp), axis=0) + avg = np.average(all_data, axis=0) + std = np.std(all_data, axis=0) + return avg, std + + +def _brute_aparam_pt(data, ndim): + adata = [ii["aparam"] for ii in data] + all_data = [] + for ii in adata: + tmp = np.reshape(ii, [-1, ndim]) + if len(all_data) == 0: + all_data = np.array(tmp) + else: + all_data = np.concatenate((all_data, tmp), axis=0) + avg = np.average(all_data, axis=0) + std = np.std(all_data, axis=0) + return avg, std + + +class TestEnerFittingStat(unittest.TestCase): + def test(self) -> None: + descrpt = DescrptSeA(6.0, 5.8, [46, 92], neuron=[25, 50, 100], axis_neuron=16) + fitting = EnergyFittingNet( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + numb_fparam=3, + numb_aparam=3, + ) + avgs = [0, 10, 100] + stds = [2, 0.4, 0.00001] + sys_natoms = [10, 100] + sys_nframes = [5, 2] + all_data = _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds) + frefa, frefs = _brute_fparam_pt(all_data, len(avgs)) + arefa, arefs = _brute_aparam_pt(all_data, len(avgs)) + fitting.compute_input_stats(all_data, protection=1e-2) + frefs_inv = 1.0 / frefs + arefs_inv = 1.0 / arefs + frefs_inv[frefs_inv > 100] = 100 + arefs_inv[arefs_inv > 100] = 100 + np.testing.assert_almost_equal(frefa, fitting.fparam_avg) + np.testing.assert_almost_equal(frefs_inv, fitting.fparam_inv_std) + np.testing.assert_almost_equal(arefa, fitting.aparam_avg) + np.testing.assert_almost_equal(arefs_inv, fitting.aparam_inv_std) diff --git a/source/tests/common/dpmodel/test_padding_atoms.py b/source/tests/common/dpmodel/test_padding_atoms.py new file mode 100644 index 0000000000..d4ea39f598 --- /dev/null +++ b/source/tests/common/dpmodel/test_padding_atoms.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from copy import ( + deepcopy, +) + +import numpy as np + +from deepmd.dpmodel.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting import ( + PropertyFittingNet, +) +from deepmd.dpmodel.model.property_model import ( + PropertyModel, +) + + +class TestCaseSingleFrameWithoutNlist: + def setUp(self) -> None: + # nf=2, nloc == 3 + self.nloc = 3 + self.nt = 2 + self.coord = np.array( + [ + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + [ + [1, 0, 1], + [0, 1, 1], + [1, 1, 0], + ], + ], + dtype=np.float64, + ) + self.atype = np.array([[0, 0, 1], [1, 1, 0]], dtype=int).reshape([2, self.nloc]) + self.cell = 2.0 * np.eye(3).reshape([1, 9]) + self.cell = np.array([self.cell, self.cell]).reshape(2, 9) + self.sel = [16, 8] + self.rcut = 2.2 + self.rcut_smth = 0.4 + self.atol = 1e-12 + + +class TestPaddingAtoms(unittest.TestCase, TestCaseSingleFrameWithoutNlist): + def setUp(self): + TestCaseSingleFrameWithoutNlist.setUp(self) + + def test_padding_atoms_consistency(self): + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = PropertyFittingNet( + self.nt, + ds.get_dim_out(), + mixed_types=ds.mixed_types(), + intensive=True, + ) + type_map = ["foo", "bar"] + model = PropertyModel(ds, ft, type_map=type_map) + var_name = model.get_var_name() + args = [self.coord, self.atype, self.cell] + result = model.call(*args) + # test intensive + np.testing.assert_allclose( + result[f"{var_name}_redu"], + np.mean(result[f"{var_name}"], axis=1), + atol=self.atol, + ) + # test padding atoms + padding_atoms_list = [1, 5, 10] + for padding_atoms in padding_atoms_list: + coord = deepcopy(self.coord) + atype = deepcopy(self.atype) + atype_padding = np.pad( + atype, + pad_width=((0, 0), (0, padding_atoms)), + mode="constant", + constant_values=-1, + ) + coord_padding = np.pad( + coord, + pad_width=((0, 0), (0, padding_atoms), (0, 0)), + mode="constant", + constant_values=0, + ) + args = [coord_padding, atype_padding, self.cell] + result_padding = model.call(*args) + np.testing.assert_allclose( + result[f"{var_name}_redu"], + result_padding[f"{var_name}_redu"], + atol=self.atol, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/common/test_argument_parser.py b/source/tests/common/test_argument_parser.py index 4e39df8659..df0d06e4b5 100644 --- a/source/tests/common/test_argument_parser.py +++ b/source/tests/common/test_argument_parser.py @@ -15,7 +15,6 @@ from typing import ( TYPE_CHECKING, Any, - Union, ) from deepmd.main import ( @@ -30,7 +29,7 @@ from typing_extensions import TypedDict # python<=3.7 class DATA(TypedDict): - type: Union[type, tuple[type]] + type: type | tuple[type] value: Any TEST_DICT = dict[str, DATA] @@ -322,6 +321,32 @@ def test_parser_test(self) -> None: self.run_test(command="test", mapping=ARGS) + def test_parser_test_train_data(self) -> None: + """Test test subparser with train-data.""" + ARGS = { + "--model": {"type": str, "value": "MODEL.PB"}, + "--train-data": { + "type": (str, type(None)), + "value": "INPUT.JSON", + "dest": "train_json", + }, + } + + self.run_test(command="test", mapping=ARGS) + + def test_parser_test_valid_data(self) -> None: + """Test test subparser with valid-data.""" + ARGS = { + "--model": {"type": str, "value": "MODEL.PB"}, + "--valid-data": { + "type": (str, type(None)), + "value": "INPUT.JSON", + "dest": "valid_json", + }, + } + + self.run_test(command="test", mapping=ARGS) + def test_parser_compress(self) -> None: """Test compress subparser.""" ARGS = { diff --git a/source/tests/common/test_deepmd_data.py b/source/tests/common/test_deepmd_data.py new file mode 100644 index 0000000000..c7a26491cd --- /dev/null +++ b/source/tests/common/test_deepmd_data.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import tempfile +import unittest +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.utils.data import ( + DeepmdData, +) + + +class TestDeepmdDataTypeMap(unittest.TestCase): + def setUp(self) -> None: + self.tmpdir = tempfile.TemporaryDirectory() + self.root = Path(self.tmpdir.name) + self.set_dir = self.root / "set.000" + self.set_dir.mkdir() + + # minimal required dataset + atom_types = np.array([0, 1, 0, 1], dtype=np.int32) + np.savetxt(self.root / "type.raw", atom_types, fmt="%d") + np.savetxt( + self.root / "type_map.raw", + np.array(["O", "H", "Si"], dtype=object), + fmt="%s", + ) + + coord = np.zeros((1, atom_types.size * 3), dtype=np.float32) + box = np.eye(3, dtype=np.float32).reshape(1, 9) + np.save(self.set_dir / "coord.npy", coord) + np.save(self.set_dir / "box.npy", box) + + def tearDown(self) -> None: + self.tmpdir.cleanup() + + def test_remap_with_unused_types(self) -> None: + data = DeepmdData(str(self.root), type_map=["H", "O", "Si"]) + + expected_atom_types = np.array([1, 0, 1, 0], dtype=np.int32) + np.testing.assert_array_equal(data.atom_type, expected_atom_types) + self.assertEqual(data.type_map, ["H", "O", "Si"]) + + loaded = data._load_set(self.set_dir) + expected_sorted = expected_atom_types[data.idx_map] + np.testing.assert_array_equal(loaded["type"], np.tile(expected_sorted, (1, 1))) diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 09e59bf711..6c9f1e43a2 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -68,6 +68,7 @@ input_files_multi = ( p_examples / "water_multi_task" / "pytorch_example" / "input_torch.json", p_examples / "water_multi_task" / "pytorch_example" / "input_torch_sharefit.json", + p_examples / "water_multi_task" / "pytorch_example" / "input_torch_with_alias.json", p_examples / "hessian" / "multi_task" / "input.json", ) diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 7ecb5ea5a5..c1e6644be6 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -8,6 +8,9 @@ ABC, abstractmethod, ) +from collections.abc import ( + Callable, +) from enum import ( Enum, ) @@ -16,10 +19,7 @@ ) from typing import ( Any, - Callable, ClassVar, - Optional, - Union, ) from uuid import ( uuid4, @@ -80,18 +80,18 @@ class CommonTest(ABC): """Arguments data.""" additional_data: ClassVar[dict] = {} """Additional data that will not be checked.""" - tf_class: ClassVar[Optional[type]] + tf_class: ClassVar[type | None] """TensorFlow model class.""" - dp_class: ClassVar[Optional[type]] + dp_class: ClassVar[type | None] """Native DP model class.""" - pt_class: ClassVar[Optional[type]] + pt_class: ClassVar[type | None] """PyTorch model class.""" - jax_class: ClassVar[Optional[type]] + jax_class: ClassVar[type | None] """JAX model class.""" - pd_class: ClassVar[Optional[type]] + pd_class: ClassVar[type | None] """Paddle model class.""" - array_api_strict_class: ClassVar[Optional[type]] - args: ClassVar[Optional[Union[Argument, list[Argument]]]] + array_api_strict_class: ClassVar[type | None] + args: ClassVar[Argument | list[Argument] | None] """Arguments that maps to the `data`.""" skip_dp: ClassVar[bool] = False """Whether to skip the native DP model.""" @@ -354,9 +354,6 @@ def test_tf_consistent_with_ref(self) -> None: data1.pop("@version") data2.pop("@version") - if tf_obj.__class__.__name__.startswith("Polar"): - data1["@variables"].pop("bias_atom_e") - np.testing.assert_equal(data1, data2) for rr1, rr2 in zip(ret1, ret2): np.testing.assert_allclose( diff --git a/source/tests/consistent/descriptor/common.py b/source/tests/consistent/descriptor/common.py index 617312145e..8af1c7ea64 100644 --- a/source/tests/consistent/descriptor/common.py +++ b/source/tests/consistent/descriptor/common.py @@ -57,7 +57,15 @@ class DescriptorTest: """Useful utilities for descriptor tests.""" - def build_tf_descriptor(self, obj, natoms, coords, atype, box, suffix): + def build_tf_descriptor( + self, + obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, + suffix: str, + ) -> tuple[list[Any], dict[Any, np.ndarray]]: t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_coord") t_type = tf.placeholder(tf.int32, [None], name="i_type") t_natoms = tf.placeholder(tf.int32, natoms.shape, name="i_natoms") @@ -83,7 +91,13 @@ def build_tf_descriptor(self, obj, natoms, coords, atype, box, suffix): } def eval_dp_descriptor( - self, dp_obj: Any, natoms, coords, atype, box, mixed_types: bool = False + self, + dp_obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, + mixed_types: bool = False, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( coords.reshape(1, -1, 3), @@ -102,7 +116,13 @@ def eval_dp_descriptor( return dp_obj(ext_coords, ext_atype, nlist=nlist, mapping=mapping) def eval_pt_descriptor( - self, pt_obj: Any, natoms, coords, atype, box, mixed_types: bool = False + self, + pt_obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, + mixed_types: bool = False, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts_pt( torch.from_numpy(coords).to(PT_DEVICE).reshape(1, -1, 3), @@ -124,7 +144,13 @@ def eval_pt_descriptor( ] def eval_jax_descriptor( - self, jax_obj: Any, natoms, coords, atype, box, mixed_types: bool = False + self, + jax_obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, + mixed_types: bool = False, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( jnp.array(coords).reshape(1, -1, 3), @@ -146,7 +172,13 @@ def eval_jax_descriptor( ] def eval_pd_descriptor( - self, pd_obj: Any, natoms, coords, atype, box, mixed_types: bool = False + self, + pd_obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, + mixed_types: bool = False, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts_pd( paddle.to_tensor(coords).to(PD_DEVICE).reshape([1, -1, 3]), @@ -170,10 +202,10 @@ def eval_pd_descriptor( def eval_array_api_strict_descriptor( self, array_api_strict_obj: Any, - natoms, - coords, - atype, - box, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, mixed_types: bool = False, ) -> Any: ext_coords, ext_atype, mapping = extend_coord_with_ghosts( diff --git a/source/tests/consistent/descriptor/test_dpa1.py b/source/tests/consistent/descriptor/test_dpa1.py index db5fe4dae0..a0a3d9c833 100644 --- a/source/tests/consistent/descriptor/test_dpa1.py +++ b/source/tests/consistent/descriptor/test_dpa1.py @@ -2,7 +2,6 @@ import unittest from typing import ( Any, - Optional, ) import numpy as np @@ -127,12 +126,13 @@ def data(self) -> dict: "use_tebd_bias": use_tebd_bias, "type_map": ["O", "H"] if use_econf_tebd else None, "seed": 1145141919810, + "trainable": False, } def is_meaningless_zero_attention_layer_tests( self, attn_layer: int, - temperature: Optional[float], + temperature: float | None, ) -> bool: return attn_layer == 0 and (temperature is not None) diff --git a/source/tests/consistent/descriptor/test_dpa2.py b/source/tests/consistent/descriptor/test_dpa2.py index ef840bf9d7..6864d91f26 100644 --- a/source/tests/consistent/descriptor/test_dpa2.py +++ b/source/tests/consistent/descriptor/test_dpa2.py @@ -181,7 +181,7 @@ def data(self) -> dict: "smooth": smooth, "exclude_types": exclude_types, "env_protection": 0.0, - "trainable": True, + "trainable": False, "use_econf_tebd": use_econf_tebd, "use_tebd_bias": use_tebd_bias, "type_map": ["O", "H"] if use_econf_tebd else None, diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py index b99117b9e7..367ff29a3b 100644 --- a/source/tests/consistent/descriptor/test_dpa3.py +++ b/source/tests/consistent/descriptor/test_dpa3.py @@ -130,7 +130,7 @@ def data(self) -> dict: "exclude_types": exclude_types, "env_protection": 0.0, "use_loc_mapping": use_loc_mapping, - "trainable": True, + "trainable": False, } @property @@ -171,14 +171,7 @@ def skip_pd(self) -> bool: n_multi_edge_message, precision, ) = self.param - return ( - not INSTALLED_PD - or precision == "bfloat16" - or edge_init_use_dist - or use_exp_switch - or use_dynamic_sel - or use_loc_mapping - ) # not supported yet + return CommonTest.skip_pd @property def skip_dp(self) -> bool: diff --git a/source/tests/consistent/descriptor/test_se_atten_v2.py b/source/tests/consistent/descriptor/test_se_atten_v2.py index 459fac037b..129343c177 100644 --- a/source/tests/consistent/descriptor/test_se_atten_v2.py +++ b/source/tests/consistent/descriptor/test_se_atten_v2.py @@ -2,7 +2,6 @@ import unittest from typing import ( Any, - Optional, ) import numpy as np @@ -132,7 +131,7 @@ def is_meaningless_zero_attention_layer_tests( attn_layer: int, attn_dotr: bool, normalize: bool, - temperature: Optional[float], + temperature: float | None, ) -> bool: return attn_layer == 0 and (attn_dotr or normalize or temperature is not None) diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index 396ee2d492..010944d109 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -61,6 +61,7 @@ (True, False), # resnet_dt ("float64", "float32"), # precision (True, False), # mixed_types + (None, [0]), # sel_type ) class TestDipole(CommonTest, DipoleFittingTest, unittest.TestCase): @property @@ -69,13 +70,37 @@ def data(self) -> dict: resnet_dt, precision, mixed_types, + sel_type, ) = self.param - return { + data = { "neuron": [5, 5, 5], "resnet_dt": resnet_dt, "precision": precision, + "sel_type": sel_type, "seed": 20240217, } + return data + + def pass_data_to_cls(self, cls, data) -> Any: + """Pass data to the class.""" + if cls not in (self.tf_class,): + sel_type = data.pop("sel_type", None) + if sel_type is not None: + all_types = list(range(self.ntypes)) + exclude_types = [t for t in all_types if t not in sel_type] + data["exclude_types"] = exclude_types + return cls(**data, **self.additional_data) + + @property + def skip_tf(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + sel_type, + ) = self.param + # mixed_types + sel_type is not supported + return CommonTest.skip_tf or (mixed_types and sel_type is not None) @property def skip_pt(self) -> bool: @@ -83,6 +108,7 @@ def skip_pt(self) -> bool: resnet_dt, precision, mixed_types, + sel_type, ) = self.param return CommonTest.skip_pt @@ -112,6 +138,7 @@ def additional_data(self) -> dict: resnet_dt, precision, mixed_types, + sel_type, ) = self.param return { "ntypes": self.ntypes, @@ -125,6 +152,7 @@ def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: resnet_dt, precision, mixed_types, + sel_type, ) = self.param return self.build_tf_fitting( obj, @@ -141,6 +169,7 @@ def eval_pt(self, pt_obj: Any) -> Any: resnet_dt, precision, mixed_types, + sel_type, ) = self.param return ( pt_obj( @@ -159,6 +188,7 @@ def eval_dp(self, dp_obj: Any) -> Any: resnet_dt, precision, mixed_types, + sel_type, ) = self.param return dp_obj( self.inputs, @@ -200,6 +230,7 @@ def rtol(self) -> float: resnet_dt, precision, mixed_types, + sel_type, ) = self.param if precision == "float64": return 1e-10 @@ -215,6 +246,7 @@ def atol(self) -> float: resnet_dt, precision, mixed_types, + sel_type, ) = self.param if precision == "float64": return 1e-10 @@ -222,3 +254,39 @@ def atol(self) -> float: return 1e-4 else: raise ValueError(f"Unknown precision: {precision}") + + def test_tf_consistent_with_ref(self) -> None: + """Test whether TF and reference are consistent.""" + # Special handle for sel_types + if self.skip_tf: + self.skipTest("Unsupported backend") + ref_backend = self.get_reference_backend() + if ref_backend == self.RefBackend.TF: + self.skipTest("Reference is self") + ret1, data1 = self.get_reference_ret_serialization(ref_backend) + ret1 = self.extract_ret(ret1, ref_backend) + self.reset_unique_id() + tf_obj = self.tf_class.deserialize(data1, suffix=self.unique_id) + ret2, data2 = self.get_tf_ret_serialization_from_cls(tf_obj) + ret2 = self.extract_ret(ret2, self.RefBackend.TF) + if tf_obj.__class__.__name__.startswith(("Polar", "Dipole", "DOS")): + # tf, pt serialization mismatch + common_keys = set(data1.keys()) & set(data2.keys()) + data1 = {k: data1[k] for k in common_keys} + data2 = {k: data2[k] for k in common_keys} + + # not comparing version + data1.pop("@version") + data2.pop("@version") + + if tf_obj.__class__.__name__.startswith("Polar"): + data1["@variables"].pop("bias_atom_e") + for ii, networks in enumerate(data2["nets"]["networks"]): + if networks is None: + data1["nets"]["networks"][ii] = None + np.testing.assert_equal(data1, data2) + for rr1, rr2 in zip(ret1, ret2): + np.testing.assert_allclose( + rr1.ravel()[: rr2.size], rr2.ravel(), rtol=self.rtol, atol=self.atol + ) + assert rr1.dtype == rr2.dtype, f"{rr1.dtype} != {rr2.dtype}" diff --git a/source/tests/consistent/fitting/test_ener.py b/source/tests/consistent/fitting/test_ener.py index f5a79acabe..ad70bd0bfa 100644 --- a/source/tests/consistent/fitting/test_ener.py +++ b/source/tests/consistent/fitting/test_ener.py @@ -70,7 +70,7 @@ (True, False), # resnet_dt ("float64", "float32", "bfloat16"), # precision (True, False), # mixed_types - (0, 1), # numb_fparam + ((0, None), (1, None), (1, [1.0])), # (numb_fparam, default_fparam) ((0, False), (1, False), (1, True)), # (numb_aparam, use_aparam_as_mask) ([], [-12345.6, None]), # atom_ener ) @@ -81,7 +81,7 @@ def data(self) -> dict: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -91,6 +91,7 @@ def data(self) -> dict: "precision": precision, "numb_fparam": numb_fparam, "numb_aparam": numb_aparam, + "default_fparam": default_fparam, "seed": 20240217, "atom_ener": atom_ener, "use_aparam_as_mask": use_aparam_as_mask, @@ -102,7 +103,7 @@ def skip_pt(self) -> bool: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -116,7 +117,7 @@ def skip_array_api_strict(self) -> bool: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -129,13 +130,25 @@ def skip_pd(self) -> bool: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param # Paddle do not support "bfloat16" in some kernels, # so skip this in CI test - return not INSTALLED_PD or precision == "bfloat16" + return not INSTALLED_PD or precision == "bfloat16" or default_fparam is not None + + @property + def skip_tf(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + (numb_fparam, default_fparam), + (numb_aparam, use_aparam_as_mask), + atom_ener, + ) = self.param + return not INSTALLED_TF or default_fparam is not None tf_class = EnerFittingTF dp_class = EnerFittingDP @@ -165,7 +178,7 @@ def additional_data(self) -> dict: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -180,7 +193,7 @@ def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -199,7 +212,7 @@ def eval_pt(self, pt_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -209,7 +222,7 @@ def eval_pt(self, pt_obj: Any) -> Any: torch.from_numpy(self.atype.reshape(1, -1)).to(device=PT_DEVICE), fparam=( torch.from_numpy(self.fparam).to(device=PT_DEVICE) - if numb_fparam + if (numb_fparam and default_fparam is None) # test default_fparam else None ), aparam=( @@ -228,14 +241,14 @@ def eval_dp(self, dp_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param return dp_obj( self.inputs, self.atype.reshape(1, -1), - fparam=self.fparam if numb_fparam else None, + fparam=self.fparam if (numb_fparam and default_fparam is None) else None, aparam=self.aparam if numb_aparam else None, )["energy"] @@ -244,7 +257,7 @@ def eval_jax(self, jax_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -252,7 +265,9 @@ def eval_jax(self, jax_obj: Any) -> Any: jax_obj( jnp.asarray(self.inputs), jnp.asarray(self.atype.reshape(1, -1)), - fparam=jnp.asarray(self.fparam) if numb_fparam else None, + fparam=jnp.asarray(self.fparam) + if (numb_fparam and default_fparam is None) + else None, aparam=jnp.asarray(self.aparam) if numb_aparam else None, )["energy"] ) @@ -262,7 +277,7 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -270,7 +285,9 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: array_api_strict_obj( array_api_strict.asarray(self.inputs), array_api_strict.asarray(self.atype.reshape(1, -1)), - fparam=array_api_strict.asarray(self.fparam) if numb_fparam else None, + fparam=array_api_strict.asarray(self.fparam) + if (numb_fparam and default_fparam is None) + else None, aparam=array_api_strict.asarray(self.aparam) if numb_aparam else None, )["energy"] ) @@ -280,7 +297,7 @@ def eval_pd(self, pd_obj: Any) -> Any: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -317,7 +334,7 @@ def rtol(self) -> float: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param @@ -337,7 +354,7 @@ def atol(self) -> float: resnet_dt, precision, mixed_types, - numb_fparam, + (numb_fparam, default_fparam), (numb_aparam, use_aparam_as_mask), atom_ener, ) = self.param diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index d6b7348a91..7ae9af6891 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -158,6 +158,7 @@ def test_deep_eval(self) -> None: prefix + backend.suffixes[suffix_idx], reference_data ) deep_eval = DeepEval(prefix + backend.suffixes[suffix_idx]) + self.assertIsInstance(deep_eval.get_model_def_script(), dict) if deep_eval.get_dim_fparam() > 0: fparam = np.ones((nframes, deep_eval.get_dim_fparam())) else: diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index 7cf71000db..778ae519c6 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -77,8 +77,8 @@ def build_tf_model( ] elif ret_key == "polar": ret_list = [ - ret["polar"], ret["global_polar"], + ret["polar"], ] else: raise NotImplementedError diff --git a/source/tests/consistent/model/test_dipole.py b/source/tests/consistent/model/test_dipole.py index bb381129a3..339dcae7c3 100644 --- a/source/tests/consistent/model/test_dipole.py +++ b/source/tests/consistent/model/test_dipole.py @@ -73,6 +73,7 @@ def data(self) -> dict: pt_class = DipoleModelPT jax_class = DipoleModelJAX args = model_args() + atol = 1e-8 def get_reference_backend(self): """Get the reference backend. @@ -89,7 +90,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return True # need to fix tf consistency + return not INSTALLED_TF @property def skip_jax(self) -> bool: @@ -203,3 +204,15 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ret[1].ravel(), ) raise ValueError(f"Unknown backend: {backend}") + + def test_atom_exclude_types(self): + if self.skip_pt: + self.skipTest("Unsupported backend") + if self.skip_tf: + self.skipTest("Unsupported backend") + _ret, data = self.get_reference_ret_serialization(self.RefBackend.PT) + data["atom_exclude_types"] = [1] + self.reset_unique_id() + tf_obj = self.tf_class.deserialize(data, suffix=self.unique_id) + pt_obj = self.pt_class.deserialize(data) + self.assertEqual(tf_obj.get_sel_type(), pt_obj.get_sel_type()) diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 83e33e499a..ef72e9096b 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -90,7 +90,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return True # need to fix tf consistency + return not INSTALLED_TF @property def skip_jax(self) -> bool: diff --git a/source/tests/consistent/model/test_polar.py b/source/tests/consistent/model/test_polar.py index 5295bc4705..1405814f03 100644 --- a/source/tests/consistent/model/test_polar.py +++ b/source/tests/consistent/model/test_polar.py @@ -73,6 +73,7 @@ def data(self) -> dict: pt_class = PolarModelPT jax_class = PolarModelJAX args = model_args() + atol = 1e-8 def get_reference_backend(self): """Get the reference backend. @@ -89,7 +90,7 @@ def get_reference_backend(self): @property def skip_tf(self): - return True # need to fix tf consistency + return not INSTALLED_TF @property def skip_jax(self) -> bool: @@ -197,3 +198,15 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: ret[1].ravel(), ) raise ValueError(f"Unknown backend: {backend}") + + def test_atom_exclude_types(self): + if self.skip_pt: + self.skipTest("Unsupported backend") + if self.skip_tf: + self.skipTest("Unsupported backend") + _ret, data = self.get_reference_ret_serialization(self.RefBackend.PT) + data["atom_exclude_types"] = [1] + self.reset_unique_id() + tf_obj = self.tf_class.deserialize(data, suffix=self.unique_id) + pt_obj = self.pt_class.deserialize(data) + self.assertEqual(tf_obj.get_sel_type(), pt_obj.get_sel_type()) diff --git a/source/tests/infer/case.py b/source/tests/infer/case.py index 8b8481a194..828974c6e6 100644 --- a/source/tests/infer/case.py +++ b/source/tests/infer/case.py @@ -26,9 +26,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import yaml @@ -125,6 +122,13 @@ def __init__(self, data: dict) -> None: else: self.descriptor = None + if "fit_ll" in data: + self.fit_ll = np.array(data["fit_ll"], dtype=np.float64).reshape( + self.nloc, -1 + ) + else: + self.fit_ll = None + class Case: """Test case. @@ -149,7 +153,7 @@ def __init__(self, filename: str) -> None: self.model_def_script = config.get("model_def_script") @lru_cache - def get_model(self, suffix: str, out_file: Optional[str] = None) -> str: + def get_model(self, suffix: str, out_file: str | None = None) -> str: """Get the model file with the specified suffix. Parameters diff --git a/source/tests/infer/deepdipole_pt.pth b/source/tests/infer/deepdipole_pt.pth new file mode 100644 index 0000000000..4c93a1b864 Binary files /dev/null and b/source/tests/infer/deepdipole_pt.pth differ diff --git a/source/tests/infer/deeppot-testcase.yaml b/source/tests/infer/deeppot-testcase.yaml index 9523b8d1ea..772a06a89b 100644 --- a/source/tests/infer/deeppot-testcase.yaml +++ b/source/tests/infer/deeppot-testcase.yaml @@ -350,6 +350,45 @@ results: 1.391094495316195001e+00, 7.036614101584164338e-01, ] + fit_ll: + [ + -1.930622006643730598e-02, + 7.105172146387829235e-01, + 8.063835335367619539e-01, + -8.414936892447275607e-01, + 1.076881365346436414e+00, + -5.058153291569045251e-01, + -3.104797373867691779e-02, + 7.915138025598530414e-01, + 8.704498369678651537e-01, + -9.394329433114724237e-01, + 1.081177674358831053e+00, + -5.122829163516022799e-01, + 5.307913125575804136e-03, + 7.644783775007328863e-01, + 8.548853566716824171e-01, + -9.264496186379944653e-01, + 1.087178488222722672e+00, + -4.893627623467682874e-01, + -1.098746804357388085e-01, + 8.092546382430507723e-01, + 8.757043853926992361e-01, + -9.036627000544070754e-01, + 1.064706190677472852e+00, + -5.670533963064982030e-01, + -1.270062329805081158e-01, + 8.618261193779762630e-01, + 8.979592934126284787e-01, + -9.939941754957831721e-01, + 1.072078883192923771e+00, + -5.780043831847785363e-01, + -8.617331266742107865e-02, + 8.388158674801169390e-01, + 8.904977456468012864e-01, + -9.751383339999978306e-01, + 1.075378146084344344e+00, + -5.508880199511664300e-01, + ] - coord: [ 12.83, diff --git a/source/tests/infer/test_get_model.py b/source/tests/infer/test_get_model.py new file mode 100644 index 0000000000..4c52dda0a1 --- /dev/null +++ b/source/tests/infer/test_get_model.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +from deepmd.infer.deep_eval import ( + DeepEval, +) + +from ..consistent.common import ( + parameterized, +) +from .case import ( + get_cases, +) + + +@parameterized( + ( + "se_e2_a", + "fparam_aparam", + ), # key + (".pb", ".pth"), # model extension +) +class TestGetModelMethod(unittest.TestCase): + """Test the new get_model method functionality.""" + + @classmethod + def setUpClass(cls) -> None: + key, extension = cls.param + cls.case = get_cases()[key] + cls.model_name = cls.case.get_model(extension) + cls.dp = DeepEval(cls.model_name) + + @classmethod + def tearDownClass(cls) -> None: + cls.dp = None + + def test_get_model_method_exists(self): + """Test that get_model method exists.""" + self.assertTrue( + hasattr(self.dp, "get_model"), "DeepEval should have get_model method" + ) + + def test_get_model_returns_valid_object(self): + """Test that get_model returns a valid model object.""" + model = self.dp.get_model() + self.assertIsNotNone(model, "get_model should return a non-None object") + + def test_get_model_backend_specific(self): + """Test that get_model returns the expected type for each backend.""" + key, extension = self.param + model = self.dp.get_model() + + if extension == ".pth": + # For PyTorch .pth models (TorchScript), should return torch.jit.ScriptModule + import torch + + self.assertIsInstance( + model, + torch.jit.ScriptModule, + "PyTorch .pth model should return TorchScript ScriptModule instance", + ) + # TorchScript modules are also nn.Module instances + self.assertIsInstance( + model, + torch.nn.Module, + "PyTorch .pth model should be a torch.nn.Module instance", + ) + # Check if it has common model methods + self.assertTrue( + hasattr(model, "get_type_map"), + "PyTorch model should have get_type_map method", + ) + self.assertTrue( + hasattr(model, "get_rcut"), + "PyTorch model should have get_rcut method", + ) + elif extension == ".pb": + # For TensorFlow models, should return graph + try: + # Should be a TensorFlow graph or have graph-like properties + self.assertTrue( + hasattr(model, "get_operations") + or str(type(model)).find("Graph") >= 0, + "TensorFlow model should be a graph or graph-like object", + ) + except ImportError: + # If TensorFlow not available, skip this assertion + pass + + def test_get_model_consistency(self): + """Test that get_model always returns the same object.""" + model1 = self.dp.get_model() + model2 = self.dp.get_model() + # Should return the same object (not necessarily equal, but same reference) + self.assertIs( + model1, model2, "get_model should return consistent object reference" + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/infer/test_models.py b/source/tests/infer/test_models.py index c0543e91cd..7f7b7cc21c 100644 --- a/source/tests/infer/test_models.py +++ b/source/tests/infer/test_models.py @@ -164,6 +164,24 @@ def test_descriptor(self) -> None: expected_descpt = result.descriptor np.testing.assert_almost_equal(descpt.ravel(), expected_descpt.ravel()) + def test_fitting_last_layer(self) -> None: + _, extension = self.param + if extension == ".pb": + self.skipTest("fitting_last_layer not supported for TensorFlow models") + for ii, result in enumerate(self.case.results): + if result.fit_ll is None: + continue + fit_ll = self.dp.eval_fitting_last_layer( + result.coord, result.box, result.atype + ) + expected_fit_ll = result.fit_ll + np.testing.assert_almost_equal(fit_ll.ravel(), expected_fit_ll.ravel()) + fit_ll = self.dp.eval_fitting_last_layer( + result.coord, result.box, result.atype + ) + expected_fit_ll = result.fit_ll + np.testing.assert_almost_equal(fit_ll.ravel(), expected_fit_ll.ravel()) + def test_2frame_atm(self) -> None: for ii, result in enumerate(self.case.results): coords2 = np.concatenate((result.coord, result.coord)) diff --git a/source/tests/jax/test_dp_hessian_model.py b/source/tests/jax/test_dp_hessian_model.py index 798b893651..89c066e980 100644 --- a/source/tests/jax/test_dp_hessian_model.py +++ b/source/tests/jax/test_dp_hessian_model.py @@ -7,25 +7,23 @@ from deepmd.dpmodel.common import ( to_numpy_array, ) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.jax.env import ( + jnp, +) +from deepmd.jax.fitting.fitting import ( + EnergyFittingNet, +) +from deepmd.jax.model.ener_model import ( + EnergyModel, +) -if sys.version_info >= (3, 10): - from deepmd.jax.common import ( - to_jax_array, - ) - from deepmd.jax.descriptor.se_e2_a import ( - DescrptSeA, - ) - from deepmd.jax.env import ( - jnp, - ) - from deepmd.jax.fitting.fitting import ( - EnergyFittingNet, - ) - from deepmd.jax.model.ener_model import ( - EnergyModel, - ) - - dtype = jnp.float64 +dtype = jnp.float64 @unittest.skipIf( diff --git a/source/tests/jax/test_make_hessian_model.py b/source/tests/jax/test_make_hessian_model.py index 185660e2be..8666ff4ad4 100644 --- a/source/tests/jax/test_make_hessian_model.py +++ b/source/tests/jax/test_make_hessian_model.py @@ -10,30 +10,28 @@ from deepmd.dpmodel.output_def import ( OutputVariableCategory, ) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.jax.env import ( + jax, + jnp, +) +from deepmd.jax.fitting.fitting import ( + EnergyFittingNet, +) +from deepmd.jax.model import ( + EnergyModel, +) + +from ..seed import ( + GLOBAL_SEED, +) -if sys.version_info >= (3, 10): - from deepmd.jax.common import ( - to_jax_array, - ) - from deepmd.jax.descriptor.se_e2_a import ( - DescrptSeA, - ) - from deepmd.jax.env import ( - jax, - jnp, - ) - from deepmd.jax.fitting.fitting import ( - EnergyFittingNet, - ) - from deepmd.jax.model import ( - EnergyModel, - ) - - from ..seed import ( - GLOBAL_SEED, - ) - - dtype = jnp.float64 +dtype = jnp.float64 def finite_hessian(f, x, delta=1e-6): diff --git a/source/tests/jax/test_padding_atoms.py b/source/tests/jax/test_padding_atoms.py new file mode 100644 index 0000000000..b63b464721 --- /dev/null +++ b/source/tests/jax/test_padding_atoms.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import sys +import unittest +from copy import ( + deepcopy, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.jax.common import ( + to_jax_array, +) +from deepmd.jax.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.jax.env import ( + jnp, +) +from deepmd.jax.fitting.fitting import ( + PropertyFittingNet, +) +from deepmd.jax.model.property_model import ( + PropertyModel, +) + +dtype = jnp.float64 + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestCaseSingleFrameWithoutNlist: + def setUp(self) -> None: + # nf=2, nloc == 3 + self.nloc = 3 + self.nt = 2 + self.coord = np.array( + [ + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + [ + [1, 0, 1], + [0, 1, 1], + [1, 1, 0], + ], + ], + dtype=np.float64, + ) + self.atype = np.array([[0, 0, 1], [1, 1, 0]], dtype=int).reshape([2, self.nloc]) + self.cell = 2.0 * np.eye(3).reshape([1, 9]) + self.cell = np.array([self.cell, self.cell]).reshape(2, 9) + self.sel = [16, 8] + self.rcut = 2.2 + self.rcut_smth = 0.4 + self.atol = 1e-12 + + +@unittest.skipIf( + sys.version_info < (3, 10), + "JAX requires Python 3.10 or later", +) +class TestPaddingAtoms(unittest.TestCase, TestCaseSingleFrameWithoutNlist): + def setUp(self): + TestCaseSingleFrameWithoutNlist.setUp(self) + + def test_padding_atoms_consistency(self): + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = PropertyFittingNet( + self.nt, + ds.get_dim_out(), + mixed_types=ds.mixed_types(), + intensive=True, + ) + type_map = ["foo", "bar"] + model = PropertyModel(ds, ft, type_map=type_map) + var_name = model.get_var_name() + args = [to_jax_array(ii) for ii in [self.coord, self.atype, self.cell]] + result = model.call(*args) + # test intensive + np.testing.assert_allclose( + to_numpy_array(result[f"{var_name}_redu"]), + np.mean(to_numpy_array(result[f"{var_name}"]), axis=1), + atol=self.atol, + ) + # test padding atoms + padding_atoms_list = [1, 5, 10] + for padding_atoms in padding_atoms_list: + coord = deepcopy(self.coord) + atype = deepcopy(self.atype) + atype_padding = np.pad( + atype, + pad_width=((0, 0), (0, padding_atoms)), + mode="constant", + constant_values=-1, + ) + coord_padding = np.pad( + coord, + pad_width=((0, 0), (0, padding_atoms), (0, 0)), + mode="constant", + constant_values=0, + ) + args = [ + to_jax_array(ii) for ii in [coord_padding, atype_padding, self.cell] + ] + result_padding = model.call(*args) + np.testing.assert_allclose( + to_numpy_array(result[f"{var_name}_redu"]), + to_numpy_array(result_padding[f"{var_name}_redu"]), + atol=self.atol, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pd/common.py b/source/tests/pd/common.py index d73544c5f1..f96777c9a3 100644 --- a/source/tests/pd/common.py +++ b/source/tests/pd/common.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import pathlib -from typing import ( - Optional, - Union, -) import numpy as np import paddle @@ -51,10 +47,10 @@ def run_dp(cmd: str) -> int: def eval_model( model, - coords: Union[np.ndarray, paddle.Tensor], - cells: Optional[Union[np.ndarray, paddle.Tensor]], - atom_types: Union[np.ndarray, paddle.Tensor, list[int]], - spins: Optional[Union[np.ndarray, paddle.Tensor]] = None, + coords: np.ndarray | paddle.Tensor, + cells: np.ndarray | paddle.Tensor | None, + atom_types: np.ndarray | paddle.Tensor | list[int], + spins: np.ndarray | paddle.Tensor | None = None, atomic: bool = False, infer_batch_size: int = 2, denoise: bool = False, @@ -79,7 +75,12 @@ def eval_model( if spins is not None: assert isinstance(spins, paddle.Tensor), err_msg assert isinstance(atom_types, paddle.Tensor) or isinstance(atom_types, list) - atom_types = paddle.to_tensor(atom_types, dtype=paddle.int32, place=DEVICE) + if isinstance(atom_types, paddle.Tensor): + atom_types = ( + atom_types.clone().detach().to(dtype=paddle.int32, device=DEVICE) + ) + else: + atom_types = paddle.to_tensor(atom_types, dtype=paddle.int32, place=DEVICE) elif isinstance(coords, np.ndarray): if cells is not None: assert isinstance(cells, np.ndarray), err_msg @@ -101,28 +102,57 @@ def eval_model( else: natoms = len(atom_types[0]) - coord_input = paddle.to_tensor( - coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE - ) - spin_input = None - if spins is not None: - spin_input = paddle.to_tensor( - spins.reshape([-1, natoms, 3]), + if isinstance(coords, paddle.Tensor): + coord_input = ( + coords.reshape([-1, natoms, 3]) + .clone() + .detach() + .to(dtype=GLOBAL_PD_FLOAT_PRECISION, device=DEVICE) + ) + else: + coord_input = paddle.to_tensor( + coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE, ) + spin_input = None + if spins is not None: + if isinstance(spins, paddle.Tensor): + spin_input = ( + spins.reshape([-1, natoms, 3]) + .clone() + .detach() + .to(dtype=GLOBAL_PD_FLOAT_PRECISION, device=DEVICE) + ) + else: + spin_input = paddle.to_tensor( + spins.reshape([-1, natoms, 3]), + dtype=GLOBAL_PD_FLOAT_PRECISION, + place=DEVICE, + ) has_spin = getattr(model, "has_spin", False) if callable(has_spin): has_spin = has_spin() - type_input = paddle.to_tensor(atom_types, dtype=paddle.int64, place=DEVICE) + if isinstance(atom_types, paddle.Tensor): + type_input = atom_types.clone().detach().to(dtype=paddle.int64, device=DEVICE) + else: + type_input = paddle.to_tensor(atom_types, dtype=paddle.int64, place=DEVICE) box_input = None if cells is None: pbc = False else: pbc = True - box_input = paddle.to_tensor( - cells.reshape([-1, 3, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE - ) + if isinstance(cells, paddle.Tensor): + box_input = ( + cells.reshape([-1, 3, 3]) + .clone() + .detach() + .to(dtype=GLOBAL_PD_FLOAT_PRECISION, device=DEVICE) + ) + else: + box_input = paddle.to_tensor( + cells.reshape([-1, 3, 3]), dtype=GLOBAL_PD_FLOAT_PRECISION, place=DEVICE + ) num_iter = int((nframes + infer_batch_size - 1) / infer_batch_size) for ii in range(num_iter): diff --git a/source/tests/pd/conftest.py b/source/tests/pd/conftest.py index 530cb18907..6320d59e7c 100644 --- a/source/tests/pd/conftest.py +++ b/source/tests/pd/conftest.py @@ -6,4 +6,5 @@ @pytest.fixture(scope="package", autouse=True) def clear_cuda_memory(request): yield - paddle.device.cuda.empty_cache() + if paddle.device.get_device() != "cpu": + paddle.device.empty_cache() diff --git a/source/tests/pd/model/test_atomic_model_atomic_stat.py b/source/tests/pd/model/test_atomic_model_atomic_stat.py index 93aa7b8905..e80ef42c6e 100644 --- a/source/tests/pd/model/test_atomic_model_atomic_stat.py +++ b/source/tests/pd/model/test_atomic_model_atomic_stat.py @@ -5,7 +5,7 @@ Path, ) from typing import ( - Optional, + NoReturn, ) import h5py @@ -80,11 +80,11 @@ def forward( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} @@ -114,10 +114,10 @@ def forward( class TestAtomicModelStat(unittest.TestCase, TestCaseSingleFrameWithNlist): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.merged_output_stat = [ { @@ -171,7 +171,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_output_stat(self): + def test_output_stat(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -237,10 +237,12 @@ def cvt_ret(x): expected_ret1["foo"] = ret0["foo"] + foo_bias[at] expected_ret1["bar"] = ret0["bar"] + bar_bias[at] for kk in ["foo", "bar"]: - np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) + np.testing.assert_almost_equal( + ret1[kk], expected_ret1[kk], err_msg=f"{kk} not equal" + ) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) @@ -284,10 +286,10 @@ def raise_error(): class TestAtomicModelStatMergeGlobalAtomic( unittest.TestCase, TestCaseSingleFrameWithNlist ): - def tearDown(self): + def tearDown(self) -> None: self.tempdir.cleanup() - def setUp(self): + def setUp(self) -> None: TestCaseSingleFrameWithNlist.setUp(self) self.merged_output_stat = [ { @@ -341,7 +343,7 @@ def setUp(self): pass self.stat_file_path = DPPath(h5file, "a") - def test_output_stat(self): + def test_output_stat(self) -> None: nf, nloc, nnei = self.nlist.shape ds = DescrptDPA1( self.rcut, @@ -401,7 +403,7 @@ def cvt_ret(x): np.testing.assert_almost_equal(ret1[kk], expected_ret1[kk]) # 3. test bias load from file - def raise_error(): + def raise_error() -> NoReturn: raise RuntimeError md0.compute_or_load_out_stat(raise_error, stat_file_path=self.stat_file_path) diff --git a/source/tests/pd/model/test_atomic_model_global_stat.py b/source/tests/pd/model/test_atomic_model_global_stat.py index abd7928a0f..053596bce2 100644 --- a/source/tests/pd/model/test_atomic_model_global_stat.py +++ b/source/tests/pd/model/test_atomic_model_global_stat.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import h5py import numpy as np @@ -92,11 +89,11 @@ def forward( self, descriptor: paddle.Tensor, atype: paddle.Tensor, - gr: Optional[paddle.Tensor] = None, - g2: Optional[paddle.Tensor] = None, - h2: Optional[paddle.Tensor] = None, - fparam: Optional[paddle.Tensor] = None, - aparam: Optional[paddle.Tensor] = None, + gr: paddle.Tensor | None = None, + g2: paddle.Tensor | None = None, + h2: paddle.Tensor | None = None, + fparam: paddle.Tensor | None = None, + aparam: paddle.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pd/model/test_deeppot.py b/source/tests/pd/model/test_deeppot.py new file mode 100644 index 0000000000..24696dea86 --- /dev/null +++ b/source/tests/pd/model/test_deeppot.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import unittest +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) + +import numpy as np +import paddle + +from deepmd.infer.deep_pot import DeepPot as DeepPotUni +from deepmd.pd.entrypoints.main import ( + freeze, + get_trainer, +) +from deepmd.pd.infer.deep_eval import ( + DeepPot, +) + + +class TestDeepPot(unittest.TestCase): + def setUp(self) -> None: + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + self.config["training"]["training_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + self.config["training"]["validation_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + self.input_json = "test_dp_test.json" + with open(self.input_json, "w") as fp: + json.dump(self.config, fp, indent=4) + + trainer = get_trainer(deepcopy(self.config)) + trainer.run() + + ori_place = paddle.device.get_device() + paddle.device.set_device("cpu") + input_dict, label_dict, _ = trainer.get_data(is_train=False) + paddle.device.set_device(ori_place) + trainer.wrapper(**input_dict, label=label_dict, cur_lr=1.0) + self.model = "model.pd" + + def tearDown(self) -> None: + for f in os.listdir("."): + if f in ["lcurve.out", self.input_json]: + os.remove(f) + + def test_dp_test(self) -> None: + dp = DeepPot(str(self.model)) + cell = np.array( + [ + 5.122106549439247480e00, + 4.016537340154059388e-01, + 6.951654033828678081e-01, + 4.016537340154059388e-01, + 6.112136112297989143e00, + 8.178091365465004481e-01, + 6.951654033828678081e-01, + 8.178091365465004481e-01, + 6.159552512682983760e00, + ] + ).reshape(1, 3, 3) + coord = np.array( + [ + 2.978060152121375648e00, + 3.588469695887098077e00, + 2.792459820604495491e00, + 3.895592322591093115e00, + 2.712091020667753760e00, + 1.366836847133650501e00, + 9.955616170888935690e-01, + 4.121324820711413039e00, + 1.817239061889086571e00, + 3.553661462345699906e00, + 5.313046969500791583e00, + 6.635182659098815883e00, + 6.088601018589653080e00, + 6.575011420004332585e00, + 6.825240650611076099e00, + ] + ).reshape(1, -1, 3) + atype = np.array([0, 0, 0, 1, 1]).reshape(1, -1) + + ret = dp.eval(coord, cell, atype, atomic=True) + e, f, v, ae, av = ret[0], ret[1], ret[2], ret[3], ret[4] + self.assertEqual(e.shape, (1, 1)) + self.assertEqual(f.shape, (1, 5, 3)) + self.assertEqual(v.shape, (1, 9)) + self.assertEqual(ae.shape, (1, 5, 1)) + self.assertEqual(av.shape, (1, 5, 9)) + + self.assertEqual(dp.get_type_map(), ["O", "H"]) + self.assertEqual(dp.get_ntypes(), 2) + self.assertEqual(dp.get_dim_fparam(), 0) + self.assertEqual(dp.get_dim_aparam(), 0) + self.assertEqual(dp.deep_eval.model_type, DeepPot) + + def test_uni(self) -> None: + dp = DeepPotUni("model.pd") + self.assertIsInstance(dp, DeepPot) + # its methods has been tested in test_dp_test + + def test_eval_typeebd(self) -> None: + dp = DeepPot(str(self.model)) + eval_typeebd = dp.eval_typeebd() + self.assertEqual( + eval_typeebd.shape, (len(self.config["model"]["type_map"]) + 1, 8) + ) + np.testing.assert_allclose(eval_typeebd[-1], np.zeros_like(eval_typeebd[-1])) + + +@unittest.skip(reason="Freezed model(.json) do not support getting attributes") +class TestDeepPotFrozen(TestDeepPot): + def setUp(self) -> None: + super().setUp() + frozen_model = "frozen_model.json" + freeze( + model=self.model, + output=frozen_model, + head=None, + do_atomic_virial=True, + ) + self.model = frozen_model + + # Note: this can not actually disable cuda device to be used + # only can be used to test whether devices are mismatched + @unittest.skipIf(not paddle.device.is_compiled_with_cuda(), "CUDA not available") + @unittest.mock.patch("deepmd.pd.utils.env.DEVICE", paddle.CPUPlace()) + @unittest.mock.patch("deepmd.pd.infer.deep_eval.DEVICE", paddle.CPUPlace()) + def test_dp_test_cpu(self) -> None: + self.test_dp_test() diff --git a/source/tests/pd/model/test_descriptor.py b/source/tests/pd/model/test_descriptor.py index dc78856851..293dbd9efc 100644 --- a/source/tests/pd/model/test_descriptor.py +++ b/source/tests/pd/model/test_descriptor.py @@ -13,6 +13,9 @@ Path, ) +from deepmd.common import ( + expand_sys_str, +) from deepmd.pd.model.descriptor import ( prod_env_mat, ) @@ -31,9 +34,6 @@ from deepmd.pd.utils.nlist import ( extend_input_and_build_neighbor_list, ) -from deepmd.tf.common import ( - expand_sys_str, -) from deepmd.tf.env import ( op_module, ) diff --git a/source/tests/pd/model/test_descriptor_dpa1.py b/source/tests/pd/model/test_descriptor_dpa1.py index bfcf4ba6ee..f4cf6a8005 100644 --- a/source/tests/pd/model/test_descriptor_dpa1.py +++ b/source/tests/pd/model/test_descriptor_dpa1.py @@ -368,7 +368,7 @@ def translate_se_atten_and_type_embd_dicts_to_dpa1( source_dict, type_embd_dict, ): - all_keys = list(target_dict.keys()) + all_keys = [key for key in target_dict.keys() if "buffer_" not in key] record = [False for ii in all_keys] for kk, vv in source_dict.items(): tk = "se_atten." + kk diff --git a/source/tests/pd/model/test_descriptor_dpa2.py b/source/tests/pd/model/test_descriptor_dpa2.py index 12017bb840..b8e48580d0 100644 --- a/source/tests/pd/model/test_descriptor_dpa2.py +++ b/source/tests/pd/model/test_descriptor_dpa2.py @@ -190,7 +190,7 @@ def translate_type_embd_dicts_to_dpa2( source_dict, type_embd_dict, ): - all_keys = list(target_dict.keys()) + all_keys = [key for key in target_dict.keys() if "buffer_" not in key] record = [False for ii in all_keys] for kk, vv in source_dict.items(): record[all_keys.index(kk)] = True diff --git a/source/tests/pd/model/test_dynamic_sel.py b/source/tests/pd/model/test_dynamic_sel.py new file mode 100644 index 0000000000..a605d97f85 --- /dev/null +++ b/source/tests/pd/model/test_dynamic_sel.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import paddle + +from deepmd.dpmodel.descriptor.dpa3 import ( + RepFlowArgs, +) +from deepmd.pd.model.descriptor import ( + DescrptDPA3, +) +from deepmd.pd.utils import ( + env, +) +from deepmd.pd.utils.env import ( + PRECISION_DICT, +) + +from ...seed import ( + GLOBAL_SEED, +) +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) + +dtype = env.GLOBAL_PD_FLOAT_PRECISION + + +class TestDescrptDPA3DynamicSel(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + + def test_consistency( + self, + ) -> None: + rng = np.random.default_rng(100) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for ( + ua, + rus, + ruri, + acr, + nme, + prec, + ect, + optim, + ) in itertools.product( + [True, False], # update_angle + ["res_residual"], # update_style + ["norm", "const"], # update_residual_init + [0, 1], # a_compress_rate + [1, 2], # n_multi_edge_message + ["float64"], # precision + [False], # use_econf_tebd + [True, False], # optim_update + ): + dtype = PRECISION_DICT[prec] + # rtol, atol = get_tols(prec) + rtol, atol = 1e-5, 1e-7 + if prec == "float64": + atol = 1e-8 # marginal GPU test cases... + + repflow = RepFlowArgs( + n_dim=20, + e_dim=10, + a_dim=10, + nlayers=3, + e_rcut=self.rcut, + e_rcut_smth=self.rcut_smth, + e_sel=nnei, + a_rcut=self.rcut - 0.1, + a_rcut_smth=self.rcut_smth, + a_sel=nnei, + a_compress_rate=acr, + n_multi_edge_message=nme, + axis_neuron=4, + update_angle=ua, + update_style=rus, + update_residual_init=ruri, + optim_update=optim, + smooth_edge_update=True, + sel_reduce_factor=1.0, # test consistent when sel_reduce_factor == 1.0 + ) + + # dpa3 new impl + dd0 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + repflow.use_dynamic_sel = True + + # dpa3 new impl + dd1 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + dd0.repflows.mean = paddle.to_tensor(davg, dtype=dtype).to( + device=env.DEVICE + ) + dd0.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype).to( + device=env.DEVICE + ) + rd0, _, _, _, _ = dd0( + paddle.to_tensor(self.coord_ext, dtype=dtype).to(device=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype=paddle.int64).to( + device=env.DEVICE + ), + paddle.to_tensor(self.nlist, dtype=paddle.int64).to(device=env.DEVICE), + paddle.to_tensor(self.mapping, dtype=paddle.int64).to( + device=env.DEVICE + ), + ) + # serialization + dd1.repflows.mean = paddle.to_tensor(davg, dtype=dtype).to( + device=env.DEVICE + ) + dd1.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype).to( + device=env.DEVICE + ) + rd1, _, _, _, _ = dd1( + paddle.to_tensor(self.coord_ext, dtype=dtype).to(device=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype=paddle.int64).to( + device=env.DEVICE + ), + paddle.to_tensor(self.nlist, dtype=paddle.int64).to(device=env.DEVICE), + paddle.to_tensor(self.mapping, dtype=paddle.int64).to( + device=env.DEVICE + ), + ) + np.testing.assert_allclose( + rd0.numpy(), + rd1.numpy(), + rtol=rtol, + atol=atol, + ) diff --git a/source/tests/pd/model/test_embedding_net.py b/source/tests/pd/model/test_embedding_net.py index 2dcc9f821b..80d2b77dc5 100644 --- a/source/tests/pd/model/test_embedding_net.py +++ b/source/tests/pd/model/test_embedding_net.py @@ -18,6 +18,9 @@ Path, ) +from deepmd.common import ( + expand_sys_str, +) from deepmd.pd.model.descriptor import ( DescrptSeA, ) @@ -34,9 +37,6 @@ from deepmd.pd.utils.nlist import ( extend_input_and_build_neighbor_list, ) -from deepmd.tf.common import ( - expand_sys_str, -) from deepmd.tf.descriptor import DescrptSeA as DescrptSeA_tf from ..test_finetune import ( diff --git a/source/tests/pd/model/test_force_grad.py b/source/tests/pd/model/test_force_grad.py index d7b569ef38..eb6975afec 100644 --- a/source/tests/pd/model/test_force_grad.py +++ b/source/tests/pd/model/test_force_grad.py @@ -5,9 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import paddle @@ -31,7 +28,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ): super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pd/model/test_model.py b/source/tests/pd/model/test_model.py index ce91fd3f21..e619171e44 100644 --- a/source/tests/pd/model/test_model.py +++ b/source/tests/pd/model/test_model.py @@ -19,6 +19,9 @@ Path, ) +from deepmd.common import ( + expand_sys_str, +) from deepmd.dpmodel.utils.learning_rate import LearningRateExp as MyLRExp from deepmd.pd.loss import ( EnergyStdLoss, @@ -32,9 +35,6 @@ from deepmd.pd.utils.env import ( DEVICE, ) -from deepmd.tf.common import ( - expand_sys_str, -) from deepmd.tf.descriptor import DescrptSeA as DescrptSeA_tf from deepmd.tf.fit import ( EnerFitting, @@ -400,7 +400,7 @@ def test_consistency(self) -> None: .detach() .numpy(), ) - self.assertIsNone(model_predict_1.get("atom_virial", None)) + # self.assertIsNone(model_predict_1.get("atom_virial", None)) np.testing.assert_allclose( head_dict["atom_virial"], p_atomic_virial.reshape(head_dict["atom_virial"].shape) diff --git a/source/tests/pd/model/test_rotation.py b/source/tests/pd/model/test_rotation.py index 94e3442631..48f5ae8983 100644 --- a/source/tests/pd/model/test_rotation.py +++ b/source/tests/pd/model/test_rotation.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import paddle @@ -29,7 +26,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ): super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pd/model/test_saveload_dpa1.py b/source/tests/pd/model/test_saveload_dpa1.py index 54a82e479a..0b8878534f 100644 --- a/source/tests/pd/model/test_saveload_dpa1.py +++ b/source/tests/pd/model/test_saveload_dpa1.py @@ -13,6 +13,9 @@ DataLoader, ) +from deepmd.common import ( + expand_sys_str, +) from deepmd.pd.loss import ( EnergyStdLoss, ) @@ -32,9 +35,6 @@ from deepmd.pd.utils.stat import ( make_stat_input, ) -from deepmd.tf.common import ( - expand_sys_str, -) def get_dataset(config): diff --git a/source/tests/pd/model/test_saveload_se_e2_a.py b/source/tests/pd/model/test_saveload_se_e2_a.py index c1c2ba2cdd..0eabbdb12e 100644 --- a/source/tests/pd/model/test_saveload_se_e2_a.py +++ b/source/tests/pd/model/test_saveload_se_e2_a.py @@ -13,6 +13,9 @@ DataLoader, ) +from deepmd.common import ( + expand_sys_str, +) from deepmd.pd.loss import ( EnergyStdLoss, ) @@ -32,9 +35,6 @@ from deepmd.pd.utils.stat import ( make_stat_input, ) -from deepmd.tf.common import ( - expand_sys_str, -) def get_dataset(config): diff --git a/source/tests/pd/test_multitask.py b/source/tests/pd/test_multitask.py index 0b85816f7b..f345630429 100644 --- a/source/tests/pd/test_multitask.py +++ b/source/tests/pd/test_multitask.py @@ -11,6 +11,7 @@ ) import numpy as np +import paddle from deepmd.pd.entrypoints.main import ( get_trainer, @@ -232,8 +233,15 @@ def setUp(self) -> None: self.config["model"], self.shared_links = preprocess_shared_params( self.config["model"] ) + if not paddle.device.is_compiled_with_cuda(): + self.FLAGS_use_stride_kernel = paddle.get_flags("FLAGS_use_stride_kernel")[ + "FLAGS_use_stride_kernel" + ] + paddle.set_flags({"FLAGS_use_stride_kernel": False}) def tearDown(self) -> None: + if not paddle.device.is_compiled_with_cuda(): + paddle.set_flags({"FLAGS_use_stride_kernel": self.FLAGS_use_stride_kernel}) MultiTaskTrainTest.tearDown(self) @@ -271,9 +279,17 @@ def setUp(self) -> None: self.config["model"], self.shared_links = preprocess_shared_params( self.config["model"] ) + self.config["learning_rate"]["start_lr"] = 1e-5 self.share_fitting = True + if not paddle.device.is_compiled_with_cuda(): + self.FLAGS_use_stride_kernel = paddle.get_flags("FLAGS_use_stride_kernel")[ + "FLAGS_use_stride_kernel" + ] + paddle.set_flags({"FLAGS_use_stride_kernel": False}) def tearDown(self) -> None: + if not paddle.device.is_compiled_with_cuda(): + paddle.set_flags({"FLAGS_use_stride_kernel": self.FLAGS_use_stride_kernel}) MultiTaskTrainTest.tearDown(self) diff --git a/source/tests/pd/test_sampler.py b/source/tests/pd/test_sampler.py index 2af5a9c05c..e920660704 100644 --- a/source/tests/pd/test_sampler.py +++ b/source/tests/pd/test_sampler.py @@ -13,15 +13,15 @@ DataLoader, ) +from deepmd.common import ( + expand_sys_str, +) from deepmd.pd.utils.dataloader import ( DpLoaderSet, get_weighted_sampler, ) -from deepmd.tf.common import ( - expand_sys_str, -) from deepmd.tf.utils import random as tf_random -from deepmd.tf.utils.data_system import ( +from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -55,6 +55,7 @@ def setUp(self): tf_random.seed(10) self.dp_dataset = DeepmdDataSystem(self.systems, self.batch_size, 1, self.rcut) + @unittest.skip("Temporarily skipped due to changes in numpy 2.4") def test_sampler_debug_info(self): dataloader = DataLoader( self.my_dataset, diff --git a/source/tests/pd/test_training.py b/source/tests/pd/test_training.py index 8958dcb165..692a8fb32f 100644 --- a/source/tests/pd/test_training.py +++ b/source/tests/pd/test_training.py @@ -89,7 +89,11 @@ def test_dp_train(self) -> None: state_dict_trained[state_key].numpy(), state_dict_finetuned_empty[state_key].numpy(), ) - if "fitting_net" not in state_key: + if ( + ("fitting_net" not in state_key) + or ("fparam" in state_key) + or ("aparam" in state_key) + ): np.testing.assert_allclose( state_dict_trained[state_key].numpy(), state_dict_finetuned_random[state_key].numpy(), @@ -132,11 +136,14 @@ def test_trainable(self) -> None: def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pd"): - os.remove(f) + if os.path.exists(f): + os.remove(f) if f in ["lcurve.out"]: - os.remove(f) + if os.path.exists(f): + os.remove(f) if f in ["stat_files"]: - shutil.rmtree(f) + if os.path.exists(f): + shutil.rmtree(f) class TestEnergyModelSeA(unittest.TestCase, DPTrainTest): @@ -150,9 +157,25 @@ def setUp(self) -> None: self.config["model"] = deepcopy(model_se_e2_a) self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 - # import paddle enable_prim(True) - # assert paddle.framework.core._is_eager_prim_enabled() + + def tearDown(self) -> None: + DPTrainTest.tearDown(self) + + +class TestEnergyModelGradientAccumulation(unittest.TestCase, DPTrainTest): + def setUp(self) -> None: + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + data_file = [str(Path(__file__).parent / "water/data/data_0")] + self.config["training"]["training_data"]["systems"] = data_file + self.config["training"]["validation_data"]["systems"] = data_file + self.config["model"] = deepcopy(model_se_e2_a) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + self.config["training"]["acc_freq"] = 4 + enable_prim(True) def tearDown(self) -> None: DPTrainTest.tearDown(self) @@ -174,6 +197,7 @@ def setUp(self) -> None: self.config["training"]["save_freq"] = 1 self.set_path = Path(__file__).parent / "water/data/data_0" / "set.000" shutil.copyfile(self.set_path / "energy.npy", self.set_path / "fparam.npy") + self.config["model"]["data_stat_nbatch"] = 100 def tearDown(self) -> None: (self.set_path / "fparam.npy").unlink(missing_ok=True) diff --git a/source/tests/pt/common.py b/source/tests/pt/common.py index 8709c8b4f9..0592be1423 100644 --- a/source/tests/pt/common.py +++ b/source/tests/pt/common.py @@ -1,9 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import pathlib -from typing import ( - Optional, - Union, -) import numpy as np import torch @@ -51,10 +47,10 @@ def run_dp(cmd: str) -> int: def eval_model( model, - coords: Union[np.ndarray, torch.Tensor], - cells: Optional[Union[np.ndarray, torch.Tensor]], - atom_types: Union[np.ndarray, torch.Tensor, list[int]], - spins: Optional[Union[np.ndarray, torch.Tensor]] = None, + coords: np.ndarray | torch.Tensor, + cells: np.ndarray | torch.Tensor | None, + atom_types: np.ndarray | torch.Tensor | list[int], + spins: np.ndarray | torch.Tensor | None = None, atomic: bool = False, infer_batch_size: int = 2, denoise: bool = False, @@ -79,7 +75,12 @@ def eval_model( if spins is not None: assert isinstance(spins, torch.Tensor), err_msg assert isinstance(atom_types, torch.Tensor) or isinstance(atom_types, list) - atom_types = torch.tensor(atom_types, dtype=torch.int32, device=DEVICE) + if isinstance(atom_types, torch.Tensor): + atom_types = ( + atom_types.clone().detach().to(dtype=torch.int32, device=DEVICE) + ) + else: + atom_types = torch.tensor(atom_types, dtype=torch.int32, device=DEVICE) elif isinstance(coords, np.ndarray): if cells is not None: assert isinstance(cells, np.ndarray), err_msg @@ -101,28 +102,59 @@ def eval_model( else: natoms = len(atom_types[0]) - coord_input = torch.tensor( - coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE - ) - spin_input = None - if spins is not None: - spin_input = torch.tensor( - spins.reshape([-1, natoms, 3]), + if isinstance(coords, torch.Tensor): + coord_input = ( + coords.reshape([-1, natoms, 3]) + .clone() + .detach() + .to(dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE) + ) + else: + coord_input = torch.tensor( + coords.reshape([-1, natoms, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE, ) + spin_input = None + if spins is not None: + if isinstance(spins, torch.Tensor): + spin_input = ( + spins.reshape([-1, natoms, 3]) + .clone() + .detach() + .to(dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE) + ) + else: + spin_input = torch.tensor( + spins.reshape([-1, natoms, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) has_spin = getattr(model, "has_spin", False) if callable(has_spin): has_spin = has_spin() - type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE) + if isinstance(atom_types, torch.Tensor): + type_input = atom_types.clone().detach().to(dtype=torch.long, device=DEVICE) + else: + type_input = torch.tensor(atom_types, dtype=torch.long, device=DEVICE) box_input = None if cells is None: pbc = False else: pbc = True - box_input = torch.tensor( - cells.reshape([-1, 3, 3]), dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE - ) + if isinstance(cells, torch.Tensor): + box_input = ( + cells.reshape([-1, 3, 3]) + .clone() + .detach() + .to(dtype=GLOBAL_PT_FLOAT_PRECISION, device=DEVICE) + ) + else: + box_input = torch.tensor( + cells.reshape([-1, 3, 3]), + dtype=GLOBAL_PT_FLOAT_PRECISION, + device=DEVICE, + ) num_iter = int((nframes + infer_batch_size - 1) / infer_batch_size) for ii in range(num_iter): diff --git a/source/tests/pt/model/models/se_e2_a.json b/source/tests/pt/model/models/se_e2_a.json new file mode 100644 index 0000000000..d3b83b5a5c --- /dev/null +++ b/source/tests/pt/model/models/se_e2_a.json @@ -0,0 +1,77 @@ +{ + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_e2_a", + "sel": [ + 23, + 46 + ], + "rcut_smth": 0.50, + "rcut": 4.00, + "neuron": [ + 2, + 4 + ], + "resnet_dt": false, + "axis_neuron": 4, + "type_one_side": true, + "seed": 1, + "_comment": " that's all" + }, + "fitting_net": { + "neuron": [ + 120 + ], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" + }, + "data_stat_nbatch": 20, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "_comment": " that's all" + }, + "training": { + "training_data": { + "systems": [ + "../data/data_0", + "../data/data_1", + "../data/data_2" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../data/data_3" + ], + "batch_size": 1, + "numb_btch": 3, + "_comment": "that's all" + }, + "numb_steps": 100000, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 10000, + "_comment": "that's all" + }, + "_comment": "that's all" +} diff --git a/source/tests/pt/model/models/se_e2_a.pbtxt b/source/tests/pt/model/models/se_e2_a.pbtxt new file mode 100644 index 0000000000..c928fa375b --- /dev/null +++ b/source/tests/pt/model/models/se_e2_a.pbtxt @@ -0,0 +1,17035 @@ +node { + name: "train_attr/min_nbor_dist" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + } + double_val: 0.9103434477930997 + } + } + } +} +node { + name: "train_attr/training_script" + op: "Const" + attr { + key: "dtype" + value { + type: DT_STRING + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_STRING + tensor_shape { + } + string_val: "{\"model\":{\"type_map\":[\"O\",\"H\"],\"descriptor\":{\"type\":\"se_e2_a\",\"sel\":[23,46],\"rcut_smth\":0.5,\"rcut\":4.0,\"neuron\":[2,4],\"resnet_dt\":false,\"axis_neuron\":4,\"type_one_side\":true,\"precision\":\"float64\",\"seed\":1,\"activation_function\":\"tanh\",\"trainable\":true,\"exclude_types\":[],\"env_protection\":0.0,\"set_davg_zero\":false},\"fitting_net\":{\"neuron\":[120],\"resnet_dt\":true,\"precision\":\"float64\",\"seed\":1,\"type\":\"ener\",\"numb_fparam\":0,\"numb_aparam\":0,\"default_fparam\":null,\"dim_case_embd\":0,\"activation_function\":\"tanh\",\"trainable\":true,\"rcond\":null,\"atom_ener\":[],\"use_aparam_as_mask\":false},\"data_stat_nbatch\":10,\"data_stat_protect\":0.01,\"data_bias_nsample\":10,\"pair_exclude_types\":[],\"atom_exclude_types\":[],\"preset_out_bias\":null,\"srtab_add_bias\":true,\"type\":\"standard\",\"model_branch_alias\":[],\"info\":{}},\"learning_rate\":{\"type\":\"exp\",\"decay_steps\":5000,\"start_lr\":0.001,\"stop_lr\":3.51e-08,\"scale_by_worker\":\"linear\",\"decay_rate\":null},\"loss\":{\"type\":\"ener\",\"start_pref_e\":0.02,\"limit_pref_e\":1,\"start_pref_f\":1000,\"limit_pref_f\":1,\"start_pref_v\":0,\"limit_pref_v\":0,\"start_pref_h\":0.0,\"limit_pref_h\":0.0,\"start_pref_ae\":0.0,\"limit_pref_ae\":0.0,\"start_pref_pf\":0.0,\"limit_pref_pf\":0.0,\"enable_atom_ener_coeff\":false,\"start_pref_gf\":0.0,\"limit_pref_gf\":0.0,\"numb_generalized_coord\":0,\"use_huber\":false,\"huber_delta\":0.01},\"training\":{\"training_data\":{\"systems\":[\"../data/single\"],\"batch_size\":\"auto\",\"rglob_patterns\":null,\"auto_prob\":\"prob_sys_size\",\"sys_probs\":null},\"validation_data\":{\"systems\":[\"../data/single\"],\"batch_size\":1,\"numb_btch\":3,\"rglob_patterns\":null,\"auto_prob\":\"prob_sys_size\",\"sys_probs\":null},\"numb_steps\":100,\"seed\":10,\"disp_file\":\"lcurve.out\",\"disp_freq\":100,\"save_freq\":1000,\"save_ckpt\":\"model.ckpt\",\"max_ckpt_keep\":5,\"change_bias_after_training\":false,\"disp_training\":true,\"time_training\":true,\"disp_avg\":false,\"profiling\":false,\"profiling_file\":\"timeline.json\",\"enable_profiler\":false,\"tensorboard\":false,\"tensorboard_log_dir\":\"log\",\"tensorboard_freq\":1,\"acc_freq\":1,\"opt_type\":\"Adam\"}}" + } + } + } +} +node { + name: "model_type" + op: "Const" + attr { + key: "dtype" + value { + type: DT_STRING + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_STRING + tensor_shape { + } + string_val: "original_model" + } + } + } +} +node { + name: "t_box" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "shape" + value { + shape { + dim { + size: -1 + } + } + } + } +} +node { + name: "t_coord" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "shape" + value { + shape { + dim { + size: -1 + } + } + } + } +} +node { + name: "t_type" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "shape" + value { + shape { + dim { + size: -1 + } + } + } + } +} +node { + name: "t_natoms" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "shape" + value { + shape { + dim { + size: 4 + } + } + } + } +} +node { + name: "t_mesh" + op: "Placeholder" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "shape" + value { + shape { + dim { + size: -1 + } + } + } + } +} +node { + name: "model_attr/tmap" + op: "Const" + attr { + key: "dtype" + value { + type: DT_STRING + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_STRING + tensor_shape { + } + string_val: "O H" + } + } + } +} +node { + name: "model_attr/model_type" + op: "Const" + attr { + key: "dtype" + value { + type: DT_STRING + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_STRING + tensor_shape { + } + string_val: "ener" + } + } + } +} +node { + name: "model_attr/model_version" + op: "Const" + attr { + key: "dtype" + value { + type: DT_STRING + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_STRING + tensor_shape { + } + string_val: "1.1" + } + } + } +} +node { + name: "model_attr/model_attr/t_out_bias" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 1 + } + dim { + size: 2 + } + dim { + size: 1 + } + } + tensor_content: "\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "model_attr/model_attr/t_out_bias/read" + op: "Identity" + input: "model_attr/model_attr/t_out_bias" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@model_attr/model_attr/t_out_bias" + } + } + } +} +node { + name: "model_attr/model_attr/t_out_std" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 1 + } + dim { + size: 2 + } + dim { + size: 1 + } + } + tensor_content: "\000\000\000\000\000\000\360?\000\000\000\000\000\000\360?" + } + } + } +} +node { + name: "model_attr/model_attr/t_out_std/read" + op: "Identity" + input: "model_attr/model_attr/t_out_std" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@model_attr/model_attr/t_out_std" + } + } + } +} +node { + name: "strided_slice/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice/stack" + input: "strided_slice/stack_1" + input: "strided_slice/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "mul/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "mul" + op: "Mul" + input: "strided_slice" + input: "mul/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape/shape" + op: "Pack" + input: "Reshape/shape/0" + input: "mul" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape" + op: "Reshape" + input: "t_coord" + input: "Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_1/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_1/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_1/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_1" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_1/stack" + input: "strided_slice_1/stack_1" + input: "strided_slice_1/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_1/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_1/shape" + op: "Pack" + input: "Reshape_1/shape/0" + input: "strided_slice_1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_1" + op: "Reshape" + input: "t_type" + input: "Reshape_1/shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "descrpt_attr/rcut" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + } + double_val: 4.0 + } + } + } +} +node { + name: "descrpt_attr/ntypes" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "descrpt_attr/sel" + op: "Const" + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\027\000\000\000.\000\000\000" + } + } + } +} +node { + name: "descrpt_attr/original_sel" + op: "Const" + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\027\000\000\000.\000\000\000" + } + } + } +} +node { + name: "descrpt_attr/t_avg" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + dim { + size: 276 + } + } + tensor_content: "\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\177\320&\365X\250?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\221\322-?\366\245?\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "descrpt_attr/t_avg/read" + op: "Identity" + input: "descrpt_attr/t_avg" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@descrpt_attr/t_avg" + } + } + } +} +node { + name: "descrpt_attr/t_std" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + dim { + size: 276 + } + } + tensor_content: "\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\327\177:\3328\216\306?;\342AW\370\371\272?;\342AW\370\371\272?;\342AW\370\371\272?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?\210<\245\222\211\262\302?yQ&\335_\200\266?yQ&\335_\200\266?yQ&\335_\200\266?" + } + } + } +} +node { + name: "descrpt_attr/t_std/read" + op: "Identity" + input: "descrpt_attr/t_std" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@descrpt_attr/t_std" + } + } + } +} +node { + name: "strided_slice_3/stack" + op: "Const" + input: "^descrpt_attr/original_sel" + input: "^descrpt_attr/sel" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_3/stack_1" + op: "Const" + input: "^descrpt_attr/original_sel" + input: "^descrpt_attr/sel" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_3/stack_2" + op: "Const" + input: "^descrpt_attr/original_sel" + input: "^descrpt_attr/sel" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_3" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_3/stack" + input: "strided_slice_3/stack_1" + input: "strided_slice_3/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "mul_1/y" + op: "Const" + input: "^descrpt_attr/original_sel" + input: "^descrpt_attr/sel" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "mul_1" + op: "Mul" + input: "strided_slice_3" + input: "mul_1/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape_2/shape/0" + op: "Const" + input: "^descrpt_attr/original_sel" + input: "^descrpt_attr/sel" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_2/shape" + op: "Pack" + input: "Reshape_2/shape/0" + input: "mul_1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_2" + op: "Reshape" + input: "Reshape" + input: "Reshape_2/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape_3/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\t\000\000\000" + } + } + } +} +node { + name: "Reshape_3" + op: "Reshape" + input: "t_box" + input: "Reshape_3/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_4/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_4/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_4/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_4" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_4/stack" + input: "strided_slice_4/stack_1" + input: "strided_slice_4/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_4/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_4/shape" + op: "Pack" + input: "Reshape_4/shape/0" + input: "strided_slice_4" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_4" + op: "Reshape" + input: "Reshape_1" + input: "Reshape_4/shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "ProdEnvMatA" + op: "ProdEnvMatA" + input: "Reshape_2" + input: "Reshape_4" + input: "t_natoms" + input: "Reshape_3" + input: "t_mesh" + input: "descrpt_attr/t_avg/read" + input: "descrpt_attr/t_std/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "rcut_a" + value { + f: -1.0 + } + } + attr { + key: "rcut_r" + value { + f: 4.0 + } + } + attr { + key: "rcut_r_smth" + value { + f: 0.5 + } + } + attr { + key: "sel_a" + value { + list { + i: 23 + i: 46 + } + } + } + attr { + key: "sel_r" + value { + list { + i: 0 + i: 0 + } + } + } +} +node { + name: "Reshape_7/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\024\001\000\000" + } + } + } +} +node { + name: "Reshape_7" + op: "Reshape" + input: "ProdEnvMatA" + input: "Reshape_7/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "o_rmat" + op: "Identity" + input: "Reshape_7" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "o_rmat_deriv" + op: "Identity" + input: "ProdEnvMatA:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "o_rij" + op: "Identity" + input: "ProdEnvMatA:2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "o_nlist" + op: "Identity" + input: "ProdEnvMatA:3" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_5/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_5/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_5/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_5" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_5/stack" + input: "strided_slice_5/stack_1" + input: "strided_slice_5/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_8/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_8/shape/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 276 + } + } + } +} +node { + name: "Reshape_8/shape" + op: "Pack" + input: "Reshape_8/shape/0" + input: "strided_slice_5" + input: "Reshape_8/shape/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_8" + op: "Reshape" + input: "o_rmat" + input: "Reshape_8/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape_9/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\024\001\000\000" + } + } + } +} +node { + name: "Reshape_9" + op: "Reshape" + input: "Reshape_8" + input: "Reshape_9/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/Slice/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice/size" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\\\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice" + op: "Slice" + input: "Reshape_9" + input: "filter_type_all/Slice/begin" + input: "filter_type_all/Slice/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Shape" + op: "Shape" + input: "filter_type_all/Slice" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "filter_type_all/strided_slice/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "filter_type_all/strided_slice/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/strided_slice/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/strided_slice" + op: "StridedSlice" + input: "filter_type_all/Shape" + input: "filter_type_all/strided_slice/stack" + input: "filter_type_all/strided_slice/stack_1" + input: "filter_type_all/strided_slice/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "filter_type_all/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape" + op: "Reshape" + input: "filter_type_all/Slice" + input: "filter_type_all/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/Slice_1/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_1/size" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_1" + op: "Slice" + input: "filter_type_all/Reshape" + input: "filter_type_all/Slice_1/begin" + input: "filter_type_all/Slice_1/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_1/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_1" + op: "Reshape" + input: "filter_type_all/Slice_1" + input: "filter_type_all/Reshape_1/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/matrix_1_0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 1 + } + dim { + size: 2 + } + } + tensor_content: "(\350\310\014\010o\255\2772\337\220\353&\346\340?" + } + } + } +} +node { + name: "filter_type_all/matrix_1_0/read" + op: "Identity" + input: "filter_type_all/matrix_1_0" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/matrix_1_0" + } + } + } +} +node { + name: "filter_type_all/bias_1_0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "Y\2130\240\003\260\335\277m\375`\217\372\211\000\300" + } + } + } +} +node { + name: "filter_type_all/bias_1_0/read" + op: "Identity" + input: "filter_type_all/bias_1_0" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/bias_1_0" + } + } + } +} +node { + name: "filter_type_all/MatMul" + op: "MatMul" + input: "filter_type_all/Reshape_1" + input: "filter_type_all/matrix_1_0/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "filter_type_all/BiasAdd" + op: "BiasAdd" + input: "filter_type_all/MatMul" + input: "filter_type_all/bias_1_0/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "filter_type_all/Tanh" + op: "Tanh" + input: "filter_type_all/BiasAdd" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_2/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\002\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_2" + op: "Reshape" + input: "filter_type_all/Tanh" + input: "filter_type_all/Reshape_2/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/concat" + op: "ConcatV2" + input: "filter_type_all/Reshape_1" + input: "filter_type_all/Reshape_1" + input: "filter_type_all/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/add" + op: "AddV2" + input: "filter_type_all/concat" + input: "filter_type_all/Reshape_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/matrix_2_0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + dim { + size: 4 + } + } + tensor_content: "\233\353C\217\023\227\253?\213\317\355\362\250&\312?^\256\023\342g\343\333\277J\345u\2118R\320?L{\254\033W\312\244?-g\263\335\302\\\351\277\301\253*\264r\234\330?\025\244Y\240\017\014\275\277" + } + } + } +} +node { + name: "filter_type_all/matrix_2_0/read" + op: "Identity" + input: "filter_type_all/matrix_2_0" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/matrix_2_0" + } + } + } +} +node { + name: "filter_type_all/bias_2_0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 4 + } + } + tensor_content: "x\231b\217\234\201\366\277\263\010\n-E0\354\277\202\373|\243\222\300\317?\306\371\210\203\016\245\316\277" + } + } + } +} +node { + name: "filter_type_all/bias_2_0/read" + op: "Identity" + input: "filter_type_all/bias_2_0" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/bias_2_0" + } + } + } +} +node { + name: "filter_type_all/MatMul_1" + op: "MatMul" + input: "filter_type_all/add" + input: "filter_type_all/matrix_2_0/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "filter_type_all/BiasAdd_1" + op: "BiasAdd" + input: "filter_type_all/MatMul_1" + input: "filter_type_all/bias_2_0/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "filter_type_all/Tanh_1" + op: "Tanh" + input: "filter_type_all/BiasAdd_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_3/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_3" + op: "Reshape" + input: "filter_type_all/Tanh_1" + input: "filter_type_all/Reshape_3/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/concat_1/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/concat_1" + op: "ConcatV2" + input: "filter_type_all/add" + input: "filter_type_all/add" + input: "filter_type_all/concat_1/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/add_1" + op: "AddV2" + input: "filter_type_all/concat_1" + input: "filter_type_all/Reshape_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_4/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\377\377\377\377\027\000\000\000\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_4" + op: "Reshape" + input: "filter_type_all/add_1" + input: "filter_type_all/Reshape_4/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/Reshape_5/shape/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 23 + } + } + } +} +node { + name: "filter_type_all/Reshape_5/shape/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 4 + } + } + } +} +node { + name: "filter_type_all/Reshape_5/shape" + op: "Pack" + input: "filter_type_all/strided_slice" + input: "filter_type_all/Reshape_5/shape/1" + input: "filter_type_all/Reshape_5/shape/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "filter_type_all/Reshape_5" + op: "Reshape" + input: "filter_type_all/Slice" + input: "filter_type_all/Reshape_5/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/MatMul_2" + op: "BatchMatMulV2" + input: "filter_type_all/Reshape_5" + input: "filter_type_all/Reshape_4" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: true + } + } + attr { + key: "adj_y" + value { + b: false + } + } + attr { + key: "grad_x" + value { + b: false + } + } + attr { + key: "grad_y" + value { + b: false + } + } +} +node { + name: "filter_type_all/Slice_2/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\000\000\000\000\\\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_2/size" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\270\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_2" + op: "Slice" + input: "Reshape_9" + input: "filter_type_all/Slice_2/begin" + input: "filter_type_all/Slice_2/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Shape_1" + op: "Shape" + input: "filter_type_all/Slice_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "filter_type_all/strided_slice_1/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "filter_type_all/strided_slice_1/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/strided_slice_1/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/strided_slice_1" + op: "StridedSlice" + input: "filter_type_all/Shape_1" + input: "filter_type_all/strided_slice_1/stack" + input: "filter_type_all/strided_slice_1/stack_1" + input: "filter_type_all/strided_slice_1/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "filter_type_all/Reshape_6/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_6" + op: "Reshape" + input: "filter_type_all/Slice_2" + input: "filter_type_all/Reshape_6/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/Slice_3/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_3/size" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_3" + op: "Slice" + input: "filter_type_all/Reshape_6" + input: "filter_type_all/Slice_3/begin" + input: "filter_type_all/Slice_3/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_7/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_7" + op: "Reshape" + input: "filter_type_all/Slice_3" + input: "filter_type_all/Reshape_7/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/matrix_1_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 1 + } + dim { + size: 2 + } + } + tensor_content: "\317\221\332=\'\355\337?\237\224\230\311\376\r\270\277" + } + } + } +} +node { + name: "filter_type_all/matrix_1_1/read" + op: "Identity" + input: "filter_type_all/matrix_1_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/matrix_1_1" + } + } + } +} +node { + name: "filter_type_all/bias_1_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\0372\016\0074\345\321\277=\363\302\320@]\344\277" + } + } + } +} +node { + name: "filter_type_all/bias_1_1/read" + op: "Identity" + input: "filter_type_all/bias_1_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/bias_1_1" + } + } + } +} +node { + name: "filter_type_all/MatMul_3" + op: "MatMul" + input: "filter_type_all/Reshape_7" + input: "filter_type_all/matrix_1_1/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "filter_type_all/BiasAdd_2" + op: "BiasAdd" + input: "filter_type_all/MatMul_3" + input: "filter_type_all/bias_1_1/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "filter_type_all/Tanh_2" + op: "Tanh" + input: "filter_type_all/BiasAdd_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_8/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\002\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_8" + op: "Reshape" + input: "filter_type_all/Tanh_2" + input: "filter_type_all/Reshape_8/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/concat_2/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/concat_2" + op: "ConcatV2" + input: "filter_type_all/Reshape_7" + input: "filter_type_all/Reshape_7" + input: "filter_type_all/concat_2/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/add_2" + op: "AddV2" + input: "filter_type_all/concat_2" + input: "filter_type_all/Reshape_8" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/matrix_2_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + dim { + size: 4 + } + } + tensor_content: "$4\265\372=\004\335?\3271\224K\344\206\334\277\274]\\W\024\177\332?q\305\257\367V>\354?I?\345G\205\212\257?\324\253\255\002\031\204\306?\322&\3603\320\273\334?Ux\2526\2024\321\277" + } + } + } +} +node { + name: "filter_type_all/matrix_2_1/read" + op: "Identity" + input: "filter_type_all/matrix_2_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/matrix_2_1" + } + } + } +} +node { + name: "filter_type_all/bias_2_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 4 + } + } + tensor_content: "N\361\035]x\245\351\277\357\010\224A\343B\311\277\322x\017\"M\253\273?9\027K\2228e\374\277" + } + } + } +} +node { + name: "filter_type_all/bias_2_1/read" + op: "Identity" + input: "filter_type_all/bias_2_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@filter_type_all/bias_2_1" + } + } + } +} +node { + name: "filter_type_all/MatMul_4" + op: "MatMul" + input: "filter_type_all/add_2" + input: "filter_type_all/matrix_2_1/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "filter_type_all/BiasAdd_3" + op: "BiasAdd" + input: "filter_type_all/MatMul_4" + input: "filter_type_all/bias_2_1/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "filter_type_all/Tanh_3" + op: "Tanh" + input: "filter_type_all/BiasAdd_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_9/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_9" + op: "Reshape" + input: "filter_type_all/Tanh_3" + input: "filter_type_all/Reshape_9/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/concat_3/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "filter_type_all/concat_3" + op: "ConcatV2" + input: "filter_type_all/add_2" + input: "filter_type_all/add_2" + input: "filter_type_all/concat_3/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/add_3" + op: "AddV2" + input: "filter_type_all/concat_3" + input: "filter_type_all/Reshape_9" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Reshape_10/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\377\377\377\377.\000\000\000\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_10" + op: "Reshape" + input: "filter_type_all/add_3" + input: "filter_type_all/Reshape_10/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/Reshape_11/shape/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 46 + } + } + } +} +node { + name: "filter_type_all/Reshape_11/shape/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 4 + } + } + } +} +node { + name: "filter_type_all/Reshape_11/shape" + op: "Pack" + input: "filter_type_all/strided_slice_1" + input: "filter_type_all/Reshape_11/shape/1" + input: "filter_type_all/Reshape_11/shape/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "filter_type_all/Reshape_11" + op: "Reshape" + input: "filter_type_all/Slice_2" + input: "filter_type_all/Reshape_11/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "filter_type_all/MatMul_5" + op: "BatchMatMulV2" + input: "filter_type_all/Reshape_11" + input: "filter_type_all/Reshape_10" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: true + } + } + attr { + key: "adj_y" + value { + b: false + } + } + attr { + key: "grad_x" + value { + b: false + } + } + attr { + key: "grad_y" + value { + b: false + } + } +} +node { + name: "filter_type_all/AddN" + op: "AddN" + input: "filter_type_all/MatMul_2" + input: "filter_type_all/MatMul_5" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/truediv/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + } + double_val: 69.0 + } + } + } +} +node { + name: "filter_type_all/truediv" + op: "RealDiv" + input: "filter_type_all/AddN" + input: "filter_type_all/truediv/y" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/Slice_4/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\000\000\000\000\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_4/size" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\377\377\377\377\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Slice_4" + op: "Slice" + input: "filter_type_all/truediv" + input: "filter_type_all/Slice_4/begin" + input: "filter_type_all/Slice_4/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "filter_type_all/MatMul_6" + op: "BatchMatMulV2" + input: "filter_type_all/truediv" + input: "filter_type_all/Slice_4" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: true + } + } + attr { + key: "adj_y" + value { + b: false + } + } + attr { + key: "grad_x" + value { + b: false + } + } + attr { + key: "grad_y" + value { + b: false + } + } +} +node { + name: "filter_type_all/Reshape_12/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\020\000\000\000" + } + } + } +} +node { + name: "filter_type_all/Reshape_12" + op: "Reshape" + input: "filter_type_all/MatMul_6" + input: "filter_type_all/Reshape_12/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Shape_2" + op: "Shape" + input: "Reshape_8" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "strided_slice_9/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_9/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_9/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_9" + op: "StridedSlice" + input: "Shape_2" + input: "strided_slice_9/stack" + input: "strided_slice_9/stack_1" + input: "strided_slice_9/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_10/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_10/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_10/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_10" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_10/stack" + input: "strided_slice_10/stack_1" + input: "strided_slice_10/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_12/shape/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 16 + } + } + } +} +node { + name: "Reshape_12/shape" + op: "Pack" + input: "strided_slice_9" + input: "strided_slice_10" + input: "Reshape_12/shape/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_12" + op: "Reshape" + input: "filter_type_all/Reshape_12" + input: "Reshape_12/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "concat_1/concat" + op: "Identity" + input: "Reshape_12" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "o_descriptor" + op: "Identity" + input: "concat_1/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "fitting_attr/dfparam" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "fitting_attr/daparam" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "fitting_attr/t_bias_atom_e" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\232\231\231\031\220dW\300\231\231\231\031\220dg\300" + } + } + } +} +node { + name: "fitting_attr/t_bias_atom_e/read" + op: "Identity" + input: "fitting_attr/t_bias_atom_e" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@fitting_attr/t_bias_atom_e" + } + } + } +} +node { + name: "strided_slice_13/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_13/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_13/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_13" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_13/stack" + input: "strided_slice_13/stack_1" + input: "strided_slice_13/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_14/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_14/shape/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 16 + } + } + } +} +node { + name: "Reshape_14/shape" + op: "Pack" + input: "Reshape_14/shape/0" + input: "strided_slice_13" + input: "Reshape_14/shape/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_14" + op: "Reshape" + input: "o_descriptor" + input: "Reshape_14/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_14/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_14/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_14/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_14" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_14/stack" + input: "strided_slice_14/stack_1" + input: "strided_slice_14/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_15/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_15/shape" + op: "Pack" + input: "Reshape_15/shape/0" + input: "strided_slice_14" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_15" + op: "Reshape" + input: "t_type" + input: "Reshape_15/shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_15/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_15/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_15/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_15" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_15/stack" + input: "strided_slice_15/stack_1" + input: "strided_slice_15/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Slice_1/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "Slice_1/size/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Slice_1/size" + op: "Pack" + input: "Slice_1/size/0" + input: "strided_slice_15" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Slice_1" + op: "Slice" + input: "Reshape_15" + input: "Slice_1/begin" + input: "Slice_1/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "GreaterEqual/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "GreaterEqual" + op: "GreaterEqual" + input: "Slice_1" + input: "GreaterEqual/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "Cast" + op: "Cast" + input: "GreaterEqual" + attr { + key: "DstT" + value { + type: DT_DOUBLE + } + } + attr { + key: "SrcT" + value { + type: DT_BOOL + } + } + attr { + key: "Truncate" + value { + b: false + } + } +} +node { + name: "Reshape_16/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_16" + op: "Reshape" + input: "Slice_1" + input: "Reshape_16/shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "clip_by_value/Minimum/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "clip_by_value/Minimum" + op: "Minimum" + input: "Reshape_16" + input: "clip_by_value/Minimum/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "clip_by_value/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "clip_by_value" + op: "Maximum" + input: "clip_by_value/Minimum" + input: "clip_by_value/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_16/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_16/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 3 + } + } + } +} +node { + name: "strided_slice_16/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_16" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_16/stack" + input: "strided_slice_16/stack_1" + input: "strided_slice_16/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Slice_2/begin" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\000\000\000\000\000\000\000\000\000\000\000\000" + } + } + } +} +node { + name: "Slice_2/size/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Slice_2/size/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Slice_2/size" + op: "Pack" + input: "Slice_2/size/0" + input: "strided_slice_16" + input: "Slice_2/size/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Slice_2" + op: "Slice" + input: "Reshape_14" + input: "Slice_2/begin" + input: "Slice_2/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "Reshape_17/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\020\000\000\000" + } + } + } +} +node { + name: "Reshape_17" + op: "Reshape" + input: "Slice_2" + input: "Reshape_17/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "layer_0_type_0/matrix" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 16 + } + dim { + size: 120 + } + } + tensor_content: "e0\352\312O\'\224?0\307E\334p\265\236\277\262\272\257\354\370}\262\277\331.)u\')\226\277\014)A\3328\362\304?\261\335OyT.\235\277\254j>Z\037\023\240?\031\031\3618\273\307\302\277\200\366^8Js\254?\265W\253\210fN\260\277\354\245&\276`\206\213?/\024\014\244F\'\266\277S\261<\247c\023\200?WtW\327\211\300\300?^eh\271\235\213\260?\334\263\016\204\210(\243\277\355\345\016\253\365\005\267\277\216n\263\201\377\372\312?h+{sV\343\235\277J\312\025m\003\376\245?W\n]\247\031\353\255?g\354\332\177)b\252?zS\333\355\364\376\222?\305\010N\224\375\210\237\2771\224O\032\034\013\242\277e\013\241Gtn\220? N-\365\343\353h?\336Y,\'E\367\306\277\277&\372cV\301\267\277\024\365Z7\205H\246\277$5a\231\026\241\262\277\200\223$\\\353\177\246\277 jY~\'\215\214\277\t\256\201\371:Z\267\277\243\2735\347\336\272\261\277 =;x\237U\232\277\254\352g\221\231p\267?a\315m\n\362+\230?\336\233o\323\367I\201\277\006[\315\312\364\325\261\277o\031\353\2634\363\254\277\216\366\336\240\347\344\304?\224\303w\225\027c\273\277\001z\265~i\341\234\277\215\207\035+\337f\257?\013\3462\022\307\026\251\277\033g+T\025\014\276\277\022d%G\022(u\277\332)\002\262a$\220?\003\330\016#\345<\271?\355\335\337\026Z\264\235\277\202\306I\"i\220\237\277c?\201yO\217\303?Nl&E\201\272\221\2775\211Q\265\240k\312\277\337\204\370\247V\233\246?k\031\255\371\002\254\214?\204\032z\345\314;\241\277\364\377(R\364\t@\277HV\342\211l\262g?\025i\372\353y\200\267?^!(\177\333\200z?|\302\377\312\202,\274\277\301G\026 ]n\231?[\370\3605\225*\252\277\t@\303\252\360\035p?\005\306\004\244c\336\247\277].\244\014\205\362\276?\n\021\005i\222O\221?\376J\002\344B\215\306?\331xZn\332\257}\277*\261?\204\214\010J\333\347y\277\252AL\257\031;\255\277s\361S\260s#q?\256\225\202\n\n:\223\277\252\206\021Oz<\275?P\362i\376\232[\263?W\366p\314{=\300\277Qj\360\362\224\267u\277\361;\r\226 G\212\277\347D\226\263\334\262\177\277&)\317\215O\007\303?\345p\272-\021\327f?\371\214[Yi\244\301?\006\200(\327r\021\264?\267\256\335V\002\345\235?\217?\262\217\254`\274\277\370I\346\245\354\200\300?\335>\264\304\205B\222?\210d\"\232\2237\303\277\224rO\321su\253\277\271\017Pm1\227\262?\206\220I\225\272\213\305?\252\361)F[=\221?\000\250\303\314\t\030\240?x\010\370\"\0242\303?]p\223&\373\'\246?\250\330\230\027\237\314\241?\033\020\214?Ce\305\277\233\3564%\334d|?\244\235\324\352\377\247\232?[SS\325\315\334\215\277\027\317=\306\000\245\263\277`\337\337\2348\351\260\277\n\3470\035\352\357\263?\230\265\344t\007\302\277?\224\272\'}x\305\240?\303)n\032@J\370\276\365b\037\027\214Xf?\016\r\276\311CU\301\277aFR\020\345\231\271?\220\313.\213\377\344\273\277\321\322\2679\305\231\270?\355\351\025\372\244\204\247\277p:!\246sy\254\2772\300*\224QN\275?\362\2602/\330\220\230\277\341\323:r\275\242\303\277\303\340\020\367\332\220j\277\225\313\232\220\255\027\223\277\252f\202\273\212\242\271\277D\212yh=\236\227\277\235\261\322;8\334\273?n\360\360\032\356\253\234\277\016\036\211\257\255n\237?\240\303\371h\rY\272\277`x\020\377\346o\246?\001\231\261\234\2613\274\277/\036]\361E\301\262\277\346\202\"\016\001\305\213?\231K\010\002\257\267\230\277\337\2725\230\335\366\273?|\034&\347G\362\273\277z\331h\2154\333\265\277p\202&\2029\300\250\2776A W\257/\271\277\313x\352\350\033F\247?Y\364\373d\203\222e\277\306EW(\341\251\250\277\334H\272QM\007\257?\211Ln\277\201\211\264?\007\000\037\2516\257\236?yO\"7\353\023\265\277\351\204mi\243\035\300?\3010\370\212\270S\272?\315\267\214\265s\014\266\277;\250\357\3339%\261?7\322\333\2060A\260?\230`\264y\367\277\277?\203\230\202\372\247\256\243\277B\205\000gn\330\273?\275\004&$\265\206[\277?\220\206\235,\276u\277/\346{^\203\373\254\277\314k\2623\000\344\266?!2ny\237\t\247?O\267&\n\352u\237?\256\002\316\226\3601\271\277c\233}\002\334\032\240?u\031\006\373B\357\257\277\326\006{\345\350\360n\277\260\330\024\211\244\326\250\277\2638KdN\320\253\277\360/\240\031WE\250?K\010^\232\376\272\253\277g/\223T\020\253\250\277\200\327`_*j\313\277N\016\346i\033~\254?>\364\256\000\256J\271?\325\033\255)\t\245\260?]|M\014\035f\262?\232c\351\322\322\316\260\277\337\010\032~],\273\277\236\230\"\206ZR\267?q\031\267\313\323\303\272\277\256\3711\026\213[\302?\234Ib\202\207\216\266\277\244\271\204\303v\221\250\2777\304\207\321\2355\260?5\315\031\234\365\023\215\277\272\2547\316\322\306\237?ZO}\214\003\215\220\277\036\251-\347\255\331\220?\237\203\256=\013\036\220\277\3219\200\025R\004\257?\300\230\375\021\261\223\240?\304\323\022J\263\240\264\277\004\035|\030&+\277\277\366<\227Hbx\231?\022\342I\'Y\265\242\277\035\310\277\367=9\260?N\244\340\321c\345\215\277I\265\307F+j\243?\317wF\306.z\204\277\205\201\307!\032I\246?\347\361=NU\330\271?\217\nk\206\226\001\254\277\343\330\2130n\250\241\277\370@?\017\026T\274\277\205\271\265\370\235A\225\277\251MjX\275\n\262?-\030\206<\315\007\205\277\224\234,\013\022\374\262?\365\327\343\013=8\211?\324\321F\232\000\010\203\277@W\320\rG\234\215\277|\254\216\036\363-p?\007\272\266\3145&\247?\350\354\221\254\324\326\243\2778Rk\301\377\\\270\277\005\027\256\334\246_i?\264\356\026\304\203:\270\277s\226\350Z\305\326\250\277k\342\262z\242\237\250\277=\345\336z7\253\223?\333t\000$Z\253\300?q\346\211w$\\\275?\236\263\343\251,\003\267\277q\267N%\370\311\276\277\327\271q\345\370W\232?\203\331\030\246\016_\306?\342\251\203\206n\354\240?\247\3327\262\341\370\254?\017`\375\017\315\232\241\277&\340x\221\356\001\225\277\301\353\211\032\341?\242?\221\240r\256\177h\256\277\225\266M\023\027\321\260?\\\3269\351hE\223\2778)\326\303h\303\263?\357\215lKk\032\223\277\374\245\017\033\300\372\253?\267K\001\004\370%\276\277\315x(\346\363_\312?\316[\231\231Q9\302?\n\254\022P\014x\264\277\" d\376\334\307\244?\216\240\255\344W\336\244?Lk\243 \005t\213\277\304]\332#\352\261\232?\256\010\236\325\242\342\256\277\222\rD\037J\313\241?4\367gq\360\240\241?d\245?\310\277\374\310\277\204YZ\265\304\334\240?\032\266[G}\247\277?\217F\321\033\355\n\241\277\274\025\320\244\325\237\253?\377\310\260\270\333\007\204\277\227\237\331\263S\314\241\277\362\311g\344W\241\201\277\224\1775\005\315V\211?\273\252\305ai\372\273\277\317\317VOR8\270?\271D\335\237\273X\227?\203\315\204\327\302J\306\277<\300B6\000k\214\277-Z\016\225\344q\257\277\313\327\236\036=\230\277?leD\212O\352\260?\362\0270\332Y\345\260?\017\022\371\203\275\210\305\277\007\2531\014(\207\220?\343?\217ey\362\241\277\225\206\323Xvm\257\277\214\320\347\2472\367\250?|My\226\237\353\236\277F\245o\237\204\016\252\277\243h\351eE\317b\277rU\300C\034\252\210?\266j\376m\234c\252?\234\325\031\205m\227\315\277\210\343]\333\275\306\276?\320\320\250K\301\032\251\277\355\364\321\367\031\031\271\277\223BUAV\201\264\277R\002\274\355\323\312\240?O\240 \326\320~\251\277\365\312\300}\036/\277\277\262\263#y\304w\254\2776\336\213\345n^{\277\301\365\313\022\274\370\306?\017\257\365C\374\212\244?\222/\324\352?\306\261\277\362\'\025d{\005\243?\2657\224\340\"}~?]\237\263\322\031\272\260\277\330?t\004\203\330\265?\221\024=\344\316\373\261\277\323\373\272M}\r\260?\3649j\336\374\266\300\277\'0\360A\017@\276\277\3166\013\037\205a\243\277\241i\245q\000\017\254?N\311\223\370\024\177\275\277\276\336\225\350\210\324\270\277-2\374\013\2504\240\277\370\301\337\274\235\233\251\277c\355\347\307\214w\217\277\330ILQWy\265?\274}\251L\237\250\256\277\271\224\013Yf\324\242\277\222C\224-\206E\277\277K\200\rr0\371\244\277[\010\252\3168m\235?\212m/\341\'\252\241\277Q&=\035A\177\301\277\353u\306\010|7\202\302\203\277-\361`\354\254:\241\277U5\017\3446`\246\277\365\027$\241%\331\265\277\214[\n<\250\207\266\277\343\244;\324C\233\275?\253\031*\221\000\352s\277\r\207\340\310\320{\305\277\3740<\342\232\305~\277\225(n\365\256eR\277/\241\333\252.1\230\277)N\251\223\237;\272?\262\r\231B6\177\266\277\336\237\230s\302\361\262?\211\202\036r\2505\265?\342\303\226\037\356\026\246?\256\022_\200\022y\241?\317R\276X\022\210\263\277\365\272H\236\017\016\245\277\037*\234%\031\267\302\277\026rG\200\331\375\260\277\213<\224\235\014o\222\277xhR<7\267\251\277\024\306\370\252I\336\240\277-[\222@\330\361\214?\020\241\035\266\250\266\224\277\201\234\317\306\352\311\235\277\230&\200\213\3559\255\277#G3\367yQ\240?\253\366s\361\017K\260?\034\322d\302c4\251?H^\206\344\343\211\263\277\234\260\000\376W\323\260?\302\n0\306?\373\264?\3416\306\021R\372\310\277\213aq\206\005\003\244\277\\\223K\335<|\273?\325\027K^\263\246\266?\353\005]\244\355\243\266?\010m\236\361\235\362\301\277\241:\265\017\343\300\263\277\241\312\000\014\351\026\245?z\222\247\262\025K\266\277\231\t\241\345\327\371\301\277\344G\033\030\236[\241\277\265\343\310\026Z\225\230\2775[\026OT\n\261?n3\3622\327{\305?\276H\240\251b\227\262\277G\031\200\241\034\341\240\277=+\002\321:t\303?\277\306\320\262\364\034\266\277\303\177\270~\251\035\263\277G\260c\206\311\004\246\277{\367\235<\3255\205?\216\026\336M\002\210\255\277p\0141S-\233\224\277\300\3668\244\326\347\224\277\306\206\360\345\370:\274?\020-\371\242\231\211\254\277s\003L@!\376y\277`\262\013\243\375\377\227?\242\271\343\032\252\347\307\277\345\326\273\r\375\254\261\277#\255\275\276\237\312x?\315\001\226\020C\362\241\277qR\263T;\\\300\277\347\r\234\224\375\320\263?Z\224\007\360\310y\301?\215n\325\212\305\253\266\277\327#\344=o\262\267\277\301i~\257\255J\225\277\321\370\002\310/X\227?\366\307\225\037nr\267?T\017=K\364\243\262\277\312\330\016\302\032\215\232\277\236\003*0\312\311\252?|_\027j)_\273\277\203[\021\325\340b\315\277\266X\270\267\010\321\233?\213K\201\210\030e\272\277\315E\252a\362\200\255?T\211A\004`\253\260\277I6\016)\271\300\243\277\230Q\234\313c~p\277a\025\372\202\010\007\255\277 \212\272\324\006$\300?P\275dTti\307?s\356\260\277a\207\256?+}a\336ZXs\277\345\247\230\n\243\346\302?\327\t\2516\341v\261\277\367\260]i<\314\271\277?zx\346\033\222\246\2777\200\234 k\241\242\2770\272\217f\240\035\237?\037T\370\2647\367\177?\211c&\255\241\214\317\277\374r\003\231\340\335\303\277T%O\311\263\273\237?\357\336\2146\222\352\255?>\246\310\227\323\026\262?\322+\371\277\027\355\220\277=\273\365\201\212\036\256?VB\206h\210<\304\2773\024;\273(]\301?\371\2547A`j\262?]\227\367\345@\314\241\2772i\\q\366\343\201\277[\367\362\307G\002\245?\344\021\255\275\265q\262\277*\034\370\021\032#\267?\274\202\262\366\215\251\264\277\350,\2352\212\373\252\277\261{\322\276O\\\241?\333\270\240W\377\004\264?\343\306\201\231\265\332\263\277&=7\004\021\217\267?\366\241:\275P+\230\277f\376\023\243\324\266p?\214\033#\312\264\214\261?p\014vL\033X\253\277\310A>\211\035K\221\277\323\270m8\272\315`\277\365\224\204\230)\361\256\277H\344\232\214E\036\253?\251I\274r/\207\220?\263.\236\253,\t\252?\361\330b\234z>x?d\212\035\026?y\244\277\027\217D\267\371\221\312?\377\222\216\261.8\241?\026\227K8\344!\262?\216\346\213\014=\"\265?\251\320d\233/\027\260?-Q\270\245\032\365\301?>\\\035\374\030\273\252\277\205%\257\216eI\227\277\277\320\226\204\026\243\221\277\224d~\214\355\376\256\277\350\204g\315\202\013\300\277\350\"8Xw#\260\277\n\330\334\344\243;\230\277\220\241<\001\227S\212?\253\020\275o\234\245\260\277\324\017\355\242\335\204\271\277\243\302R\275\255\233\266\277\336\270\002:d)\310\277\355\226v\356\350\036t?\366\235\256\363q-\272?;L%\270\030\317\260\277\237\021\362y\037\023\275?\362\272\250\2018\000\261\277=A\n\270\342Q\217\277\310\267&S\021B\305?[\342\373\366H\023\262?X\nH\361\004h\277?\265(\245\272[~\270?\024[\324\275sd\251\277\322\344\t[\200\020\220?1\246\335\266\215\315\255\277\305\267\267\"\311\017\263\277\250\212\217\023\300\017g\277\017\220\256\350\244\366\264\277i6\335\303u\324\247\277{\314?\267c\377\217?\251>\270\032\005\350\276\277m\213\254\225\013~\200\277\250\026\2215\252O\267?\355\321.\007h+\230\277%c]|\305\271t?\230\363\006\245\303\243\277\277\036O\222O*\365\246?n/]\363\r\030\300?\216\253\323\003\320+\303\277\267\221b\213fV\222\277^\367uvo{\202?y\256\004\357\331f\266?W>}8Z\212\270\277\271\3624\242\311 \314\277h\222\006&Qu\314\277\037\315n\004\302\222\252?\317\016y\265\023.\276?w\370\261Q\221\341\275\2772\334N\355el\240\277\020~\271|;\350\267\277\372\2034eF/\214?Y\346Sv\236U\262\277y\327!\030%\260\226\277\022u\347\231\033\246\302?\326\331\345\014>y\261?|\205\013\r\253\270\255?{\307N\334\257Sy\277U\005\240\0018\277\300?\354\357\331\300\354;~\277<\003\033\315\233\276\232?u|\231W\345\017x?\013=\321\337\311l\254?&\322\r\201=\3368?\217@F\272\357\276\246?O\207\022\270\235\037\274\277\353\373\210\270\211\006\302\277\217)@\024\004\222k?n\305?}\341\304\311\277G\255\271\274\032\347\251?\003G4\033\037W\300?P\306\342\254\206b\273\277\203\272f=h\200\246?P\265w+\335\025\260?\202z\271\372\245;\235\277\002I\213\315\253\303\221?\007L&\221nO\250\277\024\210\030{\006\326\242?\237\276\364n\243\2778\266\323\272\010\222\300? 3x\246\272\320\240?\177\030\031h7*\200\277\262\2044\033\271\373\226\277Hyn\220\353tp?\245.we\331\t\264\277\301\365\2510\373\343\260\2778\210\034\332\010|\306?\n\302\2605\260\354x\277\235\345\036\342z!\247?7\317P|3\264\264?\272C0\316\202/\260??\346W\357\376\207\302?b\004J\323\347\035\244?I\000\'x!\033\226?7\270\020\262)p\300\277\253\337\316,\311z\277?\361\227\007\006m\355\261\277a\261)\270T\032\262\277\233\324\001\277\321k\210\277\335\236pp\200#\224\277ex\013/!\313s\277\214\274.F5i\245?\037\202v\257s2\244?5\300\347\274\240\226\241?T\022\275\\\007m\263\277j\036`\030\345\341\230?\272\357\272f\016\303\253?%qb\367\372\003\305?\331\301\013&\232\333\246?UU\206\0052P\221?p\342a=\310\224\224?-np\023\003\305\205\277\n|\354\232\355\262\241?b\271\223\200H\t\302?Y*\310\007\340\244\300?\366\343\340W\2223\244?\215\276\214\215w\331\251?G\200\344\211\277\251\265\277\3111\212J,h\321?2RE\257\300\373\260\277\014\344><<\326\261?9@\255.\330\260\246?\354\354D[\340\022\244?\n5\352\262\331I\300\277\302Z(r3\315\253?\'\250\244u\377\301}?K\240\307t\020x\274?\243\317\232t\006H\266\277\366vh\2642\013\217\277\213\376\263\354\204Q\230\277\310\325\250\241\275\330\246\277G (\277\242e\252?\350\242\002\352\317\262\243?\005\304\275\335\245\301\232\277\0235\247\"Uh\250?\026\203)1\262\222\302\277\003\315\347\307:H\302?\326\233[\301M\200\222?`&\275\312_\252\233\277\360\333=\216\336\323\276\277 .\"}\013{\227\277\337\226\322&\371\r\261?\257\350\346\210\221`\234?G\252j\214\324.\305?\237\237h\316[\346\267\277\246l\030qF>\235?\333BW\251\0052\237\277\377,\253\320\273zL?*\224\330\227\307\211\213?\014\014\325?^\306\300\277Y\211\236c\205\033\240\277\252\312>\'\002_\234\277\342s\372*\0274\276\277\360o\361\002\221s\215\277j\023\276\304\373\017\220\277\t\010{gR\233\222\277\302}\304\325\004\220\261\27703\327\352A|\260?\311\327\301\016x \202?\224\217\024\237\004\315\265\277\0204 \221;\373\203??\\\006\374\347n\247\277D\355\263e\355Q\305?$p\216a\336_\274?\245H\034Ch\332\274\277\333\022\n\203\213?\310\277\355\321\235\274\357g\261\277\211\231\"tj\363\256\277\320\222:\347\2216\257?\035\256UW)j\251\277X\2127\312\3379\304?\005\342T\3146.~\277\205\022\234\212P\341\200\277>\337\\Yb+V?\276\261\260\363\235\366\224\277\252\322\2571\025R\231?\031C\2347\230\332\305?\037\224\377%63\211?!\014\353\033\233\216\275?}b\336(Z\236\273\277\253=\326Ex\370\234\277\364\224\210\202\327N\307\277\341\233\216TFM\266\277\026<\276g\377\372\235\277\030!\350\013\213\232\275?\227\323\216D\215F\257?\243\024\205\224\035O\314\277*\"%\205\317\263\310?\264\313\242\325\311\344\243\277\002\227\241\271\371\200\300\277\306\256\346.`\204\233?\t\355Q\206\211\320\304?\275\253\334\235\3605i\277[\211\350\363\005\316s\277\203\221N\356O\t\220?2!eQ\021\241\251\277\275|\227\025\003c\262?\267\262\324\361\356\177\261\277\206\351\245\302\225n\262?\244Q\373)\001\270\273\277\227\364G_DS\266\277\351\263\252+R\306\235?\231Q\300Ju\240\262?:\024\263a\315\014\252\277\263\372\245\266F\271\276\277}9\344\264\306\347v?\027\340\325\217\220*\212?w\265\201\t\261E\237?\273\027|%$\314\222?7\026\314\000B\n\267?WT\257@n\233\305?\002\275\327K\205F\242?\362\337\374.D\257\265\277\355\271\025\272x+\260\277\265L\246\365\336\215\261?wA\367IJ\n\260?\216\265Z\020\0238]?\232\222s\231\307j\263\277o\217\244\244s\273\241\277\026yR\232U\211%\277\272\'\270A\307\215\262\277N\353\355\231\345\001\246\277z{\267L|Y\265?\3420\355\021\263&\251?G\267\256\311\250f\270?%O\3204\"t\307\277\277\"42\206)\300?\2635U\324\256\372\244\277%\234\217\222\240\254\243\277y\277<\354F\263\220\277T\367\"(\025\316\301?\002[\326M\014\372\262?\\fs?\315i\240?8p~\304\250\033\215?\204\320ZC]\364b\277X\233X\360\250H\250?N5\342\352\205\321\235\277^EcT\265-\237?\023m9\353\023\200\251\277\374\372\204\331\251a\230\277\200L\023\341\253\377\222\277\250\371\2678`\345\254?\270\0179\212\353q\300?/\376\177\355\221\311g\277\314h\177\027\273?\305?\024^1_\356\240\304?\307~\327/7\037\203\277\252]\374\254J<\265??\337\333_\311I\267\277\326\315D^U\272\203\277R`\200\303\200\326\262\277l<\326N\371,\222?\370a&\205*\021\266\277*\271\347\264B!\307\277i\002\247\313\374V\250\277\243\277R](\243\271\277\324p\201\305\313\375\204\277\340\353N\2455C\302?\257za\313Vf\250\277\035\262\352g\2551\314?Z\363\032\234\2279\277?\325!^\336\267\032\241\277\021\014\334\'\2401\253?\215;;fG\004\256\2771/&\363P\303\213\277\034\265\263O\321\022\255?:\255\333\356\313\201\301\277\331i,*\020L{\277\026p\020vG\006\270\277\312\245\206\367\224}\231?\321O\022\253\002K\256\2775\032.M\307\263\242?k\324\305\261\303\\\262\277\323f\035\021\034+\261\277\240u\021\231\215o\227?:Lw\230\232l\230\277\340\306\006V\276\033\235\277\230\245x\2747Q\254?\326Di\027\3427\267\277L\'\270\245\326\016\301?\306\"/\364\374\373\210\277/\223z\005\3229\225\277\226o\375\r\354{\305?\261\273Ps\205\232\231?\214nQ\n5\234\225?\314d\010\274G6\247\277$\362S&\226\r\301\277\247\335`\322\303Q\301?\2730\225\240S\316\255\277T\\\337*\006s\223?7T\001\352\275\325\307?#P\025\253|\362\262?w\003\253X\025\013\206\277&\323Xz\032\205\260\277&E\311\205\0261\227\277\225\233{\356MT\260?N\302\003\333\300\332\234\277 \006\001\211\262\255\206\277\260zz\030\3633\261\277\205\211\023\031\037\037\270\277\244\3434\373\271\006\260?\376\260\2174\335[\305?\306\2277\274\024\350\273\277\001\'\251VK\033|?\214\242\341%W\263\313\277FEA\341G\237\244?#\230\022a%\355\261\277\031\355\326n\257\207\256?%\354\006\205-\252\261\277a`\211\234\224\316\300?fj\314RKA\223?\320|\331\322U\237\243\277\267\340+\007b\035\244?Y\2333\257\200\305\276\277ar\223\233T\376\234\277\2001\274\301\023%\303\277\344>9\304\251\026\274\277\212-w\357Z\356\300?\201\216\307\357\"\020\226\277az\307\276\217\217\257?\240E\310\307\336\310\270\277\251\272;I\271\010\222?T\021\357\360\223%\230?\347\234U/q\345\250\277\345\255\304Q\321\200\254\277\332\235y\036\234\214\237\277\245\010G\333\366\212\261?&\200\340\364\362\343\247\277\263\022%\353\327\324\262\277yC\333\034@~\300\277\tL\361\270E\362\315\277\023D\371:\220\366\255\277\021\224Zo\324\325\227\277\2268z$\353Q\276?\307\267\\\261W\256\311\277\303;lG\2322\212\277\317\020\2200H\364\271?\205fo\347!\310\273?\237\264\274\323\310\310\225\277\2166\315\312\303\361\\\277u\375\340[\326{\251\277\374^i\n\226\203|?ZS\325\213|\206\257?h\313\217\240\225\203\224\277\2615\202\262\204X\231?\337\025\025,HI\302\277*\360\316\301\022\366\200\277\357k\326\300i\"\221\277\007\364\202\327\300\032\262\277\324I\002d\347\373\221?`\256\367A\020\035\230\277\251K7[\201\266\224\277\211\316\302\210guh\277<{\265\363I\017H?V\337Ob\351\257\252\277\341\306\211\216\272\377\224\277`\214eS\222\213\307?\005\237\366\245\345u\266?8\353/\234U\206z\277JX\034rA\r\274\277\340\376\222\212\372*\300?[\240\361\206\330\272\246?\373\210\342\300\247\361\247?\002\305\354;\014\037\273\277M\272\240)\001:\222\277\340\240R\355\tH\251\277\365\215\\\220\313\277\235?\001OPCT\033\262\277<\3054\340\274\317\237?D\266\303\212$M\305\2777i\272\227\036\331\254? 2\336\220%M\233?\024\006O\310\204db?\215\022\234\006^\342\246\277\025t\324\203\356\3336\277\373lR\n&\361\260\277\027\275v\3110m\267?\365/y\036\362\333\202\277M\2506\225\347\025\307?S\204\225.(e\302\2778P&\201NN\273\27741\316\215\356\362\300?\216\217\314\001H\\\257?\003v\360D\227|\306\277\272\325v\033\247\025\260\277D\210\021O\267\375\262\277G\223\221\243\247\310w\277\0342\320\347}f\306\277I\001\rV\322\254\220\277\'4\247\234\317\326\237\277\310HW2\010P\252\277G\357\023\0147I\264\277M\361G\350\334\352\304?<\315\302\014>\366\261\277d\203._h \240?\372s\034\230\266\215\262\277\335!*R\337\233\300\277\006\360\202\267\247q\263?\336\267\000\213\360_\272\277\315b\376}\371L\207\277G>\257\"<\315|\277n\340\274\273\275[\223?\030K\326mI\351\255?\261SA[c\021\250\277\332\207\246YC\212\236?\256\225=\242\370=\225\277\211\001\311\357\317\030\255\277\320\244\351_vE\276?\250\027\277F\004\362\233?\305g\020`\213\266\213?\351P\212W(P\273\277\205\207&\304\033\304\307?\216\3100\332\007\022\243?I\327\220\252\300\314\221?n\362\252}\372\276\257?%\316\"|\332\327\223?\356\234\270N\211\333\243?^L\325s\367J\237?x\246\327Qw\317\255?\223\326*\354\202\212\261?\376C\352\214\364\363R\277k\217\312\3674\265\241\277\035j\t\361\343`\257?0e\255JO\270\272?%\352\274<\323\374v?\362f\tt\275\017\311\277\223\323\275\356\241@\210\277\322\257n\312\352N\264?h\351\233\374d\022o?\252!c\212\264\276\214\277\221i8\\K\031\274?\367\271\353\376\021\344\260?\2614\\\236\3509\242\277\270g/\354\0202\236\277\234\035\250HE\365\305\277\032\0203\304\220~\260?\271\213=W~\313\245?r\276\377\241\226\027\303?i\326\342>\207\242\267\277\035\021\377\204\273\261z\277\001d\304\215\374-\245?\252A\226\310a\177\272?\312\357\341\265\354\353\252?z\233(\206\t\213\276?\346j\343\312\017r\272?\271\253&s\007N\264?\206{B\322rd\245\277W\340\207}\374\274\301\277\244\303\2643\'4\245?<5\232\365\247\324\206?f?^\2071\247\264\277\277\340\352\006C\202\254\277\316\260@\006\301\342\246?qb\"\304V\250\231?<\223\204C9\035\273?\032\2416\304\273w\214?\371\257s\233\025\271\274\277\335\263M\217=\330\225?4g\215\223w\203\301\2773\311\252\2719\321\301\277@\032\251\307\216U\217\277\014\021So8\324\205\277\263\n\357p-Z\241?\371U\364\362\000\321\276\277\026{\006/AG~\277\375\342\363\224\331\354\300?\363).w\301\320\220?:\230\002+\270?\262\277W>\022oB\027\307\277\001D\000\361\240\326\267?\342C\023t\364Y\277?\321\225\252Zc\202b?\3038\226\242\371d\216?Kb\365>\014H\273?rU\221\351+Z\307\277\324\340\016C\247\312\272?M)\216E\024d\241\277\032p\2215\331\374b\277\241b;\036\310\221\264\277\2627i\344sK\262?Cv<\201\"\002\235?=>\341\013\344\0236?\251)]\205\031\212\301?\216\033x\275A\001\310?P\324\013\'qt\214\277\241\n>[X\236\245?\235>=\312\357\261\223?\312\3348\343xc\250\277\245\002\326\243\312\310\254\277\017x\024\215\225\247\200?\257\242\265\376\036\n\263?\267\261G\277\0012\237\277\030\332*\376\\1\276\277T\342!\nWo\277?\257\347\351Q\007R\275\277[\365\n\306\360\353\255\277L\304]4\201\355\263?\375\324\014\314h\366\235?&\331\232D:\367\237?\351f\312#\256\314\254?8@\032\2303\246\306?@?\034\316\0367\272?\271%\2356\261?pm\\\025q\241\203?\300\302\273\302J\270\224\277.\203`$;\347\244?\005\342\372e*\311\263?S\014\263=\253\033\221?2\303\206:]\004\246?\320\213n08H\232\277\3024 \352/`\257\277\266\371\357\354\262\251\250?M\023+l\004K\214?\336n\336\031\214b\243\277\\\214\265\215lQ\250\277\207\222A\353\006\373\231?\375\nnul\247\262?\320\363\177\333\356Z\263?\005Z\370\301 \013\266?v\353\036\3567\337\254\277\\\340~\274\366!\301?\025\226~\002\362f\255\277Z\202K\3548O\263\277y\232\010\361\214\013\247\277\036b\275\346\335D\235\277\240\354\234/\021$\237?\206\232\032L&\250\237\277)\244;\001\324\252\251?\014\203\026\2300\210u\277\030\rF\374\021\232\246?\214\336e\305\202\027\222?\025\030C2\237\353\270\277\035\277\213\330t\332\261\277S6\204p\252\214\241?zH\316\235\324\354\234?w\346\036|\342\212K?\246B\342\322\275\345\305\277\366\323B\2238\361\207\277\004Ew\215\320\251\236?B\211!\035\262\201\227\277fl\337\322\352\001\216?iipx\237\034\262?\353D\264K\206\242\274\277\267\233\353\\N\327\260\277R\347bZ\222\345\273\277\270\024\030\361gF\252?_\275 Sl\310\226?^j\231\362n\242\221?\372\006\234\312d\031\225\277\036\313\024\'\200\313\247\277?\344\324 `\367\246?\341pn#\353\320\266\277R\234\337\265M\306\273?To\274\364\203Av?Y\247A\014\320\345\270\277}\025\250\354\033\351\205\277YF\010\244\337\350\212?\325\313\014\242s\310\240?D\365\266\254in\301?\322=\306\014\253+\253\277\225%c\330e\023\267\277\"~ eJ\232\243?\020\341\006U\020\n\240\277\025}\271Ku\200\245?\364!\202\310\204{\300?a\004\0334\264\177\270\277&\'\325?\342\216\274?\341Mp\376[\233\261\277\221d6{D\034\235?\231k\020\357s \213\277\370)\003\244|\216\314?q,\372)a\305\240\277\200\227`\223\014:\226?w\330N\265\014F\300?\302\367#\000\241\364\264?R\340\357VP\032\270\277\312\225#\213P9`?|\2411\332\376\256\206?\034|\347\215&\006d\277\277\240U\335Ik\265\277\333\306{\362\340l\263?\322~\025\000\356<\252?\234\001\004\013p\317\225?\'\367\037J\252\010\263?=\376\234n\351\213\247?\037\253\\\005\r\201\274\277q\272#\200\347\306\301\277\323\322\263\251>\323\223\277\\\225A<\204\216\261\277\301\006\304\351\036\260s?z\274\270\215\301\206\220\277\260\031\345\372\"0\266\277\375\315\311\r\363\303\262?\023\2418\304\002\202\201\277sD\227\007\314\206\250?R\344\035\311\235l\307\277\253\240\037\242\210\017\307?\005\374\266\004\002y\270?\242\365r\026\322\330\247?\"Ow^\307;\260\277\314)@\273\251q\270\277,\273\261Z\355\275\241\277\0057\013y\321\364\302\277\354\372\201\007\177Y\257?\346\253\032\307\231e\304\277\355\361N\374\026<\236?B\0001=\311Nw\277T\326\360\037\305\222\224\277f\217d\025\230y\225?\013|\235\302\333/\270?\232\352.^\024z\252\277}\210x]M\037r\277\233#\221g\371\353\260?\267\211\336\362\317\215\267?_\307K\241i\351\257?u\'\264\3501\315\265\277\261\301\324\312#+\251?Sy\027{\021\275\273\277;=\266\266\236\350\226\277\347?\372\314/\030\263?\345\034D\272\"\254\261\277\302\261\211\004\324\333\307\277\2419y\235W\034\242\2772VO\'\3124\267\277+y\21176\r\235?\347\201\371\341\346\310J\2774z\205\223\246 \234?\224C\336\005\342\320\234\277\024\336\264\223\2331\303\277\353p\336\034I\232\270\277\350O\240\356\255V\275\277\370\314\3326\261\222\234\277\2443\177\332\250\363\237?\326\220\375\370g\207\242\277\240\375\315\220\205\323\323?\225fmJ88\264\277i\240O\345\004\373\231?7\253\1771\001\017\255\277a\240uk\323y\273\277\253X\341uX\033\270?\377e\357\356v\263\242?\003U\314\200\246>\241?X\240iM\373\033\214\277x3(\344\307\367q?a?\364\000\276E\242?\342!\367d&\027\304?]\361\203\231\013\351\257?q.\215hB\342\265\277\332\025\327\344;`\271?\370\371e\253z\344\256?\253\375z\210e\017\250?\033Fo2<\250\262?\205v\307\310\353T\225?\001~\2410\241Y\234?\247\300X\2308w\305\277/\221\227nc.\304\277V\032b\273\357\004\227\277\211k\002\013\357\002\230\277o?<\355ZO\241?\360-\301t\303\370\304?OC$\342\353\374\277\277fo\020\265\371\201\271\277Slo\032%\212\223\277\305\033,\001\263\275\213?3f|\324\031Q\276?\355\3741H\255\203\264?-&!\230\325\343\240?yVG\207\"\337\255?\316\031\264\016\252E\204\277\267\2508\347\351\250\304?M\353\210\002\356\203\267?\334cW\027\236\376\244\2771\034\310\323Y}\200\277\035|{>K\244\266?\267\0337\301\027\366\262\277z\004\223.\221/\300?\366\031\363K\t\307\263?\373Z\270]I\231\250?\364\230?\342\364 \275\2775\271\225\356\213B\300?\334 \323\027\033:\231?\345\271LA\010d\301\277f3\374u\2516\266?\000\'\246z\325i\306?\263B9{\333\224\237?\306\300R\375\262\206\270\277\333\373\363\232K\013\271\277\005D\n\276\020\231\264?\035\255=\340\2416\272\277\221bB\320|\330\245?v\207x\204\300\203\272?\253i%\335\377\ry\277m\313Fw\0107\211\277\006\244\rXT)\251?\213q\356\203\367\335\277?\337\016\277\311\344?h\277\365\251\207\236z3}?\253\237qc}l\303\277F\201@p\007^\241\277\374\177\341\241\257\035\257\277\343|l\255\205\254\274\277\026\324\365~\263o\260?\264\204O\017\206\216\252\277\242\241\335\362\300\r\254\277\210P\220Z\256\346\245? }\r\216\315W\304?\250\021\316J\262~\231?\246\253\377\001Np\256\277b\226\243j#\035\260?\234c\327\210\023\311\264\277\340\241\'\231YY\231\277\337D\033\355\227\260h\277\253\351\020\301\325\375\307\277\244;Z\262\254.\252\277\363\224\271\223\331\271\242?\206v\3277\013\342\235\277\341\251\356\223\365o\270\277\303\023s\024\3131\242?\371\334\255\375\214\346\272?\313}\236\177b\300\305?7\\rv\202d\260?\243I#\014\201\205\245\277\000\344\353]\256)\243?b\220B\322\003\257\244\277\363;n\027I\314\232\277\315V\023\220&\204\260?\276>\342\306\223\270\234\277\346\370;~\3659\256?\202\367\007g\035t\225?W\020\351\313\242\313\260\277G\214\264\251\377\222\243?a$\361e\023:\266?g\222S\235\256\371U\277\270\323\342p\251\031\273?m\333\303\263?\3054\275\200\301!\305\277\001+\017\251\2034\260?\'\267U\243\316*\267\277\221RD\313.\333\263\277M-\200\252?\337\244\277p\245\263Fj{\265\277PZK\'\333!\267\277\212~C\363\251h\266?U\003\036\241A\317\240?\232\207\251\030\377\341\251\277y\265f@e\210\244\277x;c\r>B\261?\232\346H\0001\310\264?-\257K\034\236\271\211?L\004\241\021\236\222\275\277R`l\036q\305\250?\2773/\0266\213\267\277\215\262._:\361\300\277\210\"\2477\366\327\310?\355Z\010\307en\271?\232[YD\260!\206?\371\t\363\356\020\234\274?1\364\202\213:,\223\277\003Z\013{rG\227\277\3207\002wo\203\234\277\037\204\322eV\321\222?\017\214o\235\304\347\245\277\232\246\177\314\301\331\252\277\365\217\370\r\254a\243\277U\017P\250N\357\215\277e\373\265\365@2\221?{\313\357.{\305\267\277\317\213\032 \260\276\270?\\\023R\246\005s\222?U\035\217\260p\337\313?\362\213\356h\352\033\212?+vY\211O\226\236\277\230\254\305o\0209\270\277\335\261\310\001\032\377\262\277%K\247R\354\020y?\260T%?\264\257\306\277\017E\034V4y\262\277\220W&\326\304\311r\277G\260?<)\275}?\224\350\337;\371\201p\277\n\"\346\032$\214\204\277\203K\205>0\251\232\277b\034\222\217~\307\302\277x:\3014z>\307?_@\241m\365\303\245\277\207\035d%\2770\206?\223\310e\340\326w\250?ZT\027\\\232>\260?^C\303\334\245\211\241\277\\@\273\213^j\300\277\243\346\201N\240\372\301?\300\3155\206\221\240\303?\337\300\342\031\001\206\235\277-\320\210\253\277\003y?b\2711S\276\305\260?\306F\007w\235\367\276\277\262\261\257=&\356\270\277\254\334\272\035c=\300?\301\002\306Xv\210l\277\010Q <\007\251\251\277[HG\006\227\017\226?\013\361](\211\345\273?<]\262\302\006^\241\277JU\360LWM\246?\276A\206\007r\025\304\277\20198\364#\350\301?\025\032\325G\0001\264?\360;\327\246\364O\234\277\t$\027R\302\024\223\277X\355\256\031\207n\245\277\310\032\235)i\351\263?g\203;\032\361\336\263?\005\200|\231\230\333\305\277\3349yVf\350\272\277\251\355f\240j4\300?{\213\207\rb\016\270\277A\273\26722\344\225?\303}\237\307\\\352\255?\000T\223\235\346\020\301\277v\005\317\276<\322\262\277\333\016\340\336\343s\270\277Lmnx{\265\304?\036\0374\355\n\266\252?\301pt\210<\\\200\277/\211Fi\305\337\300?{lw\313[\003\252?k\316\014\260\266\014\261?\342\347\253\347\217\302\243?\025\327\365v\265\277\264\277\246c#\'\024\r\305?i?V2\322\364\251?\307i\255L\325m\220\277\322\376\2511\277\022\201\277|\231s\337\177\243\245\277d\376\032\217\367K\242\277\244\230\317\240\241\254\232?\227\237\377\236@_\307\277t\024D\372$\322\263\277\300\004n\206`\016\243?\213\305An\347\223\225\2779\317}\320\346}\251?\373\361\342\372\270\230\273?\364j\275\001\352\362\232\277\356\355tO[\n\272?x$z\220A\261\240\277\300\3717H\031\023\275\277\367\332\241\347?\016\264?\237M\270\253,\\D?|\372\260+\316\257\277\277mAX\252\322\n\263\277\345U\261\362\355\345\265?b&\344\013\255m\274?Q\220M\306\342\263~?-\310)^W\005\274?f\200#\202\202(\215?\023\r\035\32492\267\277\321P\006\037\227\213`?\273\235\326\306V\022\314\277\345\231\346\361\213\030\250?\355\243\242\345\305\016\271?\265\252\275\036es\214?\233\370J\332$\036\212\277\346\224J*G\'\260?r\021f\214\235u\302\277?>\3059\324\276\247?\3142\301\325d)]?\226\232\321\234v:g?\360\252r\365\377>\250?\333\017^\213\017<\244?\373W\226s\263\303\302?\311\224J\001_C}\2774\016y\033\005\027\253?\224\341\231{-\336\276\277\274B\340\340\320x\266\277\000\301\014\341CK\275?\203S\307\342W\026\262?\310Ne\024\\\223\234?\350\330\323\323\212\374\245?v:\r\330U>\271?\232\324\321\274-m\264\277\343\005\241\202\350\275\246\277\343\276\232\201d\030\311\2770Y\016\006e]\272\277z\265\000\027\201\210\260?\035\213\337\266\322.\270?\007\255\235\351\254\326\224?\0022[\026\232\343\262\277\312\336\031I\300\313\302\2773o\005\346\327H\301\277\260\376C\253C\203\252\277/B\277\305F\310\263\277p\245rsnP\304\277\2715\264)\350V\272?\235\325q\021\212k\274\277Z\366fQ\242b\307\277\024\273\030\335\260\226\254?\252\346\327A\035\264\272?\367\207\"o(MG?3\266\362\026\314p\260?\334\241\361d\224\211\217\277\233\317n\224\350;\255\277Q\236\274\237\213h\225\277y\341v\206\024]\302?\240O\3118\214$x?=>r\246\205;;?\311\271\201\275\352\220\203?\332\316\002\354-\353\267?\244\332\215n\033\262/\277\033\214_\312\255&\205\277\000\217\225\207\263%\231\277\363M\262_\035J\304\277FO\227q.8\224\277.\177p\331\t\311\302?9\205\217\256wr\257\277I!r \033\235\311\277\237\241\264\235\212\274\274?\237\374\312\260\323\007\214\277\331<\007x\256\202\224?k\331\007\204\204\334\260\277\013\023\013\330\304\014\244\277\002I\216l\377\"r?\303\344MMYx\276?\271\347\233\227\223\354\272\277\331\270>h\321\244\261\277\035\333\017-M\250\223\277\"\005\231\037\003\030\257?\363:0\td\262\305?\362\375^\342\375(\306\277\222\020\372\377m\257\253\277\341_\357\263\373\263\311?0u\246\013\243\330\264\277\2353\304l\212\305\230?t\213\244\020\371\264\207\277\317\324\354\006:?\313\277\004\3378\210\037\022\235?\355\345\006\032\371\356\276?\350\017-\243\237\272\300\277\262N6\352\010[\223\277\030\204{\200\377C\265\277\323\334\221\032r\327\247\277\317\3317v_!\310?\354\241\001\341\353F\267\277!\302\242\277\356\236\276?\327Q\210Zu\255\315?\206;\376\222\377\370\221\277\251\"P+\262\373\213?\2411\357R/=\266?w|U\006\304\204\302\277@\206F\314\334Y\247?\227\357\204\020\342_\250?\3155\203\2061q\242?O\263JF|\266\213?v\302%R|W\240\277\206\342\315\023\355\322\253?a\334~\267 \256\263\277\205x\325r\233\251\253?f\225\232`\377v\232?\'\306#\255\364\373\253\277\206\227\320\252\374S\251\277\366y\363\241Y\370\261?\243\213L\345\"H\244?\002\207\273m\200\242\240\277\360\r*\004\222\214\217?\207\030\325\032\353\251\303?\001\037\274\317\241\005\300\277\246\017\310\337\024\307\223?\354\210_\333.h\251?\002\306|\236{>\241\277Vn|ag\215\247?\035\006\312\215\374f\245?\2351qI\303\310\300?:\231\306\347(\\\272\277\013\322w\2337s\213?\206\306K\302\031r\254\277l\325\356\t\027\230\234\277\237\035\n\262\225lg?\005\021\n\037\272\313\203?\301\014\217\270,\360\256?\242*\304\221\253>\275\277\017\266@xz=\264\277\251^\375Q\256n\222\277Y\351@=\032\242#\277\211(\'\024\267#\265\277\251b\270\357\244M\233?s\305\373e\264y\300?\031\2248\223\233\002\263\277\250(I>\375\025\213?\367\242\351\355\211\353\256?#\244\352\372I/\263?@nR\036\"b\237?\316\322\372k\267\227\253?\273\220\315\215\335-\262?\213$\030G\327\321\317?;\343\257\350\316\317\230?\211^\321wK\261\254?\221\216\234#!w\264?\361\016R%&d\263?\232\017\001\031r\267\260\277\214:w\215\310\371\306?\t\271\320 {t\265?)\212%q\346\367\270?\352\330yi*\034\275?UI7\271\274^\311?t\272|\301\265\273\260\277\024=@&\014\316\272\277y\024\305\254;%\222\2771\347$b\357\003F\277\346\305\260rl\267\242\277\202\323\036N\203J\275?\303\363\251\036{\002\256?\2304\204(\217\251\275?B\217\360\260\'}\266?\224\350\346\030\260\210\273? \325`\320n\334\247\277B\326\035\341\347\211\241?\020+_$C;\272\277,\230\223\206k#\264?\343\372\220\235K:\227\277c\220\251\276\315\241\247\277\271\017:\245Oy\303?\243zr\370\3624\252\277Z\300B7\036\225\302?g\027\376hT\205\301?PC\021*(Jl\277\255L7(\014\360\245\277\303+\373\223\334\307\271\277\2640\354\034WG`?\3530-\214@\341\223?\n4\376\004\345\322\270?\274\t=\314>\235\245?\005\3275\214_\311\262\277\302\223\216q\341\037\246\277\025\t/\033>s\261\277\354\345-F\333\370\265?\273\021+\262\026)\277\277\0231@\362<\267\206?\265f\006*\326/\234\277\272{\034[-\235\216?k\262)\343?\225\260\277\034\334\005\001\345\370\303?\243I\024\350\251\256\300\277\355\320\342\234(\233\311?\003BUkMv\275?\031\010?\342\341*\235?\31349\022n,\310\27794@;\322_\242?1\357\002\230C\367\274\277\003\301\351C\316{\265\277\023\346\363\255\212!\303?\230\240\3701)\332|?<*\364\276\246\265\236\277\275u{\247\031\267\222?bo\223\362\013\354n\277\230\311G\211;.\270\277#Lj\366\342\014\275\277\033_\247\247\222j\276\277\037r\007\311\2231\321\277!\274\200\350\340\360\235?\250\266T\005q\322\213\277\303\222s\256\004\220\243?d\252\317<\302y\265?{\0317O\340\276\303\277\020\214\222Gc+\277?+\253a\030\026\267\245?\017r\025@Tm\300\277\345/\310\344x-\213?\"\n\256G\366\242\272\277|\374\363)\362\264\300?\251\367\204I\021\032\226\277g:\2331\362\014\210?\331\206\330\236\323\t\262?|\246n\344\347)\257\277j\014d)x\354\302\277\t\232\275\375Z\013\230\277,`\352aP\024\303\277s\224\210\007\026\021\264?\353\350\375\021\237\224\243\2778@\215j\252=\240?i\255\326\336\323\345\233\277\262>:D\213~\266\277VD\267\365E\343\221\277\002:\374RX\322\240\277e\216\204\214f\255\227?\265\r\323`L\346\262?d_\345>}s\310\277%\365\235Wlw\216?l\357\020\204\237\233\207\277\374svc\"`\262\277w\3352wg\247\222\277\214>xjN\032\235?\001\027\367hl2\267?\267\247\235[:\302\207\277\362\035\026P\017}\201?U,\361u\250\265x?\037\261u\343\304\n\251\277\201g\347\263\'\276\237?(\326\t\034D\315\221\277WhO\350F\333\267\277zQ\255\347E\317\240?\337\nES\014\305\267\277\026\265\0242\330C\224\277" + } + } + } +} +node { + name: "layer_0_type_0/matrix/read" + op: "Identity" + input: "layer_0_type_0/matrix" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@layer_0_type_0/matrix" + } + } + } +} +node { + name: "layer_0_type_0/bias" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 120 + } + } + tensor_content: "b!m\002\033\311\364\277mj\027\32472\263?z#G\005R\'\333\277 \220\336\246\274\360\334?\376=\235P&\320\332?\210\342\tk\272y\322\277y\21169m\360\352\277\346B];B\340\302?m6\304\357\307w\240\277\316\230\250\207\374\360\377?;Pw+\007\271\335\277\317\264f73\227\230?\227\341\255vO\324\334?i\250\010z\302m\360?%d\237\227\240q\367?T\206\247\003\261\301\360\277:\232\373G3\230\347?a\003r,\243n\342?\270\271S\217\252h\317\277\230Xo\215\n\025\376\277p\347\\\334\202\315\341\277(\026P\222\213Y\264\277\223\362r1\350\r\275?\005\327\216P\025y\330?\261\n\321@\n\206\002@\334\221\215\\tc\342\277\314\005\002\247o+\371\277\346\371\342\251\300\311\367?\316\356\373[\352.\344\2776R\010\322V\320\331\277\253\"W\234$\265\355?9\265u\362\306\330\307\277Q\265\370\337/R\322?\350\357\336f\255\311\260?\311\246\017\340(T\300?\006\355\271\355\204\034\007@\\\201\362\360P\342\374?C\274\372\315\177\\\373?c\020m\201\336\214\377?\n\3350\006\267\306\354\2770\257\210+U\216\350?\025\nn\000R\242\360?v\261(\032\373\265\346\277k\256\370\237}\320\333?\243\324\220\326\"\344\357?\026\325<\030\330:\000@\362\274q\200\244\253\343?\007ESLH\366\264?-s\370\255\264\200\360?\363h\026\217\233j\363\277\334U0\254\256<\000\300\342>\037\033\246\r\364?\313\t\317\371\260\346\325\277=\\Nj\273\030\336\277)\351[;:7\340\277\023pQ\202ci\322\277f\271X\244\3307\360?KV\326\250\032\003\370\277\277\372 \223\302\243\332\277K\336\353\307V.\321?\255\026\016\260\227\315\345?\353T\305\036+\r\354\277F^\313\206>x\345?\0226\004(\340\200\363?\r\365\251r#0\370\277\367\365\312\026Sd\266?H92M\360V\353\277U!\376\237\343\034\351\277v\302\344\310\2203\363?\3769\364k\300\202\371?\231\310I\323\227_\341?\020RTg\222\232\312?\337oq\025\203\007\320?w\320\206\330\016=\336?\355\366\214\364\350\024\346\277\320\345\300\276Y\n\377\277\024\000\267\003\334\214\346\277>Z\021\307\274\373\231\277\301\262\354)\017\227\334\277\005x\0317\367\262\236\277\272\317\0051\310I\327\277\236\254\033\326\275\\\006@\221\364S\234K;\373?^\332\364t\221x\361?\003\272\246~\2703\333\277\361\261\326\300<}\373\277\224\177\255f^\303\277\277.N\317\310\206\257\300\277\303g\220\255Yb\324?\260\200\245;.\272\227?\345\201I\027\212\022\320\277\225\332\r\344K7\350?\241mp\326\326\263\363?~\322\370\351=\027\314\277\r\211pp\261\022\346?\207$\363\024}\273\364?e\005\263\244Q\336\301?\311Pv\204bL\322\277\353\"`\337\226\330\346?LC\331nu\276\350\277\307\010\331\240\223\177\373?4\362T7w\'\360?8:o7OM\343\277sbZ\272\324\220\361?\\\242!\004M\276\332\277\276\256{+\337z\355?xb\034\020\322\376\335?VB^\346l\t\373\277\2503~jJ\250\361?&\324\226\232L\344\370?\272\354\327z#6\374\277KU\031\224D\356\371?TD\306\230S\220\273?\331\020\213J\275\311\367?&t\270\033Aw\365\2773\262\263\376\270X\334\277\255\232@G\263\375\275\277\314ng#P\353\331\277\204\304\177E\'\355\341?6\324\376\027Y\344\373?" + } + } + } +} +node { + name: "layer_0_type_0/bias/read" + op: "Identity" + input: "layer_0_type_0/bias" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@layer_0_type_0/bias" + } + } + } +} +node { + name: "layer_0_type_0/MatMul" + op: "MatMul" + input: "Reshape_17" + input: "layer_0_type_0/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "layer_0_type_0/BiasAdd" + op: "BiasAdd" + input: "layer_0_type_0/MatMul" + input: "layer_0_type_0/bias/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "layer_0_type_0/Tanh" + op: "Tanh" + input: "layer_0_type_0/BiasAdd" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "layer_0_type_0/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377x\000\000\000" + } + } + } +} +node { + name: "layer_0_type_0/Reshape" + op: "Reshape" + input: "layer_0_type_0/Tanh" + input: "layer_0_type_0/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "final_layer_type_0/matrix" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 120 + } + dim { + size: 1 + } + } + tensor_content: "C\000\024%\330B\215\277\003\225\340\252\3436\265?\205\307,\022\005\237\303\277\222\247\367\300@2i?0\005T\333\324\242\266?\007\030\006\322\302*\246\277\342\207l\0263\215\274?\202\007eL#\304\246?\024\363\230zK\203\300\277,\037\234\002E\211r?4\347\234\371\200\376\265?\0237\3329\371\301\241\277\363N-\306\030}\205\277\357\233\3063-\301\240?\236\033X\000h\264C\2771(\302c\\~\302?#$\037\206\345\227\213\277l\013\365\337\374\226\274?Q\004A\214\260g\261?\225%\006\026$\343\267\277\223\220\361\275\245(\260?g\224\237\300\353\210\246\277\205=P[\001\024\307\277\204\202n\264.\332\271?r\3215N\223`x?44\346\236\370\235\233\277\241\206\214\364\357\231\202?\'\t\346\273S\245\273\277\277\t`\013\251-\311?Q/\037\234\033\025\302?(\030Oh\0310\275?P\177\"\010\256[\213?\202\313\3036 \237\267?6\326\355\377{\325\300\277\302\344_\305\266\032\301\277\255y\271\037\222\021\276\277hF\267\254n\363\247?\364W\326\351\252\360\253?\303\306/\222]\261\265\277`zfd\347Q\240\277,:h\331\203\375\222\277\247\221\320\227Ll\210?w\271I\333\233]x?\314\0335\024\270\207\243?\334\266:1l\275\301?\034/w\270\275/\264?p\307\201\207\226\023\266?O\347\201\372\222\030\261\277B\3531\210\025z\301?.m\364\035p\230\270?\210\023az\332\220\300?~\261ZZ}W\246\277\222ocN\207\350\266\277d\237\344V\236\210\265\277DHb\370\275\211\215\277\033\344\010\322\343\rA\277.]\263\207Q\003\260?Vk\310c^\355\260\277\266\000Hf\033P\227\277\331\343\244\362\177\204\265?f\262\203\234\375\244\303\277\232~\232V\007\305\267\277\363i\244\325:F\244\277\350\365\305\275\363\252\237\277\023\256\002\333 r\221?\374\254)\004\240\355\305??S5\226\366\035\243?_*\2274\266;\275?+T\371\247,\212\221\277\312\233H\215\202\016\266?\305\005\213O\203\306\255\277R\334N\313\014K\274\2777\036\3768\263v\262\277\217\240\317\305\362\230\243\277\361x\327\371\010n\261?\202<\220\363u\t\272?\226\234\350%(\337\244?\007%\353\351S\315\236?B\n!\330\2667\254?\222Um\365j\314\263?\273f5\217NX\207\277m`\010\351\242>\226?\273\257\225\266\007S\177\277\367\232\024\377K\210\303\277\276d=\231q\230\204\277\355\334!\214\304\237\273\277\330hx\352\272c\251?\314\261r\307`!\263?;\000\362\320\371\320\256?\361\312~\374\270\001\234?\223\352\014\231\006\316|?\357U\323\034.!\251?\002\322\315?:\276\307?\255\002\031\377\325X\303\277E\323\237\251H\337\311?~H}\374\321\277v?|\010\274Q\366[\272\277\344{X\025\021\270\235?\366\023\211RNq\224?.Q\330\302X8\307\277\343,\352\340\361\232\262?\214z\340\323\2615}\277\367\007\221{$_\260?\032\256K\260\251\022P\277\026\367\023L`\265\236\277\345~u\236\316n\264\277W\265K\317\014\217\315\277V\316\255\025\252d\257?\317\344\331\326\314\337\206?\034f\344\263I<\310?\262\221\305#\202b\252?x\367O\244\245\200\222\277;$\005=\211\324\271?\013f5R\250\227\245?\177>\364\2747\335{?)Wfk\230 \243\277\036=:\376\037\364\262\277\240\310%\001p\356\255\277\260Bc\331\231\021\262?\244\231\354eY\361\240?" + } + } + } +} +node { + name: "final_layer_type_0/matrix/read" + op: "Identity" + input: "final_layer_type_0/matrix" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@final_layer_type_0/matrix" + } + } + } +} +node { + name: "final_layer_type_0/bias" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 1 + } + } + double_val: -0.45491870953396435 + } + } + } +} +node { + name: "final_layer_type_0/bias/read" + op: "Identity" + input: "final_layer_type_0/bias" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@final_layer_type_0/bias" + } + } + } +} +node { + name: "final_layer_type_0/MatMul" + op: "MatMul" + input: "layer_0_type_0/Reshape" + input: "final_layer_type_0/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "final_layer_type_0/BiasAdd" + op: "BiasAdd" + input: "final_layer_type_0/MatMul" + input: "final_layer_type_0/bias/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "Shape_4" + op: "Shape" + input: "Reshape_14" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "strided_slice_17/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_17/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_17/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_17" + op: "StridedSlice" + input: "Shape_4" + input: "strided_slice_17/stack" + input: "strided_slice_17/stack_1" + input: "strided_slice_17/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_18/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_18/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 3 + } + } + } +} +node { + name: "strided_slice_18/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_18" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_18/stack" + input: "strided_slice_18/stack_1" + input: "strided_slice_18/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_18/shape" + op: "Pack" + input: "strided_slice_17" + input: "strided_slice_18" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_18" + op: "Reshape" + input: "final_layer_type_0/BiasAdd" + input: "Reshape_18/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_19/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_19/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 3 + } + } + } +} +node { + name: "strided_slice_19/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_19" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_19/stack" + input: "strided_slice_19/stack_1" + input: "strided_slice_19/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "add_1/x" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "add_1" + op: "AddV2" + input: "add_1/x" + input: "strided_slice_19" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_20/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 3 + } + } + } +} +node { + name: "strided_slice_20/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 4 + } + } + } +} +node { + name: "strided_slice_20/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_20" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_20/stack" + input: "strided_slice_20/stack_1" + input: "strided_slice_20/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Slice_3/begin/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "Slice_3/begin/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "Slice_3/begin" + op: "Pack" + input: "Slice_3/begin/0" + input: "add_1" + input: "Slice_3/begin/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Slice_3/size/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Slice_3/size/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Slice_3/size" + op: "Pack" + input: "Slice_3/size/0" + input: "strided_slice_20" + input: "Slice_3/size/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Slice_3" + op: "Slice" + input: "Reshape_14" + input: "Slice_3/begin" + input: "Slice_3/size" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "Reshape_19/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\020\000\000\000" + } + } + } +} +node { + name: "Reshape_19" + op: "Reshape" + input: "Slice_3" + input: "Reshape_19/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "layer_0_type_1/matrix" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 16 + } + dim { + size: 120 + } + } + tensor_content: "(T\3151\263^\210?\301\360F\335\312\363\245?\257S\244i\310}\267\277c\220\340\001\275\335\253?e\025\030M\364\373\201?\023\273\353\254A@\305\277\026\373,\312\030\223\264?D\001\036\021\244q\230\277\030\324\024\255t\256\245?^\300p\325i\224\301?\345K\002E\335\302\304?\3344B\317\023\371\254?\'\031^\032\262c\303?\237\374\016\275\323\240\306?.aBP\016n\260\277\352\034\001\177\314v\212?\356|\246\256\214\313\251\277\3223\263\2722\335\253\277\303J\226W\205\364\303?\257\353\252\2349J\302\277\254U\232w\205\325\243?\000H\314\216\004?\241\277d\202\263\260\371z\247\277?B\327\027\3351\246?n\305\356 C\226\256?j\320kx\2011\275?\356 \t\342v\213\304\277\231U\036)\333\370q\277a\n\254;\354\301\300?B\263}P\004\305\300?[\341\214\373\321\017\231\277\254_/x\226e\303?\335XNm\226\325\253\277;\305\213Fb%\300\277#\3062\033\314\"\213\277\2655?\342\002z\201\277\204\345\244\355\0361\230\277\200c\023|\374\270\307\277C\365\332\005W\017m\277S\2460\277\252M\270\277)\336\006\313\263=\204?\305\320\307\315\035D\265\277Qu1\375\304\320\240?{\363\0063\255\336\241\277\260[\220*\014\266\274?\022\367\001XZ\323\256\277\323-\300\321oV\252\277\323\216S_\222\001\265\2779\334e\254vE\273\277\314\026\010\312\257\206\260\277X\030`\315-e\243?\324M\345\320\366\355\226?\366r\316\246R\314\300\277\340\010\207\221\205\357\273\277\"b\234\243B\257\223?\324\274\014\214\375\034\254?\236\0133\300\255\225\233?\341a6\203?\254\241\277^\207\236#N\211\261?\366+zg\2121\214\277\273\000\254\342|b\231?l\006\373\267\341\253\206\277\200\370\223\035\002\362r\277\010n\021]\200\214\306\277\234\032\366\336\020\313\251?w\032\343\331Q\201\300?\367i\353\312\204\307\241?~\362\331V\344J\265\277\002\001\024\013y\034\256?\035\267\033\217\016\355\257?\243,\341\223%q\267?\032\226\376\2370o\253?3\277:\375W0]\277\370\207^\304\327\030\272?\\g\346)\272\250\271?\366\332\205\315z\314\314\277_o\204\364;\262\266\277g\364\177-;\254\210?\001,\020TZ\020\303?\320M\314_oP\203?\317\202\251{\000\303\251\277\274Z\231%>\255\245?\036SY\264Q \311\277k\336\013\025fk\300?\354\207\326\027\246m\244?\241\304*\007\346\216\244\2779-\'\262]{\253\277\003Lvc\267\351\247\277,A\244j}\013\240?\336\2711\237\302\220\250\277\270x\322\327\316|\273?\273\313\323\210ua\261\277\002\341\010t8=\237\277\302\217\211\330\213\245\301\277}\014\263\371\356\205\277\277D\221\330\376wo\237?\325;\2564&r\236\277r\036U\201E|\307\277\331\373\321\314e|\234?\350%b\020\034\375\223\2774/\3739w{\230\277\021B\357\271e\010\307\277xP\3342\303\225\251\277\004\341\002>\367\236\207\277\310\345\350\343\264\343\260\277\020\007\217\\Q\375\245?\246\266\334\304w|\236\277\014\025\035\334\210v\235?\3746\202\t\261\262\262?\333\354\202\"V\334\225?x\315\n\343A\235\261\277\221\022\356\020+\323\241?\311\362\2629`\334\307\277\343\013\264\257=\224\231?\274;\316qm9\207?\316C\225\360\266\245\256?0\034?\2200L\201\277*p\345\374\034\310\267?\363\370\204\271\235\202\251\277><\354\270\267\325\271?\177:\316A\262\334\264\277\350W\030\323\244\035\243\277\223l\266\307\343\220\274\277\001A\254?\262\367\244?L\257\324v\227\377\262?q!\350\321\257{\274?[\032CL\3358\306?\274\324\237o\305\322\236?E\031i\227\355^\234\277\014\315\314\222;\270\271\277\375\216\223\027\240#\270\277>\323\205g\255X\224?J\266\017\327\302\236\316\277\313\\\220\346\322zw?i\212\214\3132\243\302\277\266\2164\373m\201\253?\212(\315\000\277\323\255?\001\327\212\320N-\267?\306\336[\356g\317\256\277:\325\320\350\2644\222\277Nq\316\316\240\000\263?\354\352B`\221 \270?\330\210Rhj2\270\277v\\\332\216)\212\245?\343p\375.X\303z\2778\366lkR#\277\277\347\210\216\257\245\375\253?\351\340\223\250W\274\256\277\234 \274\005\272O\301\277\2716Qc\207\231\266\277\265\362/D\254d\207?\"\272\354\347\347\374\257?\355\256o\326N\325\244\277C\303h\235\344\034\270?\227\364wD\267g\260?\343\310\035\243Q\246\263\277\351^\2050\240j\227?F\200;\364\202\004\235?\2455\275\"\274\351\305\277\341|\231\325QZ\254\277\002\266\262\332\263\233\302?\272l9\376J\206\246\277*\305\265\027\320\367\210?\266\207\333\021\367\243\255\277\302Z:g\017\251\265?\'\365\2551L\205\246?\265~\365\274\301y\303?F\222\036>\203j\241\277~\201\016\225@O\267\277M\t\310A\216\276\277\277\355\023\241\277^8\250\277\244Dn\004\3225\274\277\3657D\225\np\300?Q\024\3632\235\334\232?\0270\353\314\0259\267\277\207\264\265\241\207\332\267?A=\336\004A\207\242?]\246\010\376L\324\242\277QD\377\347(\211\234\277C\031\237\017\253\361\242\277\337b\266\346\257\330\226?\211\213\341\235\251\324\236?PN\016\206Z\202\256\277\334\2619u\366Y\214\277zR\3139\030L\225\277\021\n\350\235\333\262\272\277\200\2573K\357n\243\277\347I\255\260\006\265\203\277\306\017\014E\266=\305?\377\266_\370e\320\244?n\317.p]\005\303?\347\226\266!9\032\260\277\227*\204\265;\237\222\277\337\236\024\216\235f\320\277\024\213m\253\352\345\264?\342\013\021*\312G\216?\337\275\020_\311\007c?/\306+e\302\242\231\277\310\303\322>\332&\264\277ZW\317\021\251z\305?\316\327\3614\326c\227\277\302}\222\031\241p\212\277\223\314\251\301\276\005\205?\3140\241J<\302\265?\031\204a\035\324\230\267\277S!\340\373V\263\262?\255r?\212\014\363\260?\035$I\226\300\304\260?\341\356k\2243\247\247?\267\337\237\000\336\302\303\277\203\270\240UI\216\302\277\344\001\305\261\371\340b?\036\022Z\227`_\261\277\342\277\274\030\017U\264?\006\306\370\301\262\220\261\2776\273\346\335\213J\243\277\324\323IR\333\305\252\277\2442H:\032\267\250\277\025Ohy\240\232\235\277U|x\261\312\'\276\277\231i\216\031\221h\311\277\r\323\327\275\240\256\302\277W \0048l\260\253?\0203\033\345\222j\222\277\365}\005\277\027q\242\27794\226tI\370\263\277Q\215\201%E/\262?4h\345FQ[\300\277Q\321\302\360\247\320\306\277\003\234\216\227\262\327\257?w\366Fx_/\244?\335\255\334|(=\270\277>\323\374\251\263\206\251?a\353\303\217L\026\275?q\014I\031~\331\276?\"\366\024\013\232(\262?\333i~Q\202v\251\277\303\312\254\225\357j\216?\341N\260,45\311\277\224\255\007&6\305\242?\t6\3006\322\251\270\277\267\216\246G~\216\260?\330\372\303\224`\362\251\277s\325}\357\022\323\252?~\2373\365U7\263\277XI&\340\207\036\304?o\356\217\307\265\227\226\277\364T\310\251\004\200\241\277\'\241\337\253S6\222?\005\t\006\302\262[\307?\272*e\001X|\266?\363>\007\356\372\350\263\277\271*aMZ\306\201\277\217\353\347,\262h\234?\260\3320\274\t\204\305\277\334\005\363\036\020\tn?M\355\214\310\244]\234\277z\n\206$\210\262\267?S6.\246>\344\243\277\350t\230\207!\033\264\277\350\316\214\336m\357\252?\330\327\001\271\n\276\261\277\272\003\325\324t\206\300?KwV!\247\343\226?\326\312\2517\204\331w?\2476\214\024\332\374\267??\365\200\304\027+\254?wHT\020\016^\207?\344\034\301\340\016\275\236\277%So\025~K\301?bI\306\rR>\247?\200e\200}7\363\260\277%)R@\353\203\256?\377\210\272l\226\206\246?\034EZ\357\204\310\262?\237_\tz\271\023\252?\203\204Q\201\350i\261?\nM\212vL`\307\277\026\211arP\312\251\277z1`3\341\014\247\277\035\023\261w\0202\262?\326*\350\262\270 \212?B\323oN\r\t\264?==O\211Z\220\271?\035FW\213\037\256\265?tl\004\346\261`\250\277\274\'{\264l\300\254?9p\370\002[%\300\2778\367\377\333D8\241\277\031\245\252\253\032)\271\277y\016\203\024 5\271?\366\363\350\000\323m\277\277\312!E\022r\310\274?\206\t\310\267\247\341\270?\237<\226B\355\205s\277\220\301\004(\3101\217?\220\"WX\017\377\253\2777\213\376\270\207\311\267?v\364[\357\312\252\246\277\237\027\375^\262[\252?\337,^\340\312\250\271\277\256\032[\334\037\275\234?\235\031\010YLY\241\277\310X\341Z=P\246\277\331\214\376\230\205Z\304\277 \371\201\3245\000\272?WM\336QEq\241?(\007\225\336\361\246\260\277HL\020S3\237p?\274V\303\014\"\323\271\277yz\350\225\305\r\242?\210:\302Z\254\277\264\277\226\202\344\027;\320\262?fb\371^\343\217d? \354\277\204\251\315\242?I{\022\231\000\322\277?\264\304\005U\375\277\221?\365)p\306I\261\262?%\373\231\355\036\236\243\277\356\002\272\\E\312\303\277\227\370\307=\036\207\242?:\326{\'\233\310\215\277\025\241\206\323\272\243\262\277\223\267c\216\026!b?\311Iwx\213\250\243?\322\311\277\373,Z\262?pC\033I\020u\202\277\006\010\312\324\254\301\223?\276\356\263:jx\270\277B_\372\321T\371\247?\262g\334\257\355\260\302?)\025\262\201\026\204\235\277\320\261\347d\237x\274\277\313\346d\216\026\217N\277\307N\013\373\037F\265\277\2045\\\'\257\232\265?\357\276^\341V^S\277Kw\303\242\021,\242?(J\342\206\205\266\241\277\270\334\002\211A\374\246\277\027\017\251\301\263\246\256?(\2725H@\207\276\277N\240]fz\004\302\2779\263\277\027p\265\236\277H\221S\305J\024\234\277\360\331\3264\203c\245\277\030P~\r\326\312\303?\\A5\r\257u\210\277\007\347O\241\231T\214\277\010\022\305&\357,\301\277\031?\325\020\031F\244\277\240\245.\216c\305\274\277\204\364\312>w6\252?}\366\217\207\350J\217\277\022\2210A\030-\224?o\354i3\024\250\254\277]\253\201\000\372\366\260?h\321\305\021\354;\276\277}\345\025\\-\357p?D\350O\323\356\353\273\2774\252+\304\314J\261?\022S 1\267\265\246\277\250\376\230d\275q\214?n\232O0\227\262?\207O\342?\256\302\222?\0019\363\267\243f\226?h\311v\251t\244\240?g\037\254\243\363P\273\277_\272\305\262\205`{\277\360s\3762\256\307A\277\317N\211\335.Y\272?\233\336\273\020\022\227\240?\362\305\260g\360\r\264\277\242\257\303\360\202\242\254\277L\213x\226\200\257\267\277\332\374d\317\343V\246?\324\234\202\263s\332\222?\366\236[L\247\366\214\277W\360\212\311W\355\215\277\"\032\333\203\361\177\270\277B\247\274\260\270\264\267\277\332\206\243\222m\030\264?^b\364\220\007\007\215?<\322\225T5\255\275?l\277Y\037}\254\245?\205\321xH|\t\256?\342\234}\204\366\267\254?\322\225\006i\233q\272\277\014t\217\310\341\005\300?\253/\237\n*\272\301\277\034\200\030\225\353\275\262?\310\357U\260\252\213\242\277pW\250bd\243\270\277D\223\344=\316\272\254\277\211\345\306\240\263\312\243?*sB\034\261\037\246?q\376!\022\300\374\225\277YjB\033N\345\274\277@\346s\214\276$I?[+\312\254\343\341\211\277\304&\367\357 R\217\277\013\233i\025j\312\255?\324\200{Y\215T\260\277AC\341\324\002a\264?B?\247\361\203,\254\277\215\037\316e\247\236\221\277\244\372\314\345|0\201?V\355\010`\362T\272?\347\217\256Z=\313\232\277W\265\261Q\235\213\243\277\014\n\316\025\010\307\226?I\274\235\rfs\262?\031\342\350\260-\365\304?\2158y\254\001\272\230\277$\315\377yT\317\306?\354\227\372rI\007\232\277\246\320\353\251\302o\241?\232LA^\273\374\246\277onS~l\341\247?\357\321~\312~\223\210\277\n\313\262\004\201\030\264?AY\032\301\000\323\274\277\357a\346__\232\220\277\325Rp1=t\211?5H\307>k\254\210\277\277\253V\372\005D\254?\352\207\260\373e9\267?d\007\323\311\311\227\246\277\232\334 @\305\223x?\214\316!\016&\237\256\277\236|0A\214\036\253\277$\243\256#\033\334\253\277\202\357\332\200D-\275?\370D@\331\"\235\247?rp\205\360\360\255\267?\242\336@\254\230>\205\277\310\202\212\22119\303?Yn\"\252Q2\236\277$&\255\227#-\240\277\247P\345\2756\007\303?\233U@\001\315\370\250\277\342\217w\265\315\177\256?\312\024\375\251A\350\273?\360T\325#\021\337\235\277b\275)\367K\205\307\277}=\246$w\311\261\277!\327(\014Jb\305\277N=\260\274\244\024\265?\023=\004@\034;\246\277\370\2159B\364\335\210\277\177\035\237\254\216\377\215\277\200\262QY+\177\263?\250\321\006y\240\347r?\264\005\333\307k\343\256\277/\247:&\035\323\265\277\262\345\'\260{\355\264\277\260\035E\t\363~\257?v*$\004\262h\315\277!L\316\374l\027\265\277\354\231\000\243~!\267?T?\313D\021K\236?7\302F\350\216j\263\277\262=\366\365r\211\256\2770\336\r\376\346\255e?Ki\365p\235\345L?V\332d\203\205n\277?\367\002T\3139\030j?\235/\336\272\262t\305?\214\215\231\224I\007\241\277\227\313\027\020\346\326\255?\345\026\262U\250\312\227\277\323\005\017[1\002\247\277\270\274\353\266\233\032\255\277\020\035T\251\3142\224\277\000\334\027\201\231\276\270?\333Bh\227\341y\251\277O?\330H4a\225?\261@\034/\311\234\227?v0\2233\305\245\234?\226\240\356\335\377\327\230\277\315U\332\235\330P\262\277\264\322\323\004\227\262$\277\225P\177v\202\315\266\277\260\362\313}\003\030\262?\\\016\254\275+\013\276\277%\230x\266O\030\253\277\231}\241t\006\307\220?5\267\030\221\272\255\264?\273~\243\365(V\241\277\2405\264\335\217\221\303?\037\'\260\273\364\023\272\277)\354\264\230\340C\263?\021\t\030k5P\237\277\202\223-w\177k\275\277-2 p\342i\203\277\215\337\243#\342\225\301\277\032\253\317T\240\205\306?\323\343\361\374\300\346\277\2779\217ju\370\350\255?\264\2659\202\310/\305\277\314_F\230#D\254?\002\263\222\0328<\265?\233\3634\376W\276b\277o>hRGb\263\277\3752Is\236\344\275\277\372\021t7\004\326\251\277\336&q\261\004\037\224?\376\262@-\212\032\260\277\t\2632Q\255\317\314?\262\356\247\311\325\255\250?:\301\357\361\243\232\306\277\301p\300S\016\236\306?MS\247\014b^\221\277U\021\217\371\352\250\262\277\200\264\"/\251r\273?\265/\264`\303\033\213\277\362*\300\336!6\262\277\272\373\360\257\330\221\267?\224\333/\027\262]Y?K\363\300T\205\336\240?C\366vK\255m\272\277#\301\350\307\025y\260\277\022@\3217<\232\274?p.R\234.\"\266\277\314\245\026\274L\234\303\277\033|G\363\261;\265?\271\334qbR=\240?c\005\225^\"\224\247\277\237D\220\314\001\035\264\277\212\344\245\362\322\356\301?\262N\352Fj\210\246\277\367\253\231\025\353T\242?I\214\033\357}g\244?-S\n\253\271\202\305\277\217\212\327\214\327\325\320?\3005\253}\305\254V\277\226\343\263fs\\\257\277\312\276MT\232y\270\277*\375\301\031\000k\202\277\010Sg\347\035F\263?=\327\327\246>\345\277\277\236\252\363\200\264\006\310?\311Zz\226\373\313\240\277)\245\313?S\032\301\277\200\351w\206\335g\301?\002\364\027\277h\261\230\277\226\207\371\305\352\267\274\277\357\024og\270(\300?4\221)\264\270\375\237\277Tf\272\3238|\253?XJi\337\364\223\267\277\360\365\004Cbb\300\277\310\247\340\273m\177\314\277\036\026e\342\267\334\221\277e\005\206J\016\363\263?\356\321\262\022\340\202\266?\374\246L\271\305\222\204\277V\250\004p\006[\301?m\304\3357\321\003\273\277\300\256bx-`\235?\332\341\014I\210\203\301?x+\021\337\302\254\250?J\013\023\375D\010\245?*\177\301\"\274Cx\277N\363\223D~5\240\277.;\"\246\036\212Q\277`#\001\203\216A\255\277@\357,\006sG\233?\322P\221z\r{\252\277K\016=\367\207\023\267\277\370\242O\231\301*\260?A\303p\017\274\255\214?q\037\n\313\013\234\273\277\360%\315\310W\235\250\277\277r]\351jd\272?\230\357\026\032F\027\301\277]\035\266\033\022?\300?+\305\"\331\030.\247\277\241@\3406\3232\225\277\240\310\tw\024\014\250\2777M\010\032Z\003\233?\010\333\254\177\006(\257\277\221<\013.)o\227?\247\267\356\354\350\027n\277q\306\200WE\004\204\277E\337\"\213!\325\220\277\350\243\2550\230\177\305\277\271\332\255i\277g\236?\016?7\264)\313\300?\361\341\232\221\321i\274?\341FX4\327\344\240?\342N\026\251I\037\277?kyt8#\216\256\277a o\250V\236\241\277\377\305\312\375\302\337\274?\030s;O\022$\256\277\324pU\355Y\264}\277}\272\357\033\375\225\216?\027O\023\237\311\250\201\277(\256\303\232\236\335\266\277\n`G\"\217`\301?U\327\366\375\267\354\303?\245\022\251\252\037\014\244?\177\246\251S\335\377\261\277Vb\nGl\\\223\277\345w\272V\274\275\253?\342\262\222\345\364\357\272?Z\207\211pm\362\301?\013\323\223\266\343\216\200\277j\257\233\347+Y\271\277\017\256\215\310+\271\220?\232c\222YlY\234?d\273\3738\207\235\241\277kN\373\331\301\003\305?\203@\355\024/e\261?\331~\224\363\2622\226?\030\201\240\271\016B\212\277\311\253C\245;\217\250\277\244\027\034\342\001\'\233\277\262|R\277\355\035\240?Q\316\322\270\357\325\231?F\2426Oi?\304?\244\300\205Tv\331~?\346\227\377\245\333&}\277\247V\024\363\231w\263\277IMC+(\204\273\277\326\340j\336\272\331\226\277\233\213\323\272[?\242\277\236\374\2051&\233\025\277I\211\006\364\270\261\247?p\242\264\226\013\220\303\277\201\000\361\311\235\026\262?\025*\213*@A\244?\252\021\257\215oC\264?*\367mj\222q\245\277\321\337\362\310\001?v\277hR?\326\304[\243?[\227\026\207 \013\270?]\033\201&S\242\202\277\014`M\226R\007\212?\235\373{\266|\000\234?{u\264\325\303b~\277\226\377\353#\010\324\307?\257\356\177\377\004\217q?\203\221\275\2469\257\260\277aY\262\233\304\2231\277\207S\276\220\331a\305?\251yh9j\331\240?\025a_\202\224\024\260?\307\031\357\337\277K\302?\242\220\2658\247}\245?\264\270\270\355\276\305\255\277C\364\264\341\244h\262?\225\001M>N\302\227\277\255\273\3108\320\031\273?\200\2464\211\337\260\277?P\376d\375\350\212\245?3|E\211X\216\241\277\034\227\320\216<\021>?\215\335\227\031\303\315|?\211N\310\017N\306\270?\203\241\337\340\317\372o?\221a\202H\007m\234\277\215\337\242\224\220\325}\277\342\234\275\211\236h\270\277x\227mDqb\233?.\n\0143\177\027\250\277\240\005Z\347B\020\275?\270?\211|\375u\241?.u\204b\000\226\246?\307j7\337\324\036\275\277\207\027C\t\t\252y\277\251[\216 \305\252\266?\334\026=\334q\027\214?v\231\2117\\\316\237?\326\025E\004!\247\272\277`\217*}\211Q\221\277\303\266\333\270^\372\237\277\211qp^\'\316\252\277\323\375\267\017\232@\306?\316\253\007\214a\314\273\27774\206\251\206U\224\277\2226\230P\245\325\256?\347\346\323\263\275o\255?\215>k\211fw\234?\003+)D5\001\245?\214u\216+\305\026\267\277\212\334\303\003H\323\256?\014\340\205\037\200\006\276\277\324\313&\2018\226\265\277v\002\\0>\302\241?5k\346\226\355\007\301\277\264\324\n\336g\317\270\277\357\374\363\027b\277\275?\236*\205\222|\244\261\277\264\257\003d\2638\222\277\211\257uTu\336R?T\260F\371{4\244?\342\211\323\232 \006\241?Zo\256\330\013m\301?\227\334_\036\212\337\233\277\241n\364\361\313M\227?PR\373[\314R\230?\320\307(\2668\310\305\277\212\246Co^\233\260\277\257&\322\266\271\237\270\277Z`\246K\241\200\251?\341\375\244z\200\212\307\277l\026\212D\313i\227?\317#\351\336\232\234\250\277\0146\324`\211\355\262?\007\371R\336\222H\274\277\020V\201^:\226\260\277Q\n\206\261&%\237\277oO\023\354\220\354\305?\207\254m\376\0374~?m\025\262\213\232u\270?\205\252\226\231\313\255\263\277\355\375\352\213i\321\206\277\001\\#w\225\302\307?D&\316p\n\326\254?\326o\251)\221r\266\277\307\271\0334\350\362\203?\205\315\017k\311\200\254\277\336#\273\303!\003\265\277\232q\232\003\"q\233\277\234\326\257\255NO\305\2778\244\r\334\374\t\251?;\337\023[H\371b?\025p\354|n\234\273\277\370\n\272\250\332V\266\277\262\264\254D\264x\267?\207\221L\2454o\264?\014\007\323\202\245`\267?%\331X\007q\226\274?5\273\235\364yA\256\277\207R\256\027!\237\260\277\r\334\255( U\265\277\233c\331\235\371g\232?\n\371\326N\270f\245\277\346\327\306\360\234\315\223\277\240Ri\241\275\274\274?&\231\214\374\226s\246?\325\223\375F\223H\263?H\334\035\2552e\310\277\365\367\'o\312\233\270\277\275be\250\333\236\236?c\215\203\364x\007\301\277u\374\251\362Y+\276?\002\025\013\214\034Gp\277\304\371\251\317\316S\252\277\207\313\036K\215T\300?\313$\263af2\222\277F\305\355\014\275\003\271\277.J\267M\336\006\225?\245b\243\332\365/\306\2776\251\213I\206j\270\277\317\035\221\310\005\205w?}7~\376\305\247\274?\022\217B\211\241\372\216\277\301\240\2433\033f\300?`\330n\267\275\'\260?\344\203\363\377\274\023\232?\004cy)>~\264?N\275\0144A\030\241\277\233\307:M\374B\221\277\332\206}\235Yb\246\277\252\300\312,\005w\305?\207\356\030\356gs\301\277)h\375A\360\351\247?6)0\366IE\245?\213*\372\3738\203\244?\207-\216\354\010\276\275?\205\301f\354\025R\240?\021\277]\353\t\037`?\304\326\223\2170\276\204\277\367\264\376\356z\004\322\277\372+\327-j\312\257?<\215\273N\025 @\277\2563\t4\207\372\223\277u\270\204\235\324\222\225?,\2273\307\037\336s\277\220\256\343\215\205\243\245?\331;\353D;\232\232?rC\365\367\024|\276?\360\254\277\276\036s\261?\2551\025\304}Ax\277\342F\314d0\302\302?{T\260\023;\232\264\277s\277\244\320\264\241\270?\031\245\204\347\316*\300?\025\374\331s\333\355\242\2770\345\036\022\206\253\232?BlM\211B\t\276?\034\352sF\\KK\277\351\207\335\374fr\220?\002\222\334d#\017\274?\\>\202\346\332Y\252\277F&\027\343\233\200\300?o\346\260\354\333\212\261?u\006|Y\215\364\247\277il\033|\212H\265?vS\300\025\353\213\264?\234\207\256W\322\021\240\277\261\325\346M\3042\255?\266\342\271tl\350\215?\272\321\275\351\262s\276?l\342\376\333c\356\265?\256x\300\2455[\321?\230\366T\257\234?\275\277\236M\0368\304#\243\277C\350\3755\277\275\246?\250(\212\227\275R\304\277\034\252w:5\345\244?\006\255YJ\226\006\235\277D\374T\334\304\370\240\277u@\375c]\002\255?\226\345 \037\260\354\262\2775\256\305\n\2348\235\277\'\324\\\224\367\233\245\277\313\206\346\340\202G\221?\350}A\220\204\025\260\277,\t;\004LJ\263?o!\231\262\203\036\212?r\267d\372\375!\226\277\250\217\022\261\201%\305\277T\207?\350X\'\255? -\273\212\256\177j\277\223\037\031\321\246\253\301\277\357\225\262\243\035\331\300\277\332\201\020\222\035a\275?p\231\033\006B;\231\277\257\225\234\212\345\235\253\2777s\254\033\204\265\261\277\340\330\022\324\246N\236\277\257n\017[\021\037\274?\377\204\032\007\200p\240\2774huZ,\235\270?N\023\200\332^Z\231?.m\364\204\321\333f?\236=\230\242]\360\260?F\033\314\224\375\276\257\277L;\243<\247<\271\277~\376\230\2702T\267?\030j\313;\375\334\177\277\326z\370\020H\313\204\277H\310\373\353\0133\245?\354,q\314\362\311\275\277\245\'\314\033w\010\302\277\366\352K\253\250\262\221\277\250Y\250Amb\252\277\272\351\365\266\262\270\244\277\341\323l\331w\373\256\277\370\203\027\211>2\263?H\262\335\307\317\355\246?\204:\231\305Ll\232\277a\355\344\t\237\036\277?\373\033\326F=\035\241?\026r\226DT\235j?{\366\305\020z\001\264?>oY\034\2661\264?S>\026\352\376+\251\277\352\225\232n\010\340\266\277}f;\241\307k\224\277\277s\312\334\264\262p\277\3161\2565\325\345\210?q\243e6\021Z\223?\203\330\014\316\007\027\262?\311\266Jk\227$\230?\324\307\001\374\000\215\257?\361\333\315\206O\032y\277Z\270\211\253\363\377\267\277P\371*\"\327|s\277$\210\014\301\314\016\250\277\242\037l7\276d\266\277\213v\276VI\373\236?\352\366?\253mG\241\277\336\334W4\n6\271?\360\201^\266\211u\250?\221\320\333\227~8\263\277\t\335}\344\021\350\307?\035\246\267\021\350\337\257?\000\320\303W\033p\266?\r\303\301\245/\027\252?\214]h\206\257\273\242\277Y:r\263~\334\263?\305+\3724\365l\233?2\322,2Wk\265?\315\230\367\014\007$\251?q>y\221\024@\232\277\352 \264>4\327\303\277\266\232?\277tc\240?\252z\262tq\356\270?~\261\324\312\327\253\220?\253{4\263\r\324\226?\031\0028\235A\016\252?f0\275\335\'\314\253?\354\010\230:= \274?3\201F\307>\345\206\277\337s\364C\r\263\302\277F\270dU\030!\260\277\357P\340\'\000\003\271?h\224\r!\\\232\236?\210\0036\212\016\333\275\277Q>\317\277\377\351\300?#\032\273\367\276\217v\277@\232\347\264.~\226\277bUR+Y*\214?U\311C\300\030\236\255?\000O\220\201\320\205\252?l\247\366\353\301\235\261?(\207\247\0039\312\270\277Z4\254\340\202\223\263?\002\264\242\272\200q\240\277\023\270H\357\223\364\273\277\035\324P\355\341>\\\277 \210\266\231\356q\275\277\275wq/\275\016\215\277\244\234\367\204\204@[\277\363C5%\250\211\200?\357\260~\230\327\202\266?\215\0347\223\014\037\267?\'+\227\rCl\271?)\266e\r)\344\244?\001w\361\232\304u\303?hGu5!I\304\277\343\2107[9\313\255?\352\326\222\034\314\222\232\277Q\315\237;\213e\272?Q\361\355\001\365\005\303\277\260u\277\300\302\365m?\330\034]\242\346]w?f\304\324o\344\343\253?\313\265\037\320\320\362\265?t\006\376\240\2771\242\277\207\353(\337J\225\275\277\234\306G\037\205\374\302?p\220\356\020x\257\266\277q\232\305\374\267,\307?\3009*\372\220\021\262?`\323\364$C\365\251?\034\346\237\307\323\364\302?\260\034\005e\024\003\246\277Hb/\371\217\372\236\277\312\021<\260\257Lc?j+\375\237\275o\303?k\306+\2258C\257?\264\232\310\217^\213\307?\020\214c\036\020\333\241\277\024\377\230\223\223c\271?\337\233\017:\253\304\274?m\334\0066\271\037\204\277\314\232vlqu\214?=Yp\257\022\246\313\277*M\222\371\322*\272\277\267\265\276\026\230\244\244?\3739\375x\325v\241\277\370\033\253t_1\273\277N\316\253LB\233\266\277\241Pk9\256s\277?\244T?\267HJ\244?1u\236h\275\027\241\277WJ\321\270$R\300\277\\\252\242\262?&\341\272\227\304\304\274?\001:\357b\272{\240\277.\026f\035w\305\261?\035\323\311\231\214\017\271\277\364D\260\261\\!\205?\300\201f)y\330\300?\365b\232I\233\035\266?W\300Q\255a\337\267?\262\332\3565\314\r\262\277\0315\306\341\2518\301\277\272\324-\277\0074\320\277\261\025\376b\013B\266?s\362\353\3419q\260?\365\353\211pW\234\250\277\246\261.\024\255\300\272?\275\267\026\362\276\211\304?1\323\365\377\227j\247?M4\\q\\\037\261?\355\315\210\'\265\301\267?\365\007\272\362\023\036\256?\214\204\364\361T,\205?rR\031\314/#t\277a\365)i\236>\221\277\213y\313+\245\032\223?\336\363+\303\316\023\201?\312\353\032\220\344Tu\277\210\3740-\211\362\242\277\262\367^\316/u\300?\357\366t|&b\221\277\340\376\326V#\005\264\277\037l 2\325\246\244\277\231U\355\220\347\n\177\277\201.\207\235[\\\300?Ii\332\302L\347\300?\025\270\360\023\257\211\255?\253\374\350\230C\377\244\277\277\311^?\274\360\243\277\200\027\353e\002\033\247\277wt\235\214\020\005\300\277\202\005\312\206\271x\254\277o\234\264\327\025\006\310\277\014+T\275\276\027\242?\275\244\275#J\366\271?P\235\2435:C\301?\275\372g\n\356\212\241\277\314\376\333F\030#\240?l\342\016u\251\260\272?\365Ja\226X\224\306\277\226\007S\321N\216\235?\310H\024\2036(\300\277L\203\353\352\265\277\227\277\032F\326\013\016\266\261\277\241\333\230pH\270q?%\321\263\302\304\016\266\277\260`\254\r\177\351\264?\324\254AI\003\r\223?S\237\327n}\353\233\277\263\3729\013\211\230\261\277\202\265\"\016\021\000\273\277\342\320%\226f`\261\277\037\037\261\036\201\254\312?`9V\263\316s\262\277\325U\372\224\261\305\267\277\256\310\0205\273\371\242?\250T\000\026\217r\312?H\204e\331\014\177z\277\000f[\016[\016\300\277\374\031\322\010\320\274\263\2775h\314\267\006\263\257\277\252\267\025)^\353\262\277\004\322ajjj\220?A\243\276I\252\027\311?6\223i7\323\362\231?@\025a\035H;\261\277@\\\031\251\013\224\260?\235\2160\365r\247\265?h\230\224\326{\n\225?\227\321\254\307\344@\301\277&\347\035\321\004\276\235?\251\346\350Q\337K\264?\200\2243\037R\342\300\277\251\010\027\257\247Q\254\277*\246\344O\355\217\303\277\322\302\272/\250\311{\277\333\217b\030\231\020\307?\2045\335S\3720\245?\352\010\3177\217Q\262\277s\355\252\244\005\233\222\277c\325\277$C\033\305\277\201U7\027v\210\300?\363f\357\215\276 \232\277\215\010\257U\327\270\231\277I&\254\273\030\267\236?\352 \201\336\356d\267?G\t\"\035\343M\244?\201\341M\221\310\242\260?\341\304~F\236\026\257\277\270Wu\3534\362\202\277\033hw\220\313\020\251?\356\327rBD\371\254\277\032\344\311?\007I\234BD^\252\277\376\t\024N\276V\221?U\314Jc/\367\241?\311&\"\340\024\254\243?\242iB\036E\324\231?\215A\360\316\304\254\240\2774H]q!Q\303\277\242\350!\262\013\275\266\277\332Gp\320\260\342\300?\330l\235\024\247\352F\2772Y.#\257o\240?\032`\236?om\235\277\277\314\276:p\330\251?\373c.\013t\226\255\277u\370`\331\261\374\301?\003zf\021\371\007\251\277\270\"\035\357\255\260\260\277\371\302>\261\020\335\251?!\326\r\"C\034\221\277\273\304\377\331M\233\300?\207\231\351\203ZU\202?5\357\025F\n\214\304\277\205\373t\332\010\244\300\277\n=\242\234L\331\266\277\r\030\253\240\364S\266\277w\026\004Y\024\036\262?N\002\344\254\324\354\265?\212\'\264\377\265\200\264\277y\035\020U \306\223\277\257dd\003\013\212\245?gD\232[=\342\270?\216\253\202\255\201\026\272?P\r\020\205\020%\206?s\271\332u\025\260\207\277$2\314\362\264\234\221?\304\253\243\355\310\252\226\277\362\240\206+\276\203\232\277Z\3570\371\241\212\212?\235\346e\034\303\226\247\277\000\266\225\014}\357f\277\016!\341\332I\232\234\277\243\222\300\344I\267\211?\306\224\230(^\341\247?T\016\356\014\351\273p\277\031Y\001\236q\214\274\277\330\252\350nb\315\272?U>\277\250\326\322\257?\263\270`\"\364D\255?\221\'\332\222\242\031\256\277\345\246\201\023\006\001\252?\335\271\265%\207]\301?Sd\373\317\031\017\233?\007\322I\344!\017\262\277\331I\352\251\257\204\212?\345U)\225\277+\200\2774\026\'\004\366\261e\277\376/\275\354\017$\270\277\256\020\277\261M \223\277K\207\224,\213Dh?\354B\353m\036\025\302\2779\023\323Y=\274\254\277\n\250\2123u\257\253?\3710\0168E\233\234\277]\261R\034\277\025\261\277\037\255oP\337\205\257\277\231\304\316d\375\035\270\277\306D\347\010 \346\241\277\205u\337\300\273d\202?t\263b\t2*\267?\032\316y\243\353\315\224?\321\236\265p\242\207\234?HKTh\255\365\210?J2\325\261 3\251?\237L#\253O\354H?l+\227\027\014\027\240?\304\344\016\226\245\020\246\277a\350\276\265\252H\277\277\007\304\253\303~c\235?\252\303\221\030h\322\260\277\371n$\032D\242\302?\202r}\271z\204\306?\026\341a\356_\372\244\277\212\035L\215\207\210\260\277n\n\354.\027X\242\277\254\031\357\333e=\241?\254\276\342\235\024\027\240\277\363\014)>Db\273\277\314\340\246<\253\004\274\277%\2246&H\017\217\277Y\237\035_,*\307\277\003Sh\347\325\300\222?\006\205B\364\203\233\304?\343:\022\244\267\344\235\277\277\267\241~E\377\267?\245\325\202st\256\276?\370\320C\345}l\271?:\031\367\034\\\"\242\277\"O!\3454X\255\277\013z\264\372\243\331\203\277\204\270E\376\251l\260?\355T&\266C\320\240?\027\017\017\20556\264\277\037\313\027\221\266$\300\277\222^@\3638\366\276\277Q\310\351\334\025L\300?Y\350\241n\207\201\263\277\"D\347\353\215\235\264?\364(\034\000\237\335\264\277n*\223\2304\242\270\277\250\344Y[\210\274\263\277\353\344 \350\372x\210\277\231ZI,\211\222\265?\3335&\026\2056\210?\3552\037\365\321\322\273\277\3547:\017\260\007\275\277\025\212t\361\322\300\216?\343\251\315\215\222\370\261?b\343\232\361\r\220\201?R\356\0018H\330|?\374\341\0315\356\007\253?\206M\242\271\376\217y\277\207\352\332\035\376N\303?Z\361ee\306(\267\277x\344\256\315\210\016\270\277_\257\244n\202$\301\277\330* ;\336Sw?\020\211\021\026J\246\300\277\343\025\267s\265w\265??s\314\365Jk\266?\301\364\354D\322a\272?C\246C@\275\030z\277\266\306\276\030\250\020\274?cX\321\266\007\207\274?]Ou.Y\343\232?w\225+E\240\325\232?\235\322m\317 \343\260?\027/\236\366\340\r\276\277\215\353\363K\335E\271\277_f\311G\001o\206?\"\321?\375^+\266\277;\330\245\026\331\300\224\277\030\2432\222u\322\311\277\244\021\363\233\002\226\244\277\213N(\323B\202\266\277\020\3317\242\300\277\256\277\r\226\321!7\253\270?\311\226Q\271y\330\265?u6lw\327\311\233?\357\351\245\353k^\260?\370\204L\337\327\310\255\277\360\326\220z\325dv?\357\322C03\307\266?\256\237\322\356\206\267\220\277\013\253\366\221\374\220\263\277\375\3447\030\262\360\230?\000\007\211\364v\367\275\277\032j;m+&\262?\010\231\363\335\260Z\263?\030\226\263\224\365\037\273?0\323\035HGD\300\277Y\022 \272\370)\254?x!H\354\244=\240?e\021M\2776\270\227\277\2339\373z\367\347\262?\027\000\245\202I\262\272?\332E\220\321p\037\244\277\375\237\002\265Q\001\234?\250\021\353Z\352\rz\277@\352X\344\024\243\301\277~\300>\367X\365\266?\240\031\276\266@\026\263?[\233\314r\334\202\270?\000\345\364\'\212\037\244\277\037\017f\027\340T\245?\036\231U\214\255\230\303?gN\332\'\335\321\230?6iPC\333\013\266?GM-\356\233\213\203\277\211r6\331\274C\275?\321\035\251L>\374\267?%\242\307WG\336\243\277\013\232\213F\346w\272\277\010cgS\004\234i?\255U\263\256H8\203\277\235\254\336\222\327\370\202\277\307(\n\205\333Q\255\277~\034uwA\344\203?\214\372.~\026y\211?r\007\3053\243.\215?\013\325\357w\357W\221?.\271\2111\252\245\267\277\356G\316{\343T\255?_\214\211d\2115\303\277\016\016\177c\215Q\243\277\324\364\244(\345\317\247\277\\E<\034\372i\256?\034\2630\350\217\t\243\277Qjd\232k7\273?\331\376\340\306\277\373\262\305\033\010\034\304?k\332\341\377\312`\255?`O*\246*^\272?\rI\026\030\264\250\271?,\023b\003\r\021\267?\266\020\001\021\320\026\234?\004\306\217_\\\245\253?\003\t\316\230O\260\302\277aN*\262\263u\275\277\346~\3005\211\200\256?\210.\331\356\342$\265\277\276O\306\361\371\333\262?\377\255\373I\007\004k\277e\256oV\241\234\314\277\340\377\262\306\274V\244\277e\331\336Y\3401\310?\007+!@\2638\301?\020_\325\316\004\220\261?\326\303\335\324\030\244\257\277\356j\0011[\017\3005\316\032\260\277\237\240\322C\237\376\273?\317|[\351\335$\246?\234d\367\327)H\302?\246\345\304\317\272\000\266?z\246\006\234\357\302\300?\227\322\225\263\340S\251\277\264\006)\374b\275{\277Y8\316\271\357\330\207\277\332\311\305\2052P\303\277\275\364Wk>\017\216?BW\251-\220Q\243\277\232?Gkt\022\301\277e\214\374\037|\263:\201\2778\203\005n\037b\257\277\032\346\264\367\375=\241\277\256\277/f*\177\304?\236\020\314\354\262\336\303\277H\230\3248y&\204?q\206\361\235@M\251?G9\252DH~\236?\'\341k\316\232\310\263?v2\313\347\335~\266\277\241\323;\303\373\034\244?\004\375\177bY\334\277?s\304<\311\240\335\272?\343\364X\237\207\205\276?\276\375\204N\t\004\226?\314\375_\302H9\302?\255F:\236\216\031\243\277G)\004\274W\312\254?\372\307G\266\245\231\301?o\014d\346U6\263?Kx\202\275u\221\270?&E\276\r\3142\247\277\325\277\246*\202\027\206\277\001teJ\265\307\276?@m\230VB\006\277\277\007\032\377\256N<\231?\304g\272\034\254\221\262?9}\006\237\372\234\261\277\352\032\345Go\002\220?\340\003\372U\030\202\263?Q\346^\336\262F\306?\266\351:\353c\211\265\277\3509\244t^9\263\277\350\345%\014\212X\305?\377\310!\265\374\335\230\277y\335\221H\274\r\277?\257\326:q\311\302l?\017\213\240\245k\005\263\277}\037b\022\361\345\264?\254\253\210&\022}\251\277\031,\013\\b\340\263?\316\255\361L\021`\255?\010\372\017\211\310\337\251\277z\331\226\371(A\243?\216\246C\020U\362\260\277]\315\372\026z\014x\2777p\2305j\312\247\277\300\220Y\321i\201\244?\331_x\n\375\333\262\277\277\023\003f\247\021\300\277\234\260S\223\273\262\253\277o\215f\n\311\037\212\277\026\320\033a\304\274n\277\025\00771\276l\241\277\r\364\035(\240\'\273?\014+`\270T\006\267\2777\026\324\2021a\251\277R\216\343t\033\177\273?\022\313\215\365\033\212\233\277\377[\305\023\034\356\267?|\321\355\240\254^\225?J\277\001d\361\304\302?\034\177{+\3310\304\277c\340\204\264\213\343\267?\247\376qXW\241\266?\341\301\335\377d:\255\277$T\"\023\033\n\307\277u\314\211*\223\236\264?\006\000\250\260\345\031\243\277\236\r\206\3252J\255\277\313\237\n\246%q\274?\026E\301\274\331\221\236?\262\360\214\212%\205\240\277N\362\274\023\323T\262\277\260\234\340\025\014\036g\277\'\307\324\351\347j\251?\212\213\236\"\215G\236\277\231\354B\343\177\374\311?\213J\345L\201M\224\277\334@\372\315uo\312\277\2757-\014L\036\246\277\2631\'\252\013\004\245\277s\203\3605m\305\251\277J\365\272\204}U\262?\215\t\031\275\210\223\271\277@|\307\304~b\266?\026\202\211\000&\033\310\277(\302\267\002\331\360\252\277g\321qU\345\351\270?\334\003\025\033\366\271\260\277\315\363\242\013zC\230\277\177\223\036q \033\020??B\332\3703(\305\277\005\300\000\366{\246\233\277M8\262<,\223\257?\3360\177\372\3605\247?\214\241\005\031\203\r\227\277\037\321\231\226|\251\232\277~\230\376y\307?\260\277\t\r\033r\236\230\271\277W}\333\206/\351\221\277\237\357\266\217\007Z\2108\"\265?-\351\331\020f\275\274?\357\244\250\257~\263\237\277}NK\210l\363\261\277|{\335j\323\021\300\277\021)\341\271\222\267\261\277\265,|\267]?\241\277ogZE\230\223\246?\351DQ\361\263I\204?\021\211\301\320\357J\\\277Q.\206\316%@\227?\313\370\375/\3728\250\277\262_w_\367\322\242?d\372\023M$\362\246?<=\352\361\372\t\252?\336h\027v\235\007\274\277q~\351\264\205\214\200?\320\231\241^)w\261?\232V9<\277\225\260\277\200\334\315\270U\301\204?&\332\023\352\322^\232\277:1\273-\031d\217?k\375n\237kH\240\277}5>`\343\022\230\277\355\240\244\373\273\344\240\277FE\234\032\342-\236?\301d\230G\244-\260?=j\317\242H\206\221\277\253\230\316\267\370}\261?\312\236\r\255M\022\320\277*\214O\370\346\354\177\277_\263P\313\345.\250?\352\331\2758\225\334\250?\362\255\313\020\201\372\270?[Z\025\337E\021m?\243\007\24384\372\277\277\241\024\024\210\262l\213?ZJ\026\004\263\277Z\204x\006qN\270\277\032D\037\301n\267\271\277?\244\306\245\256*\252?m\275\017n\301\307\264?D\376!\'\232\317\270?\230\020L\305\345\260\272?\315\3247\226\030\256\253\277\022\201\235\334\'\377\256\277\033\262\032`\034%\231?G\227YLX>\300?~\317\300\246\205\257\255?\347\221\345$\253\216\223?\371\367\326@)\364l\277\303<\365\3355{\251?*\022\332\230\207\n\224?u\001\270\301#>\300\277a\251?%\253\300ar]F?\273\323\222\322\000X\234\277=8\241\231\367\360\254?{y\371 {\373\267\277\234W\007hq\362q?W\344\267\213\303\226\256\277uKyDu\252\254?\376\340\370\221qf\245\277\263\314\251\210~*\262\277\274[\330GZA\252\277\272Y\221\277\035\266\215?#\360D\371\3461\263?0\353.\r\272\263\264?\253\270\336\327o\226\272\277V\275\337\361F\332g?\341\263\362\255\370U\246\277\\\310\256\203\2753\246?\207$\305\250\342a\267\277\354\002\2452\302K}?l\272\216\201T \266\277\266\276A\t\017\327\303\277\300\225\027Z\244\342\242\277\275D}\357\230E\215\277\363\271\273\026\n7\252\277@\234\254h\326\263\253?\230\016\330\374\326\253\303?\335i\014\253\345\251\260\277\275\241\240\217v\315\270?QB\315\231\342\375\244?\3254eU\1773\265\277\244\020\250i\276C\220\277\307\342\366\302\203\367\270?N\032\331\222J\326\250\277]\225k\002\222K\\\277\030&$,\253C\273\2771\002oKt>\266?\305\326\017\263M?\232\277M\361q)S\330\240?\306z\322\235\273\010\245\277" + } + } + } +} +node { + name: "layer_0_type_1/matrix/read" + op: "Identity" + input: "layer_0_type_1/matrix" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@layer_0_type_1/matrix" + } + } + } +} +node { + name: "layer_0_type_1/bias" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 120 + } + } + tensor_content: "\374\325\034\265c\177\366\277\350\226\303\214\246G\354\277\237\244zZRz\320?\010\346\246\257\212[\316\277\225\334\225\000\246\323\362\277\252\372A\252+\322\332?\345BA\357O%\325?\251d\2348\237\023\360?\256\004\262U\374\236\375?\366\030\225k\237A\353\277D\332=jN\357\344?h\0356\010/\300\354\277\372\205\003\321\r\301\304\277\010\033)Y\246\275\312\277\344\300\212\374\245\250\326\277\236\005\0057\232pw\277\244d\305\372h\034\371?\267\212G\351\267;\323?\320|\0016c\022\357? \353\233\033\3369\316?\375>\341\376\240l\343?\316\277i\207\325!\374\277\253Z*\247\0242\374\2775!\000\014\360k\366\277k\361\302cO\245\312?\327\"\323\2113\177\353\277X3\207\211\273\333\342\277W\245x\016j\202\374\277=\246j\352\230\257\356\277r\240\225\032\252\245\245?\306\277\307\200\373\253\350\277Zgq\234]K\367\277\261i\235\252\331\365\001\300]I\221\3067\225\333?\000\365\201${\273\363?f\320\325y\372\341\360?\010\200\231\\\3636\312?\363{\302=\301\201\272?\233@\t\262\363\322\311\277B\300\234\007N\201\373\277\353\370\036S\026\357\362\277z\312\016\340y\377\256?\221e\'\203\210\254\304?\311\217M\256\213:\350\277\373\021o\024:\302\361?-\'f\205\245{\352\277\362]Pd\215\362\262\277\226n> \266\027\004@\307vnN\300o\350\277\262?\204\346\017\364\341?\227kB3\251z\311?\211\214T\365\2317\352\277$tZ\0033\344\351?\013+\236\325\376X\326?4\275\311\346\207\205\345\277\022\253\234U\t1\235?W\311\027\224>P\353?\236\252lD\324X\246?\372\313\320H:d\337\277/\315\020\207Q\242\355\277\320\324\344M\325\236\241\277\247\264M\007E\t\000\300W\206\010\211\327n\364\277&\311e#~\260\364\277Y\272c)(b\273?91\351\341vI\315?\006z\177hG\316\345?.\346\036;\013\370\365\277\r6\227\312Y;\333\277\372g\014G\014\302\364?4\222\203K\006\203\361\277En\014;\372G\305\277w\237N\224\033\322\360\2773?Y\264b\340\316?w>\301q\242\356\330?\217\263\006\025\204\226\326?m\255\373\270\275\'\303\277\2434\335\262\366\320\353?\361\341AA2\347\360?=\303\315P4\006\361?$f\t\024\244%\302\277\326\275\207#+\363\353?\320\252mu#\240\362\277\345\271-20m\371\277pi\303\346%\246\260?\272Kc\264m\021\340?.\325\305\342zr\325\277?\305R\214l\351\355\277\211$\346f\231\032\002\300\010jYC*?\370?(\251\305\357G\037\345\277\241\347j\267C\213\337?\031\271\353\304\314V\274\277\330q\266\232\354\n\357\277\365\231\301\214\301\245\005\300I\356\252\0201\201\317??\332\213\243\230\301\362\277\205\233\320\335\276\207\333\277\355\031\267\376\205)\000\300\013(\2349\276\362\001\300\257\265=\310\315\354\341\277\257t\321?\301\3354?&\307\221\016\240c\352\277\221\356\275\007\025\033\305?}\332\312\304 \243\370\277\315 z\324 \330\361?5$\231Sk\037\320\277\320\372R\261\336\200\342?t\344F\251\3563\245?\003\206r$\321=\327?:\017\234\213P\240\340?\200N\315\217\'6\223?\341n\235u:\356\362?\276}\350\241\024i\373\277y\266\332\006\220#\363\277[jC\276\272X\267\277\2041\374\340\021\037\362\277\225<\306E\030\337\352\277\217O\305\227s\342\300?\267\254\030h%&\350\277" + } + } + } +} +node { + name: "layer_0_type_1/bias/read" + op: "Identity" + input: "layer_0_type_1/bias" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@layer_0_type_1/bias" + } + } + } +} +node { + name: "layer_0_type_1/MatMul" + op: "MatMul" + input: "Reshape_19" + input: "layer_0_type_1/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "layer_0_type_1/BiasAdd" + op: "BiasAdd" + input: "layer_0_type_1/MatMul" + input: "layer_0_type_1/bias/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "layer_0_type_1/Tanh" + op: "Tanh" + input: "layer_0_type_1/BiasAdd" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "layer_0_type_1/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377x\000\000\000" + } + } + } +} +node { + name: "layer_0_type_1/Reshape" + op: "Reshape" + input: "layer_0_type_1/Tanh" + input: "layer_0_type_1/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "final_layer_type_1/matrix" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 120 + } + dim { + size: 1 + } + } + tensor_content: "A\2101\303\313\035\302?\311\205]\235\3219\233\277\3633\207\"\035\355\250?`WB\241P>\307?u\324\316`\2433\226?[\333`A1x\262?\323\253\014Z\013\245\307?\353\rF\'\233\034\263\277k\001\372\253\272-\257\277z\215\210P\025L\271\277\351G\300y\227\002\275?#\335\r\210\245\344\263?\226V\316\332\031\205\241\277\317\275N(\312\013\311\277\276[&\256c\321\265?\\\311O\307\235\\\223\277M\274\361i\351\022\310?\365x\335\355\377\315\240\277\256\275\350}\317\363\274?\013\251x\033\200&\227?Su\350\327M\202\264?^yy\374\221j\314?\020r\323\206\260\214\267?m\211\027\212)\266\273?\2277\277\250\240M\262\277\266\217\307\321\266)t?r\276\036\352\212\277\267?\236\337\251\352\0134\247?\234c\342v\261\353\241?\212n\t\244J\030\263\277H\010v\tO\004\300\277Xk\013{@\342\260\277wP\\E\376\371\265?\341:\227\337\301\266\266?\003\310d\033\202#\257\277\314 \nS\224\201\271\277G6\303$>\235\260\277\250\344\355g\316n\224\277\0104I|\313\300\237?\232\313\t\356\210\326\242?G=\300C\255\037\301\277\330\337@\264\3503\315\277\262\265.v@0\220\277J\334\241\304+@\251\277\002bT\216\236\223\232\277EQ\274\r\016f\241\277\334\274+\2128\026\250?; 7D\371(\311\277\3755\023\261\024\343\261?-O\237\254\350\314\272\277\240\263\364\305\356\004\233\277\267\035\216&Z\325\234?[\000W\217\312\301\277?\\\221\267Y\t\327\252?\000\325\177\013\2622\251\277\263\261\264>\014S\222\277\372l\302C\234f\231\277\323P\357I\356k\253?\037\330\030;\226\"\266\277\362\272\375VD\006\263?x\363P\331\252n\244?$\257\251\245\307n\220?\206\232\360\n\316\025\261?\330\005\342\263\016\027\244?\315\234\252\264\014m\300\2776\033:|\001\242\220\277\217\310\275\2436g\267\277\351\235\342\r;\r\303\277\256\022\016,*N\265?\2055\352\236\237J\241\277\304+\006W\303\207\272\277\361\240to\013\342\261?*\264\365\\\370\343\301\277\306\253\024\210S3\312?\3571\325\304\177\037\233\277\'\254f\275\324z\254?\325>\3047\271\'\303\277\251h\327\225#\231\277?\231\177\367\001n\304\305\277`4}\324\277\025\252?Z\316SqG\245\201\277v\305\353Y\352b\240?\246\355\347\017f\273\216?\2358\276\272T\274\255\277\321\363|\036\002l\312?\032\373\345#\343\220\306\277\226\250\213z\304\270\271\277\214\300\032\250:k\302\277Ag5\2014\260\305\277WGp\312\343\352z?\t2x!&q\302?\311\245_7\221\306\245?i\027\216\353_\'\240?#\344\215\202\336\225\313?\036\020\351\244\255\005\226?\004`\247\203\311G\246\277\020q\261\244\233@\305\277M\256\031\326\026B\234?\204\264\275\364/\037q?)\210\354\205\261\205\207\277\242\334\254\340Y\357~?\375\264Rw>\213\274\277+8m\351\302\354\300?\234D\231\305\336\352\264?\035\230\324\003\3531\300\277<\227\361\205g\253\264\277w\242k\300q\201\243?\275\260!\210\0076\311\277}vMN\337\355\253?\006\010e\'@\301\273?\323\374\"\302\004#\272\277\276/\002\366\267P\212\277\374\336YW\310c\241\277\006\360B\360\\\n\270\277ON\016\325X\203\263?\206\037cb\327#\264\2774\024\026RE\366\267\277\271\247\274Y\177\323\203\277.\367\327:\206_\276\277m\361\355\200\366\244\265?" + } + } + } +} +node { + name: "final_layer_type_1/matrix/read" + op: "Identity" + input: "final_layer_type_1/matrix" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@final_layer_type_1/matrix" + } + } + } +} +node { + name: "final_layer_type_1/bias" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + dim { + size: 1 + } + } + double_val: 0.20219566093913552 + } + } + } +} +node { + name: "final_layer_type_1/bias/read" + op: "Identity" + input: "final_layer_type_1/bias" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@final_layer_type_1/bias" + } + } + } +} +node { + name: "final_layer_type_1/MatMul" + op: "MatMul" + input: "layer_0_type_1/Reshape" + input: "final_layer_type_1/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: false + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: false + } + } +} +node { + name: "final_layer_type_1/BiasAdd" + op: "BiasAdd" + input: "final_layer_type_1/MatMul" + input: "final_layer_type_1/bias/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "data_format" + value { + s: "NHWC" + } + } +} +node { + name: "Shape_5" + op: "Shape" + input: "Reshape_14" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "strided_slice_21/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_21/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_21/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_21" + op: "StridedSlice" + input: "Shape_5" + input: "strided_slice_21/stack" + input: "strided_slice_21/stack_1" + input: "strided_slice_21/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_22/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 3 + } + } + } +} +node { + name: "strided_slice_22/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 4 + } + } + } +} +node { + name: "strided_slice_22/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_22" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_22/stack" + input: "strided_slice_22/stack_1" + input: "strided_slice_22/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_20/shape" + op: "Pack" + input: "strided_slice_21" + input: "strided_slice_22" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_20" + op: "Reshape" + input: "final_layer_type_1/BiasAdd" + input: "Reshape_20/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "concat_3/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "concat_3" + op: "ConcatV2" + input: "Reshape_18" + input: "Reshape_20" + input: "concat_3/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "embedding_lookup_1/axis" + op: "Const" + attr { + key: "_class" + value { + list { + s: "loc:@fitting_attr/t_bias_atom_e" + } + } + } + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "embedding_lookup_1" + op: "GatherV2" + input: "fitting_attr/t_bias_atom_e/read" + input: "clip_by_value" + input: "embedding_lookup_1/axis" + attr { + key: "Taxis" + value { + type: DT_INT32 + } + } + attr { + key: "Tindices" + value { + type: DT_INT32 + } + } + attr { + key: "Tparams" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@fitting_attr/t_bias_atom_e" + } + } + } + attr { + key: "batch_dims" + value { + i: 0 + } + } +} +node { + name: "embedding_lookup_1/Identity" + op: "Identity" + input: "embedding_lookup_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "Shape_6" + op: "Shape" + input: "Reshape_14" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "strided_slice_24/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_24/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_24/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_24" + op: "StridedSlice" + input: "Shape_6" + input: "strided_slice_24/stack" + input: "strided_slice_24/stack_1" + input: "strided_slice_24/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_25/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_25/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 4 + } + } + } +} +node { + name: "strided_slice_25/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_25" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_25/stack" + input: "strided_slice_25/stack_1" + input: "strided_slice_25/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "Const" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "Sum" + op: "Sum" + input: "strided_slice_25" + input: "Const" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "Reshape_21/shape" + op: "Pack" + input: "strided_slice_24" + input: "Sum" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_21" + op: "Reshape" + input: "embedding_lookup_1/Identity" + input: "Reshape_21/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "add_3" + op: "AddV2" + input: "concat_3" + input: "Reshape_21" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "mul_3" + op: "Mul" + input: "add_3" + input: "Cast" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "Reshape_22/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_22" + op: "Reshape" + input: "mul_3" + input: "Reshape_22/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Shape_7" + op: "Shape" + input: "Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "strided_slice_26/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_26/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_26/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_26" + op: "StridedSlice" + input: "Shape_7" + input: "strided_slice_26/stack" + input: "strided_slice_26/stack_1" + input: "strided_slice_26/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_27/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_27/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_27/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_27" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_27/stack" + input: "strided_slice_27/stack_1" + input: "strided_slice_27/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_28/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_28/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_28/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_28" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_28/stack" + input: "strided_slice_28/stack_1" + input: "strided_slice_28/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "Reshape_23/shape/2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "Reshape_23/shape" + op: "Pack" + input: "strided_slice_26" + input: "strided_slice_27" + input: "Reshape_23/shape/2" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_23" + op: "Reshape" + input: "Reshape_22" + input: "Reshape_23/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape_24/shape" + op: "Pack" + input: "strided_slice_26" + input: "strided_slice_28" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_24" + op: "Reshape" + input: "Reshape_1" + input: "Reshape_24/shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Const_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "Const_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_29/stack/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_29/stack" + op: "Pack" + input: "strided_slice_29/stack/0" + input: "Const_1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "strided_slice_29/stack_1/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_29/stack_1" + op: "Pack" + input: "strided_slice_29/stack_1/0" + input: "strided_slice_27" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "strided_slice_29/stack_2/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_29/stack_2" + op: "Pack" + input: "strided_slice_29/stack_2/0" + input: "Const_2" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "strided_slice_29" + op: "StridedSlice" + input: "Reshape_24" + input: "strided_slice_29/stack" + input: "strided_slice_29/stack_1" + input: "strided_slice_29/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 3 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 1 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "GreaterEqual_1/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "GreaterEqual_1" + op: "GreaterEqual" + input: "strided_slice_29" + input: "GreaterEqual_1/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "zeros_like" + op: "ZerosLike" + input: "strided_slice_29" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "Select" + op: "Select" + input: "GreaterEqual_1" + input: "strided_slice_29" + input: "zeros_like" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_30/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_30/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_30/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_30" + op: "StridedSlice" + input: "model_attr/model_attr/t_out_bias/read" + input: "strided_slice_30/stack" + input: "strided_slice_30/stack_1" + input: "strided_slice_30/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "GatherV2/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "GatherV2" + op: "GatherV2" + input: "strided_slice_30" + input: "Select" + input: "GatherV2/axis" + attr { + key: "Taxis" + value { + type: DT_INT32 + } + } + attr { + key: "Tindices" + value { + type: DT_INT32 + } + } + attr { + key: "Tparams" + value { + type: DT_DOUBLE + } + } + attr { + key: "batch_dims" + value { + i: 0 + } + } +} +node { + name: "strided_slice_31/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_31/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_31/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_31" + op: "StridedSlice" + input: "model_attr/model_attr/t_out_std/read" + input: "strided_slice_31/stack" + input: "strided_slice_31/stack_1" + input: "strided_slice_31/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "GatherV2_1/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 0 + } + } + } +} +node { + name: "GatherV2_1" + op: "GatherV2" + input: "strided_slice_31" + input: "Select" + input: "GatherV2_1/axis" + attr { + key: "Taxis" + value { + type: DT_INT32 + } + } + attr { + key: "Tindices" + value { + type: DT_INT32 + } + } + attr { + key: "Tparams" + value { + type: DT_DOUBLE + } + } + attr { + key: "batch_dims" + value { + i: 0 + } + } +} +node { + name: "mul_5" + op: "Mul" + input: "Reshape_23" + input: "GatherV2_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "add_4" + op: "AddV2" + input: "mul_5" + input: "GatherV2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "ExpandDims/dim" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "ExpandDims" + op: "ExpandDims" + input: "GreaterEqual_1" + input: "ExpandDims/dim" + attr { + key: "T" + value { + type: DT_BOOL + } + } + attr { + key: "Tdim" + value { + type: DT_INT32 + } + } +} +node { + name: "Tile/multiples" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\001\000\000\000\001\000\000\000\001\000\000\000" + } + } + } +} +node { + name: "Tile" + op: "Tile" + input: "ExpandDims" + input: "Tile/multiples" + attr { + key: "T" + value { + type: DT_BOOL + } + } + attr { + key: "Tmultiples" + value { + type: DT_INT32 + } + } +} +node { + name: "Select_1" + op: "Select" + input: "Tile" + input: "add_4" + input: "Reshape_23" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "Shape_8" + op: "Shape" + input: "Reshape_22" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "Reshape_25" + op: "Reshape" + input: "Select_1" + input: "Shape_8" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_32/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_32/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_32/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_32" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_32/stack" + input: "strided_slice_32/stack_1" + input: "strided_slice_32/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "o_atom_energy/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "o_atom_energy/shape" + op: "Pack" + input: "o_atom_energy/shape/0" + input: "strided_slice_32" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "o_atom_energy" + op: "Reshape" + input: "Reshape_25" + input: "o_atom_energy/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "o_energy/reduction_indices" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "o_energy" + op: "Sum" + input: "o_atom_energy" + input: "o_energy/reduction_indices" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/Shape" + op: "Shape" + input: "Reshape_25" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/grad_ys_0/Const" + op: "Const" + attr { + key: "dtype" + value { + type: DT_DOUBLE + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_DOUBLE + tensor_shape { + } + double_val: 1.0 + } + } + } +} +node { + name: "gradients/grad_ys_0" + op: "Fill" + input: "gradients/Shape" + input: "gradients/grad_ys_0/Const" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "index_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Reshape_25_grad/Shape" + op: "Shape" + input: "Select_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Reshape_25_grad/Reshape" + op: "Reshape" + input: "gradients/grad_ys_0" + input: "gradients/Reshape_25_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Select_1_grad/zeros_like" + op: "ZerosLike" + input: "add_4" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/Select_1_grad/Select" + op: "Select" + input: "Tile" + input: "gradients/Reshape_25_grad/Reshape" + input: "gradients/Select_1_grad/zeros_like" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/Select_1_grad/Select_1" + op: "Select" + input: "Tile" + input: "gradients/Select_1_grad/zeros_like" + input: "gradients/Reshape_25_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/add_4_grad/Shape" + op: "Shape" + input: "mul_5" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/add_4_grad/Shape_1" + op: "Shape" + input: "GatherV2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/add_4_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/add_4_grad/Shape" + input: "gradients/add_4_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/add_4_grad/Sum" + op: "Sum" + input: "gradients/Select_1_grad/Select" + input: "gradients/add_4_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/add_4_grad/Reshape" + op: "Reshape" + input: "gradients/add_4_grad/Sum" + input: "gradients/add_4_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/mul_5_grad/Mul" + op: "Mul" + input: "gradients/add_4_grad/Reshape" + input: "GatherV2_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/mul_5_grad/Shape" + op: "Shape" + input: "Reshape_23" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/mul_5_grad/Shape_1" + op: "Shape" + input: "GatherV2_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/mul_5_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/mul_5_grad/Shape" + input: "gradients/mul_5_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/mul_5_grad/Sum" + op: "Sum" + input: "gradients/mul_5_grad/Mul" + input: "gradients/mul_5_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/mul_5_grad/Reshape" + op: "Reshape" + input: "gradients/mul_5_grad/Sum" + input: "gradients/mul_5_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN" + op: "AddN" + input: "gradients/Select_1_grad/Select_1" + input: "gradients/mul_5_grad/Reshape" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/Select_1_grad/Select_1" + } + } + } +} +node { + name: "gradients/Reshape_23_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -1 + } + } + } +} +node { + name: "gradients/Reshape_23_grad/Reshape" + op: "Reshape" + input: "gradients/AddN" + input: "gradients/Reshape_23_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Reshape_22_grad/Shape" + op: "Shape" + input: "mul_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Reshape_22_grad/Reshape" + op: "Reshape" + input: "gradients/Reshape_23_grad/Reshape" + input: "gradients/Reshape_22_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/mul_3_grad/Mul" + op: "Mul" + input: "gradients/Reshape_22_grad/Reshape" + input: "Cast" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/mul_3_grad/Shape" + op: "Shape" + input: "add_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/mul_3_grad/Shape_1" + op: "Shape" + input: "Cast" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/mul_3_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/mul_3_grad/Shape" + input: "gradients/mul_3_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/mul_3_grad/Sum" + op: "Sum" + input: "gradients/mul_3_grad/Mul" + input: "gradients/mul_3_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/mul_3_grad/Reshape" + op: "Reshape" + input: "gradients/mul_3_grad/Sum" + input: "gradients/mul_3_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/add_3_grad/Shape" + op: "Shape" + input: "concat_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/add_3_grad/Shape_1" + op: "Shape" + input: "Reshape_21" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/add_3_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/add_3_grad/Shape" + input: "gradients/add_3_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/add_3_grad/Sum" + op: "Sum" + input: "gradients/mul_3_grad/Reshape" + input: "gradients/add_3_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/add_3_grad/Reshape" + op: "Reshape" + input: "gradients/add_3_grad/Sum" + input: "gradients/add_3_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/concat_3_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/concat_3_grad/mod" + op: "FloorMod" + input: "concat_3/axis" + input: "gradients/concat_3_grad/Rank" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/concat_3_grad/ShapeN" + op: "ShapeN" + input: "Reshape_18" + input: "Reshape_20" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/concat_3_grad/ConcatOffset" + op: "ConcatOffset" + input: "gradients/concat_3_grad/mod" + input: "gradients/concat_3_grad/ShapeN" + input: "gradients/concat_3_grad/ShapeN:1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "shape_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/concat_3_grad/Slice" + op: "Slice" + input: "gradients/add_3_grad/Reshape" + input: "gradients/concat_3_grad/ConcatOffset" + input: "gradients/concat_3_grad/ShapeN" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/concat_3_grad/Slice_1" + op: "Slice" + input: "gradients/add_3_grad/Reshape" + input: "gradients/concat_3_grad/ConcatOffset:1" + input: "gradients/concat_3_grad/ShapeN:1" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/Reshape_18_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "gradients/Reshape_18_grad/Reshape" + op: "Reshape" + input: "gradients/concat_3_grad/Slice" + input: "gradients/Reshape_18_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Reshape_20_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "gradients/Reshape_20_grad/Reshape" + op: "Reshape" + input: "gradients/concat_3_grad/Slice_1" + input: "gradients/Reshape_20_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/final_layer_type_0/MatMul_grad/MatMul" + op: "MatMul" + input: "gradients/Reshape_18_grad/Reshape" + input: "final_layer_type_0/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/final_layer_type_1/MatMul_grad/MatMul" + op: "MatMul" + input: "gradients/Reshape_20_grad/Reshape" + input: "final_layer_type_1/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/layer_0_type_0/Reshape_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377x\000\000\000" + } + } + } +} +node { + name: "gradients/layer_0_type_0/Reshape_grad/Reshape" + op: "Reshape" + input: "gradients/final_layer_type_0/MatMul_grad/MatMul" + input: "gradients/layer_0_type_0/Reshape_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } +} +node { + name: "gradients/layer_0_type_1/Reshape_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377x\000\000\000" + } + } + } +} +node { + name: "gradients/layer_0_type_1/Reshape_grad/Reshape" + op: "Reshape" + input: "gradients/final_layer_type_1/MatMul_grad/MatMul" + input: "gradients/layer_0_type_1/Reshape_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } +} +node { + name: "gradients/layer_0_type_0/Tanh_grad/TanhGrad" + op: "TanhGrad" + input: "layer_0_type_0/Tanh" + input: "gradients/layer_0_type_0/Reshape_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/layer_0_type_1/Tanh_grad/TanhGrad" + op: "TanhGrad" + input: "layer_0_type_1/Tanh" + input: "gradients/layer_0_type_1/Reshape_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/layer_0_type_0/MatMul_grad/MatMul" + op: "MatMul" + input: "gradients/layer_0_type_0/Tanh_grad/TanhGrad" + input: "layer_0_type_0/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/layer_0_type_1/MatMul_grad/MatMul" + op: "MatMul" + input: "gradients/layer_0_type_1/Tanh_grad/TanhGrad" + input: "layer_0_type_1/matrix/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/Reshape_17_grad/Shape" + op: "Shape" + input: "Slice_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Reshape_17_grad/Reshape" + op: "Reshape" + input: "gradients/layer_0_type_0/MatMul_grad/MatMul" + input: "gradients/Reshape_17_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Reshape_19_grad/Shape" + op: "Shape" + input: "Slice_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Reshape_19_grad/Reshape" + op: "Reshape" + input: "gradients/layer_0_type_1/MatMul_grad/MatMul" + input: "gradients/Reshape_19_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_2_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "gradients/Slice_2_grad/Shape" + op: "Shape" + input: "Slice_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Slice_2_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/Slice_2_grad/stack" + op: "Pack" + input: "gradients/Slice_2_grad/Rank" + input: "gradients/Slice_2_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/Slice_2_grad/Reshape" + op: "Reshape" + input: "Slice_2/begin" + input: "gradients/Slice_2_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_2_grad/Shape_1" + op: "Shape" + input: "Reshape_14" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Slice_2_grad/sub" + op: "Sub" + input: "gradients/Slice_2_grad/Shape_1" + input: "gradients/Slice_2_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_2_grad/sub_1" + op: "Sub" + input: "gradients/Slice_2_grad/sub" + input: "Slice_2/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_2_grad/Reshape_1" + op: "Reshape" + input: "gradients/Slice_2_grad/sub_1" + input: "gradients/Slice_2_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_2_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/Slice_2_grad/concat" + op: "ConcatV2" + input: "gradients/Slice_2_grad/Reshape" + input: "gradients/Slice_2_grad/Reshape_1" + input: "gradients/Slice_2_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_2_grad/Pad" + op: "Pad" + input: "gradients/Reshape_17_grad/Reshape" + input: "gradients/Slice_2_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_3_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "gradients/Slice_3_grad/Shape" + op: "Shape" + input: "Slice_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Slice_3_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/Slice_3_grad/stack" + op: "Pack" + input: "gradients/Slice_3_grad/Rank" + input: "gradients/Slice_3_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/Slice_3_grad/Reshape" + op: "Reshape" + input: "Slice_3/begin" + input: "gradients/Slice_3_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_3_grad/Shape_1" + op: "Shape" + input: "Reshape_14" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Slice_3_grad/sub" + op: "Sub" + input: "gradients/Slice_3_grad/Shape_1" + input: "gradients/Slice_3_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_3_grad/sub_1" + op: "Sub" + input: "gradients/Slice_3_grad/sub" + input: "Slice_3/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_3_grad/Reshape_1" + op: "Reshape" + input: "gradients/Slice_3_grad/sub_1" + input: "gradients/Slice_3_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_3_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/Slice_3_grad/concat" + op: "ConcatV2" + input: "gradients/Slice_3_grad/Reshape" + input: "gradients/Slice_3_grad/Reshape_1" + input: "gradients/Slice_3_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Slice_3_grad/Pad" + op: "Pad" + input: "gradients/Reshape_19_grad/Reshape" + input: "gradients/Slice_3_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_1" + op: "AddN" + input: "gradients/Slice_2_grad/Pad" + input: "gradients/Slice_3_grad/Pad" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/Slice_2_grad/Pad" + } + } + } +} +node { + name: "gradients/Reshape_14_grad/Shape" + op: "Shape" + input: "o_descriptor" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Reshape_14_grad/Reshape" + op: "Reshape" + input: "gradients/AddN_1" + input: "gradients/Reshape_14_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Reshape_12_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\020\000\000\000" + } + } + } +} +node { + name: "gradients/Reshape_12_grad/Reshape" + op: "Reshape" + input: "gradients/Reshape_14_grad/Reshape" + input: "gradients/Reshape_12_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_12_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 3 + } + } + tensor_content: "\377\377\377\377\004\000\000\000\004\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_12_grad/Reshape" + op: "Reshape" + input: "gradients/Reshape_12_grad/Reshape" + input: "gradients/filter_type_all/Reshape_12_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/MatMul" + op: "BatchMatMulV2" + input: "filter_type_all/Slice_4" + input: "gradients/filter_type_all/Reshape_12_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: false + } + } + attr { + key: "adj_y" + value { + b: true + } + } + attr { + key: "grad_x" + value { + b: true + } + } + attr { + key: "grad_y" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/MatMul_1" + op: "BatchMatMulV2" + input: "filter_type_all/truediv" + input: "gradients/filter_type_all/Reshape_12_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: false + } + } + attr { + key: "adj_y" + value { + b: false + } + } + attr { + key: "grad_x" + value { + b: false + } + } + attr { + key: "grad_y" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/Shape" + op: "Shape" + input: "filter_type_all/truediv" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Slice_4" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -2 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice" + op: "StridedSlice" + input: "gradients/filter_type_all/MatMul_6_grad/Shape" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice/stack" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice/stack_1" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 1 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -2 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1" + op: "StridedSlice" + input: "gradients/filter_type_all/MatMul_6_grad/Shape_1" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1/stack" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1/stack_1" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 1 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice" + input: "gradients/filter_type_all/MatMul_6_grad/strided_slice_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/Sum" + op: "Sum" + input: "gradients/filter_type_all/MatMul_6_grad/MatMul" + input: "gradients/filter_type_all/MatMul_6_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_6_grad/Sum" + input: "gradients/filter_type_all/MatMul_6_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/Sum_1" + op: "Sum" + input: "gradients/filter_type_all/MatMul_6_grad/MatMul_1" + input: "gradients/filter_type_all/MatMul_6_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_6_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_6_grad/Sum_1" + input: "gradients/filter_type_all/MatMul_6_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/Shape" + op: "Shape" + input: "filter_type_all/Slice_4" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/stack" + op: "Pack" + input: "gradients/filter_type_all/Slice_4_grad/Rank" + input: "gradients/filter_type_all/Slice_4_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/Reshape" + op: "Reshape" + input: "filter_type_all/Slice_4/begin" + input: "gradients/filter_type_all/Slice_4_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/Shape_1" + op: "Shape" + input: "filter_type_all/truediv" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/sub" + op: "Sub" + input: "gradients/filter_type_all/Slice_4_grad/Shape_1" + input: "gradients/filter_type_all/Slice_4_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/sub_1" + op: "Sub" + input: "gradients/filter_type_all/Slice_4_grad/sub" + input: "filter_type_all/Slice_4/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/Slice_4_grad/sub_1" + input: "gradients/filter_type_all/Slice_4_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/concat" + op: "ConcatV2" + input: "gradients/filter_type_all/Slice_4_grad/Reshape" + input: "gradients/filter_type_all/Slice_4_grad/Reshape_1" + input: "gradients/filter_type_all/Slice_4_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_4_grad/Pad" + op: "Pad" + input: "gradients/filter_type_all/MatMul_6_grad/Reshape_1" + input: "gradients/filter_type_all/Slice_4_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_2" + op: "AddN" + input: "gradients/filter_type_all/MatMul_6_grad/Reshape" + input: "gradients/filter_type_all/Slice_4_grad/Pad" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/MatMul_6_grad/Reshape" + } + } + } +} +node { + name: "gradients/filter_type_all/truediv_grad/RealDiv" + op: "RealDiv" + input: "gradients/AddN_2" + input: "filter_type_all/truediv/y" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/MatMul" + op: "BatchMatMulV2" + input: "filter_type_all/Reshape_4" + input: "gradients/filter_type_all/truediv_grad/RealDiv" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: false + } + } + attr { + key: "adj_y" + value { + b: true + } + } + attr { + key: "grad_x" + value { + b: true + } + } + attr { + key: "grad_y" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/MatMul_1" + op: "BatchMatMulV2" + input: "filter_type_all/Reshape_5" + input: "gradients/filter_type_all/truediv_grad/RealDiv" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: false + } + } + attr { + key: "adj_y" + value { + b: false + } + } + attr { + key: "grad_x" + value { + b: false + } + } + attr { + key: "grad_y" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/Shape" + op: "Shape" + input: "filter_type_all/Reshape_5" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_4" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -2 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice" + op: "StridedSlice" + input: "gradients/filter_type_all/MatMul_2_grad/Shape" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice/stack" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice/stack_1" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 1 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -2 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1" + op: "StridedSlice" + input: "gradients/filter_type_all/MatMul_2_grad/Shape_1" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1/stack" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1/stack_1" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 1 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice" + input: "gradients/filter_type_all/MatMul_2_grad/strided_slice_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/Sum" + op: "Sum" + input: "gradients/filter_type_all/MatMul_2_grad/MatMul" + input: "gradients/filter_type_all/MatMul_2_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_2_grad/Sum" + input: "gradients/filter_type_all/MatMul_2_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/Sum_1" + op: "Sum" + input: "gradients/filter_type_all/MatMul_2_grad/MatMul_1" + input: "gradients/filter_type_all/MatMul_2_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_2_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_2_grad/Sum_1" + input: "gradients/filter_type_all/MatMul_2_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/MatMul" + op: "BatchMatMulV2" + input: "filter_type_all/Reshape_10" + input: "gradients/filter_type_all/truediv_grad/RealDiv" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: false + } + } + attr { + key: "adj_y" + value { + b: true + } + } + attr { + key: "grad_x" + value { + b: true + } + } + attr { + key: "grad_y" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/MatMul_1" + op: "BatchMatMulV2" + input: "filter_type_all/Reshape_11" + input: "gradients/filter_type_all/truediv_grad/RealDiv" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "adj_x" + value { + b: false + } + } + attr { + key: "adj_y" + value { + b: false + } + } + attr { + key: "grad_x" + value { + b: false + } + } + attr { + key: "grad_y" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/Shape" + op: "Shape" + input: "filter_type_all/Reshape_11" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_10" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -2 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice" + op: "StridedSlice" + input: "gradients/filter_type_all/MatMul_5_grad/Shape" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice/stack" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice/stack_1" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 1 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: -2 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1" + op: "StridedSlice" + input: "gradients/filter_type_all/MatMul_5_grad/Shape_1" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1/stack" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1/stack_1" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 1 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice" + input: "gradients/filter_type_all/MatMul_5_grad/strided_slice_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/Sum" + op: "Sum" + input: "gradients/filter_type_all/MatMul_5_grad/MatMul" + input: "gradients/filter_type_all/MatMul_5_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_5_grad/Sum" + input: "gradients/filter_type_all/MatMul_5_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/Sum_1" + op: "Sum" + input: "gradients/filter_type_all/MatMul_5_grad/MatMul_1" + input: "gradients/filter_type_all/MatMul_5_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: false + } + } +} +node { + name: "gradients/filter_type_all/MatMul_5_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_5_grad/Sum_1" + input: "gradients/filter_type_all/MatMul_5_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_5_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\\\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_5_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_2_grad/Reshape" + input: "gradients/filter_type_all/Reshape_5_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_4_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_4_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_2_grad/Reshape_1" + input: "gradients/filter_type_all/Reshape_4_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_11_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\270\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_11_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_5_grad/Reshape" + input: "gradients/filter_type_all/Reshape_11_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_10_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_10_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/MatMul_5_grad/Reshape_1" + input: "gradients/filter_type_all/Reshape_10_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/Shape" + op: "Shape" + input: "filter_type_all/concat_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/add_1_grad/Shape" + input: "gradients/filter_type_all/add_1_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/Sum" + op: "Sum" + input: "gradients/filter_type_all/Reshape_4_grad/Reshape" + input: "gradients/filter_type_all/add_1_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_1_grad/Sum" + input: "gradients/filter_type_all/add_1_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/Sum_1" + op: "Sum" + input: "gradients/filter_type_all/Reshape_4_grad/Reshape" + input: "gradients/filter_type_all/add_1_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_1_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/add_1_grad/Sum_1" + input: "gradients/filter_type_all/add_1_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/Shape" + op: "Shape" + input: "filter_type_all/concat_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_9" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/add_3_grad/Shape" + input: "gradients/filter_type_all/add_3_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/Sum" + op: "Sum" + input: "gradients/filter_type_all/Reshape_10_grad/Reshape" + input: "gradients/filter_type_all/add_3_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_3_grad/Sum" + input: "gradients/filter_type_all/add_3_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/Sum_1" + op: "Sum" + input: "gradients/filter_type_all/Reshape_10_grad/Reshape" + input: "gradients/filter_type_all/add_3_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_3_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/add_3_grad/Sum_1" + input: "gradients/filter_type_all/add_3_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_1_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/concat_1_grad/mod" + op: "FloorMod" + input: "filter_type_all/concat_1/axis" + input: "gradients/filter_type_all/concat_1_grad/Rank" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_1_grad/ShapeN" + op: "ShapeN" + input: "filter_type_all/add" + input: "filter_type_all/add" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_1_grad/ConcatOffset" + op: "ConcatOffset" + input: "gradients/filter_type_all/concat_1_grad/mod" + input: "gradients/filter_type_all/concat_1_grad/ShapeN" + input: "gradients/filter_type_all/concat_1_grad/ShapeN:1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "shape_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_1_grad/Slice" + op: "Slice" + input: "gradients/filter_type_all/add_1_grad/Reshape" + input: "gradients/filter_type_all/concat_1_grad/ConcatOffset" + input: "gradients/filter_type_all/concat_1_grad/ShapeN" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/concat_1_grad/Slice_1" + op: "Slice" + input: "gradients/filter_type_all/add_1_grad/Reshape" + input: "gradients/filter_type_all/concat_1_grad/ConcatOffset:1" + input: "gradients/filter_type_all/concat_1_grad/ShapeN:1" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/Reshape_3_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_3_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_1_grad/Reshape_1" + input: "gradients/filter_type_all/Reshape_3_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/concat_3_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/concat_3_grad/mod" + op: "FloorMod" + input: "filter_type_all/concat_3/axis" + input: "gradients/filter_type_all/concat_3_grad/Rank" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_3_grad/ShapeN" + op: "ShapeN" + input: "filter_type_all/add_2" + input: "filter_type_all/add_2" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_3_grad/ConcatOffset" + op: "ConcatOffset" + input: "gradients/filter_type_all/concat_3_grad/mod" + input: "gradients/filter_type_all/concat_3_grad/ShapeN" + input: "gradients/filter_type_all/concat_3_grad/ShapeN:1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "shape_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_3_grad/Slice" + op: "Slice" + input: "gradients/filter_type_all/add_3_grad/Reshape" + input: "gradients/filter_type_all/concat_3_grad/ConcatOffset" + input: "gradients/filter_type_all/concat_3_grad/ShapeN" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/concat_3_grad/Slice_1" + op: "Slice" + input: "gradients/filter_type_all/add_3_grad/Reshape" + input: "gradients/filter_type_all/concat_3_grad/ConcatOffset:1" + input: "gradients/filter_type_all/concat_3_grad/ShapeN:1" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/Reshape_9_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\004\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_9_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_3_grad/Reshape_1" + input: "gradients/filter_type_all/Reshape_9_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/Tanh_1_grad/TanhGrad" + op: "TanhGrad" + input: "filter_type_all/Tanh_1" + input: "gradients/filter_type_all/Reshape_3_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/Tanh_3_grad/TanhGrad" + op: "TanhGrad" + input: "filter_type_all/Tanh_3" + input: "gradients/filter_type_all/Reshape_9_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/MatMul_1_grad/MatMul" + op: "MatMul" + input: "gradients/filter_type_all/Tanh_1_grad/TanhGrad" + input: "filter_type_all/matrix_2_0/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/MatMul_4_grad/MatMul" + op: "MatMul" + input: "gradients/filter_type_all/Tanh_3_grad/TanhGrad" + input: "filter_type_all/matrix_2_1/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/AddN_3" + op: "AddN" + input: "gradients/filter_type_all/concat_1_grad/Slice" + input: "gradients/filter_type_all/concat_1_grad/Slice_1" + input: "gradients/filter_type_all/MatMul_1_grad/MatMul" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/concat_1_grad/Slice" + } + } + } +} +node { + name: "gradients/filter_type_all/add_grad/Shape" + op: "Shape" + input: "filter_type_all/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/add_grad/Shape" + input: "gradients/filter_type_all/add_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_grad/Sum" + op: "Sum" + input: "gradients/AddN_3" + input: "gradients/filter_type_all/add_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_grad/Sum" + input: "gradients/filter_type_all/add_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_grad/Sum_1" + op: "Sum" + input: "gradients/AddN_3" + input: "gradients/filter_type_all/add_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/add_grad/Sum_1" + input: "gradients/filter_type_all/add_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_4" + op: "AddN" + input: "gradients/filter_type_all/concat_3_grad/Slice" + input: "gradients/filter_type_all/concat_3_grad/Slice_1" + input: "gradients/filter_type_all/MatMul_4_grad/MatMul" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/concat_3_grad/Slice" + } + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/Shape" + op: "Shape" + input: "filter_type_all/concat_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_8" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/BroadcastGradientArgs" + op: "BroadcastGradientArgs" + input: "gradients/filter_type_all/add_2_grad/Shape" + input: "gradients/filter_type_all/add_2_grad/Shape_1" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/Sum" + op: "Sum" + input: "gradients/AddN_4" + input: "gradients/filter_type_all/add_2_grad/BroadcastGradientArgs" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_2_grad/Sum" + input: "gradients/filter_type_all/add_2_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/Sum_1" + op: "Sum" + input: "gradients/AddN_4" + input: "gradients/filter_type_all/add_2_grad/BroadcastGradientArgs:1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } + attr { + key: "keep_dims" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/add_2_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/add_2_grad/Sum_1" + input: "gradients/filter_type_all/add_2_grad/Shape_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/concat_grad/mod" + op: "FloorMod" + input: "filter_type_all/concat/axis" + input: "gradients/filter_type_all/concat_grad/Rank" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_grad/ShapeN" + op: "ShapeN" + input: "filter_type_all/Reshape_1" + input: "filter_type_all/Reshape_1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_grad/ConcatOffset" + op: "ConcatOffset" + input: "gradients/filter_type_all/concat_grad/mod" + input: "gradients/filter_type_all/concat_grad/ShapeN" + input: "gradients/filter_type_all/concat_grad/ShapeN:1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "shape_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_grad/Slice" + op: "Slice" + input: "gradients/filter_type_all/add_grad/Reshape" + input: "gradients/filter_type_all/concat_grad/ConcatOffset" + input: "gradients/filter_type_all/concat_grad/ShapeN" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/concat_grad/Slice_1" + op: "Slice" + input: "gradients/filter_type_all/add_grad/Reshape" + input: "gradients/filter_type_all/concat_grad/ConcatOffset:1" + input: "gradients/filter_type_all/concat_grad/ShapeN:1" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/Reshape_2_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\002\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_2_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_grad/Reshape_1" + input: "gradients/filter_type_all/Reshape_2_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/concat_2_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/concat_2_grad/mod" + op: "FloorMod" + input: "filter_type_all/concat_2/axis" + input: "gradients/filter_type_all/concat_2_grad/Rank" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_2_grad/ShapeN" + op: "ShapeN" + input: "filter_type_all/Reshape_7" + input: "filter_type_all/Reshape_7" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_2_grad/ConcatOffset" + op: "ConcatOffset" + input: "gradients/filter_type_all/concat_2_grad/mod" + input: "gradients/filter_type_all/concat_2_grad/ShapeN" + input: "gradients/filter_type_all/concat_2_grad/ShapeN:1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "shape_type" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/concat_2_grad/Slice" + op: "Slice" + input: "gradients/filter_type_all/add_2_grad/Reshape" + input: "gradients/filter_type_all/concat_2_grad/ConcatOffset" + input: "gradients/filter_type_all/concat_2_grad/ShapeN" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/concat_2_grad/Slice_1" + op: "Slice" + input: "gradients/filter_type_all/add_2_grad/Reshape" + input: "gradients/filter_type_all/concat_2_grad/ConcatOffset:1" + input: "gradients/filter_type_all/concat_2_grad/ShapeN:1" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/Reshape_8_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\002\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_8_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/add_2_grad/Reshape_1" + input: "gradients/filter_type_all/Reshape_8_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } + attr { + key: "_has_manual_control_dependencies" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/Tanh_grad/TanhGrad" + op: "TanhGrad" + input: "filter_type_all/Tanh" + input: "gradients/filter_type_all/Reshape_2_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/Tanh_2_grad/TanhGrad" + op: "TanhGrad" + input: "filter_type_all/Tanh_2" + input: "gradients/filter_type_all/Reshape_8_grad/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } +} +node { + name: "gradients/filter_type_all/MatMul_grad/MatMul" + op: "MatMul" + input: "gradients/filter_type_all/Tanh_grad/TanhGrad" + input: "filter_type_all/matrix_1_0/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/filter_type_all/MatMul_3_grad/MatMul" + op: "MatMul" + input: "gradients/filter_type_all/Tanh_2_grad/TanhGrad" + input: "filter_type_all/matrix_1_1/read" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "grad_a" + value { + b: true + } + } + attr { + key: "grad_b" + value { + b: false + } + } + attr { + key: "transpose_a" + value { + b: false + } + } + attr { + key: "transpose_b" + value { + b: true + } + } +} +node { + name: "gradients/AddN_5" + op: "AddN" + input: "gradients/filter_type_all/concat_grad/Slice" + input: "gradients/filter_type_all/concat_grad/Slice_1" + input: "gradients/filter_type_all/MatMul_grad/MatMul" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/concat_grad/Slice" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_1_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_1_grad/Reshape" + op: "Reshape" + input: "gradients/AddN_5" + input: "gradients/filter_type_all/Reshape_1_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_6" + op: "AddN" + input: "gradients/filter_type_all/concat_2_grad/Slice" + input: "gradients/filter_type_all/concat_2_grad/Slice_1" + input: "gradients/filter_type_all/MatMul_3_grad/MatMul" + attr { + key: "N" + value { + i: 3 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/concat_2_grad/Slice" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_7_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\001\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_7_grad/Reshape" + op: "Reshape" + input: "gradients/AddN_6" + input: "gradients/filter_type_all/Reshape_7_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/Shape" + op: "Shape" + input: "filter_type_all/Slice_1" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/stack" + op: "Pack" + input: "gradients/filter_type_all/Slice_1_grad/Rank" + input: "gradients/filter_type_all/Slice_1_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/Reshape" + op: "Reshape" + input: "filter_type_all/Slice_1/begin" + input: "gradients/filter_type_all/Slice_1_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/sub" + op: "Sub" + input: "gradients/filter_type_all/Slice_1_grad/Shape_1" + input: "gradients/filter_type_all/Slice_1_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/sub_1" + op: "Sub" + input: "gradients/filter_type_all/Slice_1_grad/sub" + input: "filter_type_all/Slice_1/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/Slice_1_grad/sub_1" + input: "gradients/filter_type_all/Slice_1_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/concat" + op: "ConcatV2" + input: "gradients/filter_type_all/Slice_1_grad/Reshape" + input: "gradients/filter_type_all/Slice_1_grad/Reshape_1" + input: "gradients/filter_type_all/Slice_1_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_1_grad/Pad" + op: "Pad" + input: "gradients/filter_type_all/Reshape_1_grad/Reshape" + input: "gradients/filter_type_all/Slice_1_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/Shape" + op: "Shape" + input: "filter_type_all/Slice_3" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/stack" + op: "Pack" + input: "gradients/filter_type_all/Slice_3_grad/Rank" + input: "gradients/filter_type_all/Slice_3_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/Reshape" + op: "Reshape" + input: "filter_type_all/Slice_3/begin" + input: "gradients/filter_type_all/Slice_3_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/Shape_1" + op: "Shape" + input: "filter_type_all/Reshape_6" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/sub" + op: "Sub" + input: "gradients/filter_type_all/Slice_3_grad/Shape_1" + input: "gradients/filter_type_all/Slice_3_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/sub_1" + op: "Sub" + input: "gradients/filter_type_all/Slice_3_grad/sub" + input: "filter_type_all/Slice_3/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/Slice_3_grad/sub_1" + input: "gradients/filter_type_all/Slice_3_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/concat" + op: "ConcatV2" + input: "gradients/filter_type_all/Slice_3_grad/Reshape" + input: "gradients/filter_type_all/Slice_3_grad/Reshape_1" + input: "gradients/filter_type_all/Slice_3_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_3_grad/Pad" + op: "Pad" + input: "gradients/filter_type_all/Reshape_7_grad/Reshape" + input: "gradients/filter_type_all/Slice_3_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\\\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/Slice_1_grad/Pad" + input: "gradients/filter_type_all/Reshape_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Reshape_6_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\270\000\000\000" + } + } + } +} +node { + name: "gradients/filter_type_all/Reshape_6_grad/Reshape" + op: "Reshape" + input: "gradients/filter_type_all/Slice_3_grad/Pad" + input: "gradients/filter_type_all/Reshape_6_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_7" + op: "AddN" + input: "gradients/filter_type_all/Reshape_5_grad/Reshape" + input: "gradients/filter_type_all/Reshape_grad/Reshape" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/Reshape_5_grad/Reshape" + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/Shape" + op: "Shape" + input: "filter_type_all/Slice" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/stack" + op: "Pack" + input: "gradients/filter_type_all/Slice_grad/Rank" + input: "gradients/filter_type_all/Slice_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/Reshape" + op: "Reshape" + input: "filter_type_all/Slice/begin" + input: "gradients/filter_type_all/Slice_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/Shape_1" + op: "Shape" + input: "Reshape_9" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/sub" + op: "Sub" + input: "gradients/filter_type_all/Slice_grad/Shape_1" + input: "gradients/filter_type_all/Slice_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/sub_1" + op: "Sub" + input: "gradients/filter_type_all/Slice_grad/sub" + input: "filter_type_all/Slice/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/Slice_grad/sub_1" + input: "gradients/filter_type_all/Slice_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/concat" + op: "ConcatV2" + input: "gradients/filter_type_all/Slice_grad/Reshape" + input: "gradients/filter_type_all/Slice_grad/Reshape_1" + input: "gradients/filter_type_all/Slice_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_grad/Pad" + op: "Pad" + input: "gradients/AddN_7" + input: "gradients/filter_type_all/Slice_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_8" + op: "AddN" + input: "gradients/filter_type_all/Reshape_11_grad/Reshape" + input: "gradients/filter_type_all/Reshape_6_grad/Reshape" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/Reshape_11_grad/Reshape" + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/Rank" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 2 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/Shape" + op: "Shape" + input: "filter_type_all/Slice_2" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/stack/1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/stack" + op: "Pack" + input: "gradients/filter_type_all/Slice_2_grad/Rank" + input: "gradients/filter_type_all/Slice_2_grad/stack/1" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/Reshape" + op: "Reshape" + input: "filter_type_all/Slice_2/begin" + input: "gradients/filter_type_all/Slice_2_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/Shape_1" + op: "Shape" + input: "Reshape_9" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/sub" + op: "Sub" + input: "gradients/filter_type_all/Slice_2_grad/Shape_1" + input: "gradients/filter_type_all/Slice_2_grad/Shape" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/sub_1" + op: "Sub" + input: "gradients/filter_type_all/Slice_2_grad/sub" + input: "filter_type_all/Slice_2/begin" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/Reshape_1" + op: "Reshape" + input: "gradients/filter_type_all/Slice_2_grad/sub_1" + input: "gradients/filter_type_all/Slice_2_grad/stack" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/concat/axis" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 1 + } + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/concat" + op: "ConcatV2" + input: "gradients/filter_type_all/Slice_2_grad/Reshape" + input: "gradients/filter_type_all/Slice_2_grad/Reshape_1" + input: "gradients/filter_type_all/Slice_2_grad/concat/axis" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "Tidx" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/filter_type_all/Slice_2_grad/Pad" + op: "Pad" + input: "gradients/AddN_8" + input: "gradients/filter_type_all/Slice_2_grad/concat" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tpaddings" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/AddN_9" + op: "AddN" + input: "gradients/filter_type_all/Slice_grad/Pad" + input: "gradients/filter_type_all/Slice_2_grad/Pad" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "_class" + value { + list { + s: "loc:@gradients/filter_type_all/Slice_grad/Pad" + } + } + } +} +node { + name: "gradients/Reshape_9_grad/Shape" + op: "Shape" + input: "Reshape_8" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "gradients/Reshape_9_grad/Reshape" + op: "Reshape" + input: "gradients/AddN_9" + input: "gradients/Reshape_9_grad/Shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "gradients/Reshape_8_grad/Reshape/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\024\001\000\000" + } + } + } +} +node { + name: "gradients/Reshape_8_grad/Reshape" + op: "Reshape" + input: "gradients/Reshape_9_grad/Reshape" + input: "gradients/Reshape_8_grad/Reshape/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "Shape_9" + op: "Shape" + input: "o_nlist" + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "out_type" + value { + type: DT_INT32 + } + } + experimental_type { + type_id: TFT_PRODUCT + args { + type_id: TFT_SHAPE_TENSOR + args { + type_id: TFT_INT32 + } + } + } +} +node { + name: "strided_slice_33/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_33/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_33/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_33" + op: "StridedSlice" + input: "Shape_9" + input: "strided_slice_33/stack" + input: "strided_slice_33/stack_1" + input: "strided_slice_33/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "strided_slice_34/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 0 + } + } + } +} +node { + name: "strided_slice_34/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_34/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_34" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_34/stack" + input: "strided_slice_34/stack_1" + input: "strided_slice_34/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "mul_6/y" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 276 + } + } + } +} +node { + name: "mul_6" + op: "Mul" + input: "strided_slice_34" + input: "mul_6/y" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape_26/shape" + op: "Pack" + input: "strided_slice_33" + input: "mul_6" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_26" + op: "Reshape" + input: "gradients/Reshape_8_grad/Reshape" + input: "Reshape_26/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "ProdForceSeA" + op: "ProdForceSeA" + input: "Reshape_26" + input: "o_rmat_deriv" + input: "o_nlist" + input: "t_natoms" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "n_a_sel" + value { + i: 69 + } + } + attr { + key: "n_r_sel" + value { + i: 0 + } + } +} +node { + name: "ProdVirialSeA" + op: "ProdVirialSeA" + input: "Reshape_26" + input: "o_rmat_deriv" + input: "o_rij" + input: "o_nlist" + input: "t_natoms" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "n_a_sel" + value { + i: 69 + } + } + attr { + key: "n_r_sel" + value { + i: 0 + } + } +} +node { + name: "strided_slice_35/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_35/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_35/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_35" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_35/stack" + input: "strided_slice_35/stack_1" + input: "strided_slice_35/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "mul_7/x" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "mul_7" + op: "Mul" + input: "mul_7/x" + input: "strided_slice_35" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "Reshape_27/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "Reshape_27/shape" + op: "Pack" + input: "Reshape_27/shape/0" + input: "mul_7" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "Reshape_27" + op: "Reshape" + input: "ProdForceSeA" + input: "Reshape_27/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_36/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_36/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_36/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_36" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_36/stack" + input: "strided_slice_36/stack_1" + input: "strided_slice_36/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "mul_8/x" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 3 + } + } + } +} +node { + name: "mul_8" + op: "Mul" + input: "mul_8/x" + input: "strided_slice_36" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "o_force/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "o_force/shape" + op: "Pack" + input: "o_force/shape/0" + input: "mul_8" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "o_force" + op: "Reshape" + input: "Reshape_27" + input: "o_force/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "o_virial/shape" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 2 + } + } + tensor_content: "\377\377\377\377\t\000\000\000" + } + } + } +} +node { + name: "o_virial" + op: "Reshape" + input: "ProdVirialSeA" + input: "o_virial/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +node { + name: "strided_slice_37/stack" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_37/stack_1" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 2 + } + } + } +} +node { + name: "strided_slice_37/stack_2" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + dim { + size: 1 + } + } + int_val: 1 + } + } + } +} +node { + name: "strided_slice_37" + op: "StridedSlice" + input: "t_natoms" + input: "strided_slice_37/stack" + input: "strided_slice_37/stack_1" + input: "strided_slice_37/stack_2" + attr { + key: "Index" + value { + type: DT_INT32 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "begin_mask" + value { + i: 0 + } + } + attr { + key: "ellipsis_mask" + value { + i: 0 + } + } + attr { + key: "end_mask" + value { + i: 0 + } + } + attr { + key: "new_axis_mask" + value { + i: 0 + } + } + attr { + key: "shrink_axis_mask" + value { + i: 1 + } + } +} +node { + name: "mul_9/x" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: 9 + } + } + } +} +node { + name: "mul_9" + op: "Mul" + input: "mul_9/x" + input: "strided_slice_37" + attr { + key: "T" + value { + type: DT_INT32 + } + } +} +node { + name: "o_atom_virial/shape/0" + op: "Const" + attr { + key: "dtype" + value { + type: DT_INT32 + } + } + attr { + key: "value" + value { + tensor { + dtype: DT_INT32 + tensor_shape { + } + int_val: -1 + } + } + } +} +node { + name: "o_atom_virial/shape" + op: "Pack" + input: "o_atom_virial/shape/0" + input: "mul_9" + attr { + key: "N" + value { + i: 2 + } + } + attr { + key: "T" + value { + type: DT_INT32 + } + } + attr { + key: "axis" + value { + i: 0 + } + } +} +node { + name: "o_atom_virial" + op: "Reshape" + input: "ProdVirialSeA:1" + input: "o_atom_virial/shape" + attr { + key: "T" + value { + type: DT_DOUBLE + } + } + attr { + key: "Tshape" + value { + type: DT_INT32 + } + } +} +library { +} +versions { + producer: 1994 +} diff --git a/source/tests/pt/model/test_atomic_model_atomic_stat.py b/source/tests/pt/model/test_atomic_model_atomic_stat.py index e6e873e6c1..876b363dc8 100644 --- a/source/tests/pt/model/test_atomic_model_atomic_stat.py +++ b/source/tests/pt/model/test_atomic_model_atomic_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -81,11 +80,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_atomic_model_global_stat.py b/source/tests/pt/model/test_atomic_model_global_stat.py index fbb0279c09..acba3b56ad 100644 --- a/source/tests/pt/model/test_atomic_model_global_stat.py +++ b/source/tests/pt/model/test_atomic_model_global_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -93,11 +92,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_descriptor_dpa1.py b/source/tests/pt/model/test_descriptor_dpa1.py index abf5d1af01..27b84879dc 100644 --- a/source/tests/pt/model/test_descriptor_dpa1.py +++ b/source/tests/pt/model/test_descriptor_dpa1.py @@ -249,6 +249,7 @@ def test_descriptor_block(self) -> None: # this is an old state dict, modify manually state_dict["compress_info.0"] = des.compress_info[0] state_dict["compress_data.0"] = des.compress_data[0] + state_dict["type_embd_data"] = des.type_embd_data des.load_state_dict(state_dict) coord = self.coord atype = self.atype @@ -377,5 +378,6 @@ def translate_se_atten_and_type_embd_dicts_to_dpa1( target_dict[tk] = type_embd_dict[kk] record[all_keys.index("se_atten.compress_data.0")] = True record[all_keys.index("se_atten.compress_info.0")] = True + record[all_keys.index("se_atten.type_embd_data")] = True assert all(record) return target_dict diff --git a/source/tests/pt/model/test_descriptor_dpa2.py b/source/tests/pt/model/test_descriptor_dpa2.py index 6a859a497a..3fa6b86636 100644 --- a/source/tests/pt/model/test_descriptor_dpa2.py +++ b/source/tests/pt/model/test_descriptor_dpa2.py @@ -196,5 +196,6 @@ def translate_type_embd_dicts_to_dpa2( target_dict[tk] = type_embd_dict[kk] record[all_keys.index("repinit.compress_data.0")] = True record[all_keys.index("repinit.compress_info.0")] = True + record[all_keys.index("repinit.type_embd_data")] = True assert all(record) return target_dict diff --git a/source/tests/pt/model/test_force_grad.py b/source/tests/pt/model/test_force_grad.py index 44dbffa536..27bf660241 100644 --- a/source/tests/pt/model/test_force_grad.py +++ b/source/tests/pt/model/test_force_grad.py @@ -5,9 +5,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import torch @@ -31,7 +28,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pt/model/test_get_model.py b/source/tests/pt/model/test_get_model.py index e323c95ce0..2db6370059 100644 --- a/source/tests/pt/model/test_get_model.py +++ b/source/tests/pt/model/test_get_model.py @@ -60,6 +60,12 @@ def test_model_attr(self) -> None: ] }, ) + full_type_list = np.arange(len(atomic_model.type_map), dtype=int) + atom_exclude_types = np.setdiff1d( + full_type_list, + self.model.get_sel_type(), + ).tolist() + self.assertEqual(atom_exclude_types, [1]) self.assertEqual(atomic_model.atom_exclude_types, [1]) self.assertEqual(atomic_model.pair_exclude_types, [[1, 2]]) diff --git a/source/tests/pt/model/test_linear_atomic_model_stat.py b/source/tests/pt/model/test_linear_atomic_model_stat.py index 90758526b9..fdc4e6dac6 100644 --- a/source/tests/pt/model/test_linear_atomic_model_stat.py +++ b/source/tests/pt/model/test_linear_atomic_model_stat.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import h5py import numpy as np @@ -73,11 +70,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} @@ -125,11 +122,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} @@ -233,16 +230,11 @@ def test_linear_atomic_model_stat_with_bias(self) -> None: linear_model.compute_or_load_out_stat( self.merged_output_stat, stat_file_path=self.stat_file_path ) - # bias applied to sub atomic models. ener_bias = np.array([1.0, 3.0]).reshape(2, 1) - linear_ret = [] - for idx, md in enumerate(linear_model.models): - ret = md.forward_common_atomic(*args) - ret = to_numpy_array(ret["energy"]) - linear_ret.append(ret_no_bias[idx] + ener_bias[at]) - np.testing.assert_almost_equal((ret_no_bias[idx] + ener_bias[at]), ret) + ret = to_numpy_array(linear_model.forward_common_atomic(*args)["energy"]) + np.testing.assert_almost_equal((ret0 + ener_bias[at]), ret) # linear model not adding bias again ret1 = linear_model.forward_common_atomic(*args) ret1 = to_numpy_array(ret1["energy"]) - np.testing.assert_almost_equal(np.mean(np.stack(linear_ret), axis=0), ret1) + np.testing.assert_almost_equal(ret, ret1) diff --git a/source/tests/pt/model/test_polar_atomic_model_stat.py b/source/tests/pt/model/test_polar_atomic_model_stat.py index 19fb2a4d3f..c8b60f1164 100644 --- a/source/tests/pt/model/test_polar_atomic_model_stat.py +++ b/source/tests/pt/model/test_polar_atomic_model_stat.py @@ -6,7 +6,6 @@ ) from typing import ( NoReturn, - Optional, ) import h5py @@ -46,11 +45,11 @@ def forward( self, descriptor: torch.Tensor, atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, + gr: torch.Tensor | None = None, + g2: torch.Tensor | None = None, + h2: torch.Tensor | None = None, + fparam: torch.Tensor | None = None, + aparam: torch.Tensor | None = None, ): nf, nloc, _ = descriptor.shape ret = {} diff --git a/source/tests/pt/model/test_rotation.py b/source/tests/pt/model/test_rotation.py index 713a957e99..f2abc6f9fd 100644 --- a/source/tests/pt/model/test_rotation.py +++ b/source/tests/pt/model/test_rotation.py @@ -4,9 +4,6 @@ from pathlib import ( Path, ) -from typing import ( - Optional, -) import numpy as np import torch @@ -29,7 +26,7 @@ class CheckSymmetry(DeepmdData): def __init__( self, sys_path: str, - type_map: Optional[list[str]] = None, + type_map: list[str] | None = None, ) -> None: super().__init__(sys_path=sys_path, type_map=type_map) self.add("energy", 1, atomic=False, must=False, high_prec=True) diff --git a/source/tests/pt/model/test_saveload_dpa1.py b/source/tests/pt/model/test_saveload_dpa1.py index d09d156d4e..73d6adc0b1 100644 --- a/source/tests/pt/model/test_saveload_dpa1.py +++ b/source/tests/pt/model/test_saveload_dpa1.py @@ -69,7 +69,6 @@ def setUp(self) -> None: batch_size=None, num_workers=0, # setting to 0 diverges the behavior of its iterator; should be >=1 drop_last=False, - pin_memory=True, ) def cycle_iterator(iterable): diff --git a/source/tests/pt/model/test_saveload_se_e2_a.py b/source/tests/pt/model/test_saveload_se_e2_a.py index 12cd854664..d9bd7e0c95 100644 --- a/source/tests/pt/model/test_saveload_se_e2_a.py +++ b/source/tests/pt/model/test_saveload_se_e2_a.py @@ -69,7 +69,6 @@ def setUp(self) -> None: batch_size=None, num_workers=0, # setting to 0 diverges the behavior of its iterator; should be >=1 drop_last=False, - pin_memory=True, ) def cycle_iterator(iterable): diff --git a/source/tests/pt/model/water/data/data_0/set.000/aparam.npy b/source/tests/pt/model/water/data/data_0/set.000/aparam.npy new file mode 100644 index 0000000000..3058d9d660 Binary files /dev/null and b/source/tests/pt/model/water/data/data_0/set.000/aparam.npy differ diff --git a/source/tests/pt/model/water/data/data_0/set.000/fparam.npy b/source/tests/pt/model/water/data/data_0/set.000/fparam.npy new file mode 100644 index 0000000000..770b4a5d66 Binary files /dev/null and b/source/tests/pt/model/water/data/data_0/set.000/fparam.npy differ diff --git a/source/tests/pt/model/water/data/data_1/set.000/box.npy b/source/tests/pt/model/water/data/data_1/set.000/box.npy new file mode 100644 index 0000000000..6ad2de625b Binary files /dev/null and b/source/tests/pt/model/water/data/data_1/set.000/box.npy differ diff --git a/source/tests/pt/model/water/data/data_1/set.000/coord.npy b/source/tests/pt/model/water/data/data_1/set.000/coord.npy new file mode 100644 index 0000000000..8bd448b125 Binary files /dev/null and b/source/tests/pt/model/water/data/data_1/set.000/coord.npy differ diff --git a/source/tests/pt/model/water/data/data_1/set.000/energy.npy b/source/tests/pt/model/water/data/data_1/set.000/energy.npy new file mode 100644 index 0000000000..d03db103f5 Binary files /dev/null and b/source/tests/pt/model/water/data/data_1/set.000/energy.npy differ diff --git a/source/tests/pt/model/water/data/data_1/set.000/force.npy b/source/tests/pt/model/water/data/data_1/set.000/force.npy new file mode 100644 index 0000000000..10b2ab83a2 Binary files /dev/null and b/source/tests/pt/model/water/data/data_1/set.000/force.npy differ diff --git a/source/tests/pt/model/water/data/data_1/type.raw b/source/tests/pt/model/water/data/data_1/type.raw new file mode 100644 index 0000000000..97e8fdfcf8 --- /dev/null +++ b/source/tests/pt/model/water/data/data_1/type.raw @@ -0,0 +1,192 @@ +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +0 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 +1 diff --git a/source/tests/pt/model/water/data/data_1/type_map.raw b/source/tests/pt/model/water/data/data_1/type_map.raw new file mode 100644 index 0000000000..e900768b1d --- /dev/null +++ b/source/tests/pt/model/water/data/data_1/type_map.raw @@ -0,0 +1,2 @@ +O +H diff --git a/source/tests/pt/model/water/data/single/set.000/aparam.npy b/source/tests/pt/model/water/data/single/set.000/aparam.npy new file mode 100644 index 0000000000..af871fd2ab Binary files /dev/null and b/source/tests/pt/model/water/data/single/set.000/aparam.npy differ diff --git a/source/tests/pt/model/water/data/single/set.000/fparam.npy b/source/tests/pt/model/water/data/single/set.000/fparam.npy new file mode 100644 index 0000000000..0a55fbe92a Binary files /dev/null and b/source/tests/pt/model/water/data/single/set.000/fparam.npy differ diff --git a/source/tests/pt/test_calculator.py b/source/tests/pt/test_calculator.py index c73bbad415..7458117ca3 100644 --- a/source/tests/pt/test_calculator.py +++ b/source/tests/pt/test_calculator.py @@ -64,17 +64,108 @@ def test_calculator(self) -> None: atomic_numbers = [1, 1, 1, 8, 8] idx_perm = [1, 0, 4, 3, 2] + # Convert tensors to numpy for ASE compatibility + cell_np = cell.numpy() + coord_np = coord.numpy() + + prec = 1e-10 + low_prec = 1e-4 + + ase_atoms0 = Atoms( + numbers=atomic_numbers, + positions=coord_np, + # positions=[tuple(item) for item in coordinate], + cell=cell_np, + calculator=self.calculator, + pbc=True, + ) + e0, f0 = ase_atoms0.get_potential_energy(), ase_atoms0.get_forces() + s0, v0 = ( + ase_atoms0.get_stress(voigt=True), + -ase_atoms0.get_stress(voigt=False) * ase_atoms0.get_volume(), + ) + + ase_atoms1 = Atoms( + numbers=[atomic_numbers[i] for i in idx_perm], + positions=coord_np[idx_perm, :], + # positions=[tuple(item) for item in coordinate], + cell=cell_np, + calculator=self.calculator, + pbc=True, + ) + e1, f1 = ase_atoms1.get_potential_energy(), ase_atoms1.get_forces() + s1, v1 = ( + ase_atoms1.get_stress(voigt=True), + -ase_atoms1.get_stress(voigt=False) * ase_atoms1.get_volume(), + ) + + assert isinstance(e0, float) + assert f0.shape == (natoms, 3) + assert v0.shape == (3, 3) + np.testing.assert_allclose(e0, e1, rtol=low_prec, atol=prec) + np.testing.assert_allclose(f0[idx_perm, :], f1, rtol=low_prec, atol=prec) + np.testing.assert_allclose(s0, s1, rtol=low_prec, atol=prec) + np.testing.assert_allclose(v0, v1, rtol=low_prec, atol=prec) + + +class TestCalculatorWithFparamAparam(unittest.TestCase): + def setUp(self) -> None: + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + self.config["training"]["numb_steps"] = 1 + self.config["model"]["fitting_net"]["numb_fparam"] = 2 + self.config["model"]["fitting_net"]["numb_aparam"] = 1 + self.config["training"]["save_freq"] = 1 + data_file = [str(Path(__file__).parent / "water/data/data_0")] + self.config["training"]["training_data"]["systems"] = data_file + self.config["training"]["validation_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + self.input_json = "test_dp_test.json" + with open(self.input_json, "w") as fp: + json.dump(self.config, fp, indent=4) + + trainer = get_trainer(deepcopy(self.config)) + trainer.run() + + with torch.device("cpu"): + input_dict, label_dict, _ = trainer.get_data(is_train=False) + _, _, more_loss = trainer.wrapper(**input_dict, label=label_dict, cur_lr=1.0) + + self.calculator = DPCalculator("model.pt") + + def test_calculator(self) -> None: + from ase import ( + Atoms, + ) + + natoms = 5 + cell = torch.eye(3, dtype=dtype, device="cpu") * 10 + generator = torch.Generator(device="cpu").manual_seed(GLOBAL_SEED) + coord = torch.rand([natoms, 3], dtype=dtype, device="cpu", generator=generator) + coord = torch.matmul(coord, cell) + fparam = torch.IntTensor([1, 2]).numpy() + aparam = torch.IntTensor([[1], [0], [2], [1], [0]]).numpy() + atomic_numbers = [1, 1, 1, 8, 8] + idx_perm = [1, 0, 4, 3, 2] + + # Convert tensors to numpy for ASE compatibility + cell_np = cell.numpy() + coord_np = coord.numpy() + prec = 1e-10 low_prec = 1e-4 ase_atoms0 = Atoms( numbers=atomic_numbers, - positions=coord, + positions=coord_np, # positions=[tuple(item) for item in coordinate], - cell=cell, + cell=cell_np, calculator=self.calculator, pbc=True, ) + ase_atoms0.info.update({"fparam": fparam, "aparam": aparam}) e0, f0 = ase_atoms0.get_potential_energy(), ase_atoms0.get_forces() s0, v0 = ( ase_atoms0.get_stress(voigt=True), @@ -83,12 +174,13 @@ def test_calculator(self) -> None: ase_atoms1 = Atoms( numbers=[atomic_numbers[i] for i in idx_perm], - positions=coord[idx_perm, :], + positions=coord_np[idx_perm, :], # positions=[tuple(item) for item in coordinate], - cell=cell, + cell=cell_np, calculator=self.calculator, pbc=True, ) + ase_atoms1.info.update({"fparam": fparam, "aparam": aparam[idx_perm, :]}) e1, f1 = ase_atoms1.get_potential_energy(), ase_atoms1.get_forces() s1, v1 = ( ase_atoms1.get_stress(voigt=True), diff --git a/source/tests/pt/test_dp_show.py b/source/tests/pt/test_dp_show.py index d27e5e69d9..d816e3d0b4 100644 --- a/source/tests/pt/test_dp_show.py +++ b/source/tests/pt/test_dp_show.py @@ -47,43 +47,53 @@ def setUp(self) -> None: def test_checkpoint(self) -> None: INPUT = "model.pt" - ATTRIBUTES = "type-map descriptor fitting-net size" + ATTRIBUTES = "type-map descriptor fitting-net size observed-type" with redirect_stderr(io.StringIO()) as f: run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}") - results = f.getvalue().split("\n")[:-1] - assert "This is a singletask model" in results[-8] - assert "The type_map is ['O', 'H', 'Au']" in results[-7] + results = [ + res for res in f.getvalue().split("\n")[:-1] if "DEEPMD WARNING" not in res + ] # filter out warnings + assert "This is a singletask model" in results[0] + assert "The type_map is ['O', 'H', 'Au']" in results[1] assert ( "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0" - ) in results[-6] + ) in results[2] assert ( "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}" - in results[-5] + in results[3] ) - assert "Parameter counts:" in results[-4] - assert "Parameters in descriptor: 19,350" in results[-3] - assert "Parameters in fitting-net: 119,091" in results[-2] - assert "Parameters in total: 138,441" in results[-1] + assert "Parameter counts:" in results[4] + assert "Parameters in descriptor: 19,350" in results[5] + assert "Parameters in fitting-net: 119,091" in results[6] + assert "Parameters in total: 138,441" in results[7] + assert "The observed types for this model:" in results[8] + assert "Number of observed types: 2" in results[9] + assert "Observed types: ['H', 'O']" in results[10] def test_frozen_model(self) -> None: INPUT = "frozen_model.pth" - ATTRIBUTES = "type-map descriptor fitting-net size" + ATTRIBUTES = "type-map descriptor fitting-net size observed-type" with redirect_stderr(io.StringIO()) as f: run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}") - results = f.getvalue().split("\n")[:-1] - assert "This is a singletask model" in results[-8] - assert "The type_map is ['O', 'H', 'Au']" in results[-7] + results = [ + res for res in f.getvalue().split("\n")[:-1] if "DEEPMD WARNING" not in res + ] # filter out warnings + assert "This is a singletask model" in results[0] + assert "The type_map is ['O', 'H', 'Au']" in results[1] assert ( "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0" - ) in results[-6] + ) in results[2] assert ( "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}" - in results[-5] + in results[3] ) - assert "Parameter counts:" in results[-4] - assert "Parameters in descriptor: 19,350" in results[-3] - assert "Parameters in fitting-net: 119,091" in results[-2] - assert "Parameters in total: 138,441" in results[-1] + assert "Parameter counts:" in results[4] + assert "Parameters in descriptor: 19,350" in results[5] + assert "Parameters in fitting-net: 119,091" in results[6] + assert "Parameters in total: 138,441" in results[7] + assert "The observed types for this model:" in results[8] + assert "Number of observed types: 2" in results[9] + assert "Observed types: ['H', 'O']" in results[10] # only covers two elements def test_checkpoint_error(self) -> None: INPUT = "model.pt" @@ -152,62 +162,81 @@ def setUp(self) -> None: def test_checkpoint(self) -> None: INPUT = "model.ckpt.pt" - ATTRIBUTES = "model-branch type-map descriptor fitting-net size" + ATTRIBUTES = "model-branch type-map descriptor fitting-net size observed-type" with redirect_stderr(io.StringIO()) as f: run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}") - results = f.getvalue().split("\n")[:-1] - assert "This is a multitask model" in results[-12] + results = [ + res + for res in f.getvalue().split("\n")[:-1] + if "DEEPMD WARNING" not in res + and "|" not in res + and "+-" not in res + and "Detailed information" not in res + ] # filter out warnings and tables + assert "This is a multitask model" in results[0] assert ( "Available model branches are ['model_1', 'model_2', 'RANDOM'], " "where 'RANDOM' means using a randomly initialized fitting net." - in results[-11] + in results[1] ) - assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[-10] - assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[-9] + assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[2] + assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[3] assert ( "model_1" and "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5" - ) in results[-8] + ) in results[4] assert ( "model_2" and "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5" - ) in results[-7] + ) in results[5] assert ( "The fitting_net parameter of branch model_1 is {'neuron': [1, 2, 3], 'seed': 678}" - in results[-6] + in results[6] ) assert ( "The fitting_net parameter of branch model_2 is {'neuron': [9, 8, 7], 'seed': 1111}" - in results[-5] + in results[7] ) - assert "Parameter counts for a single branch model:" in results[-4] - assert "Parameters in descriptor: 19,350" in results[-3] - assert "Parameters in fitting-net: 4,860" in results[-2] - assert "Parameters in total: 24,210" in results[-1] + assert "Parameter counts for a single branch model:" in results[8] + assert "Parameters in descriptor: 19,350" in results[9] + assert "Parameters in fitting-net: 4,860" in results[10] + assert "Parameters in total: 24,210" in results[11] + assert "The observed types for each branch:" in results[12] + assert "model_1: Number of observed types: 2" in results[13] + assert "model_1: Observed types: ['H', 'O']" in results[14] + assert "model_2: Number of observed types: 2" in results[15] + assert "model_2: Observed types: ['H', 'O']" in results[16] + assert "TOTAL number of observed types in the model: 2" in results[17] + assert "TOTAL observed types in the model: ['H', 'O']" in results[18] def test_frozen_model(self) -> None: INPUT = "frozen_model.pth" - ATTRIBUTES = "type-map descriptor fitting-net size" + ATTRIBUTES = "type-map descriptor fitting-net size observed-type" with redirect_stderr(io.StringIO()) as f: run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}") - results = f.getvalue().split("\n")[:-1] - assert "This is a singletask model" in results[-8] - assert "The type_map is ['O', 'H', 'B']" in results[-7] + results = [ + res for res in f.getvalue().split("\n")[:-1] if "DEEPMD WARNING" not in res + ] # filter out warnings + assert "This is a singletask model" in results[0] + assert "The type_map is ['O', 'H', 'B']" in results[1] assert ( "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5" - ) in results[-6] + ) in results[2] assert ( "The fitting_net parameter is {'neuron': [1, 2, 3], 'seed': 678}" - in results[-5] + in results[3] ) - assert "Parameter counts:" in results[-4] - assert "Parameters in descriptor: 19,350" in results[-3] - assert "Parameters in fitting-net: 4,860" in results[-2] - assert "Parameters in total: 24,210" in results[-1] + assert "Parameter counts:" in results[4] + assert "Parameters in descriptor: 19,350" in results[5] + assert "Parameters in fitting-net: 4,860" in results[6] + assert "Parameters in total: 24,210" in results[7] + assert "The observed types for this model:" in results[8] + assert "Number of observed types: 2" in results[9] + assert "Observed types: ['H', 'O']" in results[10] # only covers two elements def tearDown(self) -> None: for f in os.listdir("."): diff --git a/source/tests/pt/test_dp_test.py b/source/tests/pt/test_dp_test.py index c2915c7ee7..1c11541e50 100644 --- a/source/tests/pt/test_dp_test.py +++ b/source/tests/pt/test_dp_test.py @@ -15,12 +15,19 @@ import torch from deepmd.entrypoints.test import test as dp_test +from deepmd.entrypoints.test import test_ener as dp_test_ener +from deepmd.infer.deep_eval import ( + DeepEval, +) from deepmd.pt.entrypoints.main import ( get_trainer, ) from deepmd.pt.utils.utils import ( to_numpy_array, ) +from deepmd.utils.data import ( + DeepmdData, +) from .model.test_permutation import ( model_property, @@ -30,7 +37,9 @@ class DPTest: - def test_dp_test_1_frame(self) -> None: + def _run_dp_test( + self, use_input_json: bool, numb_test: int = 0, use_train: bool = False + ) -> None: trainer = get_trainer(deepcopy(self.config)) with torch.device("cpu"): input_dict, label_dict, _ = trainer.get_data(is_train=False) @@ -44,12 +53,17 @@ def test_dp_test_1_frame(self) -> None: model = torch.jit.script(trainer.model) tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") torch.jit.save(model, tmp_model.name) + val_sys = self.config["training"]["validation_data"]["systems"] + if isinstance(val_sys, list): + val_sys = val_sys[0] dp_test( model=tmp_model.name, - system=self.config["training"]["validation_data"]["systems"][0], + system=None if use_input_json else val_sys, datafile=None, + train_json=self.input_json if use_input_json and use_train else None, + valid_json=self.input_json if use_input_json and not use_train else None, set_prefix="set", - numb_test=0, + numb_test=numb_test, rand_seed=None, shuffle_test=False, detail_file=self.detail_file, @@ -93,6 +107,20 @@ def test_dp_test_1_frame(self) -> None: ).reshape(-1, 3), ) + def test_dp_test_1_frame(self) -> None: + self._run_dp_test(False) + + def test_dp_test_input_json(self) -> None: + self._run_dp_test(True) + + def test_dp_test_input_json_train(self) -> None: + with open(self.input_json) as f: + cfg = json.load(f) + cfg["training"]["validation_data"]["systems"] = ["non-existent"] + with open(self.input_json, "w") as f: + json.dump(cfg, f, indent=4) + self._run_dp_test(True, use_train=True) + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("model") and f.endswith(".pt"): @@ -140,6 +168,208 @@ def setUp(self) -> None: json.dump(self.config, fp, indent=4) +class TestDPTestSeARglob(unittest.TestCase): + def setUp(self) -> None: + self.detail_file = "test_dp_test_ener_rglob_detail" + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + data_file = [str(Path(__file__).parent / "water/data/single")] + self.config["training"]["training_data"]["systems"] = data_file + root_dir = str(Path(__file__).parent) + self.config["training"]["validation_data"]["systems"] = root_dir + self.config["training"]["validation_data"]["rglob_patterns"] = [ + "water/data/single" + ] + self.config["model"] = deepcopy(model_se_e2_a) + self.input_json = "test_dp_test_rglob.json" + with open(self.input_json, "w") as fp: + json.dump(self.config, fp, indent=4) + + def test_dp_test_input_json_rglob(self) -> None: + trainer = get_trainer(deepcopy(self.config)) + with torch.device("cpu"): + input_dict, _, _ = trainer.get_data(is_train=False) + input_dict.pop("spin", None) + model = torch.jit.script(trainer.model) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + torch.jit.save(model, tmp_model.name) + dp_test( + model=tmp_model.name, + system=None, + datafile=None, + valid_json=self.input_json, + set_prefix="set", + numb_test=1, + rand_seed=None, + shuffle_test=False, + detail_file=self.detail_file, + atomic=False, + ) + os.unlink(tmp_model.name) + self.assertTrue(os.path.exists(self.detail_file + ".e.out")) + + def tearDown(self) -> None: + for f in os.listdir("."): + if f.startswith("model") and f.endswith(".pt"): + os.remove(f) + if f.startswith(self.detail_file): + os.remove(f) + if f in ["lcurve.out", self.input_json]: + os.remove(f) + if f in ["stat_files"]: + shutil.rmtree(f) + + +class TestDPTestSeARglobTrain(unittest.TestCase): + def setUp(self) -> None: + self.detail_file = "test_dp_test_ener_rglob_train_detail" + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + root_dir = str(Path(__file__).parent) + self.config["training"]["training_data"]["systems"] = root_dir + self.config["training"]["training_data"]["rglob_patterns"] = [ + "water/data/single" + ] + data_file = [str(Path(__file__).parent / "water/data/single")] + self.config["training"]["validation_data"]["systems"] = data_file + self.config["model"] = deepcopy(model_se_e2_a) + self.input_json = "test_dp_test_rglob_train.json" + with open(self.input_json, "w") as fp: + json.dump(self.config, fp, indent=4) + + def test_dp_test_input_json_rglob_train(self) -> None: + trainer = get_trainer(deepcopy(self.config)) + with torch.device("cpu"): + input_dict, _, _ = trainer.get_data(is_train=False) + input_dict.pop("spin", None) + model = torch.jit.script(trainer.model) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + torch.jit.save(model, tmp_model.name) + dp_test( + model=tmp_model.name, + system=None, + datafile=None, + train_json=self.input_json, + set_prefix="set", + numb_test=1, + rand_seed=None, + shuffle_test=False, + detail_file=self.detail_file, + atomic=False, + ) + os.unlink(tmp_model.name) + self.assertTrue(os.path.exists(self.detail_file + ".e.out")) + + def tearDown(self) -> None: + for f in os.listdir("."): + if f.startswith("model") and f.endswith(".pt"): + os.remove(f) + if f.startswith(self.detail_file): + os.remove(f) + if f in ["lcurve.out", self.input_json]: + os.remove(f) + if f in ["stat_files"]: + shutil.rmtree(f) + + +class TestDPTestForceWeight(DPTest, unittest.TestCase): + def setUp(self) -> None: + self.detail_file = "test_dp_test_force_weight_detail" + input_json = str(Path(__file__).parent / "water/se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + system_dir = self._prepare_weighted_system() + data_file = [system_dir] + self.config["training"]["training_data"]["systems"] = data_file + self.config["training"]["validation_data"]["systems"] = data_file + self.config["model"] = deepcopy(model_se_e2_a) + self.system_dir = system_dir + self.input_json = "test_dp_test_force_weight.json" + with open(self.input_json, "w") as fp: + json.dump(self.config, fp, indent=4) + + def _prepare_weighted_system(self) -> str: + src = Path(__file__).parent / "water/data/single" + tmp_dir = tempfile.mkdtemp() + shutil.copytree(src, tmp_dir, dirs_exist_ok=True) + set_dir = Path(tmp_dir) / "set.000" + forces = np.load(set_dir / "force.npy") + forces[0, :3] += 1.0 + forces[0, -3:] += 10.0 + np.save(set_dir / "force.npy", forces) + natoms = forces.shape[1] // 3 + atom_pref = np.ones((forces.shape[0], natoms), dtype=forces.dtype) + atom_pref[:, 0] = 2.0 + atom_pref[:, -1] = 0.0 + np.save(set_dir / "atom_pref.npy", atom_pref) + return tmp_dir + + def test_force_weight(self) -> None: + trainer = get_trainer(deepcopy(self.config)) + with torch.device("cpu"): + trainer.get_data(is_train=False) + model = torch.jit.script(trainer.model) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + torch.jit.save(model, tmp_model.name) + dp = DeepEval(tmp_model.name) + data = DeepmdData( + self.system_dir, + set_prefix="set", + shuffle_test=False, + type_map=dp.get_type_map(), + sort_atoms=False, + ) + err = dp_test_ener( + dp, + data, + self.system_dir, + numb_test=1, + detail_file=None, + has_atom_ener=False, + ) + test_data = data.get_test() + coord = test_data["coord"].reshape([1, -1]) + box = test_data["box"][:1] + atype = test_data["type"][0] + ret = dp.eval( + coord, + box, + atype, + fparam=None, + aparam=None, + atomic=False, + efield=None, + mixed_type=False, + spin=None, + ) + force_pred = ret[1].reshape([1, -1]) + force_true = test_data["force"][:1] + weight = test_data["atom_pref"][:1] + diff = force_pred - force_true + mae_unweighted = np.sum(np.abs(diff)) / diff.size + rmse_unweighted = np.sqrt(np.sum(diff * diff) / diff.size) + denom = weight.sum() + mae_weighted = np.sum(np.abs(diff) * weight) / denom + rmse_weighted = np.sqrt(np.sum(diff * diff * weight) / denom) + np.testing.assert_allclose(err["mae_f"][0], mae_unweighted) + np.testing.assert_allclose(err["rmse_f"][0], rmse_unweighted) + np.testing.assert_allclose(err["mae_fw"][0], mae_weighted) + np.testing.assert_allclose(err["rmse_fw"][0], rmse_weighted) + os.unlink(tmp_model.name) + + def tearDown(self) -> None: + super().tearDown() + shutil.rmtree(self.system_dir) + + class TestDPTestPropertySeA(unittest.TestCase): def setUp(self) -> None: self.detail_file = "test_dp_test_property_detail" diff --git a/source/tests/pt/test_eval_desc.py b/source/tests/pt/test_eval_desc.py new file mode 100644 index 0000000000..ff79a0a376 --- /dev/null +++ b/source/tests/pt/test_eval_desc.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import tempfile +import unittest +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) + +import numpy as np +import torch + +from deepmd.entrypoints.eval_desc import ( + eval_desc, +) +from deepmd.pt.entrypoints.main import ( + get_trainer, +) + +from .model.test_permutation import ( + model_se_e2_a, +) + + +class DPEvalDesc: + def test_dp_eval_desc_1_frame(self) -> None: + trainer = get_trainer(deepcopy(self.config)) + with torch.device("cpu"): + input_dict, label_dict, _ = trainer.get_data(is_train=False) + has_spin = getattr(trainer.model, "has_spin", False) + if callable(has_spin): + has_spin = has_spin() + if not has_spin: + input_dict.pop("spin", None) + input_dict["do_atomic_virial"] = True + result = trainer.model(**input_dict) + model = torch.jit.script(trainer.model) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + torch.jit.save(model, tmp_model.name) + + # Test eval_desc + eval_desc( + model=tmp_model.name, + system=self.config["training"]["validation_data"]["systems"][0], + datafile=None, + output=self.output_dir, + ) + os.unlink(tmp_model.name) + + # Check that descriptor file was created + system_name = os.path.basename( + self.config["training"]["validation_data"]["systems"][0].rstrip("/") + ) + desc_file = os.path.join(self.output_dir, f"{system_name}.npy") + self.assertTrue(os.path.exists(desc_file)) + + # Load and validate descriptor + descriptors = np.load(desc_file) + self.assertIsInstance(descriptors, np.ndarray) + # Descriptors should be 3D: (nframes, natoms, ndesc) + self.assertEqual(len(descriptors.shape), 3) # Should be 3D array + self.assertGreater(descriptors.shape[0], 0) # Should have frames + self.assertGreater(descriptors.shape[1], 0) # Should have atoms + self.assertGreater(descriptors.shape[2], 0) # Should have descriptor dimensions + + def tearDown(self) -> None: + for f in os.listdir("."): + if f.startswith("model") and f.endswith(".pt"): + os.remove(f) + if f in ["lcurve.out", self.input_json]: + os.remove(f) + if f in ["stat_files"]: + shutil.rmtree(f) + # Clean up output directory + if hasattr(self, "output_dir") and os.path.exists(self.output_dir): + shutil.rmtree(self.output_dir) + + +class TestDPEvalDescSeA(DPEvalDesc, unittest.TestCase): + def setUp(self) -> None: + self.output_dir = "test_eval_desc_output" + input_json = str(Path(__file__).parent / "water" / "se_atten.json") + with open(input_json) as f: + self.config = json.load(f) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + data_file = [str(Path(__file__).parent / "water" / "data" / "single")] + self.config["training"]["training_data"]["systems"] = data_file + self.config["training"]["validation_data"]["systems"] = data_file + self.config["model"] = deepcopy(model_se_e2_a) + self.input_json = "test_eval_desc.json" + with open(self.input_json, "w") as fp: + json.dump(self.config, fp, indent=4) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_fitting_stat.py b/source/tests/pt/test_fitting_stat.py index bc02b539a0..7807523221 100644 --- a/source/tests/pt/test_fitting_stat.py +++ b/source/tests/pt/test_fitting_stat.py @@ -1,18 +1,52 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import tempfile import unittest +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) +from typing import ( + NoReturn, +) +import h5py import numpy as np +import torch +from deepmd.pt.entrypoints.main import ( + get_trainer, +) from deepmd.pt.model.descriptor import ( DescrptSeA, ) from deepmd.pt.model.task import ( EnergyFittingNet, ) +from deepmd.pt.utils.multi_task import ( + preprocess_shared_params, +) from deepmd.pt.utils.utils import ( to_numpy_array, to_torch_tensor, ) +from deepmd.utils.argcheck import ( + normalize, +) +from deepmd.utils.compat import ( + update_deepmd_input, +) +from deepmd.utils.path import ( + DPPath, +) + +from .model.test_permutation import ( + model_se_e2_a, +) def _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds): @@ -71,16 +105,18 @@ def _brute_aparam_pt(data, ndim): class TestEnerFittingStat(unittest.TestCase): + def tearDown(self) -> None: + self.tempdir.cleanup() + + def setUp(self) -> None: + self.tempdir = tempfile.TemporaryDirectory() + h5file = str((Path(self.tempdir.name) / "testcase.h5").resolve()) + with h5py.File(h5file, "w") as f: + pass + self.stat_file_path = DPPath(h5file, "a") + def test(self) -> None: descrpt = DescrptSeA(6.0, 5.8, [46, 92], neuron=[25, 50, 100], axis_neuron=16) - fitting = EnergyFittingNet( - descrpt.get_ntypes(), - descrpt.get_dim_out(), - neuron=[240, 240, 240], - resnet_dt=True, - numb_fparam=3, - numb_aparam=3, - ) avgs = [0, 10, 100] stds = [2, 0.4, 0.00001] sys_natoms = [10, 100] @@ -88,11 +124,23 @@ def test(self) -> None: all_data = _make_fake_data_pt(sys_natoms, sys_nframes, avgs, stds) frefa, frefs = _brute_fparam_pt(all_data, len(avgs)) arefa, arefs = _brute_aparam_pt(all_data, len(avgs)) - fitting.compute_input_stats(all_data, protection=1e-2) frefs_inv = 1.0 / frefs arefs_inv = 1.0 / arefs frefs_inv[frefs_inv > 100] = 100 arefs_inv[arefs_inv > 100] = 100 + + # 1. test fitting stat is applied + fitting = EnergyFittingNet( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + numb_fparam=3, + numb_aparam=3, + ) + fitting.compute_input_stats( + all_data, protection=1e-2, stat_file_path=self.stat_file_path + ) np.testing.assert_almost_equal(frefa, to_numpy_array(fitting.fparam_avg)) np.testing.assert_almost_equal( frefs_inv, to_numpy_array(fitting.fparam_inv_std) @@ -101,3 +149,347 @@ def test(self) -> None: np.testing.assert_almost_equal( arefs_inv, to_numpy_array(fitting.aparam_inv_std) ) + del fitting + + # 2. test fitting stat writing to file is correct + concat_fparam = np.concatenate( + [ + to_numpy_array(all_data[ii]["fparam"].reshape(-1, 3)) + for ii in range(len(sys_nframes)) + ] + ) + concat_aparam = np.concatenate( + [ + to_numpy_array(all_data[ii]["aparam"].reshape(-1, 3)) + for ii in range(len(sys_nframes)) + ] + ) + fparam_stat = (self.stat_file_path / "fparam").load_numpy() + aparam_stat = (self.stat_file_path / "aparam").load_numpy() + np.testing.assert_almost_equal( + fparam_stat[:, 0], np.array([concat_fparam.shape[0]] * 3) + ) + np.testing.assert_almost_equal(fparam_stat[:, 1], np.sum(concat_fparam, axis=0)) + np.testing.assert_almost_equal( + fparam_stat[:, 2], np.sum(concat_fparam**2, axis=0) + ) + np.testing.assert_almost_equal( + aparam_stat[:, 0], np.array([concat_aparam.shape[0]] * 3) + ) + np.testing.assert_almost_equal(aparam_stat[:, 1], np.sum(concat_aparam, axis=0)) + np.testing.assert_almost_equal( + aparam_stat[:, 2], np.sum(concat_aparam**2, axis=0) + ) + + # 3. test fitting stat load from file + def raise_error() -> NoReturn: + raise RuntimeError + + fitting = EnergyFittingNet( + descrpt.get_ntypes(), + descrpt.get_dim_out(), + neuron=[240, 240, 240], + resnet_dt=True, + numb_fparam=3, + numb_aparam=3, + ) + fitting.compute_input_stats( + raise_error, protection=1e-2, stat_file_path=self.stat_file_path + ) + np.testing.assert_almost_equal(frefa, to_numpy_array(fitting.fparam_avg)) + np.testing.assert_almost_equal( + frefs_inv, to_numpy_array(fitting.fparam_inv_std) + ) + np.testing.assert_almost_equal(arefa, to_numpy_array(fitting.aparam_avg)) + np.testing.assert_almost_equal( + arefs_inv, to_numpy_array(fitting.aparam_inv_std) + ) + + +def get_weighted_fitting_stat(model_prob: list, *stat_arrays, protection: float): + n_arrays = len(stat_arrays) + assert len(model_prob) == n_arrays + + nframes = [stat.shape[0] for stat in stat_arrays] + sums = [stat.sum(axis=0) for stat in stat_arrays] + squared_sums = [(stat**2).sum(axis=0) for stat in stat_arrays] + + weighted_sum = sum(model_prob[i] * sums[i] for i in range(n_arrays)) + total_weighted_frames = sum(model_prob[i] * nframes[i] for i in range(n_arrays)) + weighted_avg = weighted_sum / total_weighted_frames + + weighted_square_sum = sum(model_prob[i] * squared_sums[i] for i in range(n_arrays)) + weighted_square_avg = weighted_square_sum / total_weighted_frames + weighted_std = np.sqrt(weighted_square_avg - weighted_avg**2) + weighted_std = np.where(weighted_std < protection, protection, weighted_std) + + return weighted_avg, weighted_std + + +class TestMultiTaskFittingStat(unittest.TestCase): + def setUp(self) -> None: + multitask_sharefit_template_json = str( + Path(__file__).parent / "water/multitask_sharefit.json" + ) + with open(multitask_sharefit_template_json) as f: + multitask_se_e2_a = json.load(f) + multitask_se_e2_a["model"]["shared_dict"]["my_descriptor"] = model_se_e2_a[ + "descriptor" + ] + self.data_file = [str(Path(__file__).parent / "water/data/data_0")] + self.data_file_without_fparam = [ + str(Path(__file__).parent / "water/data/data_1") + ] + self.data_file_single = [str(Path(__file__).parent / "water/data/single")] + self.stat_files = "se_e2_a_share_fit" + os.makedirs(self.stat_files, exist_ok=True) + + self.config = multitask_se_e2_a + self.config["training"]["data_dict"]["model_1"]["stat_file"] = ( + f"{self.stat_files}/model_1" + ) + self.config["training"]["data_dict"]["model_2"]["stat_file"] = ( + f"{self.stat_files}/model_2" + ) + self.config["model"]["shared_dict"]["my_fitting"]["numb_fparam"] = 2 + self.default_fparam = [1.0, 0.0] + self.config["model"]["shared_dict"]["my_fitting"]["default_fparam"] = ( + self.default_fparam + ) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + + self.origin_config = deepcopy(self.config) + + def test_sharefitting_with_fparam(self): + # test multitask training with fparam + self.config = deepcopy(self.origin_config) + model_prob = [0.3, 0.7] + self.config["training"]["model_prob"]["model_1"] = model_prob[0] + self.config["training"]["model_prob"]["model_2"] = model_prob[1] + + self.config["training"]["data_dict"]["model_1"]["training_data"]["systems"] = ( + self.data_file + ) + self.config["training"]["data_dict"]["model_1"]["validation_data"][ + "systems" + ] = self.data_file + self.config["training"]["data_dict"]["model_2"]["training_data"]["systems"] = ( + self.data_file_single + ) + self.config["training"]["data_dict"]["model_2"]["validation_data"][ + "systems" + ] = self.data_file_single + self.config["model"]["model_dict"]["model_1"]["data_stat_nbatch"] = 100 + + self.config["model"], self.shared_links = preprocess_shared_params( + self.config["model"] + ) + self.config = update_deepmd_input(self.config, warning=True) + self.config = normalize(self.config, multi_task=True) + trainer = get_trainer(deepcopy(self.config), shared_links=self.shared_links) + trainer.run() + + # check fparam shared + multi_state_dict = trainer.wrapper.model.state_dict() + torch.testing.assert_close( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_avg"], + multi_state_dict["model_2.atomic_model.fitting_net.fparam_avg"], + ) + torch.testing.assert_close( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_inv_std"], + multi_state_dict["model_2.atomic_model.fitting_net.fparam_inv_std"], + ) + + # check fitting stat in stat_file is correct + fparam_stat_model1 = np.load(f"{self.stat_files}/model_1/O H B/fparam") + fparam_stat_model2 = np.load(f"{self.stat_files}/model_2/O H B/fparam") + fparam_data1 = np.load(f"{self.data_file[0]}/set.000/fparam.npy") + fparam_data2 = np.load(f"{self.data_file_single[0]}/set.000/fparam.npy") + np.testing.assert_almost_equal( + fparam_stat_model1[:, 0], [fparam_data1.shape[0]] * 2 + ) + np.testing.assert_almost_equal( + fparam_stat_model1[:, 1], fparam_data1.sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model1[:, 2], (fparam_data1**2).sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model2[:, 0], [fparam_data2.shape[0]] * 2 + ) + np.testing.assert_almost_equal( + fparam_stat_model2[:, 1], fparam_data2.sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model2[:, 2], (fparam_data2**2).sum(axis=0) + ) + + # check shared fitting stat is computed correctly + weighted_avg, weighted_std = get_weighted_fitting_stat( + model_prob, fparam_data1, fparam_data2, protection=1e-2 + ) + np.testing.assert_almost_equal( + weighted_avg, + to_numpy_array( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_avg"] + ), + ) + np.testing.assert_almost_equal( + 1 / weighted_std, + to_numpy_array( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_inv_std"] + ), + ) + + def test_sharefitting_using_default_fparam(self): + # test multitask training with fparam + self.config = deepcopy(self.origin_config) + # add model3 + self.config["model"]["model_dict"]["model_3"] = deepcopy( + self.config["model"]["model_dict"]["model_2"] + ) + self.config["loss_dict"]["model_3"] = deepcopy( + self.config["loss_dict"]["model_2"] + ) + self.config["training"]["model_prob"]["model_3"] = deepcopy( + self.config["training"]["model_prob"]["model_2"] + ) + self.config["training"]["data_dict"]["model_3"] = deepcopy( + self.config["training"]["data_dict"]["model_2"] + ) + self.config["training"]["data_dict"]["model_3"]["stat_file"] = self.config[ + "training" + ]["data_dict"]["model_3"]["stat_file"].replace("model_2", "model_3") + self.config["model"]["shared_dict"]["my_fitting"]["dim_case_embd"] = 3 + + model_prob = [0.1, 0.3, 0.6] + self.config["training"]["model_prob"]["model_1"] = model_prob[0] + self.config["training"]["model_prob"]["model_2"] = model_prob[1] + self.config["training"]["model_prob"]["model_3"] = model_prob[2] + + self.config["training"]["data_dict"]["model_1"]["training_data"]["systems"] = ( + self.data_file_without_fparam + ) + self.config["training"]["data_dict"]["model_1"]["validation_data"][ + "systems" + ] = self.data_file_without_fparam + self.config["training"]["data_dict"]["model_2"]["training_data"]["systems"] = ( + self.data_file_single + ) + self.config["training"]["data_dict"]["model_2"]["validation_data"][ + "systems" + ] = self.data_file_single + self.config["training"]["data_dict"]["model_3"]["stat_file"] = ( + f"{self.stat_files}/model_3" + ) + self.config["training"]["data_dict"]["model_3"]["training_data"]["systems"] = ( + self.data_file + ) + self.config["training"]["data_dict"]["model_3"]["validation_data"][ + "systems" + ] = self.data_file + data_stat_protect = 5e-3 + self.config["model"]["model_dict"]["model_1"]["data_stat_nbatch"] = 3 + self.config["model"]["model_dict"]["model_3"]["data_stat_nbatch"] = 100 + self.config["model"]["model_dict"]["model_1"]["data_stat_protect"] = ( + data_stat_protect + ) + self.config["model"]["model_dict"]["model_2"]["data_stat_protect"] = ( + data_stat_protect + ) + self.config["model"]["model_dict"]["model_3"]["data_stat_protect"] = ( + data_stat_protect + ) + + self.config["model"], self.shared_links = preprocess_shared_params( + self.config["model"] + ) + self.config = update_deepmd_input(self.config, warning=True) + self.config = normalize(self.config, multi_task=True) + trainer = get_trainer(deepcopy(self.config), shared_links=self.shared_links) + trainer.run() + + # check fparam shared + multi_state_dict = trainer.wrapper.model.state_dict() + torch.testing.assert_close( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_avg"], + multi_state_dict["model_2.atomic_model.fitting_net.fparam_avg"], + ) + torch.testing.assert_close( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_avg"], + multi_state_dict["model_3.atomic_model.fitting_net.fparam_avg"], + ) + torch.testing.assert_close( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_inv_std"], + multi_state_dict["model_2.atomic_model.fitting_net.fparam_inv_std"], + ) + torch.testing.assert_close( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_inv_std"], + multi_state_dict["model_3.atomic_model.fitting_net.fparam_inv_std"], + ) + + # check fitting stat in stat_file is correct + fparam_stat_model1 = np.load(f"{self.stat_files}/model_1/O H B/fparam") + fparam_stat_model2 = np.load(f"{self.stat_files}/model_2/O H B/fparam") + fparam_stat_model3 = np.load(f"{self.stat_files}/model_3/O H B/fparam") + fparam_data1 = np.array([self.default_fparam]).repeat(3, axis=0) + fparam_data2 = np.load(f"{self.data_file_single[0]}/set.000/fparam.npy") + fparam_data3 = np.load(f"{self.data_file[0]}/set.000/fparam.npy") + np.testing.assert_almost_equal( + fparam_stat_model1[:, 0], [fparam_data1.shape[0]] * 2 + ) + np.testing.assert_almost_equal( + fparam_stat_model1[:, 1], fparam_data1.sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model1[:, 2], (fparam_data1**2).sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model2[:, 0], [fparam_data2.shape[0]] * 2 + ) + np.testing.assert_almost_equal( + fparam_stat_model2[:, 1], fparam_data2.sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model2[:, 2], (fparam_data2**2).sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model3[:, 0], [fparam_data3.shape[0]] * 2 + ) + np.testing.assert_almost_equal( + fparam_stat_model3[:, 1], fparam_data3.sum(axis=0) + ) + np.testing.assert_almost_equal( + fparam_stat_model3[:, 2], (fparam_data3**2).sum(axis=0) + ) + + # check shared fitting stat is computed correctly + weighted_avg, weighted_std = get_weighted_fitting_stat( + model_prob, + fparam_data1, + fparam_data2, + fparam_data3, + protection=data_stat_protect, + ) + np.testing.assert_almost_equal( + weighted_avg, + to_numpy_array( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_avg"] + ), + ) + np.testing.assert_almost_equal( + 1 / weighted_std, + to_numpy_array( + multi_state_dict["model_1.atomic_model.fitting_net.fparam_inv_std"] + ), + ) + + def tearDown(self) -> None: + for f in os.listdir("."): + if f.startswith("model") and f.endswith(".pt"): + os.remove(f) + if f in ["lcurve.out", "checkpoint"]: + os.remove(f) + if f in [self.stat_files]: + shutil.rmtree(f) diff --git a/source/tests/pt/test_init_frz_model.py b/source/tests/pt/test_init_frz_model.py index 62fbeede3c..612ef45f1c 100644 --- a/source/tests/pt/test_init_frz_model.py +++ b/source/tests/pt/test_init_frz_model.py @@ -13,6 +13,9 @@ import numpy as np +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) from deepmd.pt.entrypoints.main import ( freeze, get_trainer, @@ -20,6 +23,9 @@ from deepmd.pt.infer.deep_eval import ( DeepPot, ) +from deepmd.tf.utils.convert import ( + convert_pbtxt_to_pb, +) from .common import ( run_dp, @@ -29,7 +35,7 @@ class TestInitFrzModel(unittest.TestCase): def setUp(self) -> None: input_json = str(Path(__file__).parent / "water/se_atten.json") - with open(input_json) as f: + with open(input_json, encoding="utf-8") as f: config = json.load(f) config["model"]["descriptor"]["smooth_type_embedding"] = True config["training"]["numb_steps"] = 1 @@ -58,7 +64,7 @@ def setUp(self) -> None: empty_config["model"]["fitting_net"] = {} empty_config["training"]["numb_steps"] = 0 tmp_input = tempfile.NamedTemporaryFile(delete=False, suffix=".json") - with open(tmp_input.name, "w") as f: + with open(tmp_input.name, "w", encoding="utf-8") as f: json.dump(empty_config, f, indent=4) run_dp( f"dp --pt train {tmp_input.name} --init-frz-model {self.models[-1]} --use-pretrain-script --skip-neighbor-stat" @@ -129,6 +135,43 @@ def test_dp_test(self) -> None: np.testing.assert_allclose(av1, av2, rtol=1e-10, atol=1e-10) np.testing.assert_allclose(av1, av3, rtol=1e-10, atol=1e-10) + def test_init_frz_model_pb2pth(self) -> None: + """Test initialization from frozen model converted from pb.""" + frozen_model = "frozen_model_pb2pth.pth" + # Convert pth model from pb model + convert_pbtxt_to_pb( + str(Path(__file__).parent / "model/models/se_e2_a.pbtxt"), "frozen_model.pb" + ) + convert_backend(INPUT="frozen_model.pb", OUTPUT=frozen_model) + + # Create a base model + input_json = str(Path(__file__).parent / "model/models/se_e2_a.json") + with open(input_json, encoding="utf-8") as f: + config = json.load(f) + config["training"]["save_freq"] = 1 + config["learning_rate"]["start_lr"] = 1.0 + config["training"]["training_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + config["training"]["validation_data"]["systems"] = [ + str(Path(__file__).parent / "water/data/single") + ] + config["training"]["numb_steps"] = 0 + + trainer = get_trainer(config, init_frz_model=frozen_model) + # Explicit assertions to make test success criteria clear + self.assertIsNotNone(trainer, "Trainer should be successfully initialized") + self.assertTrue( + hasattr(trainer, "model"), "Trainer should have a model attribute" + ) + self.assertTrue( + hasattr(trainer, "optimizer"), "Trainer should have an optimizer attribute" + ) + # Run the trainer (this would fail if initialization was incorrect) + trainer.run() + # Verify the model was properly initialized from the frozen model + self.assertIsNotNone(trainer.model, "Model should be properly initialized") + def tearDown(self) -> None: for f in os.listdir("."): if f.startswith("frozen_model") and f.endswith(".pth"): @@ -139,3 +182,5 @@ def tearDown(self) -> None: os.remove(f) if f in ["stat_files"]: shutil.rmtree(f) + if f.startswith("frozen_model") and f.endswith(".pb"): + os.remove(f) diff --git a/source/tests/pt/test_loss_tensor.py b/source/tests/pt/test_loss_tensor.py index 5802c0b775..67dcb568e1 100644 --- a/source/tests/pt/test_loss_tensor.py +++ b/source/tests/pt/test_loss_tensor.py @@ -57,7 +57,7 @@ def get_single_batch(dataset, index=None): if key in np_batch.keys(): np_batch[key] = np.expand_dims(np_batch[key], axis=0) pt_batch[key] = torch.as_tensor(np_batch[key], device=env.DEVICE) - if key in ["coord", "atom_dipole"]: + if key in ["coord", "atom_dipole", "atom_polarizability"]: np_batch[key] = np_batch[key].reshape(1, -1) np_batch["natoms"] = np_batch["natoms"][0] return np_batch, pt_batch diff --git a/source/tests/pt/test_model_compression_se_atten.py b/source/tests/pt/test_model_compression_se_atten.py new file mode 100644 index 0000000000..41818f81d7 --- /dev/null +++ b/source/tests/pt/test_model_compression_se_atten.py @@ -0,0 +1,1019 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import unittest + +import numpy as np + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer.deep_eval import ( + DeepEval, +) + +from .common import ( + j_loader, + run_dp, + tests_path, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 10 + + +def _file_delete(file) -> None: + if os.path.isdir(file): + shutil.rmtree(file) + elif os.path.isfile(file): + os.remove(file) + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-atten.pth") + compressed_model = str(tests_path / "dp-compressed-se-atten.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_atten descriptor with attn_layer=0 for full compression compatibility (both type embedding and geometric parts) + jdata["model"]["descriptor"] = { + "type": "se_atten_v2", + "sel": 120, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "axis_neuron": 16, + "seed": 1, + "attn": 128, + "attn_layer": 0, + "attn_dotr": True, + "attn_mask": False, + "precision": "float64", + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-atten-exclude-types.pth") + compressed_model = str(tests_path / "dp-compressed-se-atten-exclude-types.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_atten descriptor with exclude_types + jdata["model"]["descriptor"] = { + "type": "se_atten_v2", + "exclude_types": [[0, 1]], + "sel": 120, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "axis_neuron": 16, + "seed": 1, + "attn": 128, + "attn_layer": 0, + "attn_dotr": True, + "attn_mask": False, + "precision": "float64", + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_nonzero_attn_layer(): + """Initialize models with attn_layer > 0 for partial compression testing.""" + suffix = "-nonzero-attn" + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / f"dp-original-se-atten{suffix}.pth") + compressed_model = str(tests_path / f"dp-compressed-se-atten{suffix}.pth") + INPUT = str(tests_path / f"input{suffix}.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_atten descriptor with attn_layer=2 for partial compression + # Only type embedding will be compressed, geometric parts (attention layers) will not + jdata["model"]["descriptor"] = { + "type": "se_atten_v2", + "sel": 120, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "axis_neuron": 16, + "seed": 1, + "attn": 128, + "attn_layer": 2, # Non-zero attention layer for partial compression testing + "attn_dotr": True, + "attn_mask": False, + "precision": "float64", + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_skip_neighbor_stat(): + suffix = "-skip-neighbor-stat" + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / f"dp-original-se-atten{suffix}.pth") + compressed_model = str(tests_path / f"dp-compressed-se-atten{suffix}.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_atten descriptor + jdata["model"]["descriptor"] = { + "type": "se_atten_v2", + "sel": 120, + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [25, 50, 100], + "resnet_dt": False, + "axis_neuron": 16, + "seed": 1, + "attn": 128, + "attn_layer": 0, + "attn_dotr": True, + "attn_mask": False, + "precision": "float64", + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT + " --skip-neighbor-stat") + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def setUpModule() -> None: + global \ + INPUT, \ + FROZEN_MODEL, \ + COMPRESSED_MODEL, \ + INPUT_ET, \ + FROZEN_MODEL_ET, \ + COMPRESSED_MODEL_ET, \ + INPUT_NONZERO_ATTN, \ + FROZEN_MODEL_NONZERO_ATTN, \ + COMPRESSED_MODEL_NONZERO_ATTN, \ + FROZEN_MODEL_SKIP_NEIGHBOR_STAT, \ + COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT + INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() + _, FROZEN_MODEL_SKIP_NEIGHBOR_STAT, COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT = ( + _init_models_skip_neighbor_stat() + ) + INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + INPUT_NONZERO_ATTN, FROZEN_MODEL_NONZERO_ATTN, COMPRESSED_MODEL_NONZERO_ATTN = ( + _init_models_nonzero_attn_layer() + ) + + +def tearDownModule() -> None: + # Clean up files created by _init_models + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete(COMPRESSED_MODEL) + # Clean up files created by _init_models_skip_neighbor_stat + _file_delete(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) + _file_delete(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) + # Clean up files created by _init_models_exclude_types + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + # Clean up files created by _init_models_nonzero_attn_layer + _file_delete(INPUT_NONZERO_ATTN) + _file_delete(FROZEN_MODEL_NONZERO_ATTN) + _file_delete(COMPRESSED_MODEL_NONZERO_ATTN) + # Clean up other artifacts + _file_delete("out.json") + _file_delete("input_v2_compat.json") + _file_delete("checkpoint") + _file_delete("lcurve.out") + _file_delete("model.ckpt.pt") + _file_delete("model.ckpt-1.pt") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression") + + +class TestDeepPotATPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotATNoPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = None + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotATLargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_ase(self) -> None: + from ase import ( + Atoms, + ) + + from deepmd.calculator import ( + DP, + ) + + water0 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(FROZEN_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + # nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +class TestDeepPotATPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_ET) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_ET) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestSkipNeighborStat(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotATNonZeroAttnLayer(unittest.TestCase): + """Test model compression with attn_layer > 0 (partial compression).""" + + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_NONZERO_ATTN) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_NONZERO_ATTN) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + """Test model attributes are consistent between original and compressed models.""" + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 6.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 6.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + """Test single frame evaluation with partial compression.""" + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values - should be identical even with partial compression + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + """Test single frame atomic evaluation with partial compression.""" + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values - should be identical even with partial compression + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + """Test multi-frame atomic evaluation with partial compression.""" + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values - should be identical even with partial compression + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_ase(self) -> None: + """Test ASE calculator integration with partial compression.""" + from ase import ( + Atoms, + ) + + from deepmd.calculator import ( + DP, + ) + + water0 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(FROZEN_MODEL_NONZERO_ATTN), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL_NONZERO_ATTN), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + # nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_model_compression_se_e3_tebd.py b/source/tests/pt/test_model_compression_se_e3_tebd.py new file mode 100644 index 0000000000..30976f6ce4 --- /dev/null +++ b/source/tests/pt/test_model_compression_se_e3_tebd.py @@ -0,0 +1,789 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import unittest + +import numpy as np + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer.deep_eval import ( + DeepEval, +) + +from .common import ( + j_loader, + run_dp, + tests_path, +) + +if GLOBAL_NP_FLOAT_PRECISION == np.float32: + default_places = 4 +else: + default_places = 9 + + +def _file_delete(file) -> None: + if os.path.isdir(file): + shutil.rmtree(file) + elif os.path.isfile(file): + os.remove(file) + + +def _init_models(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e3-tebd.pth") + compressed_model = str(tests_path / "dp-compressed-se-e3-tebd.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_e3_tebd descriptor + jdata["model"]["descriptor"] = { + "type": "se_e3_tebd", + "sel": 120, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [2, 4, 8], + "tebd_dim": 8, + "tebd_input_mode": "strip", + "activation_function": "tanh", + "resnet_dt": False, + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_exclude_types(): + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / "dp-original-se-e3-tebd-exclude-types.pth") + compressed_model = str(tests_path / "dp-compressed-se-e3-tebd-exclude-types.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_e3_tebd descriptor with exclude_types + jdata["model"]["descriptor"] = { + "type": "se_e3_tebd", + "exclude_types": [[0, 1]], + "sel": 120, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [2, 4, 8], + "tebd_dim": 8, + "tebd_input_mode": "strip", + "activation_function": "tanh", + "resnet_dt": False, + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT) + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def _init_models_skip_neighbor_stat(): + suffix = "-skip-neighbor-stat" + data_file = str(tests_path / os.path.join("model_compression", "data")) + frozen_model = str(tests_path / f"dp-original-se-e3-tebd{suffix}.pth") + compressed_model = str(tests_path / f"dp-compressed-se-e3-tebd{suffix}.pth") + INPUT = str(tests_path / "input.json") + jdata = j_loader(str(tests_path / os.path.join("model_compression", "input.json"))) + + # Configure se_e3_tebd descriptor + jdata["model"]["descriptor"] = { + "type": "se_e3_tebd", + "sel": 120, + "rcut_smth": 0.5, + "rcut": 4.0, + "neuron": [2, 4, 8], + "tebd_dim": 8, + "tebd_input_mode": "strip", + "activation_function": "tanh", + "resnet_dt": False, + } + + jdata["training"]["training_data"]["systems"] = data_file + with open(INPUT, "w") as fp: + json.dump(jdata, fp, indent=4) + + ret = run_dp("dp --pt train " + INPUT + " --skip-neighbor-stat") + np.testing.assert_equal(ret, 0, "DP train failed!") + ret = run_dp("dp --pt freeze -o " + frozen_model) + np.testing.assert_equal(ret, 0, "DP freeze failed!") + ret = run_dp( + "dp --pt compress " + + " -i " + + frozen_model + + " -o " + + compressed_model + + " -t " + + INPUT + ) + np.testing.assert_equal(ret, 0, "DP model compression failed!") + return INPUT, frozen_model, compressed_model + + +def setUpModule() -> None: + global \ + INPUT, \ + FROZEN_MODEL, \ + COMPRESSED_MODEL, \ + INPUT_ET, \ + FROZEN_MODEL_ET, \ + COMPRESSED_MODEL_ET, \ + FROZEN_MODEL_SKIP_NEIGHBOR_STAT, \ + COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT + INPUT, FROZEN_MODEL, COMPRESSED_MODEL = _init_models() + _, FROZEN_MODEL_SKIP_NEIGHBOR_STAT, COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT = ( + _init_models_skip_neighbor_stat() + ) + INPUT_ET, FROZEN_MODEL_ET, COMPRESSED_MODEL_ET = _init_models_exclude_types() + + +def tearDownModule() -> None: + # Clean up files created by _init_models + _file_delete(INPUT) + _file_delete(FROZEN_MODEL) + _file_delete(COMPRESSED_MODEL) + # Clean up files created by _init_models_skip_neighbor_stat + _file_delete(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) + _file_delete(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) + # Clean up files created by _init_models_exclude_types + _file_delete(INPUT_ET) + _file_delete(FROZEN_MODEL_ET) + _file_delete(COMPRESSED_MODEL_ET) + # Clean up other artifacts + _file_delete("out.json") + _file_delete("input_v2_compat.json") + _file_delete("checkpoint") + _file_delete("lcurve.out") + _file_delete("model.ckpt.pt") + _file_delete("model.ckpt-1.pt") + _file_delete("model-compression/checkpoint") + _file_delete("model-compression") + + +class TestDeepPotAPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 4.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotANoPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = None + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestDeepPotALargeBoxNoPBC(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([19.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_ase(self) -> None: + from ase import ( + Atoms, + ) + + from deepmd.calculator import ( + DP, + ) + + water0 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(FROZEN_MODEL), + ) + water1 = Atoms( + "OHHOHH", + positions=self.coords.reshape((-1, 3)), + cell=self.box.reshape((3, 3)), + calculator=DP(COMPRESSED_MODEL), + ) + ee0 = water0.get_potential_energy() + ff0 = water0.get_forces() + ee1 = water1.get_potential_energy() + ff1 = water1.get_forces() + # nframes = 1 + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + + +class TestDeepPotAPBCExcludeTypes(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_ET) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_ET) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 4.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +class TestSkipNeighborStat(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.dp_original = DeepEval(FROZEN_MODEL_SKIP_NEIGHBOR_STAT) + cls.dp_compressed = DeepEval(COMPRESSED_MODEL_SKIP_NEIGHBOR_STAT) + cls.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ] + ) + cls.atype = [0, 1, 1, 0, 1, 1] + cls.box = np.array([13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0]) + + def test_attrs(self) -> None: + self.assertEqual(self.dp_original.get_ntypes(), 2) + self.assertAlmostEqual(self.dp_original.get_rcut(), 4.0, places=default_places) + self.assertEqual(self.dp_original.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_original.get_dim_fparam(), 0) + self.assertEqual(self.dp_original.get_dim_aparam(), 0) + + self.assertEqual(self.dp_compressed.get_ntypes(), 2) + self.assertAlmostEqual( + self.dp_compressed.get_rcut(), 4.0, places=default_places + ) + self.assertEqual(self.dp_compressed.get_type_map(), ["O", "H"]) + self.assertEqual(self.dp_compressed.get_dim_fparam(), 0) + self.assertEqual(self.dp_compressed.get_dim_aparam(), 0) + + def test_1frame(self) -> None: + ee0, ff0, vv0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=False + ) + ee1, ff1, vv1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=False + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_1frame_atm(self) -> None: + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + self.coords, self.box, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + self.coords, self.box, self.atype, atomic=True + ) + # check shape of the returns + nframes = 1 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + def test_2frame_atm(self) -> None: + coords2 = np.concatenate((self.coords, self.coords)) + box2 = np.concatenate((self.box, self.box)) + ee0, ff0, vv0, ae0, av0 = self.dp_original.eval( + coords2, box2, self.atype, atomic=True + ) + ee1, ff1, vv1, ae1, av1 = self.dp_compressed.eval( + coords2, box2, self.atype, atomic=True + ) + # check shape of the returns + nframes = 2 + natoms = len(self.atype) + self.assertEqual(ee0.shape, (nframes, 1)) + self.assertEqual(ff0.shape, (nframes, natoms, 3)) + self.assertEqual(vv0.shape, (nframes, 9)) + self.assertEqual(ae0.shape, (nframes, natoms, 1)) + self.assertEqual(av0.shape, (nframes, natoms, 9)) + self.assertEqual(ee1.shape, (nframes, 1)) + self.assertEqual(ff1.shape, (nframes, natoms, 3)) + self.assertEqual(vv1.shape, (nframes, 9)) + self.assertEqual(ae1.shape, (nframes, natoms, 1)) + self.assertEqual(av1.shape, (nframes, natoms, 9)) + + # check values + np.testing.assert_almost_equal(ff0, ff1, default_places) + np.testing.assert_almost_equal(ae0, ae1, default_places) + np.testing.assert_almost_equal(av0, av1, default_places) + np.testing.assert_almost_equal(ee0, ee1, default_places) + np.testing.assert_almost_equal(vv0, vv1, default_places) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_padding_atoms.py b/source/tests/pt/test_padding_atoms.py new file mode 100644 index 0000000000..8fb417c6a4 --- /dev/null +++ b/source/tests/pt/test_padding_atoms.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from copy import ( + deepcopy, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) +from deepmd.pt.model.descriptor import ( + DescrptSeA, +) +from deepmd.pt.model.model import ( + PropertyModel, +) +from deepmd.pt.model.task import ( + PropertyFittingNet, +) +from deepmd.pt.utils.utils import ( + to_torch_tensor, +) + + +class TestCaseSingleFrameWithoutNlist: + def setUp(self) -> None: + # nf=2, nloc == 3 + self.nloc = 3 + self.nt = 2 + self.coord = np.array( + [ + [ + [0, 0, 0], + [0, 1, 0], + [0, 0, 1], + ], + [ + [1, 0, 1], + [0, 1, 1], + [1, 1, 0], + ], + ], + dtype=np.float64, + ) + self.atype = np.array([[0, 0, 1], [1, 1, 0]], dtype=int).reshape([2, self.nloc]) + self.cell = 2.0 * np.eye(3).reshape([1, 9]) + self.cell = np.array([self.cell, self.cell]).reshape(2, 9) + self.sel = [16, 8] + self.rcut = 2.2 + self.rcut_smth = 0.4 + self.atol = 1e-6 + self.rtol = 1e-5 + + +class TestPaddingAtoms(unittest.TestCase, TestCaseSingleFrameWithoutNlist): + def setUp(self): + TestCaseSingleFrameWithoutNlist.setUp(self) + + def test_padding_atoms_consistency(self): + ds = DescrptSeA( + self.rcut, + self.rcut_smth, + self.sel, + ) + ft = PropertyFittingNet( + self.nt, + ds.get_dim_out(), + mixed_types=ds.mixed_types(), + intensive=True, + property_name="abc", + ) + type_map = ["foo", "bar"] + model = PropertyModel(ds, ft, type_map=type_map) + var_name = model.get_var_name() + args = [to_torch_tensor(ii) for ii in [self.coord, self.atype, self.cell]] + result = model(*args) + # test intensive + np.testing.assert_allclose( + to_numpy_array(result[var_name].cpu().detach()), + np.mean(to_numpy_array(result[f"atom_{var_name}"].cpu().detach()), axis=1), + atol=self.atol, + rtol=self.rtol, + ) + # test padding atoms + padding_atoms_list = [1, 5, 10] + for padding_atoms in padding_atoms_list: + coord = deepcopy(self.coord) + atype = deepcopy(self.atype) + atype_padding = np.pad( + atype, + pad_width=((0, 0), (0, padding_atoms)), + mode="constant", + constant_values=-1, + ) + coord_padding = np.pad( + coord, + pad_width=((0, 0), (0, padding_atoms), (0, 0)), + mode="constant", + constant_values=0, + ) + args = [ + to_torch_tensor(ii) for ii in [coord_padding, atype_padding, self.cell] + ] + result_padding = model(*args) + np.testing.assert_allclose( + to_numpy_array(result[var_name].cpu().detach()), + to_numpy_array(result_padding[var_name].cpu().detach()), + atol=self.atol, + rtol=self.rtol, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_sampler.py b/source/tests/pt/test_sampler.py index ffe459a834..3d7143b350 100644 --- a/source/tests/pt/test_sampler.py +++ b/source/tests/pt/test_sampler.py @@ -62,7 +62,6 @@ def test_sampler_debug_info(self) -> None: batch_size=None, num_workers=0, # setting to 0 diverges the behavior of its iterator; should be >=1 drop_last=False, - pin_memory=True, ) with torch.device("cpu"): batch_data = next(iter(dataloader)) diff --git a/source/tests/pt/test_tabulate.py b/source/tests/pt/test_tabulate.py index 164819408f..d6075c3a74 100644 --- a/source/tests/pt/test_tabulate.py +++ b/source/tests/pt/test_tabulate.py @@ -4,6 +4,9 @@ import numpy as np import torch +from deepmd.dpmodel.utils.network import ( + get_activation_fn, +) from deepmd.pt.utils import ( env, ) @@ -18,8 +21,27 @@ tf, ) +ACTIVATION_NAMES = { + 1: "tanh", + 2: "gelu", + 3: "relu", + 4: "relu6", + 5: "softplus", + 6: "sigmoid", + 7: "silu", +} + + +def get_activation_function(functype: int): + """Get activation function corresponding to functype.""" + if functype not in ACTIVATION_NAMES: + raise ValueError(f"Unknown functype: {functype}") + + return get_activation_fn(ACTIVATION_NAMES[functype]) + def setUpModule() -> None: + tf.reset_default_graph() tf.compat.v1.enable_eager_execution() @@ -43,92 +65,129 @@ def setUp(self) -> None: self.xbar = np.matmul(self.x, self.w) + self.b # 4 x 4 - self.y = np.tanh(self.xbar) - def test_ops(self) -> None: + """Test all activation functions using parameterized subtests.""" + for functype in ACTIVATION_NAMES.keys(): + activation_name = ACTIVATION_NAMES[functype] + activation_fn = get_activation_function(functype) + + with self.subTest(activation=activation_name, functype=functype): + self._test_single_activation(functype, activation_fn, activation_name) + + def _test_single_activation( + self, functype: int, activation_fn, activation_name: str + ) -> None: + """Test tabulation operations for a specific activation function.""" + # Compute y using the specific activation function + y = activation_fn(self.xbar) + + # Test unaggregated_dy_dx_s dy_tf = op_module.unaggregated_dy_dx_s( - tf.constant(self.y, dtype="double"), + tf.constant(y, dtype="double"), tf.constant(self.w, dtype="double"), tf.constant(self.xbar, dtype="double"), - tf.constant(1), + tf.constant(functype), ) dy_pt = unaggregated_dy_dx_s( - torch.from_numpy(self.y), + torch.from_numpy(y), self.w, torch.from_numpy(self.xbar), - 1, + functype, ) dy_tf_numpy = dy_tf.numpy() dy_pt_numpy = dy_pt.detach().cpu().numpy() - np.testing.assert_almost_equal(dy_tf_numpy, dy_pt_numpy, decimal=10) + np.testing.assert_almost_equal( + dy_tf_numpy, + dy_pt_numpy, + decimal=10, + err_msg=f"unaggregated_dy_dx_s failed for {activation_name}", + ) + # Test unaggregated_dy2_dx_s dy2_tf = op_module.unaggregated_dy2_dx_s( - tf.constant(self.y, dtype="double"), + tf.constant(y, dtype="double"), dy_tf, tf.constant(self.w, dtype="double"), tf.constant(self.xbar, dtype="double"), - tf.constant(1), + tf.constant(functype), ) dy2_pt = unaggregated_dy2_dx_s( - torch.from_numpy(self.y), + torch.from_numpy(y), dy_pt, self.w, torch.from_numpy(self.xbar), - 1, + functype, ) dy2_tf_numpy = dy2_tf.numpy() dy2_pt_numpy = dy2_pt.detach().cpu().numpy() - np.testing.assert_almost_equal(dy2_tf_numpy, dy2_pt_numpy, decimal=10) + np.testing.assert_almost_equal( + dy2_tf_numpy, + dy2_pt_numpy, + decimal=10, + err_msg=f"unaggregated_dy2_dx_s failed for {activation_name}", + ) + # Test unaggregated_dy_dx dz_tf = op_module.unaggregated_dy_dx( - tf.constant(self.y, dtype="double"), + tf.constant(y, dtype="double"), tf.constant(self.w, dtype="double"), dy_tf, tf.constant(self.xbar, dtype="double"), - tf.constant(1), + tf.constant(functype), ) dz_pt = unaggregated_dy_dx( - torch.from_numpy(self.y).to(env.DEVICE), + torch.from_numpy(y).to(env.DEVICE), self.w, dy_pt, torch.from_numpy(self.xbar).to(env.DEVICE), - 1, + functype, ) dz_tf_numpy = dz_tf.numpy() dz_pt_numpy = dz_pt.detach().cpu().numpy() - np.testing.assert_almost_equal(dz_tf_numpy, dz_pt_numpy, decimal=10) + np.testing.assert_almost_equal( + dz_tf_numpy, + dz_pt_numpy, + decimal=10, + err_msg=f"unaggregated_dy_dx failed for {activation_name}", + ) + # Test unaggregated_dy2_dx dy2_tf = op_module.unaggregated_dy2_dx( - tf.constant(self.y, dtype="double"), + tf.constant(y, dtype="double"), tf.constant(self.w, dtype="double"), dy_tf, dy2_tf, tf.constant(self.xbar, dtype="double"), - tf.constant(1), + tf.constant(functype), ) dy2_pt = unaggregated_dy2_dx( - torch.from_numpy(self.y).to(env.DEVICE), + torch.from_numpy(y).to(env.DEVICE), self.w, dy_pt, dy2_pt, torch.from_numpy(self.xbar).to(env.DEVICE), - 1, + functype, ) dy2_tf_numpy = dy2_tf.numpy() dy2_pt_numpy = dy2_pt.detach().cpu().numpy() - np.testing.assert_almost_equal(dy2_tf_numpy, dy2_pt_numpy, decimal=10) + np.testing.assert_almost_equal( + dy2_tf_numpy, + dy2_pt_numpy, + decimal=10, + err_msg=f"unaggregated_dy2_dx failed for {activation_name}", + ) if __name__ == "__main__": diff --git a/source/tests/pt/test_tabulate_fusion_se_t_tebd.py b/source/tests/pt/test_tabulate_fusion_se_t_tebd.py new file mode 100644 index 0000000000..10bf48c46f --- /dev/null +++ b/source/tests/pt/test_tabulate_fusion_se_t_tebd.py @@ -0,0 +1,1918 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest + +import torch + +from deepmd.pt.cxx_op import ( + ENABLE_CUSTOMIZED_OP, +) +from deepmd.pt.utils import ( + env, +) + +from ..consistent.common import ( + parameterized, +) + + +@parameterized((torch.float64, torch.float32)) +@unittest.skipIf(not ENABLE_CUSTOMIZED_OP, "PyTorch customized OPs are not built") +class TestTabulateFusionSeTTebdOp(unittest.TestCase): + def setUp(self) -> None: + (dtype,) = self.param + if dtype == torch.float64: + self.prec = 1e-10 + elif dtype == torch.float32: + # JZ: not sure the reason, but 1e-5 cannot pass the grad test + self.prec = 1e-3 + self.table_tensor = torch.tensor( + [ + -1.0600000163027882e02, + 7.7059358807135015e02, + -5.6954714749735385e03, + 1.2167808756610991e03, + -7.6199102434332218e01, + 1.0706136029373441e00, + -1.0600000164528124e02, + 7.7059358630452323e02, + -5.6954715659539552e03, + 1.2167808757436076e03, + -7.6199099707724926e01, + 1.0706134206080884e00, + -1.0600000163027882e02, + 7.7059358807135015e02, + -5.6954714749735385e03, + 1.2167808756610991e03, + -7.6199102434332218e01, + 1.0706136029373441e00, + -1.0600000164528124e02, + 7.7059358630452323e02, + -5.6954715659539552e03, + 1.2167808757436076e03, + -7.6199099707724926e01, + 1.0706134206080884e00, + -9.6000006759336443e01, + 6.2969719646863621e02, + -4.2053706363664551e03, + 9.0372155784831205e02, + -5.7600014239472898e01, + 8.6528676197113796e-01, + -9.6000006828502180e01, + 6.2969718981238339e02, + -4.2053709121998018e03, + 9.0372156236848912e02, + -5.7600006817493266e01, + 8.6528625106787871e-01, + -9.6000006759336443e01, + 6.2969719646863621e02, + -4.2053706363664551e03, + 9.0372155784831205e02, + -5.7600014239472898e01, + 8.6528676197113796e-01, + -9.6000006828502180e01, + 6.2969718981238339e02, + -4.2053709121998018e03, + 9.0372156236848912e02, + -5.7600006817493266e01, + 8.6528625106787871e-01, + -8.6000028021606425e01, + 5.0303296429845562e02, + -3.0008648248894533e03, + 6.4939597734382562e02, + -4.2250984019314707e01, + 6.8180015607155764e-01, + -8.6000028340480625e01, + 5.0303293978396903e02, + -3.0008656209622986e03, + 6.4939600529391078e02, + -4.2250965541906716e01, + 6.8179882734268982e-01, + -8.6000028021606425e01, + 5.0303296429845562e02, + -3.0008648248894533e03, + 6.4939597734382562e02, + -4.2250984019314707e01, + 6.8180015607155764e-01, + -8.6000028340480625e01, + 5.0303293978396903e02, + -3.0008656209622986e03, + 6.4939600529353049e02, + -4.2250965541830588e01, + 6.8179882733888086e-01, + -7.6000116148038558e01, + 3.9060139597613619e02, + -2.0515743554479322e03, + 4.4772754091167945e02, + -2.9848087537832814e01, + 5.2014755686537917e-01, + -7.6000117618125429e01, + 3.9060130821883052e02, + -2.0515765138621105e03, + 4.4772766653712006e02, + -2.9848047259266409e01, + 5.2014443989116910e-01, + -7.6000116148038558e01, + 3.9060139597613619e02, + -2.0515743554479322e03, + 4.4772754091167945e02, + -2.9848087537832814e01, + 5.2014755686537917e-01, + -7.6000117618125742e01, + 3.9060130821877993e02, + -2.0515765138659344e03, + 4.4772766652483722e02, + -2.9848047256692499e01, + 5.2014443976043645e-01, + -6.6000481290731443e01, + 2.9240425245900917e02, + -1.3271250821434478e03, + 2.9263955624337893e02, + -2.0087224005740719e01, + 3.8031147992206349e-01, + -6.6000488067863742e01, + 2.9240394960550276e02, + -1.3271304743966571e03, + 2.9264002765325057e02, + -2.0087154325946980e01, + 3.8030522013794582e-01, + -6.6000481290731443e01, + 2.9240425245900917e02, + -1.3271250821434478e03, + 2.9263955624337893e02, + -2.0087224005740719e01, + 3.8031147992206349e-01, + -6.6000488067883694e01, + 2.9240394960308691e02, + -1.3271304745319526e03, + 2.9264002727267626e02, + -2.0087154245656002e01, + 3.8030521605011575e-01, + -5.6001992867343972e01, + 2.0844745574402617e02, + -7.9715799906587699e02, + 1.7805563184427194e02, + -1.2663929104029080e01, + 2.6224978307822894e-01, + -5.6002024103130161e01, + 2.0844646075692629e02, + -7.9717003898786652e02, + 1.7805715054974732e02, + -1.2663864677938077e01, + 2.6224029170957303e-01, + -5.6001992867343972e01, + 2.0844745574402617e02, + -7.9715799906587699e02, + 1.7805563184427194e02, + -1.2663929104029080e01, + 2.6224978307822894e-01, + -5.6002024104383771e01, + 2.0844646064871867e02, + -7.9717004324410516e02, + 1.7805714044473001e02, + -1.2663862524337585e01, + 2.6224018166598279e-01, + -4.6008230210744550e01, + 1.3874976550319553e02, + -4.3134867537287749e02, + 9.7902623595157010e01, + -7.2734403121911884e00, + 1.6589123996688057e-01, + -4.6008373996710617e01, + 1.3874671965012058e02, + -4.3137141216256458e02, + 9.7906861443792735e01, + -7.2735856084076280e00, + 1.6588642735924275e-01, + -4.6008230210744550e01, + 1.3874976550319553e02, + -4.3134867537287749e02, + 9.7902623595157010e01, + -7.2734403121911884e00, + 1.6589123996688057e-01, + -4.6008374075307870e01, + 1.3874671513440606e02, + -4.3137152784492957e02, + 9.7906652364871050e01, + -7.2735401377994249e00, + 1.6588408717348646e-01, + -3.6033642533368131e01, + 8.3364086172019398e01, + -1.9942175516407502e02, + 4.6124022747838069e01, + -3.6130563858549958e00, + 9.1249773312287188e-02, + -3.6034298111245583e01, + 8.3355843868269616e01, + -1.9945266030093268e02, + 4.6135000705962462e01, + -3.6142786797647353e00, + 9.1293932043118198e-02, + -3.6033642533368131e01, + 8.3364086172019398e01, + -1.9942175516407502e02, + 4.6124022747838069e01, + -3.6130563858549958e00, + 9.1249773312287188e-02, + -3.6034302998781108e01, + 8.3355675173745269e01, + -1.9945516784358935e02, + 4.6132303200740992e01, + -3.6136582565667807e00, + 9.1261386291659793e-02, + -2.6132076703837274e01, + 4.2398929436319683e01, + -7.1037171119057973e01, + 1.3425662262407457e01, + -7.5172495708992593e-01, + 7.7522572203268742e-03, + -2.6134776894873077e01, + 4.2384732735328775e01, + -7.1030526549717337e01, + 1.3431455085299461e01, + -7.5302028721199155e-01, + 7.8186246126207160e-03, + -2.6132076703837274e01, + 4.2398929436319683e01, + -7.1037171119057973e01, + 1.3425662262405055e01, + -7.5172495708944420e-01, + 7.7522572203027138e-03, + -2.6135071381093578e01, + 4.2379566840123424e01, + -7.1067162844830236e01, + 1.3434603316099608e01, + -7.5251233833488806e-01, + 7.7734884077347950e-03, + -2.2221480705551805e01, + 3.0067218434037404e01, + -4.1779705297521097e01, + -1.9077757705724110e02, + 3.6413466026808294e02, + -1.6067397401486718e02, + -2.2225430071703467e01, + 3.0060809113889512e01, + -4.1712800191721314e01, + -1.9084786311022177e02, + 3.6410062714257685e02, + -1.6063028238785057e02, + -2.2221480705551830e01, + 3.0067218434036263e01, + -4.1779705297545611e01, + -1.9077757705723738e02, + 3.6413466026815809e02, + -1.6067397401492047e02, + -2.2226913938674084e01, + 3.0042371820589185e01, + -4.1801582285426832e01, + -1.9048619249019526e02, + 3.6373874557858261e02, + -1.6052358406417352e02, + -2.1250858373060836e01, + 2.7343847665267702e01, + -3.6044215009418814e01, + -1.7618484800469861e02, + 3.3120085405644409e02, + -1.4534825256321494e02, + -2.1254939505030809e01, + 2.7342716030835884e01, + -3.5955450545431681e01, + -1.7635550119316844e02, + 3.3127447930769307e02, + -1.4533876561022046e02, + -2.1250858373060954e01, + 2.7343847665262818e01, + -3.6044215009514119e01, + -1.7618484800464822e02, + 3.3120085405666612e02, + -1.4534825256338749e02, + -2.1257155379297881e01, + 2.7317691772612619e01, + -3.6063526926252166e01, + -1.7588696592837897e02, + 3.3079005662384850e02, + -1.4519086534447842e02, + -2.0283472228681301e01, + 2.4763027042036295e01, + -3.0876160316998963e01, + -1.6184864900381874e02, + 2.9976970905591691e02, + -1.3084395423768876e02, + -2.0287461515322455e01, + 2.4769400540137131e01, + -3.0762734380983186e01, + -1.6214886052089241e02, + 2.9998995088792128e02, + -1.3088331758129965e02, + -2.0283472228681809e01, + 2.4763027042017129e01, + -3.0876160317336627e01, + -1.6184864900359682e02, + 2.9976970905662938e02, + -1.3084395423826805e02, + -2.0290765181946348e01, + 2.4735639907973120e01, + -3.0892738413082597e01, + -1.6154574482310053e02, + 2.9934595420013272e02, + -1.3068028494926122e02, + -1.9319499689234629e01, + 2.2323824431805683e01, + -2.6243395369841849e01, + -1.4782286378121026e02, + 2.6985759662396487e02, + -1.1715474197881395e02, + -1.9323022570439292e01, + 2.2340565860680357e01, + -2.6102786429129356e01, + -1.4828764857305418e02, + 2.7027298759214750e02, + -1.1726163007473576e02, + -1.9319499689236839e01, + 2.2323824431730525e01, + -2.6243395371031539e01, + -1.4782286378021576e02, + 2.6985759662609979e02, + -1.1715474198068593e02, + -1.9327939259284843e01, + 2.2295320666731183e01, + -2.6257097174199931e01, + -1.4751677383623073e02, + 2.6942341041084092e02, + -1.1698575776762208e02, + -1.8359079763330211e01, + 2.0025118950280675e01, + -2.2113826757823226e01, + -1.3415932552431914e02, + 2.4147795894487624e02, + -1.0427314537549884e02, + -1.8361534194530734e01, + 2.0055847278170305e01, + -2.1944107342764479e01, + -1.3482982214648752e02, + 2.4214772485703989e02, + -1.0447085300268679e02, + -1.8359079763339750e01, + 2.0025118949989704e01, + -2.2113826761939308e01, + -1.3415932552009582e02, + 2.4147795895089951e02, + -1.0427314538136979e02, + -1.8368836959765495e01, + 1.9995657614892380e01, + -2.2124533894067383e01, + -1.3385233293246981e02, + 2.4103659293914149e02, + -1.0410011400771683e02, + -1.7402299525814517e01, + 1.7865597763687486e01, + -1.8455503416511757e01, + -1.2090765118569301e02, + 2.1464125749038132e02, + -9.2190581022134992e01, + -1.7402744551259310e01, + 1.7914800567904472e01, + -1.8255754666855470e01, + -1.2183089355280822e02, + 2.1563582256173194e02, + -9.2507405324257306e01, + -1.7402299525855486e01, + 1.7865597762572605e01, + -1.8455503430527756e01, + -1.2090765116826699e02, + 2.1464125750558804e02, + -9.2190581039770791e01, + -1.7413567239985614e01, + 1.7835392747330133e01, + -1.8463115133795956e01, + -1.2060260469703572e02, + 2.1419685510959093e02, + -9.2015134441585104e01, + -1.6449179896085464e01, + 1.5843762224435309e01, + -1.5236722252652665e01, + -1.0811515163854509e02, + 1.8935506712501905e02, + -8.0897437157402223e01, + -1.6446174965543889e01, + 1.5916874201410112e01, + -1.5007553197461570e01, + -1.0934291295595986e02, + 1.9075532567542470e02, + -8.1366596347119696e01, + -1.6449179896260411e01, + 1.5843762220214204e01, + -1.5236722299508587e01, + -1.0811515156878269e02, + 1.8935506715588940e02, + -8.0897437207525684e01, + -1.6462173655481337e01, + 1.5813096619069219e01, + -1.5241142983208677e01, + -1.0781563484017332e02, + 1.8891289499393798e02, + -8.0721658713418606e01, + -1.5499661595231082e01, + 1.3957945516559789e01, + -1.2426145992195885e01, + -9.5826844741964834e01, + 1.6562434781973772e02, + -7.0383233416004117e01, + -1.5491037589250178e01, + 1.4061349904707843e01, + -1.2170301483989650e01, + -9.7412966929875139e01, + 1.6751874597575440e02, + -7.1041920384880939e01, + -1.5499661595973759e01, + 1.3957945500778198e01, + -1.2426146145776961e01, + -9.5826844470313858e01, + 1.6562434784656404e02, + -7.0383233547510557e01, + -1.5514618579274794e01, + 1.3927192540790591e01, + -1.2427264674287118e01, + -9.5537423121432880e01, + 1.6519113036542510e02, + -7.0209783384625098e01, + -1.4553592409098401e01, + 1.2206343505203831e01, + -9.9929274597052196e00, + -8.4085595900823435e01, + 1.4345191724964303e02, + -6.0636862050381758e01, + -1.4536130507533649e01, + 1.2347228125716077e01, + -9.7159302678980044e00, + -8.6081002959763751e01, + 1.4592996741513730e02, + -6.1523840242331410e01, + -1.4553592412232879e01, + 1.2206343446986155e01, + -9.9929279524397305e00, + -8.4085594870780753e01, + 1.4345191706222485e02, + -6.0636862352071532e01, + -1.4570766853404239e01, + 1.2175998366492486e01, + -9.9905856922863112e00, + -8.3812185051328299e01, + 1.4303633648493073e02, + -6.0469165577726159e01, + -1.3610717065161962e01, + 1.0587059629986399e01, + -7.9068321681349163e00, + -7.2932404423885004e01, + 1.2283913327111270e02, + -5.1646910322317169e01, + -1.3579708436673444e01, + 1.0773027159520954e01, + -7.6175370796795425e00, + -7.5376833196183071e01, + 1.2597958225245242e02, + -5.2797863799745748e01, + -1.3610717078313911e01, + 1.0587059418306087e01, + -7.9068337121483454e00, + -7.2932400620636059e01, + 1.2283913169238102e02, + -5.1646910832841897e01, + -1.3630368323321786e01, + 1.0557789879027116e01, + -7.9007777139483810e00, + -7.2682825476758552e01, + 1.2245259140017740e02, + -5.1489446559796768e01, + -1.2670671078399982e01, + 9.0981634949263963e00, + -6.1383490362855788e00, + -6.2406844162279825e01, + 1.0378677653422224e02, + -4.3402055519687693e01, + -1.2619333100308433e01, + 9.3364634226935799e00, + -5.8491811509717584e00, + -6.5316414528433455e01, + 1.0763857666200300e02, + -4.4841832720191050e01, + -1.2670671133253135e01, + 9.0981627374157021e00, + -6.1383537481895356e00, + -6.2406830503476570e01, + 1.0378676818216074e02, + -4.3402055529436716e01, + -1.2693036794620980e01, + 9.0708908225804148e00, + -6.1281713411274001e00, + -6.2191660620037396e01, + 1.0344456594081470e02, + -4.3260806640248063e01, + -1.1732979767504439e01, + 7.7377614739662697e00, + -4.6587775146685351e00, + -5.2547655563671029e01, + 8.6296103981829802e01, + -3.5891515805495345e01, + -1.1651721415208119e01, + 8.0340005825064456e00, + -4.3852919661646119e00, + -5.5898160750405737e01, + 9.0851291378134590e01, + -3.7622755083739385e01, + -1.1732979994779518e01, + 7.7377588120662892e00, + -4.6587914600219875e00, + -5.2547607987974565e01, + 8.6296066930227624e01, + -3.5891510429190419e01, + -1.1758218632638741e01, + 7.7137968422318544e00, + -4.6438239588320966e00, + -5.2381405657406454e01, + 8.6019170302439520e01, + -3.5774653697918737e01, + -1.0797063195543267e01, + 6.5040766534586290e00, + -3.4402783696562169e00, + -4.3393478931462226e01, + 7.0370032342568010e01, + -2.9105535302381853e01, + -1.0672637254876815e01, + 6.8603244928014488e00, + -3.1995767859681346e00, + -4.7101348454718874e01, + 7.5530774605740319e01, + -3.1094453979913311e01, + -1.0797064129672576e01, + 6.5040675030570139e00, + -3.4403181344841500e00, + -4.3393319126804485e01, + 7.0369884883020177e01, + -2.9105501594155889e01, + -1.0825134802124644e01, + 6.4853446725127366e00, + -3.4195560956016346e00, + -4.3296381389022351e01, + 7.0187483762520671e01, + -2.9024415860031247e01, + -9.8622468030169337e00, + 5.3955359781222549e00, + -2.4558741324534137e00, + -3.4983728078555984e01, + 5.6014425934291204e01, + -2.3035887876475471e01, + -9.6769173769353625e00, + 5.8079540801032961e00, + -2.2635143148159220e00, + -3.8890523502249145e01, + 6.1563046720547966e01, + -2.5198820521877391e01, + -9.8622505990399034e00, + 5.3955054149765509e00, + -2.4559821583353774e00, + -3.4983216045684472e01, + 5.6013889382190079e01, + -2.3035736114340502e01, + -9.8926597117464805e00, + 5.3849440641688187e00, + -2.4279562878572039e00, + -3.4983707025980287e01, + 5.5966629574570753e01, + -2.3006306589550750e01, + -8.9277749780883457e00, + 4.4108678323349286e00, + -1.6793815271288624e00, + -2.7359655656676122e01, + 4.3239544183593061e01, + -1.7676416286664047e01, + -8.6587749152265552e00, + 4.8674392165289442e00, + -1.5450097170494306e00, + -3.1230915545542118e01, + 4.8829474992442343e01, + -1.9874755288141955e01, + -8.9277901202336185e00, + 4.4107699183102085e00, + -1.6796551456533098e00, + -2.7358123514289456e01, + 4.3237769027728554e01, + -1.7675844947587926e01, + -8.9590559763951383e00, + 4.4128957610428623e00, + -1.6423658138809611e00, + -2.7493743583145054e01, + 4.3380518846300511e01, + -1.7719639183506050e01, + -7.9928164326293913e00, + 3.5492331091008302e00, + -1.0852462622393610e00, + -2.0565792757352423e01, + 3.2061909496398073e01, + -1.3023704651715642e01, + -7.6125412569887647e00, + 4.0287966748633526e00, + -1.0084592804412351e00, + -2.4116992333062022e01, + 3.7252797603904497e01, + -1.5077495076198684e01, + -7.9928747817255603e00, + 3.5489404571097585e00, + -1.0858609980296849e00, + -2.0561701094768868e01, + 3.2056747083970720e01, + -1.3021877019728107e01, + -8.0213899495838241e00, + 3.5708128515175943e00, + -1.0368753205735253e00, + -2.0877831538201836e01, + 3.2456559535389509e01, + -1.3165540198118645e01, + -7.0564174984379102e00, + 2.8104770395789380e00, + -6.4821407306458223e-01, + -1.4652118176169953e01, + 2.2507145963021038e01, + -9.0780963613608154e00, + -6.5338936679228468e00, + 3.2846161494194233e00, + -6.1760141818709846e-01, + -1.7606122820367215e01, + 2.6855555289500277e01, + -1.0803821410528570e01, + -7.0566263531717324e00, + 2.8097184139861691e00, + -6.4925197579297411e-01, + -1.4643483271177150e01, + 2.2495243692983838e01, + -9.0734373052814821e00, + -7.0742646195707266e00, + 2.8621047467298468e00, + -5.8641470402843421e-01, + -1.5178915176777426e01, + 2.3211717123277591e01, + -9.3414295847965061e00, + -6.1172231064332783e00, + 2.1957964102200167e00, + -3.4265643705632465e-01, + -9.6769153352706798e00, + 1.4613873405033004e01, + -5.8450824172251430e00, + -5.4212678780860326e00, + 2.6341589573018260e00, + -3.4085224757280796e-01, + -1.1835854891340576e01, + 1.7794701474942944e01, + -7.1075278532253687e00, + -6.1178367984533244e00, + 2.1945528943967396e00, + -3.4261268423617658e-01, + -9.6695829134679272e00, + 1.4600877298870854e01, + -5.8381668136523013e00, + -6.1072022151656586e00, + 2.2922503774685161e00, + -2.6715334266026142e-01, + -1.0408120531614587e01, + 1.5617405440391840e01, + -6.2270636615178061e00, + -5.1722074807324017e00, + 1.7098190643016411e00, + -1.4098618492175408e-01, + -5.7061337346696464e00, + 8.4331806866534098e00, + -3.3349192888568142e00, + -4.2766424379800121e00, + 2.0860564217794284e00, + -1.5548660419053545e-01, + -7.0034949575065015e00, + 1.0332245608764421e01, + -4.0873492185766374e00, + -5.1727690165421372e00, + 1.7132539127425084e00, + -1.2776576793785877e-01, + -5.7565343018918274e00, + 8.4941254548170697e00, + -3.3479852132230872e00, + -5.0998839330979591e00, + 1.8678855512825561e00, + -5.7718910331047868e-02, + -6.5095346397755423e00, + 9.5462002113817768e00, + -3.7632628689263172e00, + -4.2112469382255613e00, + 1.3675717927787789e00, + -9.4961575783498800e-03, + -2.7877417589321136e00, + 3.9953503912711956e00, + -1.5499906707437840e00, + -3.1046711877098376e00, + 1.6568346830533449e00, + -4.5990009889900242e-02, + -3.3140676307068091e00, + 4.7472200808709299e00, + -1.8492173878772247e00, + -4.1976749320353317e00, + 1.4246952243441517e00, + 8.7531923058200650e-02, + -3.0996975434049761e00, + 4.4668738099197531e00, + -1.7103055321708385e00, + -4.0163145894665320e00, + 1.5923303121893606e00, + 5.8249749369824022e-02, + -3.3748048713195491e00, + 4.7925769874900315e00, + -1.8598420111853879e00, + -3.1955533414298376e00, + 1.2168024121915868e00, + 9.9474205814620603e-02, + -8.6811124876189694e-01, + 1.1994338853723501e00, + -4.4837238870567747e-01, + -1.9098914522594992e00, + 1.3654451552507061e00, + 2.9537044429980407e-03, + -9.3701125207094127e-01, + 1.2575365835116745e00, + -4.7248060681970733e-01, + -3.0285770502890443e00, + 1.6166340190704305e00, + 4.8662683065338386e-01, + -1.2308607057515726e00, + 1.6114560066217587e00, + -6.5896729332189652e-01, + -2.8078044229222514e00, + 1.4555130910035559e00, + 9.0876948497501955e-02, + -1.0566809618626720e00, + 1.3938154223720176e00, + -5.2279617091852160e-01, + -1.9963264755188566e00, + 1.3672906754961440e00, + 2.0801988470625002e-01, + 2.0083818728351077e-02, + -1.5135587406137185e-02, + -1.4175240342178652e-02, + -6.9344786794476854e-01, + 1.2280621078720415e00, + 1.2333381103148277e-02, + -1.0895386066093759e-02, + 2.1764282171790141e-02, + -1.0106900291744604e-02, + -1.2036881930169383e00, + 2.0482931230000392e00, + -1.2689218008973949e-01, + -5.0580690719339239e-01, + 3.4047786101030464e-01, + -7.0959386937004015e-02, + -1.4470760938303664e00, + 1.4285049373060201e00, + 5.5764887956399375e-02, + -2.9461990750009881e-02, + 2.3005167601875431e-02, + -1.0760396189439407e-02, + -4.3024292433642597e-01, + 1.7121633497582587e00, + 3.5705413032693957e-02, + -9.9216800479772127e-01, + 1.5115432403429119e00, + -6.3985596276149748e-01, + 5.4770961684437192e-01, + 1.2565653391084903e00, + 9.1639130181564755e-03, + -6.8547618650262643e-01, + 1.2037212931265591e00, + -5.1526772142324506e-01, + 4.8142431677326969e-01, + 1.2842025505965851e00, + -3.1103960497811806e-01, + -3.8667287940463613e-01, + 9.2663039525338942e-01, + -4.1330437951972537e-01, + 1.9976512094478704e-02, + 1.4898674304290889e00, + -2.1940405767858565e-03, + -8.0791207141984167e-01, + 1.3979310081478775e00, + -5.9845265079421794e-01, + 1.1971451112382212e00, + 1.6539633089946477e00, + -2.7009878691796618e-01, + -2.8868139196850624e00, + 4.7294193613612734e00, + -1.9578020397520424e00, + 1.8164162541717044e00, + 1.4570111710269262e00, + 2.2385898037164991e-02, + -3.1195681762439769e00, + 4.9723722392038878e00, + -2.0423972644796100e00, + 1.5812403987207633e00, + 1.1421043858413655e00, + -4.4319666868952730e-02, + -2.3144705949527720e00, + 3.7448930479898297e00, + -1.5426803544433196e00, + 1.4992161878806018e00, + 1.6612039136364238e00, + -2.2870713891204597e-02, + -3.4442115437939465e00, + 5.5057190995408973e00, + -2.2657208348376137e00, + 2.4658130352390710e00, + 1.5819912227884063e00, + -1.3204477532594588e-01, + -5.7752803465671017e00, + 9.0677018990478242e00, + -3.6843468204828174e00, + 3.1062201217160963e00, + 1.8205810727868250e00, + 7.3942159732456811e-02, + -7.3418038323250947e00, + 1.1309154676354810e01, + -4.5733470083866452e00, + 2.5667672162869133e00, + 1.3762236869878626e00, + 5.4823291778512563e-02, + -5.5558964069977943e00, + 8.5620133672289516e00, + -3.4575259608624478e00, + 2.9333361085351610e00, + 1.9771000784477066e00, + 2.1600903596218385e-02, + -7.7786452012965430e00, + 1.2026327126407146e01, + -4.8722408979121159e00, + 3.5238342146994350e00, + 1.8411341262124141e00, + 1.0485737443151430e-01, + -1.0316470080846322e01, + 1.5628354265192609e01, + -6.2547428286449396e00, + 4.3947471898784478e00, + 2.3129375587624681e00, + 1.6998863701958250e-01, + -1.3069120913924280e01, + 1.9764673064124775e01, + -7.9234176878170990e00, + 3.5464051944219954e00, + 1.7786047141550632e00, + 1.8395466553434961e-01, + -1.0256713338978345e01, + 1.5450540198835597e01, + -6.1709943751208902e00, + 4.3074781177775723e00, + 2.4284702978185178e00, + 1.2121907902830774e-01, + -1.3510697720561426e01, + 2.0490823414440431e01, + -8.2265504110307699e00, + 4.5269670710447079e00, + 2.3411415500822019e00, + 3.7814443659878427e-01, + -1.6533454371385766e01, + 2.4532574055181296e01, + -9.7222898630871342e00, + 5.6498078480438974e00, + 2.8871559084424092e00, + 3.1648740182441881e-01, + -1.9832336139347099e01, + 2.9630584562783888e01, + -1.1804975183138390e01, + 4.5317970588477650e00, + 2.3235629480266455e00, + 4.0711209040396701e-01, + -1.6523611973754900e01, + 2.4482080409856291e01, + -9.6968326211377835e00, + 5.6107427774726322e00, + 2.9693568967987254e00, + 2.6856229367890733e-01, + -2.0186235796983127e01, + 3.0228033555488111e01, + -1.2057362656117963e01, + 5.5230828784340904e00, + 3.0159142144119913e00, + 7.5032702265793638e-01, + -2.4452361306480910e01, + 3.5745746299744695e01, + -1.4059387633540990e01, + 6.8467243986091164e00, + 3.5205846294935204e00, + 5.5323452910250115e-01, + -2.7424447720726722e01, + 4.0542113968978946e01, + -1.6058340606199877e01, + 5.5241079122419858e00, + 3.0111097413061287e00, + 7.6043241689918206e-01, + -2.4453330947201032e01, + 3.5733842835424838e01, + -1.4052622761934279e01, + 6.8330970703372866e00, + 3.5730950345697865e00, + 5.0442967447855436e-01, + -2.7630302835415993e01, + 4.0921397061842079e01, + -1.6223699529825666e01, + 6.5233214752268127e00, + 3.8455313715589599e00, + 1.2738445662734672e00, + -3.4142511056048967e01, + 4.9288751118195229e01, + -1.9258816488331760e01, + 7.9798691992574877e00, + 4.2304633704347614e00, + 9.4916911879724064e-01, + -3.6082800915305256e01, + 5.2740474636382487e01, + -2.0757970588732530e01, + 6.5235391967368317e00, + 3.8442392655293900e00, + 1.2772689685023881e00, + -3.4144245582802192e01, + 4.9286600694030149e01, + -1.9257235266278844e01, + 7.9780164759860508e00, + 4.2581364755189171e00, + 9.0490824102641643e-01, + -3.6146890048111374e01, + 5.2902251888236343e01, + -2.0834714063750525e01, + 7.5301209868737518e00, + 4.8266093670811516e00, + 1.9906532239804082e00, + -4.5696171225139402e01, + 6.5222794336738914e01, + -2.5330008845677121e01, + 9.0592048208341964e00, + 5.0524444639807982e00, + 1.5639083038511417e00, + -4.6227354827270197e01, + 6.6742768625790532e01, + -2.6090733281390481e01, + 7.5301672757177256e00, + 4.8262668988539703e00, + 1.9917837214882572e00, + -4.5697152262800707e01, + 6.5222641787790508e01, + -2.5329699752317662e01, + 9.0617089689058279e00, + 5.0627200474303731e00, + 1.5306087886050987e00, + -4.6201245261995687e01, + 6.6753711704174307e01, + -2.6103836713323240e01, + 8.5439978438576958e00, + 5.9605352581937785e00, + 2.9388171122244109e00, + -5.9213652478598007e01, + 8.3623964589400401e01, + -3.2288651007290504e01, + 1.0100238105795977e01, + 6.0156046860821641e00, + 2.4311227628788585e00, + -5.8189717323516248e01, + 8.2972590004142106e01, + -3.2212869674305303e01, + 8.5440076687321067e00, + 5.9604459430021439e00, + 2.9391801366526531e00, + -5.9214078468041464e01, + 8.3624068891376510e01, + -3.2288610777657510e01, + 1.0103667533796683e01, + 6.0158650887345448e00, + 2.4107760944314816e00, + -5.8125625048064265e01, + 8.2906979417176174e01, + -3.2191629006406409e01, + 9.5650113177877785e00, + 7.2498153679976820e00, + 4.1551371399277919e00, + -7.4795843598083408e01, + 1.0457037732454131e02, + -4.0151433068943419e01, + 1.1116968561077568e01, + 7.1347098863330896e00, + 3.5688140741297674e00, + -7.2151486218593305e01, + 1.0165680693075836e02, + -3.9206269356622016e01, + 9.5650133940644455e00, + 7.2497924894015711e00, + 4.1552503042122613e00, + -7.4796005009548836e01, + 1.0457044971811401e02, + -4.0151435976986221e01, + 1.1120034079668221e01, + 7.1303147700774092e00, + 3.5594873892317103e00, + -7.2082067018068685e01, + 1.0156598726189708e02, + -3.9171834664292227e01, + 1.0593064483227742e01, + 8.6969028070512202e00, + 5.6755396034912966e00, + -9.2539537763180832e01, + 1.2813560149579646e02, + -4.8933613418447223e01, + 1.2119543877083460e01, + 8.4137603187360543e00, + 4.9925034366798311e00, + -8.8194505075704640e01, + 1.2287993196505218e02, + -4.7096724506223822e01, + 1.0593064919257221e01, + 8.6968970567044934e00, + 5.6755738143875760e00, + -9.2539593640863643e01, + 1.2813563331215474e02, + -4.8933618162805772e01, + 1.2121921818513506e01, + 8.4078642204619420e00, + 4.9908632634858190e00, + -8.8134432374832016e01, + 1.2279086550380391e02, + -4.7060844505587738e01, + 1.1627957207938659e01, + 1.0303707615441018e01, + 7.5344011042552923e00, + -1.1253294830348190e02, + 1.5438372244089408e02, + -5.8647453529357783e01, + 1.3114510015623049e01, + 9.8513572940713416e00, + 6.7213349376406626e00, + -1.0635738219113546e02, + 1.4665751311861146e02, + -5.5881528760137869e01, + 1.1627957298834614e01, + 1.0303706197478814e01, + 7.5344111366673712e00, + -1.1253296638384563e02, + 1.5438373415898508e02, + -5.8647455853629580e01, + 1.3116237925845430e01, + 9.8455331102145145e00, + 6.7243141059359051e00, + -1.0631074264006560e02, + 1.4658112805680690e02, + -5.5849452095162235e01, + 1.2669386535689361e01, + 1.2071287030293307e01, + 9.7633555455962835e00, + -1.3485075345900265e02, + 1.8336444946299886e02, + -6.9300787627414508e01, + 1.4105804414673191e01, + 1.1444289269702800e01, + 8.7789794745243590e00, + -1.2666835962860844e02, + 1.7298274034188972e02, + -6.5547771558832267e01, + 1.2669386554490638e01, + 1.2071286687068984e01, + 9.7633584027450482e00, + -1.3485075900242089e02, + 1.8336445335820781e02, + -6.9300788508071975e01, + 1.4107018463574896e01, + 1.1439185153305873e01, + 8.7843335749580440e00, + -1.2663444344319166e02, + 1.7292158897636148e02, + -6.5521162694327174e01, + 1.3716937488160630e01, + 1.3999597459400730e01, + 1.2389915672436279e01, + -1.5954894249539399e02, + 2.1510813446746886e02, + -8.0895567204040049e01, + 1.5095682313349364e01, + 1.3189272906323732e01, + 1.1192627051714643e01, + -1.4915916817312757e02, + 2.0184825850919157e02, + -7.6081293415969839e01, + 1.3716937492019641e01, + 1.3999597377767842e01, + 1.2389916464009524e01, + -1.5954894412085929e02, + 2.1510813567394996e02, + -8.0895567498068928e01, + 1.5096520030681436e01, + 1.3185064407456906e01, + 1.1198910160279951e01, + -1.4913565617175487e02, + 2.0180124290250004e02, + -7.6060129778156622e01, + 1.4770075388032444e01, + 1.6087303167766446e01, + 1.5436222950666867e01, + -1.8666021493779203e02, + 2.4962122089688103e02, + -9.3426463524457304e01, + 1.6085379191481852e01, + 1.5083589447287226e01, + 1.3991739427782750e01, + -1.7386892459375579e02, + 2.3325385095807121e02, + -8.7470099643500802e01, + 1.4770075388818769e01, + 1.6087303148664304e01, + 1.5436223164442264e01, + -1.8666021539675981e02, + 2.4962122125116741e02, + -9.3426463615076329e01, + 1.6085951551006787e01, + 1.5080238931969067e01, + 1.3998101278449143e01, + -1.7385331837944693e02, + 2.3321864790104019e02, + -8.7453697552144448e01, + 1.5828143941097450e01, + 1.8331670220961666e01, + 1.8918268274003861e01, + -2.1619095210442941e02, + 2.8688297635978756e02, + -1.0687973526499771e02, + 1.7075534787366465e01, + 1.7125200136366264e01, + 1.7207074959934751e01, + -2.0084388544719391e02, + 2.6720765911058965e02, + -9.9705133726570395e01, + 1.5828143941256627e01, + 1.8331670216557445e01, + 1.8918268330404022e01, + -2.1619095222989833e02, + 2.8688297645950814e02, + -1.0687973529137253e02, + 1.7075923730873765e01, + 1.7122590193964911e01, + 1.7213058024904747e01, + -2.0083402645820061e02, + 2.6718180837697332e02, + -9.9692640534772679e01, + 1.6890371426423382e01, + 2.0728579569842751e01, + 2.2845917469463828e01, + -2.4812083435502871e02, + 3.2684448823688496e02, + -1.2123263616047282e02, + 1.8066449820492846e01, + 1.9312661524160735e01, + 2.0870036016187061e01, + -2.3013589616073858e02, + 3.0372498377642154e02, + -1.1277999824352135e02, + 1.6890371426455424e01, + 2.0728579568840633e01, + 2.2845917484032956e01, + -2.4812083438838550e02, + 3.2684448826399682e02, + -1.2123263616782057e02, + 1.8066713333743454e01, + 1.9310657703202459e01, + 2.0875423564416035e01, + -2.3013008228413184e02, + 3.0370630494679148e02, + -1.1277060230387309e02, + 1.7955886187113396e01, + 2.3272683588860026e01, + 2.7223982220959247e01, + -2.8240595076334000e02, + 3.6943078590316281e02, + -1.3645364576977221e02, + 1.9058236733002300e01, + 2.1644988962398710e01, + 2.5012267757287322e01, + -2.6180071928343307e02, + 3.4282650121799617e02, + -1.2669036882336400e02, + 1.7955886187119816e01, + 2.3272683588634656e01, + 2.7223982224651898e01, + -2.8240595077199526e02, + 3.6943078591032139e02, + -1.3645364577174797e02, + 1.9058414960148450e01, + 2.1643466247439289e01, + 2.5016983354038196e01, + -2.6179767020610126e02, + 3.4281320617581565e02, + -1.2668337355331974e02, + 1.9023741366983238e01, + 2.5957710504548576e01, + 3.2054387652193789e01, + -3.1898571318422574e02, + 4.1454655650462962e02, + -1.5250373535684176e02, + 2.0050906563887416e01, + 2.4121527381838824e01, + 2.9665428981325245e01, + -2.9589665055055406e02, + 3.8453661583827250e02, + -1.4143340987287985e02, + 1.9023741366984520e01, + 2.5957710504498362e01, + 3.2054387653114766e01, + -3.1898571318642672e02, + 4.1454655650647550e02, + -1.5250373535735841e02, + 2.0051026978020587e01, + 2.4120379273875816e01, + 2.9669474257430963e01, + -2.9589543070583102e02, + 3.8452729731205977e02, + -1.4142824748467820e02, + 2.0092947487287756e01, + 2.8776895490568755e01, + 3.7339233558876920e01, + -9.8781982607414882e00, + 7.0916635282296292e-01, + -1.2340880155534291e-02, + 2.1044418341890132e01, + 2.6741847681518077e01, + 3.4861073630499796e01, + -9.1700568642165461e00, + 6.5220324713443967e-01, + -1.1045071585279443e-02, + 2.0092947487288011e01, + 2.8776895490557653e01, + 3.7339233559103448e01, + -9.8781982608033179e00, + 7.0916635282857932e-01, + -1.2340880155703077e-02, + 2.1044499630877905e01, + 2.6740987496092696e01, + 3.4864491165514394e01, + -9.1707199731434574e00, + 6.5223741134844682e-01, + -1.1045188698410773e-02, + 2.1162510215379026e01, + 3.1723491960797684e01, + 4.3084295875067085e01, + -4.1033675985379521e00, + -6.6095139594000130e-01, + 6.0977735530407223e-02, + 2.2038706806958309e01, + 2.9505670300337073e01, + 4.0630600131872811e01, + -2.7905442844326718e00, + -8.3885972791335117e-01, + 6.8309956404426039e-02, + 2.1162510215379076e01, + 3.1723491960795304e01, + 4.3084295875120795e01, + -4.1033675985539224e00, + -6.6095139593840913e-01, + 6.0977735530354210e-02, + 2.2038761643178379e01, + 2.9505029336592230e01, + 4.0633451796171073e01, + -2.7913314472201640e00, + -8.3878528163749511e-01, + 6.8307595298566767e-02, + 3.1719012432820758e01, + 6.7480322661109355e01, + 1.3318978565899991e02, + -1.6791944323404795e01, + -1.0181217992701848e00, + 1.2989592638281225e-01, + 3.2009499874031789e01, + 6.5013296175889408e01, + 1.3669799889514238e02, + -1.7009031615065428e01, + -1.0689880784706638e00, + 1.3388972346122466e-01, + 3.1719012432820758e01, + 6.7480322661109355e01, + 1.3318978565899991e02, + -1.6791944323404795e01, + -1.0181217992701848e00, + 1.2989592638281225e-01, + 3.2009500887769519e01, + 6.5013269472322307e01, + 1.3669829238273672e02, + -1.7009116366540379e01, + -1.0689798256828462e00, + 1.3388945486998777e-01, + 4.1931127118492086e01, + 1.1600186087954401e02, + 3.1751764022286790e02, + -4.6438894455748802e01, + -8.7599401950869438e-01, + 2.2297105562740663e-01, + 4.2002297497564768e01, + 1.1479764873768737e02, + 3.2393143797302810e02, + -4.7847299173836262e01, + -7.8150712905299369e-01, + 2.2131248436241077e-01, + 4.1931127118492086e01, + 1.1600186087954401e02, + 3.1751764022286790e02, + -4.6438894455748802e01, + -8.7599401950869438e-01, + 2.2297105562740663e-01, + 4.2002297514594851e01, + 1.1479764793294436e02, + 3.2393145467669495e02, + -4.7847304068128608e01, + -7.8150664807362491e-01, + 2.2131246858403722e-01, + 5.1984670105634827e01, + 1.7926303194781252e02, + 6.2846495111925287e02, + -1.0034649475039414e02, + 2.4606292097951082e-01, + 3.3256752105517051e-01, + 5.2000554052128159e01, + 1.7883235795593501e02, + 6.3273302895025176e02, + -1.0138733878813618e02, + 3.2804187851642969e-01, + 3.3055293107858102e-01, + 5.1984670105634827e01, + 1.7926303194781252e02, + 6.2846495111925287e02, + -1.0034649475039414e02, + 2.4606292097951082e-01, + 3.3256752105517051e-01, + 5.2000554052402805e01, + 1.7883235793562420e02, + 6.3273302962903426e02, + -1.0138733898825184e02, + 3.2804189825766372e-01, + 3.3055293042886030e-01, + 6.1996666427075382e01, + 2.5724136589119979e02, + 1.0913830717468406e03, + -1.8317243758181812e02, + 2.5193786568880601e00, + 4.6277932792022042e-01, + 6.2000133522892554e01, + 2.5710536851489377e02, + 1.0934673032018356e03, + -1.8370056934287794e02, + 2.5630609198690104e00, + 4.6162176037505448e-01, + 6.1996666427075382e01, + 2.5724136589119979e02, + 1.0913830717468406e03, + -1.8317243758181812e02, + 2.5193786568880601e00, + 4.6277932792022042e-01, + 6.2000133522896938e01, + 2.5710536851442714e02, + 1.0934673032246803e03, + -1.8370056934963364e02, + 2.5630609205366826e00, + 4.6162176035304603e-01, + 7.1999279107664492e01, + 3.4965254984584158e02, + 1.7356304176273381e03, + -3.0063395678020430e02, + 6.2079056750108883e00, + 6.1505333334154833e-01, + 7.2000032172982571e01, + 3.4961232791697932e02, + 1.7365043785874466e03, + -3.0086002522613632e02, + 6.2270725229979789e00, + 6.1452738833821030e-01, + 7.1999279107664492e01, + 3.4965254984584158e02, + 1.7356304176273381e03, + -3.0063395678020430e02, + 6.2079056750108883e00, + 6.1505333334154833e-01, + 7.2000032172982642e01, + 3.4961232791696904e02, + 1.7365043785881401e03, + -3.0086002522634379e02, + 6.2270725230187063e00, + 6.1452738833751985e-01, + 8.1999844359310714e01, + 4.5636323545227941e02, + 2.5918884526432239e03, + -4.5885344883307727e02, + 1.1616256691917803e01, + 7.8948404417119522e-01, + 8.2000007751936337e01, + 4.5635184072744744e02, + 2.5922210189842476e03, + -4.5894061525528980e02, + 1.1623761628208563e01, + 7.8927378661620728e-01, + 8.1999844359310714e01, + 4.5636323545227941e02, + 2.5918884526432239e03, + -4.5885344883307727e02, + 1.1616256691917803e01, + 7.8948404417119522e-01, + 8.2000007751936337e01, + 4.5635184072744744e02, + 2.5922210189842476e03, + -4.5894061525528980e02, + 1.1623761628208563e01, + 7.8927378661620728e-01, + ], + dtype=dtype, + device=env.DEVICE, + ).reshape(8, 174) # 1392 + self.table_info_tensor = torch.tensor( + [ + -2.1000000000000000e01, + 2.1000000000000000e01, + 1.0500000000000000e02, + 1.0000000000000000e00, + 1.0000000000000000e01, + -1.0000000000000000e00, + ], + dtype=dtype, + device="cpu", + ) + self.em_x_tensor = torch.tensor( + [ + 9.3816147034272368e-01, + -1.6703373029862567e-01, + -4.4294526064601734e-02, + -2.8798505489184573e-01, + -1.6703373029862567e-01, + 9.2489218226366088e-01, + -2.8928196536572048e-01, + -4.7833509099876154e-01, + -4.4294526064601734e-02, + -2.8928196536572048e-01, + 5.7034320185695120e-01, + 1.8771147911830000e-01, + -2.8798505489184573e-01, + -4.7833509099876154e-01, + 1.8771147911830000e-01, + 4.0174654365823070e-01, + 8.4370316144902313e-01, + -3.7813146789689916e-02, + -3.6989397568296523e-01, + -4.0554075086539937e-01, + -3.7813146789689916e-02, + 6.5766402633747112e-01, + -4.2312966361682885e-01, + 1.2685067374257861e-01, + -3.6989397568296523e-01, + -4.2312966361682885e-01, + 6.0171576901660107e-01, + 9.8283160997298613e-02, + -4.0554075086539937e-01, + 1.2685067374257861e-01, + 9.8283160997298613e-02, + 2.1324148100625978e-01, + 9.7843596341516559e-01, + -1.0492833888237871e-01, + -1.0538688914576379e-01, + -2.0453551592353389e-01, + -1.0492833888237871e-01, + 7.7943976693565231e-01, + -1.5898500035781410e-01, + 9.4834209331437741e-02, + -1.0538688914576379e-01, + -1.5898500035781410e-01, + 7.4778071691708869e-01, + -6.1895255142095873e-01, + -2.0453551592353389e-01, + 9.4834209331437741e-02, + -6.1895255142095873e-01, + 6.0844713798743799e-01, + 1.0079020879244640e00, + -2.3855984150631487e-01, + -3.4608276043004524e-02, + -4.7448768267289088e-01, + -2.3855984150631487e-01, + 4.9732018171028253e-01, + -3.1320787082485729e-01, + -1.4528004145602180e-01, + -3.4608276043004524e-02, + -3.1320787082485729e-01, + 4.7696729363954582e-01, + 1.1723268074231248e-01, + -4.7448768267289088e-01, + -1.4528004145602180e-01, + 1.1723268074231248e-01, + 4.0511515406019899e-01, + ], + dtype=dtype, + device=env.DEVICE, + ).reshape(4, 16) # 3072 + self.em_tensor = self.em_x_tensor.reshape(4, 4, 4) + self.table_info_tensor.requires_grad = False + self.table_tensor.requires_grad = False + self.em_x_tensor.requires_grad = True + self.em_tensor.requires_grad = True + + # These hint the shapes, but not used + self.last_layer_size = 4 + self.nnei_i = 4 + self.nnei_j = 4 + + self.expected_descriptor_tensor = torch.tensor( + [ + -2.079966061311932624e00, + -7.692086176232872941e-01, + -1.330535805782471437e00, + -1.535049115492904015e00, + -3.387998402723848113e00, + -2.128016321720361592e00, + -3.277106691660685023e00, + -3.038592571524457941e00, + -3.249069620732517372e00, + -1.970139560184089200e00, + -3.099042807698376478e00, + -2.871870824895611740e00, + -3.505396485462592526e00, + -2.262865452602483174e00, + -3.422716794166150045e00, + -3.179802052814081126e00, + -3.387998402723848113e00, + -2.128016321720361592e00, + -3.277106691660685023e00, + -3.038592571524457941e00, + -2.097644546971316970e00, + -7.853817517282856153e-01, + -1.357716376463428354e00, + -1.553790419377535814e00, + -3.506551159314416299e00, + -2.264198868197136782e00, + -3.424124719375644776e00, + -3.181191625971385051e00, + -3.660169689351369549e00, + -2.443224371314374199e00, + -3.605953674751581683e00, + -3.365824124673929774e00, + -3.249069620732517372e00, + -1.970139560184089200e00, + -3.099042807698376478e00, + -2.871870824895611740e00, + -3.506551159314416299e00, + -2.264198868197136782e00, + -3.424124719375644776e00, + -3.181191625971385051e00, + -2.530401019792652750e00, + -1.199448737241302965e00, + -2.045857780281324256e00, + -2.028209748272746360e00, + -2.967997665475758939e00, + -1.658223979429474859e00, + -2.714263642699213275e00, + -2.536766176574209553e00, + -3.505396485462592526e00, + -2.262865452602483174e00, + -3.422716794166150045e00, + -3.179802052814081126e00, + -3.660169689351369549e00, + -2.443224371314374199e00, + -3.605953674751581683e00, + -3.365824124673929774e00, + -2.967997665475758939e00, + -1.658223979429474859e00, + -2.714263642699213275e00, + -2.536766176574209553e00, + -2.720389222561622855e00, + -1.393795402757894042e00, + -2.345287611062552902e00, + -2.246069186600956336e00, + -2.203412666452015500e00, + -8.832922757908097111e-01, + -1.522624038797550927e00, + -1.666955513084583895e00, + -3.241305103472522209e00, + -1.961375736752476673e00, + -3.088893233422514317e00, + -2.862569539842237276e00, + -3.574866904199356821e00, + -2.343386670122550530e00, + -3.506415378108253655e00, + -3.263388258370226502e00, + -3.603390794940432240e00, + -2.376642828344087466e00, + -3.540123958621633538e00, + -3.297678764698442677e00, + -3.241305103472522209e00, + -1.961375736752476673e00, + -3.088893233422514317e00, + -2.862569539842237276e00, + -2.429999277303452665e00, + -1.100023053295984932e00, + -1.884632599239140616e00, + -1.915409451705063049e00, + -3.617225590985869044e00, + -2.392818241566808002e00, + -3.556322178869617989e00, + -3.314299042117806238e00, + -3.041076611338295876e00, + -1.738238794499435524e00, + -2.817792369079202164e00, + -2.623522438271889357e00, + -3.574866904199356821e00, + -2.343386670122550530e00, + -3.506415378108253655e00, + -3.263388258370226502e00, + -3.617225590985869044e00, + -2.392818241566808002e00, + -3.556322178869617989e00, + -3.314299042117806238e00, + -2.494628058196212361e00, + -1.163767319481099349e00, + -1.988521437257453117e00, + -1.987826554800365741e00, + -3.075717632972006399e00, + -1.776439224176796783e00, + -2.866012767430887731e00, + -2.664752088723873147e00, + -3.603390794940432240e00, + -2.376642828344087466e00, + -3.540123958621633538e00, + -3.297678764698442677e00, + -3.041076611338295876e00, + -1.738238794499435524e00, + -2.817792369079202164e00, + -2.623522438271889357e00, + -3.075717632972006399e00, + -1.776439224176796783e00, + -2.866012767430887731e00, + -2.664752088723873147e00, + -2.937692581645545253e00, + -1.625281447224701248e00, + -2.670609348993408627e00, + -2.500890937682363369e00, + -2.025708717938600856e00, + -7.199180117105589138e-01, + -1.247903242801984502e00, + -1.477847643583491077e00, + -3.319906298979989234e00, + -2.050399877432793616e00, + -3.190625595262678527e00, + -2.956812988881384996e00, + -3.320427161433143759e00, + -2.050991936909778435e00, + -3.191292681270106346e00, + -2.957438103019588649e00, + -3.426553578385282073e00, + -2.172154331535164129e00, + -3.325432621990271009e00, + -3.084944527573738160e00, + -3.319906298979989234e00, + -2.050399877432793616e00, + -3.190625595262678527e00, + -2.956812988881384996e00, + -2.284064566837665566e00, + -9.593000707401009297e-01, + -1.650522869993226838e00, + -1.754438009603198623e00, + -3.379465923571113084e00, + -2.118266849034519783e00, + -3.266349102610927257e00, + -3.028338829309447533e00, + -3.079911054493713163e00, + -1.781075057846572385e00, + -2.871812857785094764e00, + -2.669747195759931468e00, + -3.320427161433143759e00, + -2.050991936909778435e00, + -3.191292681270106346e00, + -2.957438103019588649e00, + -3.379465923571113084e00, + -2.118266849034519783e00, + -3.266349102610927257e00, + -3.028338829309447533e00, + -2.322820890996963428e00, + -9.962581166315247483e-01, + -1.712462166775041483e00, + -1.796847406967182925e00, + -3.773974160346314832e00, + -2.578147069133936320e00, + -3.733153897787112818e00, + -3.501723916740695497e00, + -3.426553578385282073e00, + -2.172154331535164129e00, + -3.325432621990271009e00, + -3.084944527573738160e00, + -3.079911054493713163e00, + -1.781075057846572385e00, + -2.871812857785094764e00, + -2.669747195759931468e00, + -3.773974160346314832e00, + -2.578147069133936320e00, + -3.733153897787112818e00, + -3.501723916740695497e00, + -2.486913710376590814e00, + -1.156110267723822860e00, + -1.976136357370776642e00, + -1.979145885578800446e00, + -1.985509025147957018e00, + -6.837428483499845555e-01, + -1.187510572446436763e00, + -1.435784454536765731e00, + -3.459824255787086411e00, + -2.210355755872039030e00, + -3.366749809979232122e00, + -3.124964584278840274e00, + -3.237455343632644311e00, + -1.957033195714938278e00, + -3.083852092710132720e00, + -2.857958563155091891e00, + -3.657182553291094873e00, + -2.439708494346079171e00, + -3.602533119650989324e00, + -3.362243802860200859e00, + -3.459824255787086411e00, + -2.210355755872039030e00, + -3.366749809979232122e00, + -3.124964584278840274e00, + -2.612858501561634128e00, + -1.282799717023056063e00, + -2.177105827922439119e00, + -2.122088127377702360e00, + -3.527506999611477845e00, + -2.288426034556684474e00, + -3.449583412378998126e00, + -3.206410310763728511e00, + -3.364731487733799042e00, + -2.101446556578417724e00, + -3.247719091239923639e00, + -3.010635929634002395e00, + -3.237455343632644311e00, + -1.957033195714938278e00, + -3.083852092710132720e00, + -2.857958563155091891e00, + -3.527506999611477845e00, + -2.288426034556684474e00, + -3.449583412378998126e00, + -3.206410310763728511e00, + -2.635737441843280848e00, + -1.306200303853801703e00, + -2.213214111467705525e00, + -2.148325634854432398e00, + -3.052719273633249131e00, + -1.751058696872038523e00, + -2.834060115075742292e00, + -2.637372572375544344e00, + -3.657182553291094873e00, + -2.439708494346079171e00, + -3.602533119650989324e00, + -3.362243802860200859e00, + -3.364731487733799042e00, + -2.101446556578417724e00, + -3.247719091239923639e00, + -3.010635929634002395e00, + -3.052719273633249131e00, + -1.751058696872038523e00, + -2.834060115075742292e00, + -2.637372572375544344e00, + -2.716586630714696859e00, + -1.389827110741064020e00, + -2.339411933160714607e00, + -2.241657223573843893e00, + ], + dtype=dtype, + device=env.DEVICE, + ).reshape(4, 4, 4, 4) + self.expected_dy_dem_x = torch.tensor( + [ + 6.020806901503564568e00, + 4.855212102289546472e00, + 5.539480671183889093e00, + 4.079184573583559903e00, + 4.855212102289546472e00, + 6.002622647924519939e00, + 4.071866106011439612e00, + 3.510130130958049044e00, + 5.539480671183889093e00, + 4.071866106011439612e00, + 5.377211408293748107e00, + 5.608493562913938746e00, + 4.079184573583559903e00, + 3.510130130958049044e00, + 5.608493562913938746e00, + 5.361908159053506928e00, + 5.856733593151207451e00, + 5.562527732431956196e00, + 3.696970497751218332e00, + 3.591262868531118535e00, + 5.562527732431956196e00, + 5.486518130714558517e00, + 3.554623757729209022e00, + 5.671203227676853587e00, + 3.696970497751218332e00, + 3.554623757729209022e00, + 5.409085428112309302e00, + 5.689207994509548172e00, + 3.591262868531118535e00, + 5.671203227676853587e00, + 5.689207994509548172e00, + 5.576190638038706737e00, + 6.061353798127412595e00, + 5.246377090750031336e00, + 5.243735237316045961e00, + 4.603911062145699340e00, + 5.246377090750031336e00, + 5.721012903047327569e00, + 4.908561477952471108e00, + 5.690718304708298980e00, + 5.243735237316045961e00, + 4.908561477952471108e00, + 5.654457931163397788e00, + 3.881456101028720962e00, + 4.603911062145699340e00, + 5.690718304708298980e00, + 3.881456101028720962e00, + 5.417053162665033561e00, + 6.074411415516496326e00, + 4.379884232572607949e00, + 5.573202499423647538e00, + 3.509692162941947302e00, + 4.379884232572607949e00, + 5.338347197725587634e00, + 3.943361243591621879e00, + 4.998151625670068654e00, + 5.573202499423647538e00, + 3.943361243591621879e00, + 5.336491405093880758e00, + 5.678289531805588375e00, + 3.509692162941947302e00, + 4.998151625670068654e00, + 5.678289531805588375e00, + 5.359752988507883842e00, + ], + dtype=dtype, + device=env.DEVICE, + ).reshape(4, 16) + + def test_forward(self) -> None: + # Call the forward function + forward_result = torch.ops.deepmd.tabulate_fusion_se_t_tebd( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor, + self.em_tensor, + self.last_layer_size, + ) + + descriptor_tensor = forward_result[0] + + # Check the shape + self.assertEqual(descriptor_tensor.shape, self.expected_descriptor_tensor.shape) + + # Check the values + torch.testing.assert_close( + descriptor_tensor, + self.expected_descriptor_tensor, + atol=self.prec, + rtol=self.prec, + ) + + def test_backward(self) -> None: + # Call the forward function + forward_result = torch.ops.deepmd.tabulate_fusion_se_t_tebd( + self.table_tensor, + self.table_info_tensor, + self.em_x_tensor, + self.em_tensor, + self.last_layer_size, + ) + + descriptor_tensor = forward_result[0] + + # Check the forward + torch.testing.assert_close( + descriptor_tensor, + self.expected_descriptor_tensor, + atol=self.prec, + rtol=self.prec, + ) + + # Create a loss and perform backward + loss = descriptor_tensor.sum() + loss.backward() + + # Check gradients + self.assertIsNotNone(self.em_x_tensor.grad) + + # No need to check em_tensor.grad, it is ignored in the op + # Since em_tensor and em_x_tensor share the same data + + # Check the shapes of the gradients + self.assertEqual(self.em_x_tensor.grad.shape, self.expected_dy_dem_x.shape) + + # Check the values of the gradients + torch.testing.assert_close( + self.em_x_tensor.grad, + self.expected_dy_dem_x, + atol=self.prec, + rtol=self.prec, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index 3df95e4b14..ff4f00f912 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -15,6 +15,7 @@ from deepmd.pt.entrypoints.main import ( get_trainer, ) +from deepmd.pt.entrypoints.main import train as train_entry from deepmd.pt.utils.finetune import ( get_finetune_rules, ) @@ -30,6 +31,8 @@ class DPTrainTest: + test_zbl_from_standard: bool = False + def test_dp_train(self) -> None: # test training from scratch trainer = get_trainer(deepcopy(self.config)) @@ -89,12 +92,44 @@ def test_dp_train(self) -> None: state_dict_trained[state_key], state_dict_finetuned_empty[state_key], ) - if "fitting_net" not in state_key: + if ( + ("fitting_net" not in state_key) + or ("fparam" in state_key) + or ("aparam" in state_key) + ): torch.testing.assert_close( state_dict_trained[state_key], state_dict_finetuned_random[state_key], ) + if self.test_zbl_from_standard: + # test fine-tuning using zbl from standard model + finetune_model = ( + self.config["training"].get("save_ckpt", "model.ckpt") + ".pt" + ) + self.config_zbl["model"], finetune_links = get_finetune_rules( + finetune_model, + self.config_zbl["model"], + ) + trainer_finetune_zbl = get_trainer( + deepcopy(self.config_zbl), + finetune_model=finetune_model, + finetune_links=finetune_links, + ) + state_dict_finetuned_zbl = trainer_finetune_zbl.wrapper.model.state_dict() + for state_key in state_dict_finetuned_zbl: + if "out_bias" not in state_key and "out_std" not in state_key: + original_key = state_key + if ".models.0." in state_key: + original_key = state_key.replace(".models.0.", ".") + if ".models.1." not in state_key: + torch.testing.assert_close( + state_dict_trained[original_key], + state_dict_finetuned_zbl[state_key], + ) + # check running + trainer_finetune_zbl.run() + # check running trainer_finetune.run() trainer_finetune_empty.run() @@ -150,8 +185,29 @@ def setUp(self) -> None: self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 + def test_yaml_input(self) -> None: + import yaml + + yaml_file = Path("input.yaml") + with open(yaml_file, "w") as fp: + yaml.safe_dump(self.config, fp) + train_entry( + input_file=str(yaml_file), + init_model=None, + restart=None, + finetune=None, + init_frz_model=None, + model_branch="main", + skip_neighbor_stat=True, + output="out.json", + ) + self.assertTrue(Path("out.json").exists()) + def tearDown(self) -> None: DPTrainTest.tearDown(self) + for ff in ["out.json", "input.yaml"]: + if Path(ff).exists(): + os.remove(ff) class TestDOSModelSeA(unittest.TestCase, DPTrainTest): @@ -204,6 +260,7 @@ def setUp(self) -> None: self.config["training"]["save_freq"] = 1 self.set_path = Path(__file__).parent / "water/data/data_0" / "set.000" shutil.copyfile(self.set_path / "energy.npy", self.set_path / "fparam.npy") + self.config["model"]["data_stat_nbatch"] = 100 def tearDown(self) -> None: (self.set_path / "fparam.npy").unlink(missing_ok=True) @@ -222,6 +279,18 @@ def setUp(self) -> None: self.config["training"]["numb_steps"] = 1 self.config["training"]["save_freq"] = 1 + self.test_zbl_from_standard = True + + input_json_zbl = str(Path(__file__).parent / "water/zbl.json") + with open(input_json_zbl) as f: + self.config_zbl = json.load(f) + data_file = [str(Path(__file__).parent / "water/data/data_0")] + self.config_zbl["training"]["training_data"]["systems"] = data_file + self.config_zbl["training"]["validation_data"]["systems"] = data_file + self.config_zbl["model"] = deepcopy(model_zbl) + self.config_zbl["training"]["numb_steps"] = 1 + self.config_zbl["training"]["save_freq"] = 1 + def tearDown(self) -> None: DPTrainTest.tearDown(self) diff --git a/source/tests/tf/test_change_bias.py b/source/tests/tf/test_change_bias.py new file mode 100644 index 0000000000..4392bbd139 --- /dev/null +++ b/source/tests/tf/test_change_bias.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import shutil +import tempfile +import unittest +from pathlib import ( + Path, +) + +from deepmd.tf.entrypoints.change_bias import ( + change_bias, +) +from deepmd.tf.train.run_options import ( + RunOptions, +) +from deepmd.tf.train.trainer import ( + DPTrainer, +) +from deepmd.tf.utils.argcheck import ( + normalize, +) +from deepmd.tf.utils.compat import ( + update_deepmd_input, +) + +from .common import ( + j_loader, + run_dp, + tests_path, +) + + +class TestChangeBias(unittest.TestCase): + def setUp(self): + """Set up test fixtures.""" + self.temp_dir = tempfile.mkdtemp() + self.temp_path = Path(self.temp_dir) + + def tearDown(self): + """Clean up test fixtures.""" + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_change_bias_frozen_model_partial_support(self): + """Test that frozen model support has limitations but provides helpful error.""" + fake_pb = self.temp_path / "model.pb" + fake_pb.write_text("fake model content") + + # Without bias_value, should suggest using bias_value or checkpoint + with self.assertRaises(NotImplementedError) as cm: + change_bias( + INPUT=str(fake_pb), + mode="change", + system=".", + ) + + self.assertIn( + "Data-based bias changing for frozen models is not yet implemented", + str(cm.exception), + ) + self.assertIn("bias-value option", str(cm.exception)) + + # With bias_value, should provide implementation guidance + with self.assertRaises(NotImplementedError) as cm: + change_bias( + INPUT=str(fake_pb), + mode="change", + bias_value=[1.0, 2.0], + system=".", + ) + + self.assertIn( + "Bias modification for frozen models (.pb) is not yet fully implemented", + str(cm.exception), + ) + self.assertIn("checkpoint_dir", str(cm.exception)) + + def test_change_bias_invalid_model_type(self): + """Test that invalid model types raise RuntimeError.""" + fake_model = self.temp_path / "model.xyz" + fake_model.write_text("fake model content") + + with self.assertRaises(RuntimeError) as cm: + change_bias( + INPUT=str(fake_model), + mode="change", + system=".", + ) + + self.assertIn( + "checkpoint file or frozen model file (.pb)", + str(cm.exception), + ) + + def test_change_bias_no_checkpoint_in_directory(self): + """Test that checkpoint files need proper checkpoint structure.""" + fake_ckpt = self.temp_path / "model.ckpt" + fake_ckpt.write_text("fake checkpoint content") + + # Create a fake data system for the test + fake_data_dir = self.temp_path / "fake_data" + fake_data_dir.mkdir() + fake_set_dir = fake_data_dir / "set.000" + fake_set_dir.mkdir() + + with self.assertRaises(RuntimeError) as cm: + change_bias( + INPUT=str(fake_ckpt), + mode="change", + system=str(fake_data_dir), + ) + + self.assertIn("No valid checkpoint found", str(cm.exception)) + + def test_change_bias_user_defined_requires_real_model(self): + """Test that user-defined bias requires a real model with proper structure.""" + fake_ckpt_dir = self.temp_path / "fake_checkpoint" + fake_ckpt_dir.mkdir() + fake_ckpt = fake_ckpt_dir / "model.ckpt" + fake_ckpt.write_text("fake checkpoint content") + (fake_ckpt_dir / "checkpoint").write_text("fake checkpoint") + # Create a minimal but complete input.json + minimal_config = { + "model": {"type_map": ["H", "O"]}, + "training": {"systems": ["."], "validation_data": {"systems": ["."]}}, + } + + (fake_ckpt_dir / "input.json").write_text(json.dumps(minimal_config)) + + # Should fail because there's no real model structure, but with different error + with self.assertRaises((RuntimeError, FileNotFoundError, Exception)) as cm: + change_bias( + INPUT=str(fake_ckpt), + mode="change", + bias_value=[1.0, 2.0], + system=".", + ) + + # The error should be about model loading, not about NotImplementedError + self.assertNotIn("not yet implemented", str(cm.exception)) + + def test_change_bias_with_real_model(self): + """Test change_bias with a real trained model and verify output.""" + # Create temporary directories for training and output + train_dir = self.temp_path / "train" + train_dir.mkdir() + checkpoint_dir = train_dir / "checkpoint" + output_file = self.temp_path / "output_model.pb" + + # Use existing test data and configuration + data_dir = tests_path / "init_frz_model" / "data" + config_file = tests_path / "init_frz_model" / "input.json" + + # Load and modify configuration for quick training + jdata = j_loader(str(config_file)) + jdata["training"]["training_data"]["systems"] = [str(data_dir)] + jdata["training"]["validation_data"]["systems"] = [str(data_dir)] + jdata["training"]["numb_steps"] = 2 # Minimal training for testing + jdata["training"]["save_freq"] = 1 + jdata["training"]["save_ckpt"] = str(checkpoint_dir / "model.ckpt") + + # Write modified config + input_json_path = train_dir / "input.json" + with open(input_json_path, "w") as f: + json.dump(jdata, f, indent=4) + + # Train the model using run_dp + ret = run_dp(f"dp train {input_json_path}") + self.assertEqual(ret, 0, "DP train failed!") + + # Verify checkpoint was created + self.assertTrue(checkpoint_dir.exists()) + checkpoint_files = list(checkpoint_dir.glob("*")) + self.assertGreater(len(checkpoint_files), 0, "No checkpoint files created") + + # Find the actual checkpoint file + checkpoint_file = checkpoint_dir / "model.ckpt" + + # Create a frozen model from the checkpoint for testing + frozen_model_path = train_dir / "frozen_model.pb" + ret = run_dp(f"dp freeze -c {checkpoint_dir} -o {frozen_model_path}") + self.assertEqual(ret, 0, "DP freeze failed!") + self.assertTrue(frozen_model_path.exists()) + + # Test change_bias function - this should provide implementation guidance for frozen models + with self.assertRaises(NotImplementedError) as cm: + change_bias( + INPUT=str(frozen_model_path), + mode="change", + system=str(data_dir), + output=str(output_file), + ) + self.assertIn( + "Data-based bias changing for frozen models is not yet implemented", + str(cm.exception), + ) + + # Now test change_bias on the real checkpoint file (this is the real test) + change_bias( + INPUT=str(checkpoint_file), + mode="change", + system=str(data_dir), + output=str(output_file), + ) + + # Verify that output model file was created + self.assertTrue(output_file.exists()) + self.assertTrue(output_file.stat().st_size > 0, "Output model file is empty") + + # Load original model to verify structure + original_run_opt = RunOptions(init_model=str(checkpoint_dir), log_level=20) + + # Load the configuration again for creating trainers + jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + jdata = normalize(jdata) + + original_trainer = DPTrainer(jdata, run_opt=original_run_opt) + + # Verify original model loads successfully + self.assertIsNotNone(original_trainer.model) + + # Verify the original model has the expected structure + original_type_map = original_trainer.model.get_type_map() + self.assertGreater(len(original_type_map), 0, "Model should have a type_map") + + # Clean up training artifacts + for artifact in ["lcurve.out", "input_v2_compat.json"]: + if os.path.exists(artifact): + os.remove(artifact) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/tf/test_nvnmd_entrypoints.py b/source/tests/tf/test_nvnmd_entrypoints.py index eaf8bfafd5..6833b4e0f8 100644 --- a/source/tests/tf/test_nvnmd_entrypoints.py +++ b/source/tests/tf/test_nvnmd_entrypoints.py @@ -515,6 +515,7 @@ def test_model_qnn_v0(self) -> None: def tearDown(self) -> None: # close nvnmd_cfg.enable = False + tf.reset_default_graph() class TestNvnmdEntrypointsV1(tf.test.TestCase): @@ -878,6 +879,7 @@ def test_wrap_qnn_v1(self) -> None: def tearDown(self) -> None: # close nvnmd_cfg.enable = False + tf.reset_default_graph() if __name__ == "__main__": diff --git a/source/tests/tf/test_out_bias_std.py b/source/tests/tf/test_out_bias_std.py new file mode 100644 index 0000000000..a5bd4a05cc --- /dev/null +++ b/source/tests/tf/test_out_bias_std.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +import unittest + +import numpy as np + +from deepmd.tf.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.tf.env import ( + tf, +) +from deepmd.tf.fit.dipole import ( + DipoleFittingSeA, +) +from deepmd.tf.fit.ener import ( + EnerFitting, +) +from deepmd.tf.model.model import ( + StandardModel, +) + + +class TestOutBiasStd(unittest.TestCase): + """Test out_bias and out_std functionality in TensorFlow backend.""" + + def setUp(self): + """Resets the default graph before each test.""" + super().setUp() + tf.reset_default_graph() + + def tearDown(self): + """Resets the default graph after each test.""" + tf.reset_default_graph() + super().tearDown() + + def test_init_out_stat_basic(self): + """Test basic init_out_stat functionality.""" + descriptor = DescrptSeA( + rcut=4.0, rcut_smth=3.5, sel=[10, 20], neuron=[8, 16, 32] + ) + fitting = EnerFitting(ntypes=2, dim_descrpt=32) + model = StandardModel( + descriptor=descriptor, fitting_net=fitting, type_map=["H", "O"] + ) + + # Test initial state + self.assertIsNone(model.out_bias) + self.assertIsNone(model.out_std) + + # Test init_out_stat + model.init_out_stat() + self.assertIsNotNone(model.out_bias) + self.assertIsNotNone(model.out_std) + self.assertEqual(model.out_bias.shape, (1, 2, 1)) # 1 output, 2 types, 1 dim + self.assertEqual(model.out_std.shape, (1, 2, 1)) + + # Check default values + np.testing.assert_array_equal(model.out_bias, np.zeros((1, 2, 1))) + np.testing.assert_array_equal(model.out_std, np.ones((1, 2, 1))) + + def test_different_fitting_dimensions(self): + """Test that different fitting types have correct dimensions.""" + descriptor = DescrptSeA( + rcut=4.0, rcut_smth=3.5, sel=[10, 20], neuron=[8, 16, 32] + ) + + # Test energy fitting (dim_out = 1) + fitting_ener = EnerFitting(ntypes=2, dim_descrpt=32) + model_ener = StandardModel( + descriptor=descriptor, fitting_net=fitting_ener, type_map=["H", "O"] + ) + model_ener.init_out_stat(suffix="_ener") + self.assertEqual(model_ener.out_bias.shape, (1, 2, 1)) + + # Test dipole fitting (dim_out = 3) + fitting_dipole = DipoleFittingSeA(ntypes=2, dim_descrpt=32, embedding_width=32) + model_dipole = StandardModel( + descriptor=descriptor, fitting_net=fitting_dipole, type_map=["H", "O"] + ) + model_dipole.init_out_stat(suffix="_dipole") + self.assertEqual(model_dipole.out_bias.shape, (1, 2, 3)) + + def test_apply_out_stat(self): + """Test that out_bias and out_std are applied during model build.""" + from deepmd.tf.env import ( + tf, + ) + + descriptor = DescrptSeA( + rcut=4.0, rcut_smth=3.5, sel=[10, 20], neuron=[8, 16, 32] + ) + fitting = EnerFitting(ntypes=2, dim_descrpt=32) + model = StandardModel( + descriptor=descriptor, fitting_net=fitting, type_map=["H", "O"] + ) + + # Set test bias and std directly + test_bias = np.array([[[1.0], [2.0]]]) # bias for type 0: 1.0, type 1: 2.0 + test_std = np.array([[[0.5], [1.5]]]) # std for type 0: 0.5, type 1: 1.5 + model.out_bias = test_bias + model.out_std = test_std + + # Create mock input data for testing + nloc = 3 + + # Mock coordinates and atom types + coord = tf.placeholder(tf.float64, [None, nloc * 3]) + atype = tf.placeholder(tf.int32, [None, nloc]) + natoms = [ + nloc, + nloc, + 1, + 2, + ] # [local atoms, total atoms, type 0 count, type 1 count] + box = tf.placeholder(tf.float64, [None, 9]) + mesh = tf.placeholder(tf.int32, [None, 6]) + + # Build the model - this should apply bias/std internally + model.build(coord, atype, natoms, box, mesh, input_dict=None) + + # Check that the bias and std variables were created + self.assertTrue(hasattr(model, "t_out_bias")) + self.assertTrue(hasattr(model, "t_out_std")) + + # Test that out_bias and out_std are preserved + np.testing.assert_array_equal(model.out_bias, test_bias) + np.testing.assert_array_equal(model.out_std, test_std) + + def test_apply_out_stat_no_bias(self): + """Test that when no bias is explicitly set, default bias (zeros) is used.""" + descriptor = DescrptSeA( + rcut=4.0, rcut_smth=3.5, sel=[10, 20], neuron=[8, 16, 32] + ) + fitting = EnerFitting(ntypes=2, dim_descrpt=32) + model = StandardModel( + descriptor=descriptor, fitting_net=fitting, type_map=["H", "O"] + ) + + # Initialize the model which should set default bias=0, std=1 + model.init_out_stat() + + # Verify that default bias and std are set correctly + bias = model.out_bias + std = model.out_std + + # Default bias should be zeros + expected_bias = np.zeros([1, 2, 1]) # [1, ntypes, dim_out] + expected_std = np.ones([1, 2, 1]) # [1, ntypes, dim_out] + + np.testing.assert_array_equal(bias, expected_bias) + np.testing.assert_array_equal(std, expected_std) + + def test_decoupled_bias_architecture(self): + """Test that out_bias and bias_atom_e are completely decoupled.""" + # Test that setting out_bias does not affect bias_atom_e and vice versa + + descriptor = DescrptSeA( + rcut=4.0, rcut_smth=3.5, sel=[10, 20], neuron=[8, 16, 32] + ) + fitting = EnerFitting(ntypes=2, dim_descrpt=32) + model = StandardModel( + descriptor=descriptor, fitting_net=fitting, type_map=["H", "O"] + ) + + # Initialize with defaults + model.init_out_stat() + + # Set out_bias directly + test_out_bias = np.array([[[1.0], [2.0]]]) + model.out_bias = test_out_bias + + # Verify out_bias is set correctly + retrieved_bias = model.out_bias + np.testing.assert_array_equal(retrieved_bias, test_out_bias) + + # Verify that out_std can be set independently + test_out_std = np.array([[[0.5], [1.5]]]) + model.out_std = test_out_std + retrieved_std = model.out_std + np.testing.assert_array_equal(retrieved_std, test_out_std) + + # Verify shapes are correct for energy models + self.assertEqual(retrieved_bias.shape, (1, 2, 1)) # [1, ntypes, dim_out] + self.assertEqual(retrieved_std.shape, (1, 2, 1)) # [1, ntypes, dim_out] + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/universal/common/cases/atomic_model/utils.py b/source/tests/universal/common/cases/atomic_model/utils.py index 054272b8c6..45b30c6454 100644 --- a/source/tests/universal/common/cases/atomic_model/utils.py +++ b/source/tests/universal/common/cases/atomic_model/utils.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) from typing import ( Any, - Callable, - Optional, ) import numpy as np @@ -41,11 +42,11 @@ class AtomicModelTestCase: """Expected whether having message passing.""" forward_wrapper: Callable[[Any], Any] """Class wrapper for forward method.""" - aprec_dict: dict[str, Optional[float]] + aprec_dict: dict[str, float | None] """Dictionary of absolute precision in each test.""" - rprec_dict: dict[str, Optional[float]] + rprec_dict: dict[str, float | None] """Dictionary of relative precision in each test.""" - epsilon_dict: dict[str, Optional[float]] + epsilon_dict: dict[str, float | None] """Dictionary of epsilons in each test.""" def test_get_type_map(self) -> None: diff --git a/source/tests/universal/common/cases/model/utils.py b/source/tests/universal/common/cases/model/utils.py index 08a369933d..8ec0edc180 100644 --- a/source/tests/universal/common/cases/model/utils.py +++ b/source/tests/universal/common/cases/model/utils.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +from collections.abc import ( + Callable, +) from copy import ( deepcopy, ) from typing import ( Any, - Callable, - Optional, ) import numpy as np @@ -54,11 +55,11 @@ class ModelTestCase: """Class wrapper for forward method.""" forward_wrapper_cpu_ref: Callable[[Any], Any] """Convert model to CPU method.""" - aprec_dict: dict[str, Optional[float]] + aprec_dict: dict[str, float | None] """Dictionary of absolute precision in each test.""" - rprec_dict: dict[str, Optional[float]] + rprec_dict: dict[str, float | None] """Dictionary of relative precision in each test.""" - epsilon_dict: dict[str, Optional[float]] + epsilon_dict: dict[str, float | None] """Dictionary of epsilons in each test.""" def test_get_type_map(self) -> None: diff --git a/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py b/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py index 7b579ae82c..7aa94e33ee 100644 --- a/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py +++ b/source/tests/universal/dpmodel/atomc_model/test_atomic_model.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import unittest +import numpy as np + from deepmd.dpmodel.atomic_model import ( DPAtomicModel, DPZBLLinearEnergyAtomicModel, @@ -72,6 +74,13 @@ ) +def make_sel_type_from_atom_exclude_types(type_map, atom_exclude_types): + """Get sel_type from complement of atom_exclude_types.""" + full_type_list = np.arange(len(type_map), dtype=int) + sel_type = np.setdiff1d(full_type_list, atom_exclude_types, assume_unique=True) + return sel_type.tolist() + + @parameterized( des_parameterized=( ( @@ -85,6 +94,7 @@ (DescriptorParamHybridMixedTTebd, DescrptHybrid), ), # descrpt_class_param & class ((FittingParamEnergy, EnergyFittingNet),), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), fit_parameterized=( ( @@ -97,6 +107,7 @@ ( *[(param_func, EnergyFittingNet) for param_func in FittingParamEnergyList], ), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") @@ -128,9 +139,7 @@ def setUpClass(cls) -> None: **cls.input_dict_ft, ) cls.module = DPAtomicModel( - ds, - ft, - type_map=cls.expected_type_map, + ds, ft, type_map=cls.expected_type_map, atom_exclude_types=cls.param[2] ) cls.output_def = cls.module.atomic_output_def().get_data() cls.expected_has_message_passing = ds.has_message_passing() @@ -138,6 +147,14 @@ def setUpClass(cls) -> None: cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + def test_sel_type_from_atom_exclude_types(self): + self.assertEqual( + make_sel_type_from_atom_exclude_types( + self.expected_type_map, self.param[2] + ), + self.expected_sel_type, + ) + @parameterized( des_parameterized=( @@ -152,6 +169,7 @@ def setUpClass(cls) -> None: (DescriptorParamHybridMixedTTebd, DescrptHybrid), ), # descrpt_class_param & class ((FittingParamDos, DOSFittingNet),), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), fit_parameterized=( ( @@ -164,6 +182,7 @@ def setUpClass(cls) -> None: ( *[(param_func, DOSFittingNet) for param_func in FittingParamDosList], ), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") @@ -195,9 +214,7 @@ def setUpClass(cls) -> None: **cls.input_dict_ft, ) cls.module = DPAtomicModel( - ds, - ft, - type_map=cls.expected_type_map, + ds, ft, type_map=cls.expected_type_map, atom_exclude_types=cls.param[2] ) cls.output_def = cls.module.atomic_output_def().get_data() cls.expected_has_message_passing = ds.has_message_passing() @@ -205,6 +222,14 @@ def setUpClass(cls) -> None: cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + def test_sel_type_from_atom_exclude_types(self): + self.assertEqual( + make_sel_type_from_atom_exclude_types( + self.expected_type_map, self.param[2] + ), + self.expected_sel_type, + ) + @parameterized( des_parameterized=( @@ -216,6 +241,7 @@ def setUpClass(cls) -> None: (DescriptorParamHybridMixed, DescrptHybrid), ), # descrpt_class_param & class ((FittingParamDipole, DipoleFitting),), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), fit_parameterized=( ( @@ -226,6 +252,7 @@ def setUpClass(cls) -> None: ( *[(param_func, DipoleFitting) for param_func in FittingParamDipoleList], ), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") @@ -258,9 +285,7 @@ def setUpClass(cls) -> None: **cls.input_dict_ft, ) cls.module = DPAtomicModel( - ds, - ft, - type_map=cls.expected_type_map, + ds, ft, type_map=cls.expected_type_map, atom_exclude_types=cls.param[2] ) cls.output_def = cls.module.atomic_output_def().get_data() cls.expected_has_message_passing = ds.has_message_passing() @@ -268,6 +293,14 @@ def setUpClass(cls) -> None: cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + def test_sel_type_from_atom_exclude_types(self): + self.assertEqual( + make_sel_type_from_atom_exclude_types( + self.expected_type_map, self.param[2] + ), + self.expected_sel_type, + ) + @parameterized( des_parameterized=( @@ -279,6 +312,7 @@ def setUpClass(cls) -> None: (DescriptorParamHybridMixed, DescrptHybrid), ), # descrpt_class_param & class ((FittingParamPolar, PolarFitting),), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), fit_parameterized=( ( @@ -289,6 +323,7 @@ def setUpClass(cls) -> None: ( *[(param_func, PolarFitting) for param_func in FittingParamPolarList], ), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") @@ -321,9 +356,7 @@ def setUpClass(cls) -> None: **cls.input_dict_ft, ) cls.module = DPAtomicModel( - ds, - ft, - type_map=cls.expected_type_map, + ds, ft, type_map=cls.expected_type_map, atom_exclude_types=cls.param[2] ) cls.output_def = cls.module.atomic_output_def().get_data() cls.expected_has_message_passing = ds.has_message_passing() @@ -331,6 +364,14 @@ def setUpClass(cls) -> None: cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + def test_sel_type_from_atom_exclude_types(self): + self.assertEqual( + make_sel_type_from_atom_exclude_types( + self.expected_type_map, self.param[2] + ), + self.expected_sel_type, + ) + @parameterized( des_parameterized=( @@ -415,6 +456,7 @@ def setUpClass(cls) -> None: (DescriptorParamHybridMixedTTebd, DescrptHybrid), ), # descrpt_class_param & class ((FittingParamProperty, PropertyFittingNet),), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), fit_parameterized=( ( @@ -428,6 +470,7 @@ def setUpClass(cls) -> None: for param_func in FittingParamPropertyList ], ), # fitting_class_param & class + ([], [0]), # atom_exclude_types ), ) @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") @@ -460,12 +503,18 @@ def setUpClass(cls) -> None: **cls.input_dict_ft, ) cls.module = DPAtomicModel( - ds, - ft, - type_map=cls.expected_type_map, + ds, ft, type_map=cls.expected_type_map, atom_exclude_types=cls.param[2] ) cls.output_def = cls.module.atomic_output_def().get_data() cls.expected_has_message_passing = ds.has_message_passing() cls.expected_sel_type = ft.get_sel_type() cls.expected_dim_fparam = ft.get_dim_fparam() cls.expected_dim_aparam = ft.get_dim_aparam() + + def test_sel_type_from_atom_exclude_types(self): + self.assertEqual( + make_sel_type_from_atom_exclude_types( + self.expected_type_map, self.param[2] + ), + self.expected_sel_type, + ) diff --git a/source/tests/universal/dpmodel/fitting/test_fitting.py b/source/tests/universal/dpmodel/fitting/test_fitting.py index 90b0668d20..29c5fcd4da 100644 --- a/source/tests/universal/dpmodel/fitting/test_fitting.py +++ b/source/tests/universal/dpmodel/fitting/test_fitting.py @@ -52,6 +52,7 @@ def FittingParamEnergy( "numb_fparam": numb_param, "numb_aparam": numb_param, "dim_case_embd": numb_param, + "default_fparam": [1.0] * numb_param if numb_param > 0 else None, } return input_dict