diff --git a/.github/labeler.yml b/.github/labeler.yml index 195d2cd217..b0a85679de 100644 --- a/.github/labeler.yml +++ b/.github/labeler.yml @@ -1,15 +1,39 @@ Python: -- deepmd/**/* -- deepmd_cli/**/* -- source/tests/**/* -Docs: doc/**/* -Examples: examples/**/* -Core: source/lib/**/* -CUDA: source/lib/src/gpu/**/* -ROCM: source/lib/src/gpu/**/* -OP: source/op/**/* -C++: source/api_cc/**/* -C: source/api_c/**/* -LAMMPS: source/lmp/**/* -Gromacs: source/gmx/**/* -i-Pi: source/ipi/**/* +- changed-files: + - any-glob-to-any-file: + - deepmd/**/* + - deepmd_utils/**/* + - source/tests/**/* +Docs: +- changed-files: + - any-glob-to-any-file: doc/**/* +Examples: +- changed-files: + - any-glob-to-any-file: examples/**/* +Core: +- changed-files: + - any-glob-to-any-file: source/lib/**/* +CUDA: +- changed-files: + - any-glob-to-any-file: source/lib/src/gpu/**/* +ROCM: +- changed-files: + - any-glob-to-any-file: source/lib/src/gpu/**/* +OP: +- changed-files: + - any-glob-to-any-file: source/op/**/* +C++: +- changed-files: + - any-glob-to-any-file: source/api_cc/**/* +C: +- changed-files: + - any-glob-to-any-file: source/api_c/**/* +LAMMPS: +- changed-files: + - any-glob-to-any-file: source/lmp/**/* +Gromacs: +- changed-files: + - any-glob-to-any-file: source/gmx/**/* +i-Pi: +- changed-files: + - any-glob-to-any-file: source/ipi/**/* diff --git a/.github/release.yml b/.github/release.yml new file mode 100644 index 0000000000..382e5db00e --- /dev/null +++ b/.github/release.yml @@ -0,0 +1,34 @@ +changelog: + exclude: + authors: + - app/pre-commit-ci + - app/dependabot + categories: + - title: Breaking Changes + labels: + - "breaking change" + - title: New Features + labels: + - "new feature" + - title: Enhancement + labels: + - enhancement + - title: Documentation + labels: + # automatically added + - Docs + # for docs outside the doc directory + - "other docs" + exclude: + labels: + - build + - bug + - title: Build and release + labels: + - build + - title: Bug fixings + labels: + - bug + - title: Other Changes + labels: + - "*" diff --git a/.github/workflows/build_cc.yml b/.github/workflows/build_cc.yml index 964a11ce37..f029517d80 100644 --- a/.github/workflows/build_cc.yml +++ b/.github/workflows/build_cc.yml @@ -21,7 +21,7 @@ jobs: dp_variant: clang steps: - uses: actions/checkout@v4 - - uses: actions/setup-python@v4 + - uses: actions/setup-python@v5 with: python-version: '3.11' cache: 'pip' @@ -37,7 +37,7 @@ jobs: 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-12-0 cuda-nvcc-12-0 + && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2 if: matrix.variant == 'cuda120' env: DEBIAN_FRONTEND: noninteractive diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml index 84c8ac4b74..23076e9bf5 100644 --- a/.github/workflows/build_wheel.yml +++ b/.github/workflows/build_wheel.yml @@ -33,6 +33,13 @@ 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 # macos-x86-64 - os: macos-latest python: 311 @@ -68,8 +75,11 @@ jobs: CIBW_ARCHS: all CIBW_BUILD: cp${{ matrix.python }}-${{ matrix.platform_id }} DP_VARIANT: ${{ matrix.dp_variant }} - - uses: actions/upload-artifact@v3 + CUDA_VERSION: ${{ matrix.cuda_version }} + DP_PKG_NAME: ${{ matrix.dp_pkg_name }} + - uses: actions/upload-artifact@v4 with: + name: cibw-cp${{ matrix.python }}-${{ matrix.platform_id }}-cu${{ matrix.cuda_version }}-${{ strategy.job-index }} path: ./wheelhouse/*.whl build_sdist: name: Build source distribution @@ -78,7 +88,7 @@ jobs: - uses: actions/checkout@v4 with: fetch-depth: 0 - - uses: actions/setup-python@v4 + - uses: actions/setup-python@v5 name: Install Python with: python-version: '3.11' @@ -87,8 +97,9 @@ jobs: - name: Build sdist run: python -m build --sdist - - uses: actions/upload-artifact@v3 + - uses: actions/upload-artifact@v4 with: + name: cibw-sdist path: dist/*.tar.gz upload_pypi: @@ -99,22 +110,31 @@ jobs: id-token: write if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags/v') steps: - - uses: actions/download-artifact@v3 + - uses: actions/download-artifact@v4 with: - name: artifact + pattern: cibw-* path: dist + merge-multiple: true - uses: pypa/gh-action-pypi-publish@release/v1 build_docker: # use the already built wheels to build docker needs: [build_wheels] runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - variant: "" + cuda_version: "12" + - variant: "_cu11" + cuda_version: "11" steps: - uses: actions/checkout@v4 - - uses: actions/download-artifact@v3 + - uses: actions/download-artifact@v4 with: - name: artifact path: source/install/docker/dist + merge-multiple: true - name: Log in to the Container registry uses: docker/login-action@343f7c4344506bcbf9b4de18042ae17996df046d with: @@ -124,27 +144,30 @@ jobs: - name: Extract metadata (tags, labels) for Docker id: meta - uses: docker/metadata-action@96383f45573cb7f253c731d3b3ab81c87ef81934 + uses: docker/metadata-action@dbef88086f6cef02e264edb7dbf63250c17cef6c with: images: ghcr.io/deepmodeling/deepmd-kit - name: Build and push Docker image - uses: docker/build-push-action@0565240e2d4ab88bba5387d719585280857ece09 + uses: docker/build-push-action@4a13e500e55cf31b7a5d59a38ab2040ab0f42f56 with: context: source/install/docker - push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' }} - tags: ${{ steps.meta.outputs.tags }} + push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' && github.actor != 'dependabot[bot]' }} + tags: ${{ steps.meta.outputs.tags }}${{ matrix.variant }} labels: ${{ steps.meta.outputs.labels }} + build-args: | + VARIANT=${{ matrix.variant }} + CUDA_VERSION=${{ matrix.cuda_version }} build_pypi_index: needs: [build_wheels, build_sdist] runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@v3 + - uses: actions/download-artifact@v4 with: - name: artifact path: dist/packages - - uses: actions/setup-python@v4 + merge-multiple: true + - uses: actions/setup-python@v5 name: Install Python with: python-version: '3.11' @@ -153,7 +176,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@v2 + uses: actions/upload-pages-artifact@v3 with: path: dist deploy_pypi_index: @@ -169,11 +192,11 @@ jobs: steps: - name: Deploy to GitHub Pages id: deployment - uses: actions/deploy-pages@v2 + uses: actions/deploy-pages@v4 pass: name: Pass testing build wheels - needs: [build_wheels, build_sdist] + needs: [build_wheels, build_sdist, build_docker, build_pypi_index] runs-on: ubuntu-latest if: always() steps: diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml new file mode 100644 index 0000000000..a9a162432c --- /dev/null +++ b/.github/workflows/codeql.yml @@ -0,0 +1,58 @@ +name: "CodeQL" + +on: + push: + pull_request: + schedule: + - cron: '45 2 * * 2' + +jobs: + analyze: + name: Analyze + runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }} + timeout-minutes: ${{ (matrix.language == 'swift' && 120) || 360 }} + permissions: + actions: read + contents: read + security-events: write + + strategy: + fail-fast: false + matrix: + language: [ 'c-cpp', 'javascript-typescript', 'python' ] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: 'pip' + if: matrix.language == 'c-cpp' + - name: "Setup dependencies" + if: matrix.language == 'c-cpp' + 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-12-2 cuda-nvcc-12-2 + python -m pip install tensorflow + env: + DEBIAN_FRONTEND: noninteractive + # Initializes the CodeQL tools for scanning. + - name: Initialize CodeQL + uses: github/codeql-action/init@v3 + with: + languages: ${{ matrix.language }} + queries: security-extended,security-and-quality + - name: "Run, Build Application using script" + run: source/install/build_cc.sh + env: + DP_VARIANT: cuda + DOWNLOAD_TENSORFLOW: "FALSE" + if: matrix.language == 'c-cpp' + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@v3 + with: + category: "/language:${{matrix.language}}" diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml index 2c8ba30ba1..877c780f1f 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@v4 + - uses: actions/labeler@v5 with: repo-token: "${{ secrets.GITHUB_TOKEN }}" diff --git a/.github/workflows/package_c.yml b/.github/workflows/package_c.yml index ada205be00..5594c79181 100644 --- a/.github/workflows/package_c.yml +++ b/.github/workflows/package_c.yml @@ -8,23 +8,37 @@ jobs: build_c: name: Build C library runs-on: ubuntu-22.04 + strategy: + matrix: + include: + - tensorflow_build_version: "2.15" + 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: Package C library run: ./source/install/docker_package_c.sh + env: + TENSORFLOW_VERSION: ${{ matrix.tensorflow_version }} + TENSORFLOW_BUILD_VERSION: ${{ matrix.tensorflow_build_version }} + - run: cp libdeepmd_c.tar.gz ${{ matrix.filename }} + if: matrix.filename != 'libdeepmd_c.tar.gz' # for download and debug - name: Upload artifact - uses: actions/upload-artifact@v3 + uses: actions/upload-artifact@v4 with: - name: libdeepmd_c - path: ./libdeepmd_c.tar.gz + name: libdeepmd_c-${{ strategy.job-index }}-${{ matrix.filename }} + path: ${{ matrix.filename }} - name: Test C library run: ./source/install/docker_test_package_c.sh - name: Release uses: softprops/action-gh-release@v1 if: startsWith(github.ref, 'refs/tags/') with: - files: libdeepmd_c.tar.gz + files: ${{ matrix.filename }} test_c: name: Test building from C library needs: [build_c] @@ -32,9 +46,10 @@ jobs: steps: - uses: actions/checkout@v4 - name: Download artifact - uses: actions/download-artifact@v3 + uses: actions/download-artifact@v4 with: - name: libdeepmd_c + pattern: libdeepmd_c-* + merge-multiple: true - run: tar -vxzf ./libdeepmd_c.tar.gz - name: Test C library run: ./source/install/build_from_c.sh diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index a98afa7a94..ef6fade8e5 100644 --- a/.github/workflows/test_cc.yml +++ b/.github/workflows/test_cc.yml @@ -8,7 +8,7 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - - uses: actions/setup-python@v4 + - uses: actions/setup-python@v5 with: python-version: '3.11' cache: 'pip' diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index ca72a32277..e74c0abde2 100644 --- a/.github/workflows/test_cuda.yml +++ b/.github/workflows/test_cuda.yml @@ -11,14 +11,14 @@ jobs: runs-on: nvidia # https://github.com/deepmodeling/deepmd-kit/pull/2884#issuecomment-1744216845 container: - image: nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04 + image: nvidia/cuda:12.2.0-devel-ubuntu22.04 options: --gpus all if: github.repository_owner == 'deepmodeling' && github.event.label.name == 'Test CUDA' || github.event_name == 'workflow_dispatch' 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@v4 + - uses: actions/setup-python@v5 with: python-version: '3.11' # cache: 'pip' @@ -31,18 +31,19 @@ jobs: 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-11-8 libcudnn8=8.9.5.*-1+cuda11.8 + && sudo apt-get -y install cuda-12-2 libcudnn8=8.9.5.*-1+cuda12.2 if: false # skip as we use nvidia image - name: Set PyPI mirror for Aliyun cloud machine run: python -m pip config --user set global.index-url https://mirrors.aliyun.com/pypi/simple/ - run: python -m pip install -U "pip>=21.3.1,!=23.0.0" - - run: python -m pip install -v -e .[gpu,test,lmp,cu11] "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz" + - run: python -m pip install "tensorflow>=2.15.0rc0" + - run: python -m pip install -v -e .[gpu,test,lmp,cu12] "ase @ https://gitlab.com/ase/ase/-/archive/8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f/ase-8c5aa5fd6448c5cfb517a014dccf2b214a9dfa8f.tar.gz" env: DP_BUILD_TESTING: 1 DP_VARIANT: cuda - CUDA_PATH: /usr/local/cuda-11.8 + CUDA_PATH: /usr/local/cuda-12.2 - run: dp --version - - run: python -m pytest -s --cov=deepmd --cov=deepmd_cli source/tests --durations=0 + - run: python -m pytest -s --cov=deepmd --cov=deepmd_utils source/tests --durations=0 - run: source/install/test_cc_local.sh env: OMP_NUM_THREADS: 1 @@ -52,7 +53,7 @@ jobs: CMAKE_GENERATOR: Ninja DP_VARIANT: cuda DP_USE_MPICH2: 1 - CUDA_PATH: /usr/local/cuda-11.8 + CUDA_PATH: /usr/local/cuda-12.2 - run: | export LD_LIBRARY_PATH=$GITHUB_WORKSPACE/dp_test/lib:$CUDA_PATH/lib64:$LD_LIBRARY_PATH export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH @@ -63,7 +64,7 @@ jobs: TF_INTRA_OP_PARALLELISM_THREADS: 1 TF_INTER_OP_PARALLELISM_THREADS: 1 LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp - CUDA_PATH: /usr/local/cuda-11.8 + CUDA_PATH: /usr/local/cuda-12.2 - uses: codecov/codecov-action@v3 with: gcov: true diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 0ac29a7d9b..1bd78bfae0 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -18,7 +18,7 @@ jobs: steps: - uses: actions/checkout@v4 - - uses: actions/setup-python@v4 + - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python }} cache: 'pip' @@ -38,7 +38,7 @@ jobs: HOROVOD_WITH_TENSORFLOW: 1 HOROVOD_WITHOUT_GLOO: 1 - run: dp --version - - run: pytest --cov=deepmd --cov=deepmd_cli source/tests --durations=0 + - run: pytest --cov=deepmd --cov=deepmd_utils source/tests --durations=0 - uses: codecov/codecov-action@v3 with: gcov: true diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e168af2c8d..d4e89f1129 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -23,23 +23,22 @@ repos: - id: check-toml # Python - repo: https://github.com/PyCQA/isort - rev: 5.12.0 + rev: 5.13.2 hooks: - id: isort files: \.py$ exclude: ^source/3rdparty - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.1.1 + rev: v0.1.13 hooks: - id: ruff args: ["--fix"] exclude: ^source/3rdparty -- repo: https://github.com/psf/black-pre-commit-mirror - rev: 23.10.0 - hooks: - - id: black-jupyter - exclude: ^source/3rdparty + types_or: [python, pyi, jupyter] + - id: ruff-format + exclude: ^source/3rdparty + types_or: [python, pyi, jupyter] # numpydoc - repo: https://github.com/Carreau/velin rev: 0.0.12 @@ -54,7 +53,7 @@ repos: - id: blacken-docs # C++ - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v17.0.3 + rev: v17.0.6 hooks: - id: clang-format exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc @@ -65,7 +64,7 @@ repos: - id: csslint # Shell - repo: https://github.com/scop/pre-commit-shfmt - rev: v3.7.0-1 + rev: v3.7.0-4 hooks: - id: shfmt # CMake diff --git a/README.md b/README.md index 5914abe607..81fdead098 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,7 @@ A full [document](doc/train/train-input-auto.rst) on options in the training inp - [Deep potential long-range](doc/model/dplr.md) - [Deep Potential - Range Correction (DPRc)](doc/model/dprc.md) - [Linear model](doc/model/linear.md) + - [Interpolation or combination with a pairwise potential](doc/model/pairtab.md) - [Training](doc/train/index.md) - [Training a model](doc/train/training.md) - [Advanced options](doc/train/training-advanced.md) @@ -134,8 +135,7 @@ A full [document](doc/train/train-input-auto.rst) on options in the training inp - [Node.js interface](doc/inference/nodejs.md) - [Integrate with third-party packages](doc/third-party/index.rst) - [Use deep potential with ASE](doc/third-party/ase.md) - - [Run MD with LAMMPS](doc/third-party/lammps.md) - - [LAMMPS commands](doc/third-party/lammps-command.md) + - [Run MD with LAMMPS](doc/third-party/lammps-command.md) - [Run path-integral MD with i-PI](doc/third-party/ipi.md) - [Run MD with GROMACS](doc/third-party/gromacs.md) - [Interfaces out of DeePMD-kit](doc/third-party/out-of-deepmd-kit.md) diff --git a/backend/dynamic_metadata.py b/backend/dynamic_metadata.py index 0502684f47..ab955c3cf8 100644 --- a/backend/dynamic_metadata.py +++ b/backend/dynamic_metadata.py @@ -27,7 +27,7 @@ def dynamic_metadata( _, _, find_libpython_requires, extra_scripts, tf_version = get_argument_from_env() if field == "scripts": return { - "dp": "deepmd_cli.main:main", + "dp": "deepmd_utils.main:main", **extra_scripts, } elif field == "optional-dependencies": @@ -44,7 +44,8 @@ def dynamic_metadata( "sphinx>=3.1.1", "sphinx_rtd_theme>=1.0.0rc1", "sphinx_markdown_tables", - "myst-nb", + "myst-nb>=1.0.0rc0", + "myst-parser>=0.19.2", "breathe", "exhale", "numpydoc", @@ -56,7 +57,7 @@ def dynamic_metadata( "sphinxcontrib-bibtex", ], "lmp": [ - "lammps~=2023.8.2.1.0", + "lammps~=2023.8.2.2.0", *find_libpython_requires, ], "ipi": [ diff --git a/backend/find_tensorflow.py b/backend/find_tensorflow.py index aa75d5ecb4..08a73f7252 100644 --- a/backend/find_tensorflow.py +++ b/backend/find_tensorflow.py @@ -87,6 +87,24 @@ def find_tensorflow() -> Tuple[Optional[str], List[str]]: # TypeError if submodule_search_locations are None # IndexError if submodule_search_locations is an empty list except (AttributeError, TypeError, IndexError): + 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 + requires.extend( + [ + "tensorflow-cpu>=2.15.0rc0; platform_machine=='x86_64' and platform_system == 'Linux'", + ] + ) + elif cuda_version in SpecifierSet(">=11,<12"): + # CUDA 11.8 + requires.extend( + [ + "tensorflow-cpu>=2.5.0rc0,<2.15; platform_machine=='x86_64' and platform_system == 'Linux'", + ] + ) + else: + raise RuntimeError("Unsupported CUDA version") requires.extend(get_tf_requirement()["cpu"]) # setuptools will re-find tensorflow after installing setup_requires tf_install_dir = None @@ -114,9 +132,9 @@ def get_tf_requirement(tf_version: str = "") -> dict: extra_requires = [] extra_select = {} - if not (tf_version == "" or tf_version in SpecifierSet(">=2.12")): + if not (tf_version == "" or tf_version in SpecifierSet(">=2.12", prereleases=True)): extra_requires.append("protobuf<3.20") - if tf_version == "" or tf_version in SpecifierSet(">=1.15"): + if tf_version == "" or tf_version in SpecifierSet(">=1.15", prereleases=True): extra_select["mpi"] = [ "horovod", "mpi4py", @@ -129,6 +147,8 @@ def get_tf_requirement(tf_version: str = "") -> dict: "cpu": [ "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'", *extra_requires, ], "gpu": [ @@ -138,9 +158,9 @@ def get_tf_requirement(tf_version: str = "") -> dict: ], **extra_select, } - elif tf_version in SpecifierSet("<1.15") or tf_version in SpecifierSet( - ">=2.0,<2.1" - ): + elif tf_version in SpecifierSet( + "<1.15", prereleases=True + ) or tf_version in SpecifierSet(">=2.0,<2.1", prereleases=True): return { "cpu": [ f"tensorflow=={tf_version}", diff --git a/codecov.yml b/codecov.yml index 24dd9e3a23..3654859423 100644 --- a/codecov.yml +++ b/codecov.yml @@ -20,7 +20,7 @@ component_management: name: Python paths: - deepmd/** - - deepmd_cli/** + - deepmd_utils/** - component_id: module_op name: OP paths: diff --git a/deepmd/__init__.py b/deepmd/__init__.py index b02817b6fc..0190bbc124 100644 --- a/deepmd/__init__.py +++ b/deepmd/__init__.py @@ -32,7 +32,7 @@ set_mkl() try: - from deepmd_cli._version import version as __version__ + from deepmd_utils._version import version as __version__ except ImportError: from .__about__ import ( __version__, diff --git a/deepmd/calculator.py b/deepmd/calculator.py index 8636ff30d2..b9c0a81006 100644 --- a/deepmd/calculator.py +++ b/deepmd/calculator.py @@ -45,6 +45,8 @@ class DP(Calculator): type_dict : Dict[str, int], optional mapping of element types and their numbers, best left None and the calculator will infer this information from model, by default None + neighbor_list : ase.neighborlist.NeighborList, optional + The neighbor list object. If None, then build the native neighbor list. Examples -------- @@ -83,10 +85,11 @@ def __init__( model: Union[str, "Path"], label: str = "DP", type_dict: Optional[Dict[str, int]] = None, + neighbor_list=None, **kwargs, ) -> None: Calculator.__init__(self, label=label, **kwargs) - self.dp = DeepPotential(str(Path(model).resolve())) + self.dp = DeepPotential(str(Path(model).resolve()), neighbor_list=neighbor_list) if type_dict: self.type_dict = type_dict else: diff --git a/deepmd/common.py b/deepmd/common.py index 472508bb08..54e3d0a6f8 100644 --- a/deepmd/common.py +++ b/deepmd/common.py @@ -1,53 +1,65 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Collection of functions and classes used throughout the whole package.""" -import json import warnings from functools import ( wraps, ) -from pathlib import ( - Path, -) from typing import ( TYPE_CHECKING, Any, Callable, - Dict, - List, - Optional, - TypeVar, Union, ) -import numpy as np import tensorflow -import yaml from tensorflow.python.framework import ( tensor_util, ) from deepmd.env import ( - GLOBAL_NP_FLOAT_PRECISION, GLOBAL_TF_FLOAT_PRECISION, op_module, tf, ) -from deepmd.utils.path import ( - DPPath, +from deepmd_utils.common import ( + add_data_requirement, + data_requirement, + expand_sys_str, + get_np_precision, + j_loader, + j_must_have, + make_default_mesh, + select_idx_map, ) if TYPE_CHECKING: - _DICT_VAL = TypeVar("_DICT_VAL") - _OBJ = TypeVar("_OBJ") - try: - from typing import Literal # python >3.6 - except ImportError: - from typing_extensions import Literal # type: ignore - _ACTIVATION = Literal[ - "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu", "gelu_tf" - ] - _PRECISION = Literal["default", "float16", "float32", "float64"] + from deepmd_utils.common import ( + _ACTIVATION, + _PRECISION, + ) + +__all__ = [ + # from deepmd_utils.common + "data_requirement", + "add_data_requirement", + "select_idx_map", + "make_default_mesh", + "j_must_have", + "j_loader", + "expand_sys_str", + "get_np_precision", + # from self + "PRECISION_DICT", + "gelu", + "gelu_tf", + "ACTIVATION_FN_DICT", + "get_activation_func", + "get_precision", + "safe_cast_tensor", + "cast_precision", + "clear_session", +] # define constants PRECISION_DICT = { @@ -115,10 +127,6 @@ def gelu_wrapper(x): return (lambda x: gelu_wrapper(x))(x) -# TODO this is not a good way to do things. This is some global variable to which -# TODO anyone can write and there is no good way to keep track of the changes -data_requirement = {} - ACTIVATION_FN_DICT = { "relu": tf.nn.relu, "relu6": tf.nn.relu6, @@ -132,164 +140,6 @@ def gelu_wrapper(x): } -def add_data_requirement( - key: str, - ndof: int, - atomic: bool = False, - must: bool = False, - high_prec: bool = False, - type_sel: Optional[bool] = None, - repeat: int = 1, - default: float = 0.0, - dtype: Optional[np.dtype] = None, -): - """Specify data requirements for training. - - Parameters - ---------- - key : str - type of data stored in corresponding `*.npy` file e.g. `forces` or `energy` - ndof : int - number of the degrees of freedom, this is tied to `atomic` parameter e.g. forces - have `atomic=True` and `ndof=3` - atomic : bool, optional - specifies whwther the `ndof` keyworrd applies to per atom quantity or not, - by default False - must : bool, optional - specifi if the `*.npy` data file must exist, by default False - high_prec : bool, optional - if true load data to `np.float64` else `np.float32`, by default False - type_sel : bool, optional - select only certain type of atoms, by default None - repeat : int, optional - if specify repaeat data `repeat` times, by default 1 - default : float, optional, default=0. - default value of data - dtype : np.dtype, optional - the dtype of data, overwrites `high_prec` if provided - """ - data_requirement[key] = { - "ndof": ndof, - "atomic": atomic, - "must": must, - "high_prec": high_prec, - "type_sel": type_sel, - "repeat": repeat, - "default": default, - "dtype": dtype, - } - - -def select_idx_map(atom_types: np.ndarray, select_types: np.ndarray) -> np.ndarray: - """Build map of indices for element supplied element types from all atoms list. - - Parameters - ---------- - atom_types : np.ndarray - array specifing type for each atoms as integer - select_types : np.ndarray - types of atoms you want to find indices for - - Returns - ------- - np.ndarray - indices of types of atoms defined by `select_types` in `atom_types` array - - Warnings - -------- - `select_types` array will be sorted before finding indices in `atom_types` - """ - sort_select_types = np.sort(select_types) - idx_map = [] - for ii in sort_select_types: - idx_map.append(np.where(atom_types == ii)[0]) - return np.concatenate(idx_map) - - -def make_default_mesh(pbc: bool, mixed_type: bool) -> np.ndarray: - """Make mesh. - - Only the size of mesh matters, not the values: - * 6 for PBC, no mixed types - * 0 for no PBC, no mixed types - * 7 for PBC, mixed types - * 1 for no PBC, mixed types - - Parameters - ---------- - pbc : bool - if True, the mesh will be made for periodic boundary conditions - mixed_type : bool - if True, the mesh will be made for mixed types - - Returns - ------- - np.ndarray - mesh - """ - mesh_size = int(pbc) * 6 + int(mixed_type) - default_mesh = np.zeros(mesh_size, dtype=np.int32) - return default_mesh - - -# TODO maybe rename this to j_deprecated and only warn about deprecated keys, -# TODO if the deprecated_key argument is left empty function puppose is only custom -# TODO error since dict[key] already raises KeyError when the key is missing -def j_must_have( - jdata: Dict[str, "_DICT_VAL"], key: str, deprecated_key: List[str] = [] -) -> "_DICT_VAL": - """Assert that supplied dictionary conaines specified key. - - Returns - ------- - _DICT_VAL - value that was store unde supplied key - - Raises - ------ - RuntimeError - if the key is not present - """ - if key not in jdata.keys(): - for ii in deprecated_key: - if ii in jdata.keys(): - warnings.warn(f"the key {ii} is deprecated, please use {key} instead") - return jdata[ii] - else: - raise RuntimeError(f"json database must provide key {key}") - else: - return jdata[key] - - -def j_loader(filename: Union[str, Path]) -> Dict[str, Any]: - """Load yaml or json settings file. - - Parameters - ---------- - filename : Union[str, Path] - path to file - - Returns - ------- - Dict[str, Any] - loaded dictionary - - Raises - ------ - TypeError - if the supplied file is of unsupported type - """ - filepath = Path(filename) - if filepath.suffix.endswith("json"): - with filepath.open() as fp: - return json.load(fp) - elif filepath.suffix.endswith(("yml", "yaml")): - with filepath.open() as fp: - return yaml.safe_load(fp) - else: - raise TypeError("config file must be json, or yaml/yml") - - def get_activation_func( activation_fn: Union["_ACTIVATION", None], ) -> Union[Callable[[tf.Tensor], tf.Tensor], None]: @@ -340,57 +190,6 @@ def get_precision(precision: "_PRECISION") -> Any: return PRECISION_DICT[precision] -# TODO port completely to pathlib when all callers are ported -def expand_sys_str(root_dir: Union[str, Path]) -> List[str]: - """Recursively iterate over directories taking those that contain `type.raw` file. - - Parameters - ---------- - root_dir : Union[str, Path] - starting directory - - Returns - ------- - List[str] - list of string pointing to system directories - """ - root_dir = DPPath(root_dir) - matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()] - if (root_dir / "type.raw").is_file(): - matches.append(str(root_dir)) - return matches - - -def get_np_precision(precision: "_PRECISION") -> np.dtype: - """Get numpy precision constant from string. - - Parameters - ---------- - precision : _PRECISION - string name of numpy constant or default - - Returns - ------- - np.dtype - numpy presicion constant - - Raises - ------ - 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") - - def safe_cast_tensor( input: tf.Tensor, from_precision: tf.DType, to_precision: tf.DType ) -> tf.Tensor: diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 2de0b63245..721bb0d534 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -41,6 +41,8 @@ GraphWithoutTensorError, ) from deepmd.utils.graph import ( + get_extra_embedding_net_suffix, + get_extra_embedding_net_variables_from_graph_def, get_pattern_nodes_from_graph_def, get_tensor_by_name_from_graph, ) @@ -204,7 +206,7 @@ def __init__( self.type_one_side = type_one_side self.spin = spin self.stripped_type_embedding = stripped_type_embedding - self.extra_embeeding_net_variables = None + self.extra_embedding_net_variables = None self.layer_size = len(neuron) # extend sel_a for spin system @@ -470,11 +472,13 @@ def enable_compression( ) if self.stripped_type_embedding: + one_side_suffix = get_extra_embedding_net_suffix(type_one_side=True) + two_side_suffix = get_extra_embedding_net_suffix(type_one_side=False) ret_two_side = get_pattern_nodes_from_graph_def( - graph_def, f"filter_type_all{suffix}/.+_two_side_ebd" + graph_def, f"filter_type_all{suffix}/.+{two_side_suffix}" ) ret_one_side = get_pattern_nodes_from_graph_def( - graph_def, f"filter_type_all{suffix}/.+_one_side_ebd" + graph_def, f"filter_type_all{suffix}/.+{one_side_suffix}" ) if len(ret_two_side) == 0 and len(ret_one_side) == 0: raise RuntimeError( @@ -487,19 +491,19 @@ def enable_compression( elif len(ret_two_side) != 0: self.final_type_embedding = get_two_side_type_embedding(self, graph) self.matrix = get_extra_side_embedding_net_variable( - self, graph_def, "two_side", "matrix", suffix + self, graph_def, two_side_suffix, "matrix", suffix ) self.bias = get_extra_side_embedding_net_variable( - self, graph_def, "two_side", "bias", suffix + self, graph_def, two_side_suffix, "bias", suffix ) self.extra_embedding = make_data(self, self.final_type_embedding) else: self.final_type_embedding = get_type_embedding(self, graph) self.matrix = get_extra_side_embedding_net_variable( - self, graph_def, "one_side", "matrix", suffix + self, graph_def, one_side_suffix, "matrix", suffix ) self.bias = get_extra_side_embedding_net_variable( - self, graph_def, "one_side", "bias", suffix + self, graph_def, one_side_suffix, "bias", suffix ) self.extra_embedding = make_data(self, self.final_type_embedding) @@ -778,16 +782,16 @@ def _pass_filter( type_i = -1 if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor: inputs_i = descrpt2r4(inputs_i, natoms) + self.atype_nloc = tf.reshape( + tf.slice(atype, [0, 0], [-1, natoms[0]]), [-1] + ) # when nloc != nall, pass nloc to mask if len(self.exclude_types): - atype_nloc = tf.reshape( - tf.slice(atype, [0, 0], [-1, natoms[0]]), [-1] - ) # when nloc != nall, pass nloc to mask mask = self.build_type_exclude_mask( self.exclude_types, self.ntypes, self.sel_a, self.ndescrpt, - atype_nloc, + self.atype_nloc, tf.shape(inputs_i)[0], ) inputs_i *= mask @@ -952,7 +956,7 @@ def _filter_lower( extra_embedding_index = self.nei_type_vec else: padding_ntypes = type_embedding.shape[0] - atype_expand = tf.reshape(self.atype, [-1, 1]) + atype_expand = tf.reshape(self.atype_nloc, [-1, 1]) idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei]) idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei]) idx = idx_i + idx_j @@ -961,20 +965,21 @@ def _filter_lower( if not self.compress: if self.type_one_side: - one_side_type_embedding_suffix = "_one_side_ebd" net_output = embedding_net( type_embedding, self.filter_neuron, self.filter_precision, activation_fn=activation_fn, resnet_dt=self.filter_resnet_dt, - name_suffix=one_side_type_embedding_suffix, + name_suffix=get_extra_embedding_net_suffix( + self.type_one_side + ), stddev=stddev, bavg=bavg, seed=self.seed, trainable=trainable, uniform_seed=self.uniform_seed, - initial_variables=self.extra_embeeding_net_variables, + initial_variables=self.extra_embedding_net_variables, mixed_prec=self.mixed_prec, ) net_output = tf.nn.embedding_lookup( @@ -997,27 +1002,21 @@ def _filter_lower( [-1, two_side_type_embedding.shape[-1]], ) - atype_expand = tf.reshape(self.atype, [-1, 1]) - idx_i = tf.tile(atype_expand * padding_ntypes, [1, self.nnei]) - idx_j = tf.reshape(self.nei_type_vec, [-1, self.nnei]) - idx = idx_i + idx_j - index_of_two_side = tf.reshape(idx, [-1]) - self.extra_embedding_index = index_of_two_side - - two_side_type_embedding_suffix = "_two_side_ebd" net_output = embedding_net( two_side_type_embedding, self.filter_neuron, self.filter_precision, activation_fn=activation_fn, resnet_dt=self.filter_resnet_dt, - name_suffix=two_side_type_embedding_suffix, + name_suffix=get_extra_embedding_net_suffix( + self.type_one_side + ), stddev=stddev, bavg=bavg, seed=self.seed, trainable=trainable, uniform_seed=self.uniform_seed, - initial_variables=self.extra_embeeding_net_variables, + initial_variables=self.extra_embedding_net_variables, mixed_prec=self.mixed_prec, ) net_output = tf.nn.embedding_lookup(net_output, idx) @@ -1327,6 +1326,15 @@ def init_variables( self.dstd = new_dstd if self.original_sel is None: self.original_sel = sel + if self.stripped_type_embedding: + self.extra_embedding_net_variables = ( + get_extra_embedding_net_variables_from_graph_def( + graph_def, + suffix, + get_extra_embedding_net_suffix(self.type_one_side), + self.layer_size, + ) + ) @property def explicit_ntypes(self) -> bool: diff --git a/deepmd/descriptor/se_a_mask.py b/deepmd/descriptor/se_a_mask.py index 780b34d294..cc2e6b4fc8 100644 --- a/deepmd/descriptor/se_a_mask.py +++ b/deepmd/descriptor/se_a_mask.py @@ -417,3 +417,16 @@ def prod_force_virial( atom_virial = tf.zeros([1, natoms[1], 9], dtype=force.dtype) return force, virial, atom_virial + + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + return local_jdata diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index 8e4c3c3ef6..1ceda23065 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -42,9 +42,10 @@ ) from deepmd.utils.graph import ( get_attention_layer_variables_from_graph_def, + get_extra_embedding_net_suffix, + get_extra_embedding_net_variables_from_graph_def, get_pattern_nodes_from_graph_def, get_tensor_by_name_from_graph, - get_tensor_by_type, ) from deepmd.utils.network import ( embedding_net, @@ -391,11 +392,12 @@ def enable_compression( raise RuntimeError("can not compress model when attention layer is not 0.") ret = get_pattern_nodes_from_graph_def( - graph_def, f"filter_type_all{suffix}/.+_two_side_ebd" + graph_def, + f"filter_type_all{suffix}/.+{get_extra_embedding_net_suffix(type_one_side=False)}", ) if len(ret) == 0: raise RuntimeError( - "can not find variables of embedding net `*_two_side_ebd` from graph_def, maybe it is not a compressible model." + f"can not find variables of embedding net `*{get_extra_embedding_net_suffix(type_one_side=False)}` from graph_def, maybe it is not a compressible model." ) self.compress = True @@ -420,11 +422,12 @@ def enable_compression( ) self.final_type_embedding = get_two_side_type_embedding(self, graph) + type_side_suffix = get_extra_embedding_net_suffix(type_one_side=False) self.matrix = get_extra_side_embedding_net_variable( - self, graph_def, "two_side", "matrix", suffix + self, graph_def, type_side_suffix, "matrix", suffix ) self.bias = get_extra_side_embedding_net_variable( - self, graph_def, "two_side", "bias", suffix + self, graph_def, type_side_suffix, "bias", suffix ) self.two_embd = make_data(self, self.final_type_embedding) @@ -1125,14 +1128,15 @@ def _filter_lower( two_side_type_embedding, [-1, two_side_type_embedding.shape[-1]], ) - two_side_type_embedding_suffix = "_two_side_ebd" embedding_of_two_side_type_embedding = embedding_net( two_side_type_embedding, self.filter_neuron, self.filter_precision, activation_fn=activation_fn, resnet_dt=self.filter_resnet_dt, - name_suffix=two_side_type_embedding_suffix, + name_suffix=get_extra_embedding_net_suffix( + type_one_side=False + ), stddev=stddev, bavg=bavg, seed=self.seed, @@ -1292,18 +1296,6 @@ def init_variables( """ super().init_variables(graph=graph, graph_def=graph_def, suffix=suffix) - if self.stripped_type_embedding: - self.two_side_embeeding_net_variables = {} - for i in range(1, self.layer_size + 1): - matrix_pattern = f"filter_type_all{suffix}/matrix_{i}_two_side_ebd" - self.two_side_embeeding_net_variables[ - matrix_pattern - ] = self._get_two_embed_variables(graph_def, matrix_pattern) - bias_pattern = f"filter_type_all{suffix}/bias_{i}_two_side_ebd" - self.two_side_embeeding_net_variables[ - bias_pattern - ] = self._get_two_embed_variables(graph_def, bias_pattern) - self.attention_layer_variables = get_attention_layer_variables_from_graph_def( graph_def, suffix=suffix ) @@ -1322,18 +1314,15 @@ def init_variables( f"attention_layer_{i}{suffix}/layer_normalization_{i}/gamma" ] - def _get_two_embed_variables(self, graph_def, pattern: str): - node = get_pattern_nodes_from_graph_def(graph_def, pattern)[pattern] - dtype = tf.as_dtype(node.dtype).as_numpy_dtype - tensor_shape = tf.TensorShape(node.tensor_shape).as_list() - if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): - tensor_value = np.frombuffer( - node.tensor_content, - dtype=tf.as_dtype(node.dtype).as_numpy_dtype, + if self.stripped_type_embedding: + self.two_side_embeeding_net_variables = ( + get_extra_embedding_net_variables_from_graph_def( + graph_def, + suffix, + get_extra_embedding_net_suffix(type_one_side=False), + self.layer_size, + ) ) - else: - tensor_value = get_tensor_by_type(node, dtype) - return np.reshape(tensor_value, tensor_shape) def build_type_exclude_mask( self, diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py index 087eb10f73..cc28e52930 100644 --- a/deepmd/entrypoints/doc.py +++ b/deepmd/entrypoints/doc.py @@ -1,20 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Module that prints train input arguments docstrings.""" - -from deepmd.utils.argcheck import ( - gen_doc, - gen_json, +from deepmd_utils.entrypoints.doc import ( + doc_train_input, ) __all__ = ["doc_train_input"] - - -def doc_train_input(*, out_type: str = "rst", **kwargs): - """Print out trining input arguments to console.""" - if out_type == "rst": - doc_str = gen_doc(make_anchor=True) - elif out_type == "json": - doc_str = gen_json() - else: - raise RuntimeError("Unsupported out type %s" % out_type) - print(doc_str) diff --git a/deepmd/entrypoints/gui.py b/deepmd/entrypoints/gui.py index 8b6b9e0a09..72de65f1c2 100644 --- a/deepmd/entrypoints/gui.py +++ b/deepmd/entrypoints/gui.py @@ -1,31 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""DP-GUI entrypoint.""" +from deepmd_utils.entrypoints.gui import ( + start_dpgui, +) - -def start_dpgui(*, port: int, bind_all: bool, **kwargs): - """Host DP-GUI server. - - Parameters - ---------- - port : int - The port to serve DP-GUI on. - bind_all : bool - Serve on all public interfaces. This will expose your DP-GUI instance - to the network on both IPv4 and IPv6 (where available). - **kwargs - additional arguments - - Raises - ------ - ModuleNotFoundError - The dpgui package is not installed - """ - try: - from dpgui import ( - start_dpgui, - ) - except ModuleNotFoundError as e: - raise ModuleNotFoundError( - "To use DP-GUI, please install the dpgui package:\npip install dpgui" - ) from e - start_dpgui(port=port, bind_all=bind_all) +__all__ = ["start_dpgui"] diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 782136b542..2c6ac26a7f 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -32,7 +32,7 @@ from deepmd.nvnmd.entrypoints.train import ( train_nvnmd, ) -from deepmd_cli.main import ( +from deepmd_utils.main import ( get_ll, main_parser, parse_args, diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index 9469b7df90..227aa13644 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -404,9 +404,7 @@ def get_nbor_stat(jdata, rcut, one_type: bool = False): None, ) tmp_data.get_batch() - assert ( - tmp_data.get_type_map() - ), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " + assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " if train_data is None: train_data = tmp_data else: diff --git a/deepmd/env.py b/deepmd/env.py index 9b7f86f0d5..f290dc0a90 100644 --- a/deepmd/env.py +++ b/deepmd/env.py @@ -28,6 +28,11 @@ ) import deepmd.lib +from deepmd_utils.env import ( + GLOBAL_ENER_FLOAT_PRECISION, + GLOBAL_NP_FLOAT_PRECISION, + global_float_prec, +) if TYPE_CHECKING: from types import ( @@ -475,24 +480,7 @@ def _get_package_constants( op_grads_module = get_module("op_grads") # FLOAT_PREC -dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "high").lower() -if dp_float_prec in ("high", ""): - # default is high - GLOBAL_TF_FLOAT_PRECISION = tf.float64 - GLOBAL_NP_FLOAT_PRECISION = np.float64 - GLOBAL_ENER_FLOAT_PRECISION = np.float64 - global_float_prec = "double" -elif dp_float_prec == "low": - GLOBAL_TF_FLOAT_PRECISION = tf.float32 - GLOBAL_NP_FLOAT_PRECISION = np.float32 - GLOBAL_ENER_FLOAT_PRECISION = np.float64 - global_float_prec = "float" -else: - raise RuntimeError( - "Unsupported float precision option: %s. Supported: high," - "low. Please set precision with environmental variable " - "DP_INTERFACE_PREC." % dp_float_prec - ) +GLOBAL_TF_FLOAT_PRECISION = tf.dtypes.as_dtype(GLOBAL_NP_FLOAT_PRECISION) def global_cvt_2_tf_float(xx: tf.Tensor) -> tf.Tensor: diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index e74d4a7e6d..4c15e57124 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -514,6 +514,11 @@ def build( self.bias_atom_e[type_i] = self.bias_atom_e[type_i] self.bias_atom_e = self.bias_atom_e[:ntypes_atom] + if nvnmd_cfg.enable: + # fix the bug: CNN and QNN have different t_bias_atom_e. + if "t_bias_atom_e" in nvnmd_cfg.weight.keys(): + self.bias_atom_e = nvnmd_cfg.weight["t_bias_atom_e"] + with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): t_dfparam = tf.constant(self.numb_fparam, name="dfparam", dtype=tf.int32) t_daparam = tf.constant(self.numb_aparam, name="daparam", dtype=tf.int32) diff --git a/deepmd/fit/polar.py b/deepmd/fit/polar.py index 0a6f7d4242..8f6631866c 100644 --- a/deepmd/fit/polar.py +++ b/deepmd/fit/polar.py @@ -213,8 +213,9 @@ def compute_input_stats(self, all_stat, protection=1e-2): # add polar_bias polar_bias.append(all_stat["polarizability"][ss].reshape((1, 9))) - matrix, bias = np.concatenate(sys_matrix, axis=0), np.concatenate( - polar_bias, axis=0 + matrix, bias = ( + np.concatenate(sys_matrix, axis=0), + np.concatenate(polar_bias, axis=0), ) atom_polar, _, _, _ = np.linalg.lstsq(matrix, bias, rcond=None) for itype in range(len(self.sel_type)): diff --git a/deepmd/infer/__init__.py b/deepmd/infer/__init__.py index 14d75d0c44..c1071af35c 100644 --- a/deepmd/infer/__init__.py +++ b/deepmd/infer/__init__.py @@ -58,6 +58,7 @@ def DeepPotential( load_prefix: str = "load", default_tf_graph: bool = False, input_map: Optional[dict] = None, + neighbor_list=None, ) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepDOS, DeepWFC]: """Factory function that will inialize appropriate potential read from `model_file`. @@ -71,6 +72,8 @@ def DeepPotential( If uses the default tf graph, otherwise build a new tf graph for evaluation input_map : dict, optional The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NeighborList, optional + The neighbor list object. If None, then build the native neighbor list. Returns ------- @@ -97,6 +100,7 @@ def DeepPotential( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) elif model_type == "dos": dp = DeepDOS( @@ -111,6 +115,7 @@ def DeepPotential( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) elif model_type == "polar": dp = DeepPolar( @@ -118,6 +123,7 @@ def DeepPotential( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) elif model_type == "global_polar": dp = DeepGlobalPolar( @@ -125,6 +131,7 @@ def DeepPotential( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) elif model_type == "wfc": dp = DeepWFC( diff --git a/deepmd/infer/deep_dipole.py b/deepmd/infer/deep_dipole.py index 6020118135..aba098a9f3 100644 --- a/deepmd/infer/deep_dipole.py +++ b/deepmd/infer/deep_dipole.py @@ -27,6 +27,8 @@ class DeepDipole(DeepTensor): If uses the default tf graph, otherwise build a new tf graph for evaluation input_map : dict, optional The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NeighborList, optional + The neighbor list object. If None, then build the native neighbor list. Warnings -------- @@ -41,6 +43,7 @@ def __init__( load_prefix: str = "load", default_tf_graph: bool = False, input_map: Optional[dict] = None, + neighbor_list=None, ) -> None: # use this in favor of dict update to move attribute from class to # instance namespace @@ -58,6 +61,7 @@ def __init__( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) def get_dim_fparam(self) -> int: diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 3f5dede1ad..0ca9f21a77 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -45,6 +45,9 @@ class DeepEval: as the initial batch size. input_map : dict, optional The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. """ load_prefix: str # set by subclass @@ -56,6 +59,7 @@ def __init__( default_tf_graph: bool = False, auto_batch_size: Union[bool, int, AutoBatchSize] = False, input_map: Optional[dict] = None, + neighbor_list=None, ): self.graph = self._load_graph( model_file, @@ -86,6 +90,8 @@ def __init__( else: raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") + self.neighbor_list = neighbor_list + @property @lru_cache(maxsize=None) def model_type(self) -> str: @@ -360,3 +366,92 @@ def eval_typeebd(self) -> np.ndarray: t_typeebd = self._get_tensor("t_typeebd:0") [typeebd] = run_sess(self.sess, [t_typeebd], feed_dict={}) return typeebd + + def build_neighbor_list( + self, + coords: np.ndarray, + cell: Optional[np.ndarray], + atype: np.ndarray, + imap: np.ndarray, + neighbor_list, + ): + """Make the mesh with neighbor list for a single frame. + + Parameters + ---------- + coords : np.ndarray + The coordinates of atoms. Should be of shape [natoms, 3] + cell : Optional[np.ndarray] + The cell of the system. Should be of shape [3, 3] + atype : np.ndarray + The type of atoms. Should be of shape [natoms] + imap : np.ndarray + The index map of atoms. Should be of shape [natoms] + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList + ASE neighbor list. The following method or attribute will be + used/set: bothways, self_interaction, update, build, first_neigh, + pair_second, offset_vec. + + Returns + ------- + natoms_vec : np.ndarray + The number of atoms. This tensor has the length of Ntypes + 2 + natoms[0]: nloc + natoms[1]: nall + natoms[i]: 2 <= i < Ntypes+2, number of type i atoms for nloc + coords : np.ndarray + The coordinates of atoms, including ghost atoms. Should be of + shape [nframes, nall, 3] + atype : np.ndarray + The type of atoms, including ghost atoms. Should be of shape [nall] + mesh : np.ndarray + The mesh in nei_mode=4. + imap : np.ndarray + The index map of atoms. Should be of shape [nall] + ghost_map : np.ndarray + The index map of ghost atoms. Should be of shape [nghost] + """ + pbc = np.repeat(cell is not None, 3) + cell = cell.reshape(3, 3) + positions = coords.reshape(-1, 3) + neighbor_list.bothways = True + neighbor_list.self_interaction = False + if neighbor_list.update(pbc, cell, positions): + neighbor_list.build(pbc, cell, positions) + first_neigh = neighbor_list.first_neigh.copy() + pair_second = neighbor_list.pair_second.copy() + offset_vec = neighbor_list.offset_vec.copy() + # get out-of-box neighbors + out_mask = np.any(offset_vec != 0, axis=1) + out_idx = pair_second[out_mask] + out_offset = offset_vec[out_mask] + out_coords = positions[out_idx] + out_offset.dot(cell) + atype = np.array(atype, dtype=int) + out_atype = atype[out_idx] + + nloc = positions.shape[0] + nghost = out_idx.size + all_coords = np.concatenate((positions, out_coords), axis=0) + all_atype = np.concatenate((atype, out_atype), axis=0) + # convert neighbor indexes + ghost_map = pair_second[out_mask] + pair_second[out_mask] = np.arange(nloc, nloc + nghost) + # get the mesh + mesh = np.zeros(16 + nloc * 2 + pair_second.size, dtype=int) + mesh[0] = nloc + # ilist + mesh[16 : 16 + nloc] = np.arange(nloc) + # numnei + mesh[16 + nloc : 16 + nloc * 2] = first_neigh[1:] - first_neigh[:-1] + # jlist + mesh[16 + nloc * 2 :] = pair_second + + # natoms_vec + natoms_vec = np.zeros(self.ntypes + 2).astype(int) + natoms_vec[0] = nloc + natoms_vec[1] = nloc + nghost + for ii in range(self.ntypes): + natoms_vec[ii + 2] = np.count_nonzero(atype == ii) + # imap append ghost atoms + imap = np.concatenate((imap, np.arange(nloc, nloc + nghost))) + return natoms_vec, all_coords, all_atype, mesh, imap, ghost_map diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 118f8c98a7..c1f981ef86 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -30,6 +30,8 @@ class DeepPolar(DeepTensor): If uses the default tf graph, otherwise build a new tf graph for evaluation input_map : dict, optional The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NeighborList, optional + The neighbor list object. If None, then build the native neighbor list. Warnings -------- @@ -44,6 +46,7 @@ def __init__( load_prefix: str = "load", default_tf_graph: bool = False, input_map: Optional[dict] = None, + neighbor_list=None, ) -> None: # use this in favor of dict update to move attribute from class to # instance namespace @@ -61,6 +64,7 @@ def __init__( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) def get_dim_fparam(self) -> int: @@ -83,10 +87,16 @@ class DeepGlobalPolar(DeepTensor): The prefix in the load computational graph default_tf_graph : bool If uses the default tf graph, otherwise build a new tf graph for evaluation + neighbor_list : ase.neighborlist.NeighborList, optional + The neighbor list object. If None, then build the native neighbor list. """ def __init__( - self, model_file: str, load_prefix: str = "load", default_tf_graph: bool = False + self, + model_file: str, + load_prefix: str = "load", + default_tf_graph: bool = False, + neighbor_list=None, ) -> None: self.tensors.update( { @@ -101,6 +111,7 @@ def __init__( model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph, + neighbor_list=None, ) def eval( diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index fc9a6a76ed..81cfdde7a8 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -51,6 +51,9 @@ class DeepPot(DeepEval): as the initial batch size. input_map : dict, optional The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. Examples -------- @@ -78,6 +81,7 @@ def __init__( default_tf_graph: bool = False, auto_batch_size: Union[bool, int, AutoBatchSize] = True, input_map: Optional[dict] = None, + neighbor_list=None, ) -> None: # add these tensors on top of what is defined by DeepTensor Class # use this in favor of dict update to move attribute from class to @@ -112,6 +116,7 @@ def __init__( default_tf_graph=default_tf_graph, auto_batch_size=auto_batch_size, input_map=input_map, + neighbor_list=neighbor_list, ) # load optional tensors @@ -479,8 +484,30 @@ def _prepare_feed_dict( aparam = np.reshape(aparam, [nframes, natoms * fdim]) # make natoms_vec and default_mesh - natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) - assert natoms_vec[0] == natoms + if self.neighbor_list is None: + natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) + assert natoms_vec[0] == natoms + mesh = make_default_mesh(pbc, mixed_type) + ghost_map = None + else: + if nframes > 1: + raise NotImplementedError( + "neighbor_list does not support multiple frames" + ) + ( + natoms_vec, + coords, + atom_types, + mesh, + imap, + ghost_map, + ) = self.build_neighbor_list( + coords, + cells if cells is not None else None, + atom_types, + imap, + self.neighbor_list, + ) # evaluate feed_dict_test = {} @@ -501,12 +528,12 @@ def _prepare_feed_dict( raise RuntimeError if self.has_efield: feed_dict_test[self.t_efield] = np.reshape(efield, [-1]) - feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type) + feed_dict_test[self.t_mesh] = mesh if self.has_fparam: feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1]) if self.has_aparam: feed_dict_test[self.t_aparam] = np.reshape(aparam, [-1]) - return feed_dict_test, imap, natoms_vec + return feed_dict_test, imap, natoms_vec, ghost_map def _eval_inner( self, @@ -522,10 +549,13 @@ def _eval_inner( natoms, nframes = self._get_natoms_and_nframes( coords, atom_types, mixed_type=mixed_type ) - feed_dict_test, imap, natoms_vec = self._prepare_feed_dict( + feed_dict_test, imap, natoms_vec, ghost_map = self._prepare_feed_dict( coords, cells, atom_types, fparam, aparam, efield, mixed_type=mixed_type ) + nloc = natoms_vec[0] + nall = natoms_vec[1] + t_out = [self.t_energy, self.t_force, self.t_virial] if atomic: t_out += [self.t_ae, self.t_av] @@ -548,6 +578,13 @@ def _eval_inner( ) else: natoms_real = natoms + if ghost_map is not None: + # add the value of ghost atoms to real atoms + force = np.reshape(force, [nframes, -1, 3]) + np.add.at(force[0], ghost_map, force[0, nloc:]) + if atomic: + av = np.reshape(av, [nframes, -1, 9]) + np.add.at(av[0], ghost_map, av[0, nloc:]) # reverse map of the outputs force = self.reverse_map(np.reshape(force, [nframes, -1, 3]), imap) @@ -556,11 +593,15 @@ def _eval_inner( av = self.reverse_map(np.reshape(av, [nframes, -1, 9]), imap) energy = np.reshape(energy, [nframes, 1]) - force = np.reshape(force, [nframes, natoms, 3]) + force = np.reshape(force, [nframes, nall, 3]) + if nloc < nall: + force = force[:, :nloc, :] virial = np.reshape(virial, [nframes, 9]) if atomic: ae = np.reshape(ae, [nframes, natoms_real, 1]) - av = np.reshape(av, [nframes, natoms, 9]) + av = np.reshape(av, [nframes, nall, 9]) + if nloc < nall: + av = av[:, :nloc, :] return energy, force, virial, ae, av else: return energy, force, virial @@ -640,10 +681,11 @@ def _eval_descriptor_inner( natoms, nframes = self._get_natoms_and_nframes( coords, atom_types, mixed_type=mixed_type ) - feed_dict_test, imap, natoms_vec = self._prepare_feed_dict( + feed_dict_test, imap, natoms_vec, ghost_map = self._prepare_feed_dict( coords, cells, atom_types, fparam, aparam, efield, mixed_type=mixed_type ) (descriptor,) = run_sess( self.sess, [self.t_descriptor], feed_dict=feed_dict_test ) + imap = imap[:natoms] return self.reverse_map(np.reshape(descriptor, [nframes, natoms, -1]), imap) diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 268523e959..a803eb0c6b 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -39,6 +39,8 @@ class DeepTensor(DeepEval): If uses the default tf graph, otherwise build a new tf graph for evaluation input_map : dict, optional The input map for tf.import_graph_def. Only work with default tf graph + neighbor_list : ase.neighborlist.NeighborList, optional + The neighbor list object. If None, then build the native neighbor list. """ tensors: ClassVar[Dict[str, str]] = { @@ -63,6 +65,7 @@ def __init__( load_prefix: str = "load", default_tf_graph: bool = False, input_map: Optional[dict] = None, + neighbor_list=None, ) -> None: """Constructor.""" DeepEval.__init__( @@ -71,6 +74,7 @@ def __init__( load_prefix=load_prefix, default_tf_graph=default_tf_graph, input_map=input_map, + neighbor_list=neighbor_list, ) # check model type model_type = self.tensors["t_tensor"][2:-2] @@ -209,8 +213,29 @@ def eval( ) # make natoms_vec and default_mesh - natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) - assert natoms_vec[0] == natoms + if self.neighbor_list is None: + natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) + assert natoms_vec[0] == natoms + mesh = make_default_mesh(pbc, mixed_type) + else: + if nframes > 1: + raise NotImplementedError( + "neighbor_list does not support multiple frames" + ) + ( + natoms_vec, + coords, + atom_types, + mesh, + imap, + _, + ) = self.build_neighbor_list( + coords, + cells if cells is not None else None, + atom_types, + imap, + self.neighbor_list, + ) # evaluate feed_dict_test = {} @@ -223,7 +248,7 @@ def eval( ) feed_dict_test[self.t_coord] = np.reshape(coords, [-1]) feed_dict_test[self.t_box] = np.reshape(cells, [-1]) - feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type) + feed_dict_test[self.t_mesh] = mesh if atomic: assert ( @@ -333,8 +358,30 @@ def eval_full( ) # make natoms_vec and default_mesh - natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) - assert natoms_vec[0] == natoms + if self.neighbor_list is None: + natoms_vec = self.make_natoms_vec(atom_types, mixed_type=mixed_type) + assert natoms_vec[0] == natoms + mesh = make_default_mesh(pbc, mixed_type) + ghost_map = None + else: + if nframes > 1: + raise NotImplementedError( + "neighbor_list does not support multiple frames" + ) + ( + natoms_vec, + coords, + atom_types, + mesh, + imap, + ghost_map, + ) = self.build_neighbor_list( + coords, + cells if cells is not None else None, + atom_types, + imap, + self.neighbor_list, + ) # evaluate feed_dict_test = {} @@ -347,7 +394,7 @@ def eval_full( ) feed_dict_test[self.t_coord] = np.reshape(coords, [-1]) feed_dict_test[self.t_box] = np.reshape(cells, [-1]) - feed_dict_test[self.t_mesh] = make_default_mesh(pbc, mixed_type) + feed_dict_test[self.t_mesh] = mesh t_out = [self.t_global_tensor, self.t_force, self.t_virial] if atomic: @@ -361,21 +408,39 @@ def eval_full( at = v_out[3] # atom tensor av = v_out[4] # atom virial + nloc = natoms_vec[0] + nall = natoms_vec[1] + + if ghost_map is not None: + # add the value of ghost atoms to real atoms + force = np.reshape(force, [nframes * nout, -1, 3]) + # TODO: is there some way not to use for loop? + for ii in range(nframes * nout): + np.add.at(force[ii], ghost_map, force[ii, nloc:]) + if atomic: + av = np.reshape(av, [nframes * nout, -1, 9]) + for ii in range(nframes * nout): + np.add.at(av[ii], ghost_map, av[ii, nloc:]) + # please note here the shape are wrong! - force = self.reverse_map(np.reshape(force, [nframes * nout, natoms, 3]), imap) + force = self.reverse_map(np.reshape(force, [nframes * nout, nall, 3]), imap) if atomic: at = self.reverse_map( np.reshape(at, [nframes, len(sel_at), nout]), sel_imap ) - av = self.reverse_map(np.reshape(av, [nframes * nout, natoms, 9]), imap) + av = self.reverse_map(np.reshape(av, [nframes * nout, nall, 9]), imap) # make sure the shapes are correct here gt = np.reshape(gt, [nframes, nout]) - force = np.reshape(force, [nframes, nout, natoms, 3]) + force = np.reshape(force, [nframes, nout, nall, 3]) + if nloc < nall: + force = force[:, :, :nloc, :] virial = np.reshape(virial, [nframes, nout, 9]) if atomic: at = np.reshape(at, [nframes, len(sel_at), self.output_dim]) - av = np.reshape(av, [nframes, nout, natoms, 9]) + av = np.reshape(av, [nframes, nout, nall, 9]) + if nloc < nall: + av = av[:, :, :nloc, :] return gt, force, virial, at, av else: return gt, force, virial diff --git a/deepmd/loggers/__init__.py b/deepmd/loggers/__init__.py index 39aa76139d..71057e3056 100644 --- a/deepmd/loggers/__init__.py +++ b/deepmd/loggers/__init__.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Module taking care of logging duties.""" +"""Alias of deepmd_utils.loggers for backward compatibility.""" -from .loggers import ( +from deepmd_utils.loggers.loggers import ( set_log_handles, ) diff --git a/deepmd/loggers/loggers.py b/deepmd/loggers/loggers.py index 015581f6bd..74ca7de63e 100644 --- a/deepmd/loggers/loggers.py +++ b/deepmd/loggers/loggers.py @@ -1,277 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Logger initialization for package.""" - -import logging -import os -from typing import ( - TYPE_CHECKING, - Optional, +"""Alias of deepmd_utils.loggers.loggers for backward compatibility.""" +from deepmd_utils.loggers.loggers import ( + set_log_handles, ) -if TYPE_CHECKING: - from pathlib import ( - Path, - ) - - from mpi4py import ( - MPI, - ) - - _MPI_APPEND_MODE = MPI.MODE_CREATE | MPI.MODE_APPEND - -logging.getLogger(__name__) - __all__ = ["set_log_handles"] - -# logger formater -FFORMATTER = logging.Formatter( - "[%(asctime)s] %(app_name)s %(levelname)-7s %(name)-45s %(message)s" -) -CFORMATTER = logging.Formatter( - # "%(app_name)s %(levelname)-7s |-> %(name)-45s %(message)s" - "%(app_name)s %(levelname)-7s %(message)s" -) -FFORMATTER_MPI = logging.Formatter( - "[%(asctime)s] %(app_name)s rank:%(rank)-2s %(levelname)-7s %(name)-45s %(message)s" -) -CFORMATTER_MPI = logging.Formatter( - # "%(app_name)s rank:%(rank)-2s %(levelname)-7s |-> %(name)-45s %(message)s" - "%(app_name)s rank:%(rank)-2s %(levelname)-7s %(message)s" -) - - -class _AppFilter(logging.Filter): - """Add field `app_name` to log messages.""" - - def filter(self, record): - record.app_name = "DEEPMD" - return True - - -class _MPIRankFilter(logging.Filter): - """Add MPI rank number to log messages, adds field `rank`.""" - - def __init__(self, rank: int) -> None: - super().__init__(name="MPI_rank_id") - self.mpi_rank = str(rank) - - def filter(self, record): - record.rank = self.mpi_rank - return True - - -class _MPIMasterFilter(logging.Filter): - """Filter that lets through only messages emited from rank==0.""" - - def __init__(self, rank: int) -> None: - super().__init__(name="MPI_master_log") - self.mpi_rank = rank - - def filter(self, record): - if self.mpi_rank == 0: - return True - else: - return False - - -class _MPIFileStream: - """Wrap MPI.File` so it has the same API as python file streams. - - Parameters - ---------- - filename : Path - disk location of the file stream - MPI : MPI - MPI communicator object - mode : str, optional - file write mode, by default _MPI_APPEND_MODE - """ - - def __init__( - self, filename: "Path", MPI: "MPI", mode: str = "_MPI_APPEND_MODE" - ) -> None: - self.stream = MPI.File.Open(MPI.COMM_WORLD, filename, mode) - self.stream.Set_atomicity(True) - self.name = "MPIfilestream" - - def write(self, msg: str): - """Write to MPI shared file stream. - - Parameters - ---------- - msg : str - message to write - """ - b = bytearray() - b.extend(map(ord, msg)) - self.stream.Write_shared(b) - - def close(self): - """Synchronize and close MPI file stream.""" - self.stream.Sync() - self.stream.Close() - - -class _MPIHandler(logging.FileHandler): - """Emulate `logging.FileHandler` with MPI shared File that all ranks can write to. - - Parameters - ---------- - filename : Path - file path - MPI : MPI - MPI communicator object - mode : str, optional - file access mode, by default "_MPI_APPEND_MODE" - """ - - def __init__( - self, - filename: "Path", - MPI: "MPI", - mode: str = "_MPI_APPEND_MODE", - ) -> None: - self.MPI = MPI - super().__init__(filename, mode=mode, encoding=None, delay=False) - - def _open(self): - return _MPIFileStream(self.baseFilename, self.MPI, self.mode) - - def setStream(self, stream): - """Stream canot 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 -): - """Set desired level for package loggers and add file handlers. - - Parameters - ---------- - level : int - logging level - log_path : Optional[str] - path to log file, if None logs will be send only to console. If the parent - directory does not exist it will be automatically created, by default None - mpi_log : Optional[str], optional - mpi log type. Has three options. `master` will output logs to file and console - only from rank==0. `collect` will write messages from all ranks to one file - opened under rank==0 and to console. `workers` will open one log file for each - worker designated by its rank, console behaviour is the same as for `collect`. - If this argument is specified, package 'mpi4py' must be already installed. - by default None - - Raises - ------ - RuntimeError - If the argument `mpi_log` is specified, package `mpi4py` is not installed. - - References - ---------- - https://groups.google.com/g/mpi4py/c/SaNzc8bdj6U - https://stackoverflow.com/questions/35869137/avoid-tensorflow-print-on-standard-error - https://stackoverflow.com/questions/56085015/suppress-openmp-debug-messages-when-running-tensorflow-on-cpu - - Notes - ----- - Logging levels: - - +---------+--------------+----------------+----------------+----------------+ - | | our notation | python logging | tensorflow cpp | OpenMP | - +=========+==============+================+================+================+ - | debug | 10 | 10 | 0 | 1/on/true/yes | - +---------+--------------+----------------+----------------+----------------+ - | info | 20 | 20 | 1 | 0/off/false/no | - +---------+--------------+----------------+----------------+----------------+ - | warning | 30 | 30 | 2 | 0/off/false/no | - +---------+--------------+----------------+----------------+----------------+ - | error | 40 | 40 | 3 | 0/off/false/no | - +---------+--------------+----------------+----------------+----------------+ - - """ - # silence logging for OpenMP when running on CPU if level is any other than debug - if level <= 10: - os.environ["KMP_WARNINGS"] = "FALSE" - - # set TF cpp internal logging level - os.environ["TF_CPP_MIN_LOG_LEVEL"] = str(int((level / 10) - 1)) - - # get root logger - root_log = logging.getLogger("deepmd") - root_log.propagate = False - - root_log.setLevel(level) - - # check if arguments are present - MPI = None - if mpi_log: - try: - from mpi4py import ( - MPI, - ) - except ImportError as e: - raise RuntimeError( - "You cannot specify 'mpi_log' when mpi4py not installed" - ) from e - - # * add console handler ************************************************************ - ch = logging.StreamHandler() - if MPI: - rank = MPI.COMM_WORLD.Get_rank() - if mpi_log == "master": - ch.setFormatter(CFORMATTER) - ch.addFilter(_MPIMasterFilter(rank)) - else: - ch.setFormatter(CFORMATTER_MPI) - ch.addFilter(_MPIRankFilter(rank)) - else: - ch.setFormatter(CFORMATTER) - - ch.setLevel(level) - ch.addFilter(_AppFilter()) - # clean old handlers before adding new one - root_log.handlers.clear() - root_log.addHandler(ch) - - # * add file handler *************************************************************** - if log_path: - # create directory - log_path.parent.mkdir(exist_ok=True, parents=True) - - fh = None - - if mpi_log == "master": - rank = MPI.COMM_WORLD.Get_rank() - if rank == 0: - fh = logging.FileHandler(log_path, mode="w") - fh.addFilter(_MPIMasterFilter(rank)) - fh.setFormatter(FFORMATTER) - elif mpi_log == "collect": - rank = MPI.COMM_WORLD.Get_rank() - fh = _MPIHandler(log_path, MPI, mode=MPI.MODE_WRONLY | MPI.MODE_CREATE) - fh.addFilter(_MPIRankFilter(rank)) - fh.setFormatter(FFORMATTER_MPI) - elif mpi_log == "workers": - rank = MPI.COMM_WORLD.Get_rank() - # if file has suffix than inser rank number before suffix - # e.g deepmd.log -> deepmd_.log - # if no suffix is present, insert rank as suffix - # e.g. deepmdlog -> deepmdlog. - if log_path.suffix: - worker_log = (log_path.parent / f"{log_path.stem}_{rank}").with_suffix( - log_path.suffix - ) - else: - worker_log = log_path.with_suffix(f".{rank}") - - fh = logging.FileHandler(worker_log, mode="w") - fh.setFormatter(FFORMATTER) - else: - fh = logging.FileHandler(log_path, mode="w") - fh.setFormatter(FFORMATTER) - - if fh: - fh.setLevel(level) - fh.addFilter(_AppFilter()) - root_log.addHandler(fh) diff --git a/deepmd/loss/dos.py b/deepmd/loss/dos.py index fa30552486..7d38f2b17a 100644 --- a/deepmd/loss/dos.py +++ b/deepmd/loss/dos.py @@ -143,16 +143,20 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): more_loss = {} if self.has_dos: l2_loss += atom_norm_ener * (pref_dos * l2_dos_loss) - more_loss["l2_dos_loss"] = l2_dos_loss + more_loss["l2_dos_loss"] = self.display_if_exist(l2_dos_loss, find_dos) if self.has_cdf: l2_loss += atom_norm_ener * (pref_cdf * l2_cdf_loss) - more_loss["l2_cdf_loss"] = l2_cdf_loss + more_loss["l2_cdf_loss"] = self.display_if_exist(l2_cdf_loss, find_dos) if self.has_ados: l2_loss += global_cvt_2_ener_float(pref_ados * l2_atom_dos_loss) - more_loss["l2_atom_dos_loss"] = l2_atom_dos_loss + more_loss["l2_atom_dos_loss"] = self.display_if_exist( + l2_atom_dos_loss, find_atom_dos + ) if self.has_acdf: l2_loss += global_cvt_2_ener_float(pref_acdf * l2_atom_cdf_loss) - more_loss["l2_atom_cdf_loss"] = l2_atom_cdf_loss + more_loss["l2_atom_cdf_loss"] = self.display_if_exist( + l2_atom_cdf_loss, find_atom_dos + ) # only used when tensorboard was set as true self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss)) diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py index 95997bad10..d7f83f09e5 100644 --- a/deepmd/loss/ener.py +++ b/deepmd/loss/ener.py @@ -291,22 +291,32 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): more_loss = {} if self.has_e: l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) - more_loss["l2_ener_loss"] = l2_ener_loss + more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy) if self.has_f: l2_loss += global_cvt_2_ener_float(pref_f * l2_force_loss) - more_loss["l2_force_loss"] = l2_force_loss + more_loss["l2_force_loss"] = self.display_if_exist( + l2_force_loss, find_force + ) if self.has_v: l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss)) - more_loss["l2_virial_loss"] = l2_virial_loss + more_loss["l2_virial_loss"] = self.display_if_exist( + l2_virial_loss, find_virial + ) if self.has_ae: l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss) - more_loss["l2_atom_ener_loss"] = l2_atom_ener_loss + more_loss["l2_atom_ener_loss"] = self.display_if_exist( + l2_atom_ener_loss, find_atom_ener + ) if self.has_pf: l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss) - more_loss["l2_pref_force_loss"] = l2_pref_force_loss + more_loss["l2_pref_force_loss"] = self.display_if_exist( + l2_pref_force_loss, find_atom_pref + ) if self.has_gf: l2_loss += global_cvt_2_ener_float(pref_gf * l2_gen_force_loss) - more_loss["l2_gen_force_loss"] = l2_gen_force_loss + more_loss["l2_gen_force_loss"] = self.display_if_exist( + l2_gen_force_loss, find_drdq + ) # only used when tensorboard was set as true self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss)) @@ -553,19 +563,25 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): more_loss = {} if self.has_e: l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) - more_loss["l2_ener_loss"] = l2_ener_loss + more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy) if self.has_fr: l2_loss += global_cvt_2_ener_float(pref_fr * l2_force_r_loss) - more_loss["l2_force_r_loss"] = l2_force_r_loss + more_loss["l2_force_r_loss"] = self.display_if_exist( + l2_force_r_loss, find_force + ) if self.has_fm: l2_loss += global_cvt_2_ener_float(pref_fm * l2_force_m_loss) - more_loss["l2_force_m_loss"] = l2_force_m_loss + more_loss["l2_force_m_loss"] = self.display_if_exist( + l2_force_m_loss, find_force + ) if self.has_v: l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss)) - more_loss["l2_virial_loss"] = l2_virial_loss + more_loss["l2_virial_loss"] = self.display_if_exist(l2_virial_loss, find_virial) if self.has_ae: l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss) - more_loss["l2_atom_ener_loss"] = l2_atom_ener_loss + more_loss["l2_atom_ener_loss"] = self.display_if_exist( + l2_atom_ener_loss, find_atom_ener + ) # only used when tensorboard was set as true self.l2_loss_summary = tf.summary.scalar("l2_loss", tf.sqrt(l2_loss)) @@ -785,8 +801,10 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): more_loss = {} l2_loss += atom_norm_ener * (pref_e * l2_ener_loss) l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss) - more_loss["l2_ener_loss"] = l2_ener_loss - more_loss["l2_ener_dipole_loss"] = l2_ener_dipole_loss + more_loss["l2_ener_loss"] = self.display_if_exist(l2_ener_loss, find_energy) + more_loss["l2_ener_dipole_loss"] = self.display_if_exist( + l2_ener_dipole_loss, find_ener_dipole + ) self.l2_loss_summary = tf.summary.scalar("l2_loss_" + suffix, tf.sqrt(l2_loss)) self.l2_loss_ener_summary = tf.summary.scalar( diff --git a/deepmd/loss/loss.py b/deepmd/loss/loss.py index 9324077691..a719a08d81 100644 --- a/deepmd/loss/loss.py +++ b/deepmd/loss/loss.py @@ -8,6 +8,8 @@ Tuple, ) +import numpy as np + from deepmd.env import ( tf, ) @@ -72,3 +74,20 @@ def eval( A dictionary that maps keys to values. It should contain key `natoms` """ + + @staticmethod + def display_if_exist(loss: tf.Tensor, find_property: float) -> tf.Tensor: + """Display NaN if labeled property is not found. + + Parameters + ---------- + loss : tf.Tensor + the loss tensor + find_property : float + whether the property is found + """ + return tf.cond( + tf.cast(find_property, tf.bool), + lambda: loss, + lambda: tf.cast(np.nan, dtype=loss.dtype), + ) diff --git a/deepmd/loss/tensor.py b/deepmd/loss/tensor.py index 74eb2b74dc..a40f95a18e 100644 --- a/deepmd/loss/tensor.py +++ b/deepmd/loss/tensor.py @@ -87,7 +87,7 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean( tf.square(self.scale * (polar - atomic_polar_hat)), name="l2_" + suffix ) - more_loss["local_loss"] = local_loss + more_loss["local_loss"] = self.display_if_exist(local_loss, find_atomic) l2_loss += self.local_weight * local_loss self.l2_loss_local_summary = tf.summary.scalar( "l2_local_loss_" + suffix, tf.sqrt(more_loss["local_loss"]) @@ -118,7 +118,7 @@ def build(self, learning_rate, natoms, model_dict, label_dict, suffix): tf.square(self.scale * (global_polar - polar_hat)), name="l2_" + suffix ) - more_loss["global_loss"] = global_loss + more_loss["global_loss"] = self.display_if_exist(global_loss, find_global) self.l2_loss_global_summary = tf.summary.scalar( "l2_global_loss_" + suffix, tf.sqrt(more_loss["global_loss"]) / global_cvt_2_tf_float(atoms), diff --git a/deepmd/model/dos.py b/deepmd/model/dos.py index 697fad9a9e..22e291a0f0 100644 --- a/deepmd/model/dos.py +++ b/deepmd/model/dos.py @@ -155,10 +155,12 @@ def build( # type embedding if any if self.typeebd is not None: - type_embedding = self.typeebd.build( + type_embedding = self.build_type_embedding( self.ntypes, reuse=reuse, suffix=suffix, + frz_model=frz_model, + ckpt_meta=ckpt_meta, ) input_dict["type_embedding"] = type_embedding input_dict["atype"] = atype_ diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index 1976c1ad51..0d8d66b305 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -203,10 +203,12 @@ def build( # type embedding if any if self.typeebd is not None and "type_embedding" not in input_dict: - type_embedding = self.typeebd.build( + type_embedding = self.build_type_embedding( self.ntypes, reuse=reuse, suffix=suffix, + ckpt_meta=ckpt_meta, + frz_model=frz_model, ) input_dict["type_embedding"] = type_embedding # spin if any diff --git a/deepmd/model/model.py b/deepmd/model/model.py index 3f24e42aec..6117b4942d 100644 --- a/deepmd/model/model.py +++ b/deepmd/model/model.py @@ -97,6 +97,9 @@ def get_class_by_input(cls, input: dict): from deepmd.model.multi import ( MultiModel, ) + from deepmd.model.pairtab import ( + PairTabModel, + ) from deepmd.model.pairwise_dprc import ( PairwiseDPRc, ) @@ -112,6 +115,8 @@ def get_class_by_input(cls, input: dict): return FrozenModel elif model_type == "linear_ener": return LinearEnergyModel + elif model_type == "pairtab": + return PairTabModel else: raise ValueError(f"unknown model type: {model_type}") @@ -331,6 +336,60 @@ def build_descrpt( self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1]) return dout + def build_type_embedding( + self, + ntypes: int, + frz_model: Optional[str] = None, + ckpt_meta: Optional[str] = None, + suffix: str = "", + reuse: Optional[Union[bool, Enum]] = None, + ) -> tf.Tensor: + """Build the type embedding part of the model. + + Parameters + ---------- + ntypes : int + The number of types + frz_model : str, optional + The path to the frozen model + ckpt_meta : str, optional + The path prefix of the checkpoint and meta files + suffix : str, optional + The suffix of the scope + reuse : bool or tf.AUTO_REUSE, optional + Whether to reuse the variables + + Returns + ------- + tf.Tensor + The type embedding tensor + """ + assert self.typeebd is not None + if frz_model is None and ckpt_meta is None: + dout = self.typeebd.build( + ntypes, + reuse=reuse, + suffix=suffix, + ) + else: + # nothing input + feed_dict = {} + return_elements = [ + f"t_typeebd{suffix}:0", + ] + if frz_model is not None: + imported_tensors = self._import_graph_def_from_frz_model( + frz_model, feed_dict, return_elements + ) + elif ckpt_meta is not None: + imported_tensors = self._import_graph_def_from_ckpt_meta( + ckpt_meta, feed_dict, return_elements + ) + else: + raise RuntimeError("should not reach here") # pragma: no cover + dout = imported_tensors[-1] + return dout + def _import_graph_def_from_frz_model( self, frz_model: str, feed_dict: dict, return_elements: List[str] ): diff --git a/deepmd/model/model_stat.py b/deepmd/model/model_stat.py index d2cc918b64..933a634ce8 100644 --- a/deepmd/model/model_stat.py +++ b/deepmd/model/model_stat.py @@ -1,68 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from collections import ( - defaultdict, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.model_stat import ( + _make_all_stat_ref, + make_stat_input, + merge_sys_stat, ) -import numpy as np - - -def _make_all_stat_ref(data, nbatches): - all_stat = defaultdict(list) - for ii in range(data.get_nsystems()): - for jj in range(nbatches): - stat_data = data.get_batch(sys_idx=ii) - for dd in stat_data: - if dd == "natoms_vec": - stat_data[dd] = stat_data[dd].astype(np.int32) - all_stat[dd].append(stat_data[dd]) - return all_stat - - -def make_stat_input(data, nbatches, merge_sys=True): - """Pack data for statistics. - - Parameters - ---------- - data - The data - nbatches : int - The number of batches - merge_sys : bool (True) - Merge system data - - Returns - ------- - all_stat: - A dictionary of list of list storing data for stat. - if merge_sys == False data can be accessed by - all_stat[key][sys_idx][batch_idx][frame_idx] - else merge_sys == True can be accessed by - all_stat[key][batch_idx][frame_idx] - """ - all_stat = defaultdict(list) - for ii in range(data.get_nsystems()): - sys_stat = defaultdict(list) - for jj in range(nbatches): - stat_data = data.get_batch(sys_idx=ii) - for dd in stat_data: - if dd == "natoms_vec": - stat_data[dd] = stat_data[dd].astype(np.int32) - sys_stat[dd].append(stat_data[dd]) - for dd in sys_stat: - if merge_sys: - for bb in sys_stat[dd]: - all_stat[dd].append(bb) - else: - all_stat[dd].append(sys_stat[dd]) - return all_stat - - -def merge_sys_stat(all_stat): - first_key = next(iter(all_stat.keys())) - nsys = len(all_stat[first_key]) - ret = defaultdict(list) - for ii in range(nsys): - for dd in all_stat: - for bb in all_stat[dd][ii]: - ret[dd].append(bb) - return ret +__all__ = [ + "make_stat_input", + "merge_sys_stat", + "_make_all_stat_ref", # used by tests +] diff --git a/deepmd/model/multi.py b/deepmd/model/multi.py index bfc67b9792..83b231c0e8 100644 --- a/deepmd/model/multi.py +++ b/deepmd/model/multi.py @@ -317,10 +317,12 @@ def build( # type embedding if any if self.typeebd is not None: - type_embedding = self.typeebd.build( + type_embedding = self.build_type_embedding( self.ntypes, reuse=reuse, suffix=suffix, + frz_model=frz_model, + ckpt_meta=ckpt_meta, ) input_dict["type_embedding"] = type_embedding input_dict["atype"] = atype_ diff --git a/deepmd/model/pairtab.py b/deepmd/model/pairtab.py new file mode 100644 index 0000000000..38934818e6 --- /dev/null +++ b/deepmd/model/pairtab.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from enum import ( + Enum, +) +from typing import ( + List, + Optional, + Union, +) + +import numpy as np + +from deepmd.env import ( + GLOBAL_TF_FLOAT_PRECISION, + MODEL_VERSION, + global_cvt_2_ener_float, + op_module, + tf, +) +from deepmd.fit.fitting import ( + Fitting, +) +from deepmd.loss.loss import ( + Loss, +) +from deepmd.model.model import ( + Model, +) +from deepmd.utils.pair_tab import ( + PairTab, +) + + +class PairTabModel(Model): + """Pairwise tabulation energy model. + + This model can be used to tabulate the pairwise energy between atoms for either + short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not + be used alone, but rather as one submodel of a linear (sum) model, such as + DP+D3. + + Do not put the model on the first model of a linear model, since the linear + model fetches the type map from the first model. + + At this moment, the model does not smooth the energy at the cutoff radius, so + one needs to make sure the energy has been smoothed to zero. + + Parameters + ---------- + tab_file : str + The path to the tabulation file. + rcut : float + The cutoff radius + sel : int or list[int] + The maxmum number of atoms in the cut-off radius + """ + + model_type = "ener" + + def __init__( + self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs + ): + super().__init__() + self.tab_file = tab_file + self.tab = PairTab(self.tab_file) + self.ntypes = self.tab.ntypes + self.rcut = rcut + if isinstance(sel, int): + self.sel = sel + elif isinstance(sel, list): + self.sel = sum(sel) + else: + raise TypeError("sel must be int or list[int]") + + def build( + self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box: tf.Tensor, + mesh: tf.Tensor, + input_dict: dict, + frz_model: Optional[str] = None, + ckpt_meta: Optional[str] = None, + suffix: str = "", + reuse: Optional[Union[bool, Enum]] = None, + ): + """Build the model. + + Parameters + ---------- + coord_ : tf.Tensor + The coordinates of atoms + atype_ : tf.Tensor + The atom types of atoms + natoms : tf.Tensor + The number of atoms + box : tf.Tensor + The box vectors + mesh : tf.Tensor + The mesh vectors + input_dict : dict + The input dict + frz_model : str, optional + The path to the frozen model + ckpt_meta : str, optional + The path prefix of the checkpoint and meta files + suffix : str, optional + The suffix of the scope + reuse : bool or tf.AUTO_REUSE, optional + Whether to reuse the variables + + Returns + ------- + dict + The output dict + """ + tab_info, tab_data = self.tab.get() + with tf.variable_scope("model_attr" + suffix, reuse=reuse): + self.tab_info = tf.get_variable( + "t_tab_info", + tab_info.shape, + dtype=tf.float64, + trainable=False, + initializer=tf.constant_initializer(tab_info, dtype=tf.float64), + ) + self.tab_data = tf.get_variable( + "t_tab_data", + tab_data.shape, + dtype=tf.float64, + trainable=False, + initializer=tf.constant_initializer(tab_data, dtype=tf.float64), + ) + t_tmap = tf.constant(" ".join(self.type_map), name="tmap", dtype=tf.string) + 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) + + with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): + t_dfparam = tf.constant(0, name="dfparam", dtype=tf.int32) + t_daparam = tf.constant(0, name="daparam", dtype=tf.int32) + with tf.variable_scope("descrpt_attr" + suffix, reuse=reuse): + t_ntypes = tf.constant(self.ntypes, name="ntypes", dtype=tf.int32) + t_rcut = tf.constant( + self.rcut, name="rcut", dtype=GLOBAL_TF_FLOAT_PRECISION + ) + coord = tf.reshape(coord_, [-1, natoms[1] * 3]) + atype = tf.reshape(atype_, [-1, natoms[1]]) + box = tf.reshape(box, [-1, 9]) + # perhaps we need a OP that only outputs rij and nlist + ( + _, + _, + rij, + nlist, + _, + _, + ) = op_module.prod_env_mat_a_mix( + coord, + atype, + natoms, + box, + mesh, + np.zeros([self.ntypes, self.sel * 4]), + np.ones([self.ntypes, self.sel * 4]), + rcut_a=-1, + rcut_r=self.rcut, + rcut_r_smth=self.rcut, + sel_a=[self.sel], + sel_r=[0], + ) + scale = tf.ones([tf.shape(coord)[0], natoms[0]], dtype=tf.float64) + tab_atom_ener, tab_force, tab_atom_virial = op_module.pair_tab( + self.tab_info, + self.tab_data, + atype, + rij, + nlist, + natoms, + scale, + sel_a=[self.sel], + sel_r=[0], + ) + energy_raw = tf.reshape( + tab_atom_ener, [-1, natoms[0]], name="o_atom_energy" + suffix + ) + energy = tf.reduce_sum( + global_cvt_2_ener_float(energy_raw), axis=1, name="o_energy" + suffix + ) + force = tf.reshape(tab_force, [-1, 3 * natoms[1]], name="o_force" + suffix) + virial = tf.reshape( + tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis=1), + [-1, 9], + name="o_virial" + suffix, + ) + atom_virial = tf.reshape( + tab_atom_virial, [-1, 9 * natoms[1]], name="o_atom_virial" + suffix + ) + model_dict = {} + model_dict["energy"] = energy + model_dict["force"] = force + model_dict["virial"] = virial + model_dict["atom_ener"] = energy_raw + model_dict["atom_virial"] = atom_virial + model_dict["coord"] = coord + model_dict["atype"] = atype + + return model_dict + + def init_variables( + self, + graph: tf.Graph, + graph_def: tf.GraphDef, + model_type: str = "original_model", + suffix: str = "", + ) -> None: + """Init the embedding net 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 + """ + # skip. table can be initialized from the file + + def get_fitting(self) -> Union[Fitting, dict]: + """Get the fitting(s).""" + # nothing needs to do + return {} + + def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + """Get the loss function(s).""" + # nothing nees to do + return + + def get_rcut(self) -> float: + """Get cutoff radius of the model.""" + return self.rcut + + def get_ntypes(self) -> int: + """Get the number of types.""" + return self.ntypes + + def data_stat(self, data: dict): + """Data staticis.""" + # nothing needs to do + + def enable_compression(self, suffix: str = "") -> None: + """Enable compression. + + Parameters + ---------- + suffix : str + suffix to name scope + """ + # nothing needs to do + + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + """Update the selection and perform neighbor statistics. + + Notes + ----- + Do not modify the input data without copying it. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + + Returns + ------- + dict + The updated local data + """ + from deepmd.entrypoints.train import ( + update_one_sel, + ) + + local_jdata_cpy = local_jdata.copy() + return update_one_sel(global_jdata, local_jdata_cpy, True) diff --git a/deepmd/model/pairwise_dprc.py b/deepmd/model/pairwise_dprc.py index 6983a31cfd..f74571febb 100644 --- a/deepmd/model/pairwise_dprc.py +++ b/deepmd/model/pairwise_dprc.py @@ -173,10 +173,12 @@ def build( atype_qmmm = gather_placeholder(atype_qmmm, forward_qmmm_map, placeholder=-1) box_qm = box - type_embedding = self.typeebd.build( + type_embedding = self.build_type_embedding( self.ntypes, reuse=reuse, suffix=suffix, + frz_model=frz_model, + ckpt_meta=ckpt_meta, ) input_dict_qm["type_embedding"] = type_embedding input_dict_qmmm["type_embedding"] = type_embedding diff --git a/deepmd/model/tensor.py b/deepmd/model/tensor.py index 9099b753a4..6a21e085f3 100644 --- a/deepmd/model/tensor.py +++ b/deepmd/model/tensor.py @@ -135,10 +135,12 @@ def build( # type embedding if any if self.typeebd is not None: - type_embedding = self.typeebd.build( + type_embedding = self.build_type_embedding( self.ntypes, reuse=reuse, suffix=suffix, + ckpt_meta=ckpt_meta, + frz_model=frz_model, ) input_dict["type_embedding"] = type_embedding input_dict["atype"] = atype_ diff --git a/deepmd/nvnmd/data/data.py b/deepmd/nvnmd/data/data.py index 29c8b84a37..9e6dd4cc89 100644 --- a/deepmd/nvnmd/data/data.py +++ b/deepmd/nvnmd/data/data.py @@ -60,6 +60,7 @@ }, "ctrl": { # NSTDM + "MAX_NNEI": 128, "NSTDM": 64, "NSTDM_M1": 32, "NSTDM_M2": 2, @@ -67,6 +68,7 @@ "NSEL": "NSTDM*NTYPE_MAX", "NSADV": "NSTDM+1", "VERSION": 0, + "SUB_VERSION": 1, }, "nbit": { # general @@ -116,6 +118,22 @@ "end": "", } +# change the configuration accordng to the max_nnei +jdata_config_v0_ni128 = jdata_config_v0.copy() +jdata_config_v0_ni256 = jdata_config_v0.copy() +jdata_config_v0_ni256["ctrl"] = { + "MAX_NNEI": 256, + "NSTDM": 128, + "NSTDM_M1": 32, + "NSTDM_M2": 4, + "NSTDM_M1X": 8, + "NSEL": "NSTDM*NTYPE_MAX", + "NSADV": "NSTDM+1", + "VERSION": 0, + "SUB_VERSION": 1, +} +jdata_config_v0_ni256["nbit"]["NBIT_NEIB"] = 9 + jdata_config_v1 = { "dscp": { # basic config from deepmd model @@ -174,6 +192,7 @@ }, "ctrl": { # NSTDM + "MAX_NNEI": 128, "NSTDM": 64, "NSTDM_M1": 32, "NSTDM_M2": 2, @@ -181,6 +200,7 @@ "NSEL": "NSTDM", "NSADV": "NSTDM+1", "VERSION": 1, + "SUB_VERSION": 1, }, "nbit": { # general @@ -230,6 +250,22 @@ "end": "", } +# change the configuration accordng to the max_nnei +jdata_config_v1_ni128 = jdata_config_v1.copy() +jdata_config_v1_ni256 = jdata_config_v1.copy() +jdata_config_v1_ni256["ctrl"] = { + "MAX_NNEI": 256, + "NSTDM": 128, + "NSTDM_M1": 32, + "NSTDM_M2": 4, + "NSTDM_M1X": 8, + "NSEL": "NSTDM", + "NSADV": "NSTDM+1", + "VERSION": 1, + "SUB_VERSION": 1, +} +jdata_config_v1_ni256["nbit"]["NBIT_NEIB"] = 9 + jdata_deepmd_input_v0 = { "model": { "descriptor": { @@ -247,6 +283,7 @@ }, "nvnmd": { "version": 0, + "max_nnei": 128, # 128 or 256 "net_size": 128, "config_file": "none", "weight_file": "none", @@ -286,6 +323,10 @@ }, } +jdata_deepmd_input_v0_ni128 = jdata_deepmd_input_v0.copy() +jdata_deepmd_input_v0_ni256 = jdata_deepmd_input_v0.copy() +jdata_deepmd_input_v0_ni256["nvnmd"]["max_nnei"] = 256 + jdata_deepmd_input_v1 = { "model": { "descriptor": { @@ -308,6 +349,7 @@ }, "nvnmd": { "version": 1, + "max_nnei": 128, # 128 or 256 "net_size": 128, "config_file": "none", "weight_file": "none", @@ -347,6 +389,10 @@ }, } +jdata_deepmd_input_v1_ni128 = jdata_deepmd_input_v1.copy() +jdata_deepmd_input_v1_ni256 = jdata_deepmd_input_v1.copy() +jdata_deepmd_input_v1_ni256["nvnmd"]["max_nnei"] = 256 + NVNMD_WELCOME = ( r" _ _ __ __ _ _ __ __ ____ ", r"| \ | | \ \ / / | \ | | | \/ | | _ \ ", diff --git a/deepmd/nvnmd/descriptor/se_a.py b/deepmd/nvnmd/descriptor/se_a.py index 67ea45924b..816f17cfa3 100644 --- a/deepmd/nvnmd/descriptor/se_a.py +++ b/deepmd/nvnmd/descriptor/se_a.py @@ -50,12 +50,17 @@ def check_switch_range(davg, dstd): else: min_dist = nvnmd_cfg.weight["train_attr.min_nbor_dist"] else: - min_dist = rmin + min_dist = None + + # fix the bug: if model initial mode is 'init_from_model', + # we need dmin to calculate smin and smax in mapt.py + if min_dist is not None: + nvnmd_cfg.dscp["dmin"] = min_dist + nvnmd_cfg.save() # if davg and dstd is None, the model initial mode is in # 'init_from_model', 'restart', 'init_from_frz_model', 'finetune' if (davg is not None) and (dstd is not None): - nvnmd_cfg.dscp["dmin"] = min_dist nvnmd_cfg.get_s_range(davg, dstd) diff --git a/deepmd/nvnmd/descriptor/se_atten.py b/deepmd/nvnmd/descriptor/se_atten.py index 727a93ca45..cfffb8a90b 100644 --- a/deepmd/nvnmd/descriptor/se_atten.py +++ b/deepmd/nvnmd/descriptor/se_atten.py @@ -49,7 +49,13 @@ def check_switch_range(davg, dstd): else: min_dist = nvnmd_cfg.weight["train_attr.min_nbor_dist"] else: - min_dist = rmin + min_dist = None + + # fix the bug: if model initial mode is 'init_from_model', + # we need dmin to calculate smin and smax in mapt.py + if min_dist is not None: + nvnmd_cfg.dscp["dmin"] = min_dist + nvnmd_cfg.save() # if davg and dstd is None, the model initial mode is in # 'init_from_model', 'restart', 'init_from_frz_model', 'finetune' @@ -58,7 +64,6 @@ def check_switch_range(davg, dstd): davg = np.zeros([ntype, ndescrpt]) if dstd is None: dstd = np.ones([ntype, ndescrpt]) - nvnmd_cfg.dscp["dmin"] = min_dist nvnmd_cfg.get_s_range(davg, dstd) diff --git a/deepmd/nvnmd/entrypoints/freeze.py b/deepmd/nvnmd/entrypoints/freeze.py index 6c356c6118..e56a0c2130 100644 --- a/deepmd/nvnmd/entrypoints/freeze.py +++ b/deepmd/nvnmd/entrypoints/freeze.py @@ -52,6 +52,7 @@ def filter_tensorVariableList(tensorVariableList) -> dict: p1 = p1 or name.startswith("filter_type_") p1 = p1 or name.startswith("layer_") p1 = p1 or name.startswith("final_layer") + p1 = p1 or name.endswith("t_bias_atom_e") p2 = "Adam" not in name p3 = "XXX" not in name if p1 and p2 and p3: @@ -75,4 +76,5 @@ def save_weight(sess, file_name: str = "nvnmd/weight.npy"): else: min_dist = 0.0 dic_key_value["train_attr.min_nbor_dist"] = min_dist + dic_key_value["t_bias_atom_e"] = dic_key_value["fitting_attr.t_bias_atom_e"] FioDic().save(file_name, dic_key_value) diff --git a/deepmd/nvnmd/entrypoints/mapt.py b/deepmd/nvnmd/entrypoints/mapt.py index eb77913983..1299d7a74e 100644 --- a/deepmd/nvnmd/entrypoints/mapt.py +++ b/deepmd/nvnmd/entrypoints/mapt.py @@ -87,9 +87,22 @@ def __init__(self, config_file: str, weight_file: str, map_file: str): jdata["weight_file"] = weight_file jdata["enable"] = True + # 0 : xyz_scatter = xyz_scatter * two_embd + xyz_scatter; + # Gs + 1, Gt + 0 + # 1 : xyz_scatter = xyz_scatter * two_embd + two_embd ; + # Gs + 0, Gt + 1 + self.Gs_Gt_mode = 1 + nvnmd_cfg.init_from_jdata(jdata) def build_map(self): + if self.Gs_Gt_mode == 0: + self.shift_Gs = 1 + self.shift_Gt = 0 + if self.Gs_Gt_mode == 1: + self.shift_Gs = 0 + self.shift_Gt = 1 + # M = nvnmd_cfg.dscp["M1"] if nvnmd_cfg.version == 0: ndim = nvnmd_cfg.dscp["ntype"] @@ -482,7 +495,7 @@ def build_s2g_grad(self): shift = 0 if nvnmd_cfg.version == 1: ndim = 1 - shift = 1 + shift = self.shift_Gs # dic_ph = {} dic_ph["s"] = tf.placeholder(tf.float64, [None, 1], "t_s") @@ -496,6 +509,13 @@ def run_s2g(self): r"""Build s-> graph and run it to get value of mapping table.""" smin = nvnmd_cfg.dscp["smin"] smax = nvnmd_cfg.dscp["smax"] + # fix the bug: if model initial mode is 'init_from_model', + # we need dmin to calculate smin and smax in mapt.py + if smin == -2: + davg, dstd = get_normalize(nvnmd_cfg.weight) + nvnmd_cfg.get_s_range(davg, dstd) + smin = nvnmd_cfg.dscp["smin"] + smax = nvnmd_cfg.dscp["smax"] tf.reset_default_graph() dic_ph = self.build_s2g_grad() @@ -567,9 +587,11 @@ def build_t2g(self): two_side_type_embedding, [-1, two_side_type_embedding.shape[-1]], ) - + # see se_atten.py in dp wbs = [get_filter_type_weight(nvnmd_cfg.weight, ll) for ll in range(1, 5)] - dic_ph["gt"] = self.build_embedding_net(two_side_type_embedding, wbs) + dic_ph["gt"] = ( + self.build_embedding_net(two_side_type_embedding, wbs) + self.shift_Gt + ) return dic_ph def run_t2g(self): diff --git a/deepmd/nvnmd/entrypoints/train.py b/deepmd/nvnmd/entrypoints/train.py index cb3dad0792..6e14b6f865 100644 --- a/deepmd/nvnmd/entrypoints/train.py +++ b/deepmd/nvnmd/entrypoints/train.py @@ -100,6 +100,7 @@ def normalized_input_qnn(jdata, PATH_QNN, CONFIG_CNN, WEIGHT_CNN, MAP_CNN): jdata_nvnmd = jdata_deepmd_input_v0["nvnmd"] jdata_nvnmd["enable"] = True jdata_nvnmd["version"] = nvnmd_cfg.version + jdata_nvnmd["max_nnei"] = nvnmd_cfg.max_nnei jdata_nvnmd["config_file"] = CONFIG_CNN jdata_nvnmd["weight_file"] = WEIGHT_CNN jdata_nvnmd["map_file"] = MAP_CNN @@ -117,6 +118,7 @@ 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], step: str, skip_neighbor_stat: bool = False, @@ -142,6 +144,7 @@ def train_nvnmd( jdata = jdata_cmd_train.copy() jdata["INPUT"] = INPUT_CNN jdata["log_path"] = LOG_CNN + jdata["init_model"] = init_model jdata["restart"] = restart jdata["skip_neighbor_stat"] = skip_neighbor_stat train(**jdata) diff --git a/deepmd/nvnmd/entrypoints/wrap.py b/deepmd/nvnmd/entrypoints/wrap.py index 455dd999df..1ba2ed7384 100644 --- a/deepmd/nvnmd/entrypoints/wrap.py +++ b/deepmd/nvnmd/entrypoints/wrap.py @@ -156,33 +156,75 @@ def wrap_head(self, nhs, nws): r"""Wrap the head information. version + nhead nheight - nweight - rcut + nwidth + rcut cut-off radius + ntype number of atomic species + nnei number of neighbors + atom_ener atom bias energy """ nbit = nvnmd_cfg.nbit ctrl = nvnmd_cfg.ctrl + dscp = nvnmd_cfg.dscp + fitn = nvnmd_cfg.fitn + weight = nvnmd_cfg.weight VERSION = ctrl["VERSION"] + SUB_VERSION = ctrl["SUB_VERSION"] + MAX_NNEI = ctrl["MAX_NNEI"] + nhead = 128 NBIT_MODEL_HEAD = nbit["NBIT_MODEL_HEAD"] NBIT_FIXD_FL = nbit["NBIT_FIXD_FL"] - rcut = nvnmd_cfg.dscp["rcut"] + rcut = dscp["rcut"] + ntype = dscp["ntype"] + SEL = dscp["SEL"] bs = "" e = Encode() # version - bs = e.dec2bin(VERSION, NBIT_MODEL_HEAD)[0] + bs + vv = VERSION + 256 * SUB_VERSION + 256 * 256 * MAX_NNEI + bs = e.dec2bin(vv, NBIT_MODEL_HEAD)[0] + bs + # nhead + bs = e.dec2bin(nhead, NBIT_MODEL_HEAD)[0] + bs # height for n in nhs: bs = e.dec2bin(n, NBIT_MODEL_HEAD)[0] + bs - # weight + # width for n in nws: bs = e.dec2bin(n, NBIT_MODEL_HEAD)[0] + bs - # dscp + # rcut RCUT = e.qr(rcut, NBIT_FIXD_FL) bs = e.dec2bin(RCUT, NBIT_MODEL_HEAD)[0] + bs + # ntype + bs = e.dec2bin(ntype, NBIT_MODEL_HEAD)[0] + bs + # nnei + if VERSION == 0: + for tt in range(ntype): + bs = e.dec2bin(SEL[tt], NBIT_MODEL_HEAD)[0] + bs + if VERSION == 1: + bs = e.dec2bin(SEL, NBIT_MODEL_HEAD)[0] + bs + # atom_ener + # fix the bug: the different energy between qnn and lammps + if "t_bias_atom_e" in weight.keys(): + atom_ener = weight["t_bias_atom_e"] + else: + atom_ener = [0] * 32 + nlayer_fit = fitn["nlayer_fit"] + if VERSION == 0: + for tt in range(ntype): + w, b, _idt = get_fitnet_weight(weight, tt, nlayer_fit - 1, nlayer_fit) + shift = atom_ener[tt] + b[0] + SHIFT = e.qr(shift, NBIT_FIXD_FL) + bs = e.dec2bin(SHIFT, NBIT_MODEL_HEAD, signed=True)[0] + bs + if VERSION == 1: + for tt in range(ntype): + w, b, _idt = get_fitnet_weight(weight, 0, nlayer_fit - 1, nlayer_fit) + shift = atom_ener[tt] + b[0] + SHIFT = e.qr(shift, NBIT_FIXD_FL) + bs = e.dec2bin(SHIFT, NBIT_MODEL_HEAD, signed=True)[0] + bs # extend hs = e.bin2hex(bs) - hs = e.extend_hex(hs, NBIT_MODEL_HEAD * 32) + hs = e.extend_hex(hs, NBIT_MODEL_HEAD * nhead) return hs def wrap_dscp(self): diff --git a/deepmd/nvnmd/utils/argcheck.py b/deepmd/nvnmd/utils/argcheck.py index 2cbff3cbdc..2b9362efb0 100644 --- a/deepmd/nvnmd/utils/argcheck.py +++ b/deepmd/nvnmd/utils/argcheck.py @@ -1,68 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from dargs import ( - Argument, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.argcheck_nvnmd import ( + nvnmd_args, ) - -def nvnmd_args(): - doc_version = ( - "configuration the nvnmd version (0 | 1), 0 for 4 types, 1 for 32 types" - ) - doc_net_size_file = ( - "configuration the number of nodes of fitting_net, just can be set as 128" - ) - doc_map_file = "A file containing the mapping tables to replace the calculation of embedding nets" - doc_config_file = "A file containing the parameters about how to implement the model in certain hardware" - doc_weight_file = "a *.npy file containing the weights of the model" - doc_enable = "enable the nvnmd training" - doc_restore_descriptor = ( - "enable to restore the parameter of embedding_net from weight.npy" - ) - doc_restore_fitting_net = ( - "enable to restore the parameter of fitting_net from weight.npy" - ) - doc_quantize_descriptor = "enable the quantizatioin of descriptor" - doc_quantize_fitting_net = "enable the quantizatioin of fitting_net" - args = [ - Argument("version", int, optional=False, default=0, doc=doc_version), - Argument("net_size", int, optional=False, default=128, doc=doc_net_size_file), - Argument("map_file", str, optional=False, default="none", doc=doc_map_file), - Argument( - "config_file", str, optional=False, default="none", doc=doc_config_file - ), - Argument( - "weight_file", str, optional=False, default="none", doc=doc_weight_file - ), - Argument("enable", bool, optional=False, default=False, doc=doc_enable), - Argument( - "restore_descriptor", - bool, - optional=False, - default=False, - doc=doc_restore_descriptor, - ), - Argument( - "restore_fitting_net", - bool, - optional=False, - default=False, - doc=doc_restore_fitting_net, - ), - Argument( - "quantize_descriptor", - bool, - optional=False, - default=False, - doc=doc_quantize_descriptor, - ), - Argument( - "quantize_fitting_net", - bool, - optional=False, - default=False, - doc=doc_quantize_fitting_net, - ), - ] - - doc_nvnmd = "The nvnmd options." - return Argument("nvnmd", dict, args, [], optional=True, doc=doc_nvnmd) +__all__ = [ + "nvnmd_args", +] diff --git a/deepmd/nvnmd/utils/config.py b/deepmd/nvnmd/utils/config.py index 96ca74c4c9..5bfd9ea54f 100644 --- a/deepmd/nvnmd/utils/config.py +++ b/deepmd/nvnmd/utils/config.py @@ -7,9 +7,15 @@ NVNMD_CITATION, NVNMD_WELCOME, jdata_config_v0, - jdata_config_v1, + jdata_config_v0_ni128, + jdata_config_v0_ni256, + jdata_config_v1_ni128, + jdata_config_v1_ni256, jdata_deepmd_input_v0, - jdata_deepmd_input_v1, + jdata_deepmd_input_v0_ni128, + jdata_deepmd_input_v0_ni256, + jdata_deepmd_input_v1_ni128, + jdata_deepmd_input_v1_ni256, ) from deepmd.nvnmd.utils.fio import ( FioDic, @@ -50,6 +56,7 @@ def init_from_jdata(self, jdata: dict = {}): return None self.version = jdata["version"] + self.max_nnei = jdata["max_nnei"] self.net_size = jdata["net_size"] self.map_file = jdata["map_file"] self.config_file = jdata["config_file"] @@ -65,7 +72,7 @@ def init_from_jdata(self, jdata: dict = {}): self.map = FioDic().load(self.map_file, {}) self.weight = FioDic().load(self.weight_file, {}) - self.init_config_by_version(self.version) + self.init_config_by_version(self.version, self.max_nnei) load_config = FioDic().load(self.config_file, self.config) self.init_from_config(load_config) # if load the file, set net_size @@ -106,7 +113,11 @@ def init_from_config(self, jdata): r"""Initialize member element one by one.""" if "ctrl" in jdata.keys(): if "VERSION" in jdata["ctrl"].keys(): - self.init_config_by_version(jdata["ctrl"]["VERSION"]) + if "MAX_NNEI" not in jdata["ctrl"].keys(): + jdata["ctrl"]["MAX_NNEI"] = 128 + self.init_config_by_version( + jdata["ctrl"]["VERSION"], jdata["ctrl"]["MAX_NNEI"] + ) # self.config = FioDic().update(jdata, self.config) self.config["dscp"] = self.init_dscp(self.config["dscp"], self.config) @@ -117,16 +128,29 @@ def init_from_config(self, jdata): self.config["nbit"] = self.init_nbit(self.config["nbit"], self.config) self.init_value() - def init_config_by_version(self, version): + def init_config_by_version(self, version, max_nnei): r"""Initialize version-dependent parameters.""" self.version = version + self.max_nnei = max_nnei log.debug("#Set nvnmd version as %d " % self.version) if self.version == 0: - self.jdata_deepmd_input = jdata_deepmd_input_v0.copy() - self.config = jdata_config_v0.copy() + if self.max_nnei == 128: + self.jdata_deepmd_input = jdata_deepmd_input_v0_ni128.copy() + self.config = jdata_config_v0_ni128.copy() + elif self.max_nnei == 256: + self.jdata_deepmd_input = jdata_deepmd_input_v0_ni256.copy() + self.config = jdata_config_v0_ni256.copy() + else: + log.error("The max_nnei only can be set as 128|256 for version 0") if self.version == 1: - self.jdata_deepmd_input = jdata_deepmd_input_v1.copy() - self.config = jdata_config_v1.copy() + if self.max_nnei == 128: + self.jdata_deepmd_input = jdata_deepmd_input_v1_ni128.copy() + self.config = jdata_config_v1_ni128.copy() + elif self.max_nnei == 256: + self.jdata_deepmd_input = jdata_deepmd_input_v1_ni256.copy() + self.config = jdata_config_v1_ni256.copy() + else: + log.error("The max_nnei only can be set as 128|256 for version 1") def init_net_size(self): r"""Initialize net_size.""" @@ -154,10 +178,15 @@ def init_dscp(self, jdata: dict, jdata_parent: dict = {}) -> dict: jdata["M1"] = jdata["neuron"][-1] jdata["M2"] = jdata["axis_neuron"] jdata["SEL"] = (jdata["sel"] + [0, 0, 0, 0])[0:4] + for s in jdata["sel"]: + if s > self.max_nnei: + log.error("The sel cannot be greater than the max_nnei") + exit(1) jdata["NNODE_FEAS"] = [1] + jdata["neuron"] jdata["nlayer_fea"] = len(jdata["neuron"]) jdata["same_net"] = 1 if jdata["type_one_side"] else 0 # neighbor + jdata["NI"] = self.max_nnei jdata["NIDP"] = int(np.sum(jdata["sel"])) jdata["NIX"] = 2 ** int(np.ceil(np.log2(jdata["NIDP"] / 1.5))) # type @@ -168,10 +197,14 @@ def init_dscp(self, jdata: dict, jdata_parent: dict = {}) -> dict: jdata["M1"] = jdata["neuron"][-1] jdata["M2"] = jdata["axis_neuron"] jdata["SEL"] = jdata["sel"] + if jdata["sel"] > self.max_nnei: + log.error("The sel cannot be greater than the max_nnei") + exit(1) jdata["NNODE_FEAS"] = [1] + jdata["neuron"] jdata["nlayer_fea"] = len(jdata["neuron"]) jdata["same_net"] = 1 if jdata["type_one_side"] else 0 # neighbor + jdata["NI"] = self.max_nnei jdata["NIDP"] = int(jdata["sel"]) jdata["NIX"] = 2 ** int(np.ceil(np.log2(jdata["NIDP"] / 1.5))) # type @@ -306,6 +339,7 @@ def get_nvnmd_jdata(self): r"""Generate `nvnmd` in input script.""" jdata = self.jdata_deepmd_input["nvnmd"] jdata["net_size"] = self.net_size + jdata["max_nnei"] = self.max_nnei jdata["config_file"] = self.config_file jdata["weight_file"] = self.weight_file jdata["map_file"] = self.map_file diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index bbcb305404..3b81740a93 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -943,6 +943,7 @@ def print_header(fp, train_results, valid_results, multi_task_mode=False): for k in train_results[fitting_key].keys(): print_str += prop_fmt % (k + "_trn") print_str += " %8s\n" % (fitting_key + "_lr") + print_str += "# If there is no available reference data, rmse_*_{val,trn} will print nan\n" fp.write(print_str) fp.flush() diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 7104eb1de4..05e7c767b8 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1,2015 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import json -import logging -from typing import ( - Callable, - List, - Optional, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.argcheck import ( + gen_args, + gen_doc, + gen_json, + list_to_doc, + normalize, + type_embedding_args, ) -from dargs import ( - Argument, - ArgumentEncoder, - Variant, - dargs, -) - -from deepmd.common import ( - ACTIVATION_FN_DICT, - PRECISION_DICT, -) -from deepmd.nvnmd.utils.argcheck import ( - nvnmd_args, -) -from deepmd.utils.plugin import ( - Plugin, -) - -log = logging.getLogger(__name__) - - -def list_to_doc(xx): - items = [] - for ii in xx: - if len(items) == 0: - items.append(f'"{ii}"') - else: - items.append(f', "{ii}"') - items.append(".") - return "".join(items) - - -def make_link(content, ref_key): - return ( - f"`{content} <{ref_key}_>`_" - if not dargs.RAW_ANCHOR - else f"`{content} <#{ref_key}>`_" - ) - - -def type_embedding_args(): - 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" - doc_activation_function = f'The activation function in the embedding 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.' - doc_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_trainable = "If the parameters in the embedding net are trainable" - - return [ - Argument("neuron", List[int], optional=True, default=[8], doc=doc_neuron), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=False, doc=doc_resnet_dt), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, default=None, doc=doc_seed), - ] - - -def spin_args(): - doc_use_spin = "Whether to use atomic spin model for each atom type" - doc_spin_norm = "The magnitude of atomic spin for each atom type with spin" - doc_virtual_len = "The distance between virtual atom representing spin and its corresponding real atom for each atom type with spin" - - return [ - Argument("use_spin", List[bool], doc=doc_use_spin), - Argument("spin_norm", List[float], doc=doc_spin_norm), - Argument("virtual_len", List[float], doc=doc_virtual_len), - ] - - -# --- Descriptor configurations: --- # - - -class ArgsPlugin: - def __init__(self) -> None: - self.__plugin = Plugin() - - def register( - self, name: str, alias: Optional[List[str]] = None - ) -> Callable[[], List[Argument]]: - """Register a descriptor argument plugin. - - Parameters - ---------- - name : str - the name of a descriptor - alias : List[str], optional - the list of aliases of this descriptor - - Returns - ------- - Callable[[], List[Argument]] - the registered descriptor argument method - - Examples - -------- - >>> some_plugin = ArgsPlugin() - >>> @some_plugin.register("some_descrpt") - def descrpt_some_descrpt_args(): - return [] - """ - # convert alias to hashed item - if isinstance(alias, list): - alias = tuple(alias) - return self.__plugin.register((name, alias)) - - def get_all_argument(self, exclude_hybrid: bool = False) -> List[Argument]: - """Get all arguments. - - Parameters - ---------- - exclude_hybrid : bool - exclude hybrid descriptor to prevent circular calls - - Returns - ------- - List[Argument] - all arguments - """ - arguments = [] - for (name, alias), metd in self.__plugin.plugins.items(): - if exclude_hybrid and name == "hybrid": - continue - arguments.append( - Argument(name=name, dtype=dict, sub_fields=metd(), alias=alias) - ) - return arguments - - -descrpt_args_plugin = ArgsPlugin() - - -@descrpt_args_plugin.register("loc_frame") -def descrpt_local_frame_args(): - 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" - doc_axis_rule = "A list of integers. The length should be 6 times of the number of types. \n\n\ -- axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.\n\n\ -- axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom.\n\n\ -- axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance.\n\n\ -- axis_rule[i*6+3]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.\n\n\ -- axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom.\n\n\ -- axis_rule[i*6+5]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance." - - return [ - Argument("sel_a", List[int], optional=False, doc=doc_sel_a), - Argument("sel_r", List[int], optional=False, doc=doc_sel_r), - Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), - Argument("axis_rule", List[int], optional=False, doc=doc_axis_rule), - ] - - -@descrpt_args_plugin.register("se_e2_a", alias=["se_a"]) -def descrpt_se_a_args(): - 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 wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' - 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`" - 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_axis_neuron = "Size of the submatrix of G (embedding matrix)." - doc_activation_function = f'The activation function in the embedding 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_type_one_side = r"If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters." - doc_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_trainable = "If the parameters in the embedding net is trainable" - doc_seed = "Random seed for parameter initialization" - 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_set_davg_zero = "Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used" - - return [ - Argument("sel", [List[int], str], optional=True, default="auto", doc=doc_sel), - Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), - Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth), - Argument( - "neuron", List[int], optional=True, default=[10, 20, 40], doc=doc_neuron - ), - Argument( - "axis_neuron", - int, - optional=True, - default=4, - alias=["n_axis_neuron"], - doc=doc_axis_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=False, doc=doc_resnet_dt), - Argument( - "type_one_side", bool, optional=True, default=False, doc=doc_type_one_side - ), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "exclude_types", - List[List[int]], - optional=True, - default=[], - doc=doc_exclude_types, - ), - Argument( - "set_davg_zero", bool, optional=True, default=False, doc=doc_set_davg_zero - ), - ] - - -@descrpt_args_plugin.register("se_e3", alias=["se_at", "se_a_3be", "se_t"]) -def descrpt_se_t_args(): - 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 wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' - 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`" - 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_activation_function = f'The activation function in the embedding 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_trainable = "If the parameters in the embedding net are trainable" - doc_seed = "Random seed for parameter initialization" - doc_set_davg_zero = "Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used" - - return [ - Argument("sel", [List[int], str], optional=True, default="auto", doc=doc_sel), - Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), - Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth), - Argument( - "neuron", List[int], optional=True, default=[10, 20, 40], doc=doc_neuron - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=False, doc=doc_resnet_dt), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "set_davg_zero", bool, optional=True, default=False, doc=doc_set_davg_zero - ), - ] - - -@descrpt_args_plugin.register("se_a_tpe", alias=["se_a_ebd"]) -def descrpt_se_a_tpe_args(): - 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." - - return [ - *descrpt_se_a_args(), - Argument("type_nchanl", int, optional=True, default=4, doc=doc_type_nchanl), - Argument("type_nlayer", int, optional=True, default=2, doc=doc_type_nlayer), - Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), - ] - - -@descrpt_args_plugin.register("se_e2_r", alias=["se_r"]) -def descrpt_se_r_args(): - 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 wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' - 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`" - 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_activation_function = f'The activation function in the embedding 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_type_one_side = r"If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters." - doc_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_trainable = "If the parameters in the embedding net are trainable" - doc_seed = "Random seed for parameter initialization" - 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_set_davg_zero = "Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used" - - return [ - Argument("sel", [List[int], str], optional=True, default="auto", doc=doc_sel), - Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), - Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth), - Argument( - "neuron", List[int], optional=True, default=[10, 20, 40], doc=doc_neuron - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=False, doc=doc_resnet_dt), - Argument( - "type_one_side", bool, optional=True, default=False, doc=doc_type_one_side - ), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "exclude_types", - List[List[int]], - optional=True, - default=[], - doc=doc_exclude_types, - ), - Argument( - "set_davg_zero", bool, optional=True, default=False, doc=doc_set_davg_zero - ), - ] - - -@descrpt_args_plugin.register("hybrid") -def descrpt_hybrid_args(): - doc_list = "A list of descriptor definitions" - - return [ - Argument( - "list", - list, - optional=False, - doc=doc_list, - repeat=True, - sub_fields=[], - sub_variants=[descrpt_variant_type_args(exclude_hybrid=True)], - fold_subdoc=True, - ) - ] - - -def descrpt_se_atten_common_args(): - 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.\ - - `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 wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' - 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`" - 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_axis_neuron = "Size of the submatrix of G (embedding matrix)." - doc_activation_function = f'The activation function in the embedding 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_type_one_side = r"If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters." - doc_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_trainable = "If the parameters in the embedding net is trainable" - doc_seed = "Random seed for parameter initialization" - 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_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 stripped_type_embedding is True" - 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" - - return [ - Argument( - "sel", [int, List[int], str], optional=True, default="auto", doc=doc_sel - ), - Argument("rcut", float, optional=True, default=6.0, doc=doc_rcut), - Argument("rcut_smth", float, optional=True, default=0.5, doc=doc_rcut_smth), - Argument( - "neuron", List[int], optional=True, default=[10, 20, 40], doc=doc_neuron - ), - Argument( - "axis_neuron", - int, - optional=True, - default=4, - alias=["n_axis_neuron"], - doc=doc_axis_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=False, doc=doc_resnet_dt), - Argument( - "type_one_side", bool, optional=True, default=False, doc=doc_type_one_side - ), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "exclude_types", - List[List[int]], - optional=True, - default=[], - doc=doc_exclude_types, - ), - Argument("attn", int, optional=True, default=128, doc=doc_attn), - Argument("attn_layer", int, optional=True, default=2, doc=doc_attn_layer), - Argument("attn_dotr", bool, optional=True, default=True, doc=doc_attn_dotr), - Argument("attn_mask", bool, optional=True, default=False, doc=doc_attn_mask), - ] - - -@descrpt_args_plugin.register("se_atten") -def descrpt_se_atten_args(): - doc_stripped_type_embedding = "Whether to strip the type embedding into a separated embedding network. Setting it to `False` will fall back to the previous version of `se_atten` which is non-compressible." - doc_smooth_type_embdding = "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" - - return [ - *descrpt_se_atten_common_args(), - Argument( - "stripped_type_embedding", - bool, - optional=True, - default=False, - doc=doc_stripped_type_embedding, - ), - Argument( - "smooth_type_embdding", - bool, - optional=True, - default=False, - doc=doc_smooth_type_embdding, - ), - Argument( - "set_davg_zero", bool, optional=True, default=True, doc=doc_set_davg_zero - ), - ] - - -@descrpt_args_plugin.register("se_atten_v2") -def descrpt_se_atten_v2_args(): - 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" - - return [ - *descrpt_se_atten_common_args(), - Argument( - "set_davg_zero", bool, optional=True, default=False, doc=doc_set_davg_zero - ), - ] - - -@descrpt_args_plugin.register("se_a_ebd_v2", alias=["se_a_tpe_v2"]) -def descrpt_se_a_ebd_v2_args(): - return descrpt_se_a_args() - - -@descrpt_args_plugin.register("se_a_mask") -def descrpt_se_a_mask_args(): - 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 wraped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".' - - 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_axis_neuron = "Size of the submatrix of G (embedding matrix)." - doc_activation_function = f'The activation function in the embedding 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_type_one_side = r"If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters." - 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_precision = f"The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_trainable = "If the parameters in the embedding net is trainable" - doc_seed = "Random seed for parameter initialization" - - return [ - Argument("sel", [List[int], str], optional=True, default="auto", doc=doc_sel), - Argument( - "neuron", List[int], optional=True, default=[10, 20, 40], doc=doc_neuron - ), - Argument( - "axis_neuron", - int, - optional=True, - default=4, - alias=["n_axis_neuron"], - doc=doc_axis_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=False, doc=doc_resnet_dt), - Argument( - "type_one_side", bool, optional=True, default=False, doc=doc_type_one_side - ), - Argument( - "exclude_types", - List[List[int]], - optional=True, - default=[], - doc=doc_exclude_types, - ), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, doc=doc_seed), - ] - - -def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: - link_lf = make_link("loc_frame", "model/descriptor[loc_frame]") - link_se_e2_a = make_link("se_e2_a", "model/descriptor[se_e2_a]") - link_se_e2_r = make_link("se_e2_r", "model/descriptor[se_e2_r]") - link_se_e3 = make_link("se_e3", "model/descriptor[se_e3]") - link_se_a_tpe = make_link("se_a_tpe", "model/descriptor[se_a_tpe]") - link_hybrid = make_link("hybrid", "model/descriptor[hybrid]") - link_se_atten = make_link("se_atten", "model/descriptor[se_atten]") - link_se_atten_v2 = make_link("se_atten_v2", "model/descriptor[se_atten_v2]") - doc_descrpt_type = "The type of the descritpor. See explanation below. \n\n\ -- `loc_frame`: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.\n\n\ -- `se_e2_a`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor.\n\n\ -- `se_e2_r`: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.\n\n\ -- `se_e3`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor.\n\n\ -- `se_a_tpe`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor.\n\n\ -- `se_atten`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism will be used by this descriptor.\n\n\ -- `se_atten_v2`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism with new modifications will be used by this descriptor.\n\n\ -- `se_a_mask`: Used by the smooth edition of Deep Potential. It can accept a variable number of atoms in a frame (Non-PBC system). *aparam* are required as an indicator matrix for the real/virtual sign of input atoms. \n\n\ -- `hybrid`: Concatenate of a list of descriptors as a new descriptor." - - return Variant( - "type", - descrpt_args_plugin.get_all_argument(exclude_hybrid=exclude_hybrid), - doc=doc_descrpt_type, - ) - - -# --- Fitting net configurations: --- # -fitting_args_plugin = ArgsPlugin() - - -@fitting_args_plugin.register("ener") -def fitting_ener(): - 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_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.' - doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_trainable = "Whether the parameters in the fitting net are trainable. This option can be\n\n\ -- bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\ -- list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of this list should be equal to len(`neuron`)+1." - doc_rcond = "The condition number used to determine the inital energy shift for each type of atoms. See `rcond` in :py:meth:`numpy.linalg.lstsq` for more details." - doc_seed = "Random seed for parameter initialization of the fitting net" - doc_atom_ener = "Specify the atomic energy in vacuum for each type" - doc_layer_name = ( - "The name of the each layer. The length of this list should be equal to n_neuron + 1. " - "If two layers, either in the same fitting or different fittings, " - "have the same name, they will share the same neural network parameters. " - "The shape of these layers should be the same. " - "If null is given for a layer, parameters will not be shared." - ) - doc_use_aparam_as_mask = ( - "Whether to use the aparam as a mask in input." - "If True, the aparam will not be used in fitting net for embedding." - "When descrpt is se_a_mask, the aparam will be used as a mask to indicate the input atom is real/virtual. And use_aparam_as_mask should be set to True." - ) - - 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( - "neuron", - List[int], - optional=True, - default=[120, 120, 120], - alias=["n_neuron"], - doc=doc_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), - Argument( - "trainable", - [List[bool], bool], - optional=True, - default=True, - doc=doc_trainable, - ), - Argument( - "rcond", [float, type(None)], optional=True, default=None, doc=doc_rcond - ), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "atom_ener", - List[Optional[float]], - optional=True, - default=[], - doc=doc_atom_ener, - ), - Argument("layer_name", List[str], optional=True, doc=doc_layer_name), - Argument( - "use_aparam_as_mask", - bool, - optional=True, - default=False, - doc=doc_use_aparam_as_mask, - ), - ] - - -@fitting_args_plugin.register("dos") -def fitting_dos(): - 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_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.' - doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_trainable = "Whether the parameters in the fitting net are trainable. This option can be\n\n\ -- bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\ -- list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of tihs list should be equal to len(`neuron`)+1." - doc_rcond = "The condition number used to determine the inital energy shift for each type of atoms. See `rcond` in :py:meth:`numpy.linalg.lstsq` for more details." - doc_seed = "Random seed for parameter initialization of the fitting net" - doc_numb_dos = ( - "The number of gridpoints on which the DOS is evaluated (NEDOS in VASP)" - ) - - 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( - "neuron", List[int], optional=True, default=[120, 120, 120], doc=doc_neuron - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("precision", str, optional=True, default="float64", doc=doc_precision), - Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), - Argument( - "trainable", - [List[bool], bool], - optional=True, - default=True, - doc=doc_trainable, - ), - Argument( - "rcond", [float, type(None)], optional=True, default=None, doc=doc_rcond - ), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument("numb_dos", int, optional=True, default=300, doc=doc_numb_dos), - ] - - -@fitting_args_plugin.register("polar") -def fitting_polar(): - 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_scale = "The output of the fitting net (polarizability matrix) will be scaled by ``scale``" - # doc_diag_shift = 'The diagonal part of the polarizability matrix will be shifted by ``diag_shift``. The shift operation is carried out after ``scale``.' - doc_fit_diag = "Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix." - doc_sel_type = "The atom types for which the atomic polarizability will be provided. If not set, all types will be selected." - doc_seed = "Random seed for parameter initialization of the fitting net" - - # YWolfeee: user can decide whether to use shift diag - doc_shift_diag = "Whether to shift the diagonal of polar, which is beneficial to training. Default is true." - - return [ - Argument( - "neuron", - List[int], - optional=True, - default=[120, 120, 120], - alias=["n_neuron"], - doc=doc_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("fit_diag", bool, optional=True, default=True, doc=doc_fit_diag), - Argument( - "scale", [List[float], float], optional=True, default=1.0, doc=doc_scale - ), - # Argument("diag_shift", [list,float], optional = True, default = 0.0, doc = doc_diag_shift), - Argument("shift_diag", bool, optional=True, default=True, doc=doc_shift_diag), - Argument( - "sel_type", - [List[int], int, None], - optional=True, - alias=["pol_type"], - doc=doc_sel_type, - ), - Argument("seed", [int, None], optional=True, doc=doc_seed), - ] - - -# def fitting_global_polar(): -# return fitting_polar() - - -@fitting_args_plugin.register("dipole") -def fitting_dipole(): - 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.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected." - doc_seed = "Random seed for parameter initialization of the fitting net" - return [ - Argument( - "neuron", - List[int], - optional=True, - default=[120, 120, 120], - alias=["n_neuron"], - doc=doc_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument( - "sel_type", - [List[int], int, None], - optional=True, - alias=["dipole_type"], - doc=doc_sel_type, - ), - Argument("seed", [int, None], optional=True, doc=doc_seed), - ] - - -# YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. -def fitting_variant_type_args(): - doc_descrpt_type = "The type of the fitting. See explanation below. \n\n\ -- `ener`: Fit an energy model (potential energy surface).\n\n\ -- `dos` : Fit a density of states model. The total density of states / site-projected density of states labels should be provided by `dos.npy` or `atom_dos.npy` in each data system. The file has number of frames lines and number of energy grid columns (times number of atoms in `atom_dos.npy`). See `loss` parameter. \n\n\ -- `dipole`: Fit an atomic dipole model. Global dipole labels or atomic dipole labels for all the selected atoms (see `sel_type`) should be provided by `dipole.npy` in each data system. The file either has number of frames lines and 3 times of number of selected atoms columns, or has number of frames lines and 3 columns. See `loss` parameter.\n\n\ -- `polar`: Fit an atomic polarizability model. Global polarizazbility labels or atomic polarizability labels for all the selected atoms (see `sel_type`) should be provided by `polarizability.npy` in each data system. The file eith has number of frames lines and 9 times of number of selected atoms columns, or has number of frames lines and 9 columns. See `loss` parameter.\n\n" - - return Variant( - "type", - fitting_args_plugin.get_all_argument(), - optional=True, - default_tag="ener", - doc=doc_descrpt_type, - ) - - -# --- Modifier configurations: --- # -def modifier_dipole_charge(): - 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/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')}" - doc_ewald_h = "The grid spacing of the FFT grid. Unit is A" - doc_ewald_beta = f"The splitting parameter of Ewald sum. Unit is A^{-1}" - - return [ - Argument("model_name", str, optional=False, doc=doc_model_name), - Argument( - "model_charge_map", List[float], optional=False, doc=doc_model_charge_map - ), - Argument("sys_charge_map", List[float], optional=False, doc=doc_sys_charge_map), - Argument("ewald_beta", float, optional=True, default=0.4, doc=doc_ewald_beta), - Argument("ewald_h", float, optional=True, default=1.0, doc=doc_ewald_h), - ] - - -def modifier_variant_type_args(): - doc_modifier_type = "The type of modifier. See explanation below.\n\n\ --`dipole_charge`: Use WFCC to model the electronic structure of the system. Correct the long-range interaction" - return Variant( - "type", - [ - Argument("dipole_charge", dict, modifier_dipole_charge()), - ], - optional=False, - doc=doc_modifier_type, - ) - - -# --- model compression configurations: --- # -def model_compression(): - 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 = ( - "The nearest distance between neighbor atoms saved in the frozen model." - ) - - return [ - Argument("model_file", str, optional=False, doc=doc_model_file), - Argument("table_config", List[float], optional=False, doc=doc_table_config), - Argument("min_nbor_dist", float, optional=False, doc=doc_min_nbor_dist), - ] - - -# --- model compression configurations: --- # -def model_compression_type_args(): - doc_compress_type = "The type of model compression, which should be consistent with the descriptor type." - - return Variant( - "type", - [Argument("se_e2_a", dict, model_compression(), alias=["se_a"])], - optional=True, - default_tag="se_e2_a", - doc=doc_compress_type, - ) - - -def model_args(exclude_hybrid=False): - 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." - doc_data_bias_nsample = "The number of training samples in a system to compute and change the energy bias." - doc_type_embedding = "The type embedding." - doc_modifier = "The modifier of model output." - doc_use_srtab = "The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly." - doc_smin_alpha = "The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided." - doc_sw_rmin = "The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided." - doc_sw_rmax = "The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided." - doc_srtab_add_bias = "Whether add energy bias from the statistics of the data to short-range tabulated atomic energy. It only takes effect when `use_srtab` is provided." - doc_compress_config = "Model compression configurations" - doc_spin = "The settings for systems with spin." - hybrid_models = [] - if not exclude_hybrid: - hybrid_models.extend( - [ - pairwise_dprc(), - linear_ener_model_args(), - ] - ) - return Argument( - "model", - dict, - [ - Argument("type_map", List[str], optional=True, doc=doc_type_map), - Argument( - "data_stat_nbatch", - int, - optional=True, - default=10, - doc=doc_data_stat_nbatch, - ), - Argument( - "data_stat_protect", - float, - optional=True, - default=1e-2, - doc=doc_data_stat_protect, - ), - Argument( - "data_bias_nsample", - int, - optional=True, - default=10, - doc=doc_data_bias_nsample, - ), - Argument("use_srtab", str, optional=True, doc=doc_use_srtab), - Argument("smin_alpha", float, optional=True, doc=doc_smin_alpha), - Argument("sw_rmin", float, optional=True, doc=doc_sw_rmin), - Argument("sw_rmax", float, optional=True, doc=doc_sw_rmax), - Argument( - "srtab_add_bias", - bool, - optional=True, - default=True, - doc=doc_srtab_add_bias, - ), - Argument( - "type_embedding", - dict, - type_embedding_args(), - [], - optional=True, - doc=doc_type_embedding, - ), - Argument( - "modifier", - dict, - [], - [modifier_variant_type_args()], - optional=True, - doc=doc_modifier, - ), - Argument( - "compress", - dict, - [], - [model_compression_type_args()], - optional=True, - doc=doc_compress_config, - fold_subdoc=True, - ), - Argument("spin", dict, spin_args(), [], optional=True, doc=doc_spin), - ], - [ - Variant( - "type", - [ - standard_model_args(), - multi_model_args(), - frozen_model_args(), - *hybrid_models, - ], - optional=True, - default_tag="standard", - ), - ], - ) - - -def standard_model_args() -> Argument: - doc_descrpt = "The descriptor of atomic environment." - doc_fitting = "The fitting of physical properties." - - ca = Argument( - "standard", - dict, - [ - Argument( - "descriptor", dict, [], [descrpt_variant_type_args()], doc=doc_descrpt - ), - Argument( - "fitting_net", - dict, - [], - [fitting_variant_type_args()], - doc=doc_fitting, - ), - ], - doc="Stardard model, which contains a descriptor and a fitting.", - ) - return ca - - -def multi_model_args() -> Argument: - doc_descrpt = "The descriptor of atomic environment. See model[standard]/descriptor for details." - doc_fitting_net_dict = "The dictionary of multiple fitting nets in multi-task mode. Each fitting_net_dict[fitting_key] is the single definition of fitting of physical properties with user-defined name `fitting_key`." - - ca = Argument( - "multi", - dict, - [ - Argument( - "descriptor", - dict, - [], - [descrpt_variant_type_args()], - doc=doc_descrpt, - fold_subdoc=True, - ), - Argument("fitting_net_dict", dict, doc=doc_fitting_net_dict), - ], - doc="Multiple-task model.", - ) - return ca - - -def pairwise_dprc() -> Argument: - qm_model_args = model_args(exclude_hybrid=True) - qm_model_args.name = "qm_model" - qm_model_args.fold_subdoc = True - qmmm_model_args = model_args(exclude_hybrid=True) - qmmm_model_args.name = "qmmm_model" - qmmm_model_args.fold_subdoc = True - ca = Argument( - "pairwise_dprc", - dict, - [ - qm_model_args, - qmmm_model_args, - ], - ) - return ca - - -def frozen_model_args() -> Argument: - doc_model_file = "Path to the frozen model file." - ca = Argument( - "frozen", - dict, - [ - Argument("model_file", str, optional=False, doc=doc_model_file), - ], - ) - return ca - - -def linear_ener_model_args() -> Argument: - doc_weights = ( - "If the type is list of float, a list of weights for each model. " - 'If "mean", the weights are set to be 1 / len(models). ' - 'If "sum", the weights are set to be 1.' - ) - models_args = model_args(exclude_hybrid=True) - models_args.name = "models" - models_args.fold_subdoc = True - models_args.set_dtype(list) - models_args.set_repeat(True) - models_args.doc = "The sub-models." - ca = Argument( - "linear_ener", - dict, - [ - models_args, - Argument( - "weights", - [list, str], - optional=False, - doc=doc_weights, - ), - ], - ) - return ca - - -# --- Learning rate configurations: --- # -def learning_rate_exp(): - 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." - doc_decay_steps = ( - "The learning rate is decaying every this number of training steps." - ) - - args = [ - Argument("start_lr", float, optional=True, default=1e-3, doc=doc_start_lr), - Argument("stop_lr", float, optional=True, default=1e-8, doc=doc_stop_lr), - Argument("decay_steps", int, optional=True, default=5000, doc=doc_decay_steps), - ] - return args - - -def learning_rate_variant_type_args(): - doc_lr = "The type of the learning rate." - - return Variant( - "type", - [Argument("exp", dict, learning_rate_exp())], - optional=True, - default_tag="exp", - doc=doc_lr, - ) - - -def learning_rate_args(): - doc_scale_by_worker = "When parallel training or batch size scaled, how to alter learning rate. Valid values are `linear`(default), `sqrt` or `none`." - doc_lr = "The definitio of learning rate" - return Argument( - "learning_rate", - dict, - [ - Argument( - "scale_by_worker", - str, - optional=True, - default="linear", - doc=doc_scale_by_worker, - ) - ], - [learning_rate_variant_type_args()], - optional=True, - doc=doc_lr, - ) - - -def learning_rate_dict_args(): - doc_learning_rate_dict = ( - "The dictionary of definitions of learning rates in multi-task mode. " - "Each learning_rate_dict[fitting_key], with user-defined name `fitting_key` in `model/fitting_net_dict`, is the single definition of learning rate.\n" - ) - ca = Argument( - "learning_rate_dict", dict, [], [], optional=True, doc=doc_learning_rate_dict - ) - return ca - - -# --- Loss configurations: --- # -def start_pref(item, label=None, abbr=None): - if label is None: - label = item - if abbr is None: - abbr = item - 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): - 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." - - -loss_args_plugin = ArgsPlugin() - - -@loss_args_plugin.register("ener") -def loss_ener(): - 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") - doc_limit_pref_f = limit_pref("force") - doc_start_pref_v = start_pref("virial", abbr="v") - doc_limit_pref_v = limit_pref("virial") - doc_start_pref_ae = start_pref("atomic energy", label="atom_ener", abbr="ae") - doc_limit_pref_ae = limit_pref("atomic energy") - doc_start_pref_pf = start_pref( - "atomic prefactor force", label="atom_pref", abbr="pf" - ) - doc_limit_pref_pf = limit_pref("atomic prefactor force") - doc_start_pref_gf = start_pref("generalized force", label="drdq", abbr="gf") - doc_limit_pref_gf = limit_pref("generalized force") - doc_numb_generalized_coord = "The dimension of generalized coordinates. Required when generalized force loss is used." - doc_relative_f = "If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label." - doc_enable_atom_ener_coeff = "If true, the energy will be computed as \\sum_i c_i E_i. c_i should be provided by file atom_ener_coeff.npy in each data system, otherwise it's 1." - return [ - Argument( - "start_pref_e", - [float, int], - optional=True, - default=0.02, - doc=doc_start_pref_e, - ), - Argument( - "limit_pref_e", - [float, int], - optional=True, - default=1.00, - doc=doc_limit_pref_e, - ), - Argument( - "start_pref_f", - [float, int], - optional=True, - default=1000, - doc=doc_start_pref_f, - ), - Argument( - "limit_pref_f", - [float, int], - optional=True, - default=1.00, - doc=doc_limit_pref_f, - ), - Argument( - "start_pref_v", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_v, - ), - Argument( - "limit_pref_v", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_v, - ), - Argument( - "start_pref_ae", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_ae, - ), - Argument( - "limit_pref_ae", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_ae, - ), - Argument( - "start_pref_pf", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_pf, - ), - Argument( - "limit_pref_pf", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_pf, - ), - Argument("relative_f", [float, None], optional=True, doc=doc_relative_f), - Argument( - "enable_atom_ener_coeff", - [bool], - optional=True, - default=False, - doc=doc_enable_atom_ener_coeff, - ), - Argument( - "start_pref_gf", - float, - optional=True, - default=0.0, - doc=doc_start_pref_gf, - ), - Argument( - "limit_pref_gf", - float, - optional=True, - default=0.0, - doc=doc_limit_pref_gf, - ), - Argument( - "numb_generalized_coord", - int, - optional=True, - default=0, - doc=doc_numb_generalized_coord, - ), - ] - - -@loss_args_plugin.register("ener_spin") -def loss_ener_spin(): - doc_start_pref_e = start_pref("energy") - doc_limit_pref_e = limit_pref("energy") - doc_start_pref_fr = start_pref("force_real_atom") - doc_limit_pref_fr = limit_pref("force_real_atom") - doc_start_pref_fm = start_pref("force_magnetic") - doc_limit_pref_fm = limit_pref("force_magnetic") - doc_start_pref_v = start_pref("virial") - doc_limit_pref_v = limit_pref("virial") - doc_start_pref_ae = start_pref("atom_ener") - doc_limit_pref_ae = limit_pref("atom_ener") - doc_start_pref_pf = start_pref("atom_pref") - doc_limit_pref_pf = limit_pref("atom_pref") - doc_relative_f = "If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label." - doc_enable_atom_ener_coeff = r"If true, the energy will be computed as \sum_i c_i E_i. c_i should be provided by file atom_ener_coeff.npy in each data system, otherwise it's 1." - return [ - Argument( - "start_pref_e", - [float, int], - optional=True, - default=0.02, - doc=doc_start_pref_e, - ), - Argument( - "limit_pref_e", - [float, int], - optional=True, - default=1.00, - doc=doc_limit_pref_e, - ), - Argument( - "start_pref_fr", - [float, int], - optional=True, - default=1000, - doc=doc_start_pref_fr, - ), - Argument( - "limit_pref_fr", - [float, int], - optional=True, - default=1.00, - doc=doc_limit_pref_fr, - ), - Argument( - "start_pref_fm", - [float, int], - optional=True, - default=10000, - doc=doc_start_pref_fm, - ), - Argument( - "limit_pref_fm", - [float, int], - optional=True, - default=10.0, - doc=doc_limit_pref_fm, - ), - Argument( - "start_pref_v", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_v, - ), - Argument( - "limit_pref_v", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_v, - ), - Argument( - "start_pref_ae", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_ae, - ), - Argument( - "limit_pref_ae", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_ae, - ), - Argument( - "start_pref_pf", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_pf, - ), - Argument( - "limit_pref_pf", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_pf, - ), - Argument("relative_f", [float, None], optional=True, doc=doc_relative_f), - Argument( - "enable_atom_ener_coeff", - [bool], - optional=True, - default=False, - doc=doc_enable_atom_ener_coeff, - ), - ] - - -@loss_args_plugin.register("dos") -def loss_dos(): - 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( - "Cumulative Distribution Function (cumulative intergral of DOS)" - ) - doc_limit_pref_cdf = limit_pref( - "Cumulative Distribution Function (cumulative intergral of DOS)" - ) - doc_start_pref_ados = start_pref("atomic DOS (site-projected DOS)") - doc_limit_pref_ados = limit_pref("atomic DOS (site-projected DOS)") - doc_start_pref_acdf = start_pref("Cumulative integral of atomic DOS") - doc_limit_pref_acdf = limit_pref("Cumulative integral of atomic DOS") - return [ - Argument( - "start_pref_dos", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_dos, - ), - Argument( - "limit_pref_dos", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_dos, - ), - Argument( - "start_pref_cdf", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_cdf, - ), - Argument( - "limit_pref_cdf", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_cdf, - ), - Argument( - "start_pref_ados", - [float, int], - optional=True, - default=1.00, - doc=doc_start_pref_ados, - ), - Argument( - "limit_pref_ados", - [float, int], - optional=True, - default=1.00, - doc=doc_limit_pref_ados, - ), - Argument( - "start_pref_acdf", - [float, int], - optional=True, - default=0.00, - doc=doc_start_pref_acdf, - ), - Argument( - "limit_pref_acdf", - [float, int], - optional=True, - default=0.00, - doc=doc_limit_pref_acdf, - ), - ] - - -# YWolfeee: Modified to support tensor type of loss args. -@loss_args_plugin.register("tensor") -def loss_tensor(): - # 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. If 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." - doc_local_weight = "The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to atomic label, i.e. `atomic_polarizability.npy` or `atomic_dipole.npy`, whose shape should be #frames x ([9 or 3] x #selected atoms). If it's larger than 0.0, this npy should be included. Both `pref` and `pref_atomic` should be provided, and either can be set to 0.0." - return [ - Argument( - "pref", [float, int], optional=False, default=None, doc=doc_global_weight - ), - Argument( - "pref_atomic", - [float, int], - optional=False, - default=None, - doc=doc_local_weight, - ), - ] - - -def loss_variant_type_args(): - 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( - "type", - loss_args_plugin.get_all_argument(), - optional=True, - default_tag="ener", - doc=doc_loss, - ) - - -def loss_args(): - 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 - ) - return ca - - -def loss_dict_args(): - doc_loss_dict = ( - "The dictionary of definitions of multiple loss functions in multi-task mode. " - "Each loss_dict[fitting_key], with user-defined name `fitting_key` in `model/fitting_net_dict`, is the single definition of loss function, whose type should be set to `tensor`, `ener` or left unset.\n" - ) - ca = Argument("loss_dict", dict, [], [], optional=True, doc=doc_loss_dict) - return ca - - -# --- Training configurations: --- # -def training_data_args(): # ! 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." - ) - doc_set_prefix = f"The prefix of the sets in the {link_sys}." - doc_batch_size = f'This key can be \n\n\ -- list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ -- int: all {link_sys} use the same batch size.\n\n\ -- string "auto": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than 32.\n\n\ -- string "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.\n\n\ -- string "mixed:N": the batch data will be sampled from all systems and merged into a mixed system with the batch size N. Only support the se_atten descriptor.\n\n\ -If MPI is used, the value should be considered as the batch size per task.' - doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ -- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ -- "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ -- "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is devided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' - doc_sys_probs = ( - "A list of float if specified. " - "Should be of the same length as `systems`, " - "specifying the probability of each system." - ) - - args = [ - Argument( - "systems", [List[str], str], optional=False, default=".", doc=doc_systems - ), - Argument("set_prefix", str, optional=True, default="set", doc=doc_set_prefix), - Argument( - "batch_size", - [List[int], int, str], - optional=True, - default="auto", - doc=doc_batch_size, - ), - Argument( - "auto_prob", - str, - optional=True, - default="prob_sys_size", - doc=doc_auto_prob_style, - alias=[ - "auto_prob_style", - ], - ), - Argument( - "sys_probs", - List[float], - optional=True, - default=None, - doc=doc_sys_probs, - alias=["sys_weights"], - ), - ] - - doc_training_data = "Configurations of training data." - return Argument( - "training_data", - dict, - optional=True, - sub_fields=args, - sub_variants=[], - doc=doc_training_data, - ) - - -def validation_data_args(): # ! 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." - ) - doc_set_prefix = f"The prefix of the sets in the {link_sys}." - doc_batch_size = f'This key can be \n\n\ -- list: the length of which is the same as the {link_sys}. The batch size of each system is given by the elements of the list.\n\n\ -- int: all {link_sys} use the same batch size.\n\n\ -- string "auto": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than 32.\n\n\ -- string "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.' - doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\ -- "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\ -- "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\ -- "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : the list of systems is devided into blocks. A block is specified by `stt_idx:end_idx:weight`, where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional to the number of batches in the system.' - doc_sys_probs = ( - "A list of float if specified. " - "Should be of the same length as `systems`, " - "specifying the probability of each system." - ) - doc_numb_btch = "An integer that specifies the number of batches to be sampled for each validation period." - - args = [ - Argument( - "systems", [List[str], str], optional=False, default=".", doc=doc_systems - ), - Argument("set_prefix", str, optional=True, default="set", doc=doc_set_prefix), - Argument( - "batch_size", - [List[int], int, str], - optional=True, - default="auto", - doc=doc_batch_size, - ), - Argument( - "auto_prob", - str, - optional=True, - default="prob_sys_size", - doc=doc_auto_prob_style, - alias=[ - "auto_prob_style", - ], - ), - Argument( - "sys_probs", - List[float], - optional=True, - default=None, - doc=doc_sys_probs, - alias=["sys_weights"], - ), - Argument( - "numb_btch", - int, - optional=True, - default=1, - doc=doc_numb_btch, - alias=[ - "numb_batch", - ], - ), - ] - - doc_validation_data = ( - "Configurations of validation data. Similar to that of training data, " - "except that a `numb_btch` argument may be configured" - ) - return Argument( - "validation_data", - dict, - optional=True, - default=None, - sub_fields=args, - sub_variants=[], - doc=doc_validation_data, - ) - - -def mixed_precision_args(): # ! 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.' - doc_compute_prec = 'The precision for mixed precision compute. " \ - "The compute precision during the mixed precision training process, "" \ - "supported options are float16 and bfloat16 currently.' - - args = [ - Argument( - "output_prec", str, optional=True, default="float32", doc=doc_output_prec - ), - Argument( - "compute_prec", str, optional=False, default="float16", doc=doc_compute_prec - ), - ] - - doc_mixed_precision = "Configurations of mixed precision." - return Argument( - "mixed_precision", - dict, - optional=True, - sub_fields=args, - sub_variants=[], - doc=doc_mixed_precision, - ) - - -def training_args(): # ! 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." - doc_disp_freq = "The frequency of printing learning curve." - doc_save_freq = "The frequency of saving check point." - doc_save_ckpt = "The path prefix of saving check point files." - doc_disp_training = "Displaying verbose information during training." - doc_time_training = "Timing durining training." - doc_profiling = "Profiling during training." - doc_profiling_file = "Output file for profiling." - doc_enable_profiler = "Enable TensorFlow Profiler (available in TensorFlow 2.3) to analyze performance. The log will be saved to `tensorboard_log_dir`." - doc_tensorboard = "Enable tensorboard" - doc_tensorboard_log_dir = "The log directory of tensorboard outputs" - doc_tensorboard_freq = "The frequency of writing tensorboard events." - doc_data_dict = ( - "The dictionary of multi DataSystems in multi-task mode. " - "Each data_dict[fitting_key], with user-defined name `fitting_key` in `model/fitting_net_dict`, " - "contains training data and optional validation data definitions." - ) - doc_fitting_weight = ( - "Each fitting_weight[fitting_key], with user-defined name `fitting_key` in `model/fitting_net_dict`, " - "is the training weight of fitting net `fitting_key`. " - "Fitting nets with higher weights will be selected with higher probabilities to be trained in one step. " - "Weights will be normalized and minus ones will be ignored. " - "If not set, each fitting net will be equally selected when training." - ) - - arg_training_data = training_data_args() - arg_validation_data = validation_data_args() - mixed_precision_data = mixed_precision_args() - - args = [ - arg_training_data, - arg_validation_data, - mixed_precision_data, - Argument( - "numb_steps", int, optional=False, doc=doc_numb_steps, alias=["stop_batch"] - ), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "disp_file", str, optional=True, default="lcurve.out", doc=doc_disp_file - ), - Argument("disp_freq", int, optional=True, default=1000, doc=doc_disp_freq), - Argument("save_freq", int, optional=True, default=1000, doc=doc_save_freq), - Argument( - "save_ckpt", str, optional=True, default="model.ckpt", doc=doc_save_ckpt - ), - Argument( - "disp_training", bool, optional=True, default=True, doc=doc_disp_training - ), - Argument( - "time_training", bool, optional=True, default=True, doc=doc_time_training - ), - Argument("profiling", bool, optional=True, default=False, doc=doc_profiling), - Argument( - "profiling_file", - str, - optional=True, - default="timeline.json", - doc=doc_profiling_file, - ), - Argument( - "enable_profiler", - bool, - optional=True, - default=False, - doc=doc_enable_profiler, - ), - Argument( - "tensorboard", bool, optional=True, default=False, doc=doc_tensorboard - ), - Argument( - "tensorboard_log_dir", - str, - optional=True, - default="log", - doc=doc_tensorboard_log_dir, - ), - Argument( - "tensorboard_freq", int, optional=True, default=1, doc=doc_tensorboard_freq - ), - Argument("data_dict", dict, optional=True, doc=doc_data_dict), - Argument("fitting_weight", dict, optional=True, doc=doc_fitting_weight), - ] - - doc_training = "The training options." - return Argument("training", dict, args, [], doc=doc_training) - - -def make_index(keys): - ret = [] - for ii in keys: - ret.append(make_link(ii, ii)) - return ", ".join(ret) - - -def gen_doc(*, make_anchor=True, make_link=True, **kwargs): - if make_link: - make_anchor = True - ptr = [] - for ii in gen_args(): - ptr.append(ii.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs)) - - key_words = [] - for ii in "\n\n".join(ptr).split("\n"): - if "argument path" in ii: - key_words.append(ii.split(":")[1].replace("`", "").strip()) - # ptr.insert(0, make_index(key_words)) - - return "\n\n".join(ptr) - - -def gen_json(**kwargs): - return json.dumps( - tuple(gen_args()), - cls=ArgumentEncoder, - ) - - -def gen_args(**kwargs) -> List[Argument]: - return [ - model_args(), - learning_rate_args(), - learning_rate_dict_args(), - loss_args(), - loss_dict_args(), - training_args(), - nvnmd_args(), - ] - - -def normalize_multi_task(data): - # single-task or multi-task mode - if data["model"].get("type", "standard") not in ("standard", "multi"): - return data - single_fitting_net = "fitting_net" in data["model"].keys() - single_training_data = "training_data" in data["training"].keys() - single_valid_data = "validation_data" in data["training"].keys() - single_loss = "loss" in data.keys() - single_learning_rate = "learning_rate" in data.keys() - multi_fitting_net = "fitting_net_dict" in data["model"].keys() - multi_training_data = "data_dict" in data["training"].keys() - multi_loss = "loss_dict" in data.keys() - multi_fitting_weight = "fitting_weight" in data["training"].keys() - multi_learning_rate = "learning_rate_dict" in data.keys() - assert (single_fitting_net == single_training_data) and ( - multi_fitting_net == multi_training_data - ), ( - "In single-task mode, 'model/fitting_net' and 'training/training_data' must be defined at the same time! " - "While in multi-task mode, 'model/fitting_net_dict', 'training/data_dict' " - "must be defined at the same time! Please check your input script. " - ) - assert not (single_fitting_net and multi_fitting_net), ( - "Single-task mode and multi-task mode can not be performed together. " - "Please check your input script and choose just one format! " - ) - assert ( - single_fitting_net or multi_fitting_net - ), "Please define your fitting net and training data! " - if multi_fitting_net: - assert not single_valid_data, ( - "In multi-task mode, 'training/validation_data' should not appear " - "outside 'training/data_dict'! Please check your input script." - ) - assert ( - not single_loss - ), "In multi-task mode, please use 'model/loss_dict' in stead of 'model/loss'! " - assert ( - "type_map" in data["model"] - ), "In multi-task mode, 'model/type_map' must be defined! " - data["model"]["type"] = "multi" - data["model"]["fitting_net_dict"] = normalize_fitting_net_dict( - data["model"]["fitting_net_dict"] - ) - data["training"]["data_dict"] = normalize_data_dict( - data["training"]["data_dict"] - ) - data["loss_dict"] = ( - normalize_loss_dict( - data["model"]["fitting_net_dict"].keys(), data["loss_dict"] - ) - if multi_loss - else {} - ) - if multi_learning_rate: - data["learning_rate_dict"] = normalize_learning_rate_dict( - data["model"]["fitting_net_dict"].keys(), data["learning_rate_dict"] - ) - elif single_learning_rate: - data[ - "learning_rate_dict" - ] = normalize_learning_rate_dict_with_single_learning_rate( - data["model"]["fitting_net_dict"].keys(), data["learning_rate"] - ) - fitting_weight = ( - data["training"]["fitting_weight"] if multi_fitting_weight else None - ) - data["training"]["fitting_weight"] = normalize_fitting_weight( - data["model"]["fitting_net_dict"].keys(), - data["training"]["data_dict"].keys(), - fitting_weight=fitting_weight, - ) - else: - assert ( - not multi_loss - ), "In single-task mode, please use 'model/loss' in stead of 'model/loss_dict'! " - assert ( - not multi_learning_rate - ), "In single-task mode, please use 'model/learning_rate' in stead of 'model/learning_rate_dict'! " - return data - - -def normalize_fitting_net_dict(fitting_net_dict): - new_dict = {} - base = Argument("base", dict, [], [fitting_variant_type_args()], doc="") - for fitting_key_item in fitting_net_dict: - data = base.normalize_value( - fitting_net_dict[fitting_key_item], trim_pattern="_*" - ) - base.check_value(data, strict=True) - new_dict[fitting_key_item] = data - return new_dict - - -def normalize_data_dict(data_dict): - new_dict = {} - base = Argument( - "base", dict, [training_data_args(), validation_data_args()], [], doc="" - ) - for data_system_key_item in data_dict: - data = base.normalize_value(data_dict[data_system_key_item], trim_pattern="_*") - base.check_value(data, strict=True) - new_dict[data_system_key_item] = data - return new_dict - - -def normalize_loss_dict(fitting_keys, loss_dict): - # check the loss dict - failed_loss_keys = [item for item in loss_dict if item not in fitting_keys] - assert ( - not failed_loss_keys - ), "Loss dict key(s) {} not have corresponding fitting keys in {}! ".format( - str(failed_loss_keys), str(list(fitting_keys)) - ) - new_dict = {} - base = Argument("base", dict, [], [loss_variant_type_args()], doc="") - for item in loss_dict: - data = base.normalize_value(loss_dict[item], trim_pattern="_*") - base.check_value(data, strict=True) - new_dict[item] = data - return new_dict - - -def normalize_learning_rate_dict(fitting_keys, learning_rate_dict): - # check the learning_rate dict - failed_learning_rate_keys = [ - item for item in learning_rate_dict if item not in fitting_keys - ] - assert ( - not failed_learning_rate_keys - ), "Learning rate dict key(s) {} not have corresponding fitting keys in {}! ".format( - str(failed_learning_rate_keys), str(list(fitting_keys)) - ) - new_dict = {} - base = Argument("base", dict, [], [learning_rate_variant_type_args()], doc="") - for item in learning_rate_dict: - data = base.normalize_value(learning_rate_dict[item], trim_pattern="_*") - base.check_value(data, strict=True) - new_dict[item] = data - return new_dict - - -def normalize_learning_rate_dict_with_single_learning_rate(fitting_keys, learning_rate): - new_dict = {} - base = Argument("base", dict, [], [learning_rate_variant_type_args()], doc="") - data = base.normalize_value(learning_rate, trim_pattern="_*") - base.check_value(data, strict=True) - for fitting_key in fitting_keys: - new_dict[fitting_key] = data - return new_dict - - -def normalize_fitting_weight(fitting_keys, data_keys, fitting_weight=None): - # check the mapping - failed_data_keys = [item for item in data_keys if item not in fitting_keys] - assert ( - not failed_data_keys - ), "Data dict key(s) {} not have corresponding fitting keys in {}! ".format( - str(failed_data_keys), str(list(fitting_keys)) - ) - empty_fitting_keys = [] - valid_fitting_keys = [] - for item in fitting_keys: - if item not in data_keys: - empty_fitting_keys.append(item) - else: - valid_fitting_keys.append(item) - if empty_fitting_keys: - log.warning( - "Fitting net(s) {} have no data and will not be used in training.".format( - str(empty_fitting_keys) - ) - ) - num_pair = len(valid_fitting_keys) - assert num_pair > 0, "No valid training data systems for fitting nets!" - - # check and normalize the fitting weight - new_weight = {} - if fitting_weight is None: - equal_weight = 1.0 / num_pair - for item in fitting_keys: - new_weight[item] = equal_weight if item in valid_fitting_keys else 0.0 - else: - failed_weight_keys = [ - item for item in fitting_weight if item not in fitting_keys - ] - assert ( - not failed_weight_keys - ), "Fitting weight key(s) {} not have corresponding fitting keys in {}! ".format( - str(failed_weight_keys), str(list(fitting_keys)) - ) - sum_prob = 0.0 - for item in fitting_keys: - if item in valid_fitting_keys: - if ( - item in fitting_weight - and isinstance(fitting_weight[item], (int, float)) - and fitting_weight[item] > 0.0 - ): - sum_prob += fitting_weight[item] - new_weight[item] = fitting_weight[item] - else: - valid_fitting_keys.remove(item) - log.warning( - f"Fitting net '{item}' has zero or invalid weight " - "and will not be used in training." - ) - new_weight[item] = 0.0 - else: - new_weight[item] = 0.0 - assert sum_prob > 0.0, "No valid training weight for fitting nets!" - # normalize - for item in new_weight: - new_weight[item] /= sum_prob - return new_weight - - -def normalize(data): - data = normalize_multi_task(data) - - base = Argument("base", dict, gen_args()) - data = base.normalize_value(data, trim_pattern="_*") - base.check_value(data, strict=True) - - return data - - -if __name__ == "__main__": - gen_doc() +__all__ = [ + "list_to_doc", + "normalize", + "gen_doc", + "gen_json", + "gen_args", + "type_embedding_args", +] diff --git a/deepmd/utils/batch_size.py b/deepmd/utils/batch_size.py index 2b3117d849..863520b3f4 100644 --- a/deepmd/utils/batch_size.py +++ b/deepmd/utils/batch_size.py @@ -1,207 +1,40 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import logging -import os -from typing import ( - Callable, - Tuple, +from packaging.version import ( + Version, ) -import numpy as np - from deepmd.env import ( + TF_VERSION, tf, ) from deepmd.utils.errors import ( OutOfMemoryError, ) +from deepmd_utils.utils.batch_size import AutoBatchSize as AutoBatchSizeBase -log = logging.getLogger(__name__) - - -class AutoBatchSize: - """This class allows DeePMD-kit to automatically decide the maximum - batch size that will not cause an OOM error. - - Notes - ----- - In some CPU environments, the program may be directly killed when OOM. In - this case, by default the batch size will not be increased for CPUs. The - environment variable `DP_INFER_BATCH_SIZE` can be set as the batch size. - - In other cases, we assume all OOM error will raise :class:`OutOfMemoryError`. - - Parameters - ---------- - initial_batch_size : int, default: 1024 - initial batch size (number of total atoms) when DP_INFER_BATCH_SIZE - is not set - factor : float, default: 2. - increased factor - - Attributes - ---------- - current_batch_size : int - current batch size (number of total atoms) - maximum_working_batch_size : int - maximum working batch size - minimal_not_working_batch_size : int - minimal not working batch size - """ - - def __init__(self, initial_batch_size: int = 1024, factor: float = 2.0) -> None: - # See also PyTorchLightning/pytorch-lightning#1638 - # TODO: discuss a proper initial batch size - self.current_batch_size = initial_batch_size - DP_INFER_BATCH_SIZE = int(os.environ.get("DP_INFER_BATCH_SIZE", 0)) - if DP_INFER_BATCH_SIZE > 0: - self.current_batch_size = DP_INFER_BATCH_SIZE - self.maximum_working_batch_size = DP_INFER_BATCH_SIZE - self.minimal_not_working_batch_size = self.maximum_working_batch_size + 1 - else: - self.maximum_working_batch_size = initial_batch_size - if tf.test.is_gpu_available(): - self.minimal_not_working_batch_size = 2**31 - else: - self.minimal_not_working_batch_size = ( - self.maximum_working_batch_size + 1 - ) - log.warning( - "You can use the environment variable DP_INFER_BATCH_SIZE to" - "control the inference batch size (nframes * natoms). " - "The default value is %d." % initial_batch_size - ) - - self.factor = factor - def execute( - self, callable: Callable, start_index: int, natoms: int - ) -> Tuple[int, tuple]: - """Excuate a method with given batch size. - - Parameters - ---------- - callable : Callable - The method should accept the batch size and start_index as parameters, - and returns executed batch size and data. - start_index : int - start index - natoms : int - natoms +class AutoBatchSize(AutoBatchSizeBase): + def is_gpu_available(self) -> bool: + """Check if GPU is available. Returns ------- - int - executed batch size * number of atoms - tuple - result from callable, None if failing to execute - - Raises - ------ - OutOfMemoryError - OOM when batch size is 1 + bool + True if GPU is available """ - if natoms > 0: - batch_nframes = self.current_batch_size // natoms - else: - batch_nframes = self.current_batch_size - try: - n_batch, result = callable(max(batch_nframes, 1), start_index) - except OutOfMemoryError as e: - # TODO: it's very slow to catch OOM error; I don't know what TF is doing here - # but luckily we only need to catch once - self.minimal_not_working_batch_size = min( - self.minimal_not_working_batch_size, self.current_batch_size - ) - if self.maximum_working_batch_size >= self.minimal_not_working_batch_size: - self.maximum_working_batch_size = int( - self.minimal_not_working_batch_size / self.factor - ) - if self.minimal_not_working_batch_size <= natoms: - raise OutOfMemoryError( - "The callable still throws an out-of-memory (OOM) error even when batch size is 1!" - ) from e - # adjust the next batch size - self._adjust_batch_size(1.0 / self.factor) - return 0, None - else: - n_tot = n_batch * natoms - self.maximum_working_batch_size = max( - self.maximum_working_batch_size, n_tot - ) - # adjust the next batch size - if ( - n_tot + natoms > self.current_batch_size - and self.current_batch_size * self.factor - < self.minimal_not_working_batch_size - ): - self._adjust_batch_size(self.factor) - return n_batch, result + return ( + Version(TF_VERSION) >= Version("1.14") + and tf.config.experimental.get_visible_devices("GPU") + ) or tf.test.is_gpu_available() - def _adjust_batch_size(self, factor: float): - old_batch_size = self.current_batch_size - self.current_batch_size = int(self.current_batch_size * factor) - log.info( - "Adjust batch size from %d to %d" - % (old_batch_size, self.current_batch_size) - ) - - def execute_all( - self, callable: Callable, total_size: int, natoms: int, *args, **kwargs - ) -> Tuple[np.ndarray]: - """Excuate a method with all given data. + def is_oom_error(self, e: Exception) -> bool: + """Check if the exception is an OOM error. Parameters ---------- - callable : Callable - The method should accept *args and **kwargs as input and return the similiar array. - total_size : int - Total size - natoms : int - The number of atoms - *args - Variable length argument list. - **kwargs - If 2D np.ndarray, assume the first axis is batch; otherwise do nothing. + e : Exception + Exception """ - - def execute_with_batch_size( - batch_size: int, start_index: int - ) -> Tuple[int, Tuple[np.ndarray]]: - end_index = start_index + batch_size - end_index = min(end_index, total_size) - return (end_index - start_index), callable( - *[ - ( - vv[start_index:end_index] - if isinstance(vv, np.ndarray) and vv.ndim > 1 - else vv - ) - for vv in args - ], - **{ - kk: ( - vv[start_index:end_index] - if isinstance(vv, np.ndarray) and vv.ndim > 1 - else vv - ) - for kk, vv in kwargs.items() - }, - ) - - index = 0 - results = [] - while index < total_size: - n_batch, result = self.execute(execute_with_batch_size, index, natoms) - if not isinstance(result, tuple): - result = (result,) - index += n_batch - if n_batch: - for rr in result: - rr.reshape((n_batch, -1)) - results.append(result) - - r = tuple([np.concatenate(r, axis=0) for r in zip(*results)]) - if len(r) == 1: - # avoid returning tuple if callable doesn't return tuple - r = r[0] - return r + # TODO: it's very slow to catch OOM error; I don't know what TF is doing here + # but luckily we only need to catch once + return isinstance(e, (tf.errors.ResourceExhaustedError, OutOfMemoryError)) diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py index 5f9c14e6d8..91bf4021ee 100644 --- a/deepmd/utils/compat.py +++ b/deepmd/utils/compat.py @@ -1,392 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Module providing compatibility between `0.x.x` and `1.x.x` input versions.""" - -import json -import warnings -from pathlib import ( - Path, -) -from typing import ( - Any, - Dict, - Optional, - Sequence, - Union, -) - -import numpy as np - -from deepmd.common import ( - j_must_have, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.compat import ( + convert_input_v0_v1, + convert_input_v1_v2, + deprecate_numb_test, + update_deepmd_input, ) - -def convert_input_v0_v1( - jdata: Dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None -) -> Dict[str, Any]: - """Convert input from v0 format to v1. - - Parameters - ---------- - jdata : Dict[str, Any] - loaded json/yaml file - warning : bool, optional - whether to show deprecation warning, by default True - dump : Optional[Union[str, Path]], optional - whether to dump converted file, by default None - - Returns - ------- - Dict[str, Any] - converted output - """ - output = {} - output["model"] = _model(jdata, jdata["use_smooth"]) - output["learning_rate"] = _learning_rate(jdata) - output["loss"] = _loss(jdata) - output["training"] = _training(jdata) - if warning: - _warning_input_v0_v1(dump) - if dump is not None: - with open(dump, "w") as fp: - json.dump(output, fp, indent=4) - return output - - -def _warning_input_v0_v1(fname: Optional[Union[str, Path]]): - 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" - ) - if fname is not None: - msg += f", and output it to file {fname}" - warnings.warn(msg) - - -def _model(jdata: Dict[str, Any], smooth: bool) -> Dict[str, Dict[str, Any]]: - """Convert data to v1 input for non-smooth model. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - smooth : bool - whether to use smooth or non-smooth descriptor version - - Returns - ------- - Dict[str, Dict[str, Any]] - dictionary with model input parameters and sub-dictionaries for descriptor and - fitting net - """ - model = {} - model["descriptor"] = ( - _smth_descriptor(jdata) if smooth else _nonsmth_descriptor(jdata) - ) - model["fitting_net"] = _fitting_net(jdata) - return model - - -def _nonsmth_descriptor(jdata: Dict[str, Any]) -> Dict[str, Any]: - """Convert data to v1 input for non-smooth descriptor. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - - Returns - ------- - Dict[str, Any] - dict with descriptor parameters - """ - descriptor = {} - descriptor["type"] = "loc_frame" - _jcopy(jdata, descriptor, ("sel_a", "sel_r", "rcut", "axis_rule")) - return descriptor - - -def _smth_descriptor(jdata: Dict[str, Any]) -> Dict[str, Any]: - """Convert data to v1 input for smooth descriptor. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - - Returns - ------- - Dict[str, Any] - dict with descriptor parameters - """ - descriptor = {} - seed = jdata.get("seed", None) - if seed is not None: - descriptor["seed"] = seed - descriptor["type"] = "se_a" - descriptor["sel"] = jdata["sel_a"] - _jcopy(jdata, descriptor, ("rcut",)) - descriptor["rcut_smth"] = jdata.get("rcut_smth", descriptor["rcut"]) - descriptor["neuron"] = j_must_have(jdata, "filter_neuron") - descriptor["axis_neuron"] = j_must_have(jdata, "axis_neuron", ["n_axis_neuron"]) - descriptor["resnet_dt"] = False - if "resnet_dt" in jdata: - descriptor["resnet_dt"] = jdata["filter_resnet_dt"] - - return descriptor - - -def _fitting_net(jdata: Dict[str, Any]) -> Dict[str, Any]: - """Convert data to v1 input for fitting net. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - - Returns - ------- - Dict[str, Any] - dict with fitting net parameters - """ - fitting_net = {} - - seed = jdata.get("seed", None) - if seed is not None: - fitting_net["seed"] = seed - fitting_net["neuron"] = j_must_have(jdata, "fitting_neuron", ["n_neuron"]) - fitting_net["resnet_dt"] = True - if "resnet_dt" in jdata: - fitting_net["resnet_dt"] = jdata["resnet_dt"] - if "fitting_resnet_dt" in jdata: - fitting_net["resnet_dt"] = jdata["fitting_resnet_dt"] - return fitting_net - - -def _learning_rate(jdata: Dict[str, Any]) -> Dict[str, Any]: - """Convert data to v1 input for learning rate section. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - - Returns - ------- - Dict[str, Any] - dict with learning rate parameters - """ - learning_rate = {} - learning_rate["type"] = "exp" - _jcopy(jdata, learning_rate, ("decay_steps", "decay_rate", "start_lr")) - return learning_rate - - -def _loss(jdata: Dict[str, Any]) -> Dict[str, Any]: - """Convert data to v1 input for loss function. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - - Returns - ------- - Dict[str, Any] - dict with loss function parameters - """ - loss: Dict[str, Any] = {} - _jcopy( - jdata, - loss, - ( - "start_pref_e", - "limit_pref_e", - "start_pref_f", - "limit_pref_f", - "start_pref_v", - "limit_pref_v", - ), - ) - if "start_pref_ae" in jdata: - loss["start_pref_ae"] = jdata["start_pref_ae"] - if "limit_pref_ae" in jdata: - loss["limit_pref_ae"] = jdata["limit_pref_ae"] - return loss - - -def _training(jdata: Dict[str, Any]) -> Dict[str, Any]: - """Convert data to v1 input for training. - - Parameters - ---------- - jdata : Dict[str, Any] - parsed input json/yaml data - - Returns - ------- - Dict[str, Any] - dict with training parameters - """ - training = {} - seed = jdata.get("seed", None) - if seed is not None: - training["seed"] = seed - - _jcopy(jdata, training, ("systems", "set_prefix", "stop_batch", "batch_size")) - training["disp_file"] = "lcurve.out" - if "disp_file" in jdata: - training["disp_file"] = jdata["disp_file"] - training["disp_freq"] = j_must_have(jdata, "disp_freq") - training["numb_test"] = j_must_have(jdata, "numb_test") - training["save_freq"] = j_must_have(jdata, "save_freq") - training["save_ckpt"] = j_must_have(jdata, "save_ckpt") - training["disp_training"] = j_must_have(jdata, "disp_training") - training["time_training"] = j_must_have(jdata, "time_training") - if "profiling" in jdata: - training["profiling"] = jdata["profiling"] - if training["profiling"]: - training["profiling_file"] = j_must_have(jdata, "profiling_file") - return training - - -def _jcopy(src: Dict[str, Any], dst: Dict[str, Any], keys: Sequence[str]): - """Copy specified keys from one dict to another. - - Parameters - ---------- - src : Dict[str, Any] - source dictionary - dst : Dict[str, Any] - destination dictionary, will be modified in place - keys : Sequence[str] - list of keys to copy - """ - for k in keys: - dst[k] = src[k] - - -def remove_decay_rate(jdata: Dict[str, Any]): - """Convert decay_rate to stop_lr. - - Parameters - ---------- - jdata : Dict[str, Any] - input data - """ - lr = jdata["learning_rate"] - if "decay_rate" in lr: - decay_rate = lr["decay_rate"] - start_lr = lr["start_lr"] - stop_step = jdata["training"]["stop_batch"] - decay_steps = lr["decay_steps"] - stop_lr = np.exp(np.log(decay_rate) * (stop_step / decay_steps)) * start_lr - lr["stop_lr"] = stop_lr - lr.pop("decay_rate") - - -def convert_input_v1_v2( - jdata: Dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None -) -> Dict[str, Any]: - tr_cfg = jdata["training"] - tr_data_keys = { - "systems", - "set_prefix", - "batch_size", - "sys_prob", - "auto_prob", - # alias included - "sys_weights", - "auto_prob_style", - } - - tr_data_cfg = {k: v for k, v in tr_cfg.items() if k in tr_data_keys} - new_tr_cfg = {k: v for k, v in tr_cfg.items() if k not in tr_data_keys} - new_tr_cfg["training_data"] = tr_data_cfg - if "training_data" in tr_cfg: - raise RuntimeError( - "Both v1 (training/systems) and v2 (training/training_data) parameters are given." - ) - - jdata["training"] = new_tr_cfg - - # remove deprecated arguments - remove_decay_rate(jdata) - - if warning: - _warning_input_v1_v2(dump) - if dump is not None: - with open(dump, "w") as fp: - json.dump(jdata, fp, indent=4) - - return jdata - - -def _warning_input_v1_v2(fname: Optional[Union[str, Path]]): - 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" - ) - if fname is not None: - msg += f", and output it to file {fname}" - warnings.warn(msg) - - -def deprecate_numb_test( - jdata: Dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None -) -> Dict[str, Any]: - """Deprecate `numb_test` since v2.1. It has taken no effect since v2.0. - - See `#1243 `_. - - Parameters - ---------- - jdata : Dict[str, Any] - loaded json/yaml file - warning : bool, optional - whether to show deprecation warning, by default True - dump : Optional[Union[str, Path]], optional - whether to dump converted file, by default None - - Returns - ------- - Dict[str, Any] - converted output - """ - try: - jdata.get("training", {}).pop("numb_test") - except KeyError: - pass - else: - if warning: - warnings.warn( - "The argument training->numb_test has been deprecated since v2.0.0. " - "Use training->validation_data->batch_size instead." - ) - - if dump is not None: - with open(dump, "w") as fp: - json.dump(jdata, fp, indent=4) - return jdata - - -def update_deepmd_input( - jdata: Dict[str, Any], warning: bool = True, dump: Optional[Union[str, Path]] = None -) -> Dict[str, Any]: - def is_deepmd_v0_input(jdata): - return "model" not in jdata.keys() - - def is_deepmd_v1_input(jdata): - return "systems" in j_must_have(jdata, "training").keys() - - if is_deepmd_v0_input(jdata): - jdata = convert_input_v0_v1(jdata, warning, None) - jdata = convert_input_v1_v2(jdata, False, None) - jdata = deprecate_numb_test(jdata, False, dump) - elif is_deepmd_v1_input(jdata): - jdata = convert_input_v1_v2(jdata, warning, None) - jdata = deprecate_numb_test(jdata, False, dump) - else: - jdata = deprecate_numb_test(jdata, warning, dump) - - return jdata +__all__ = [ + "convert_input_v0_v1", + "convert_input_v1_v2", + "deprecate_numb_test", + "update_deepmd_input", +] diff --git a/deepmd/utils/compress.py b/deepmd/utils/compress.py index c6e68dfe19..7a79dec520 100644 --- a/deepmd/utils/compress.py +++ b/deepmd/utils/compress.py @@ -43,15 +43,15 @@ def get_two_side_type_embedding(self, graph): def get_extra_side_embedding_net_variable( - self, graph_def, type_side, varialbe_name, suffix + self, graph_def, type_side_suffix, varialbe_name, suffix ): ret = {} for i in range(1, self.layer_size + 1): target = get_pattern_nodes_from_graph_def( graph_def, - f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_side}_ebd", + f"filter_type_all{suffix}/{varialbe_name}_{i}{type_side_suffix}", ) - node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}_{type_side}_ebd"] + node = target[f"filter_type_all{suffix}/{varialbe_name}_{i}{type_side_suffix}"] ret["layer_" + str(i)] = node return ret diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 423745cddf..a6f888beac 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -1,614 +1,9 @@ -#!/usr/bin/env python3 - # SPDX-License-Identifier: LGPL-3.0-or-later -import logging -from typing import ( - List, - Optional, -) - -import numpy as np - -from deepmd.env import ( - GLOBAL_ENER_FLOAT_PRECISION, - GLOBAL_NP_FLOAT_PRECISION, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.data import ( + DeepmdData, ) -from deepmd.utils import random as dp_random -from deepmd.utils.path import ( - DPPath, -) - -log = logging.getLogger(__name__) - - -class DeepmdData: - """Class for a data system. - - It loads data from hard disk, and mantains the data as a `data_dict` - - Parameters - ---------- - sys_path - Path to the data system - set_prefix - Prefix for the directories of different sets - shuffle_test - If the test data are shuffled - type_map - Gives the name of different atom types - optional_type_map - If the type_map.raw in each system is optional - modifier - Data modifier that has the method `modify_data` - trn_all_set - Use all sets as training dataset. Otherwise, if the number of sets is more than 1, the last set is left for test. - sort_atoms : bool - Sort atoms by atom types. Required to enable when the data is directly feeded to - descriptors except mixed types. - """ - - def __init__( - self, - sys_path: str, - set_prefix: str = "set", - shuffle_test: bool = True, - type_map: Optional[List[str]] = None, - optional_type_map: bool = True, - modifier=None, - trn_all_set: bool = False, - sort_atoms: bool = True, - ): - """Constructor.""" - root = DPPath(sys_path) - self.dirs = root.glob(set_prefix + ".*") - if not len(self.dirs): - 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!" - ) - 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 - # load atom type - self.atom_type = self._load_type(root) - self.natoms = len(self.atom_type) - # load atom type map - self.type_map = self._load_type_map(root) - assert ( - optional_type_map or self.type_map is not None - ), f"System {sys_path} must have type_map.raw in this mode! " - if self.type_map is not None: - assert len(self.type_map) >= max(self.atom_type) + 1 - # check pbc - self.pbc = self._check_pbc(root) - # enforce type_map if necessary - self.enforce_type_map = False - if type_map is not None and self.type_map is not None and len(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) - else: - self.enforce_type_map = True - sorter = np.argsort(type_map) - self.type_idx_map = np.array( - sorter[np.searchsorted(type_map, self.type_map, sorter=sorter)] - ) - # padding for virtual atom - self.type_idx_map = np.append( - self.type_idx_map, np.array([-1], dtype=np.int32) - ) - self.type_map = type_map - if type_map is None and self.type_map is None and self.mixed_type: - raise RuntimeError("mixed_type format must have type_map!") - # make idx map - self.sort_atoms = sort_atoms - self.idx_map = self._make_idx_map(self.atom_type) - # train dirs - self.test_dir = self.dirs[-1] - if trn_all_set: - self.train_dirs = self.dirs - else: - if len(self.dirs) == 1: - self.train_dirs = self.dirs - else: - self.train_dirs = self.dirs[:-1] - self.data_dict = {} - # add box and coord - self.add("box", 9, must=self.pbc) - self.add("coord", 3, atomic=True, must=True) - # the training times of each frame - self.add("numb_copy", 1, must=False, default=1, dtype=int) - # set counters - self.set_count = 0 - self.iterator = 0 - self.shuffle_test = shuffle_test - # set modifier - self.modifier = modifier - - def add( - self, - key: str, - ndof: int, - atomic: bool = False, - must: bool = False, - high_prec: bool = False, - type_sel: Optional[List[int]] = None, - repeat: int = 1, - default: float = 0.0, - dtype: Optional[np.dtype] = None, - ): - """Add a data item that to be loaded. - - Parameters - ---------- - key - The key of the item. The corresponding data is stored in `sys_path/set.*/key.npy` - ndof - The number of dof - atomic - The item is an atomic property. - If False, the size of the data should be nframes x ndof - If True, the size of data should be nframes x natoms x ndof - must - The data file `sys_path/set.*/key.npy` must exist. - If must is False and the data file does not exist, the `data_dict[find_key]` is set to 0.0 - high_prec - Load the data and store in float64, otherwise in float32 - type_sel - Select certain type of atoms - repeat - The data will be repeated `repeat` times. - default : float, default=0. - default value of data - dtype : np.dtype, optional - the dtype of data, overwrites `high_prec` if provided - """ - self.data_dict[key] = { - "ndof": ndof, - "atomic": atomic, - "must": must, - "high_prec": high_prec, - "type_sel": type_sel, - "repeat": repeat, - "reduce": None, - "default": default, - "dtype": dtype, - } - return self - - def reduce(self, key_out: str, key_in: str): - """Generate a new item from the reduction of another atom. - - Parameters - ---------- - key_out - The name of the reduced item - key_in - The name of the data item to be reduced - """ - assert key_in in self.data_dict, "cannot find input key" - assert self.data_dict[key_in]["atomic"], "reduced property should be atomic" - assert key_out not in self.data_dict, "output key should not have been added" - assert ( - self.data_dict[key_in]["repeat"] == 1 - ), "reduced proerties should not have been repeated" - - self.data_dict[key_out] = { - "ndof": self.data_dict[key_in]["ndof"], - "atomic": False, - "must": True, - "high_prec": True, - "type_sel": None, - "repeat": 1, - "reduce": key_in, - } - return self - - def get_data_dict(self) -> dict: - """Get the `data_dict`.""" - return self.data_dict - - def check_batch_size(self, batch_size): - """Check if the system can get a batch of data with `batch_size` frames.""" - for ii in self.train_dirs: - if self.data_dict["coord"]["high_prec"]: - tmpe = ( - (ii / "coord.npy").load_numpy().astype(GLOBAL_ENER_FLOAT_PRECISION) - ) - else: - tmpe = (ii / "coord.npy").load_numpy().astype(GLOBAL_NP_FLOAT_PRECISION) - if tmpe.ndim == 1: - tmpe = tmpe.reshape([1, -1]) - if tmpe.shape[0] < batch_size: - return ii, tmpe.shape[0] - return None - - def check_test_size(self, test_size): - """Check if the system can get a test dataset with `test_size` frames.""" - if self.data_dict["coord"]["high_prec"]: - tmpe = ( - (self.test_dir / "coord.npy") - .load_numpy() - .astype(GLOBAL_ENER_FLOAT_PRECISION) - ) - else: - tmpe = ( - (self.test_dir / "coord.npy") - .load_numpy() - .astype(GLOBAL_NP_FLOAT_PRECISION) - ) - if tmpe.ndim == 1: - tmpe = tmpe.reshape([1, -1]) - if tmpe.shape[0] < test_size: - return self.test_dir, tmpe.shape[0] - else: - return None - - 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. - - Parameters - ---------- - batch_size - size of the batch - """ - if hasattr(self, "batch_set"): - set_size = self.batch_set["coord"].shape[0] - else: - set_size = 0 - if self.iterator + batch_size > set_size: - self._load_batch_set(self.train_dirs[self.set_count % self.get_numb_set()]) - self.set_count += 1 - set_size = self.batch_set["coord"].shape[0] - iterator_1 = self.iterator + batch_size - if iterator_1 >= set_size: - iterator_1 = set_size - idx = np.arange(self.iterator, iterator_1) - self.iterator += batch_size - ret = self._get_subdata(self.batch_set, idx) - return ret - - def get_test(self, ntests: int = -1) -> dict: - """Get the test data with `ntests` frames. - - Parameters - ---------- - ntests - Size of the test data set. If `ntests` is -1, all test data will be get. - """ - if not hasattr(self, "test_set"): - self._load_test_set(self.test_dir, self.shuffle_test) - if ntests == -1: - idx = None - else: - ntests_ = ( - ntests - if ntests < self.test_set["type"].shape[0] - else self.test_set["type"].shape[0] - ) - # print('ntest', self.test_set['type'].shape[0], ntests, ntests_) - idx = np.arange(ntests_) - ret = self._get_subdata(self.test_set, idx=idx) - if self.modifier is not None: - self.modifier.modify_data(ret, self) - return ret - - def get_ntypes(self) -> int: - """Number of atom types in the system.""" - if self.type_map is not None: - return len(self.type_map) - else: - return max(self.get_atom_type()) + 1 - - def get_type_map(self) -> List[str]: - """Get the type map.""" - return self.type_map - - def get_atom_type(self) -> List[int]: - """Get atom types.""" - return self.atom_type - - def get_numb_set(self) -> int: - """Get number of training sets.""" - return len(self.train_dirs) - - 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.train_dirs[set_idx]) - ret = data["coord"].shape[0] // batch_size - if ret == 0: - ret = 1 - return ret - - def get_sys_numb_batch(self, batch_size: int) -> int: - """Get the number of batches in the data system.""" - ret = 0 - for ii in range(len(self.train_dirs)): - ret += self.get_numb_batch(batch_size, ii) - return ret - - def get_natoms(self): - """Get number of atoms.""" - return len(self.atom_type) - - def get_natoms_vec(self, ntypes: int): - """Get number of atoms and number of atoms in different types. - - Parameters - ---------- - ntypes - Number of types (may be larger than the actual number of types in the system). - - Returns - ------- - natoms - natoms[0]: number of local atoms - natoms[1]: total number of atoms held by this processor - natoms[i]: 2 <= i < Ntypes+2, number of type i atoms - """ - natoms, natoms_vec = self._get_natoms_2(ntypes) - tmp = [natoms, natoms] - tmp = np.append(tmp, natoms_vec) - return tmp.astype(np.int32) - - def avg(self, key): - """Return the average value of an item.""" - if key not in self.data_dict.keys(): - raise RuntimeError("key %s has not been added" % key) - info = self.data_dict[key] - ndof = info["ndof"] - eners = [] - for ii in self.train_dirs: - data = self._load_set(ii) - ei = data[key].reshape([-1, ndof]) - eners.append(ei) - eners = np.concatenate(eners, axis=0) - if eners.size == 0: - return 0 - 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) - natoms = new_types.shape[0] - idx = np.arange(natoms) - idx_map = np.lexsort((idx, new_types)) - return idx_map - - def _get_natoms_2(self, ntypes): - sample_type = self.atom_type - natoms = len(sample_type) - natoms_vec = np.zeros(ntypes).astype(int) - for ii in range(ntypes): - natoms_vec[ii] = np.count_nonzero(sample_type == ii) - return natoms, natoms_vec - - def _get_subdata(self, data, idx=None): - new_data = {} - for ii in data: - dd = data[ii] - if "find_" in ii: - new_data[ii] = dd - else: - if idx is not None: - new_data[ii] = dd[idx] - else: - new_data[ii] = dd - return new_data - - def _load_batch_set(self, set_name: DPPath): - if not hasattr(self, "batch_set") or self.get_numb_set() > 1: - self.batch_set = self._load_set(set_name) - if self.modifier is not None: - self.modifier.modify_data(self.batch_set, self) - self.batch_set, _ = self._shuffle_data(self.batch_set) - self.reset_get_batch() - - def reset_get_batch(self): - self.iterator = 0 - - def _load_test_set(self, set_name: DPPath, shuffle_test): - self.test_set = self._load_set(set_name) - if shuffle_test: - self.test_set, _ = self._shuffle_data(self.test_set) - - def _shuffle_data(self, data): - ret = {} - nframes = data["coord"].shape[0] - idx = np.arange(nframes) - # the training times of each frame - idx = np.repeat(idx, np.reshape(data["numb_copy"], (nframes,))) - dp_random.shuffle(idx) - for kk in data: - if ( - type(data[kk]) == np.ndarray - and len(data[kk].shape) == 2 - and data[kk].shape[0] == nframes - and "find_" not in kk - ): - ret[kk] = data[kk][idx] - else: - ret[kk] = data[kk] - return ret, idx - - def _load_set(self, set_name: DPPath): - # get nframes - 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) - else: - coord = path.load_numpy().astype(GLOBAL_NP_FLOAT_PRECISION) - if coord.ndim == 1: - coord = coord.reshape([1, -1]) - nframes = coord.shape[0] - assert coord.shape[1] == self.data_dict["coord"]["ndof"] * self.natoms - # load keys - data = {} - for kk in self.data_dict.keys(): - if self.data_dict[kk]["reduce"] is None: - data["find_" + kk], data[kk] = self._load_data( - set_name, - kk, - nframes, - self.data_dict[kk]["ndof"], - atomic=self.data_dict[kk]["atomic"], - high_prec=self.data_dict[kk]["high_prec"], - must=self.data_dict[kk]["must"], - type_sel=self.data_dict[kk]["type_sel"], - repeat=self.data_dict[kk]["repeat"], - default=self.data_dict[kk]["default"], - dtype=self.data_dict[kk]["dtype"], - ) - for kk in self.data_dict.keys(): - if self.data_dict[kk]["reduce"] is not None: - k_in = self.data_dict[kk]["reduce"] - ndof = self.data_dict[kk]["ndof"] - data["find_" + kk] = data["find_" + k_in] - tmp_in = data[k_in].astype(GLOBAL_ENER_FLOAT_PRECISION) - data[kk] = np.sum( - np.reshape(tmp_in, [nframes, self.natoms, ndof]), axis=1 - ) - - if self.mixed_type: - # nframes x natoms - atom_type_mix = self._load_type_mix(set_name) - if self.enforce_type_map: - try: - atom_type_mix_ = self.type_idx_map[atom_type_mix].astype(np.int32) - except IndexError as e: - raise IndexError( - "some types in 'real_atom_types.npy' of set {} are not contained in {} types!".format( - set_name, self.get_ntypes() - ) - ) from e - atom_type_mix = atom_type_mix_ - real_type = atom_type_mix.reshape([nframes, self.natoms]) - data["type"] = real_type - natoms = data["type"].shape[1] - # nframes x ntypes - atom_type_nums = np.array( - [(real_type == i).sum(axis=-1) for i in range(self.get_ntypes())], - dtype=np.int32, - ).T - ghost_nums = np.array( - [(real_type == -1).sum(axis=-1)], - dtype=np.int32, - ).T - assert ( - atom_type_nums.sum(axis=-1) + ghost_nums.sum(axis=-1) == natoms - ).all(), "some types in 'real_atom_types.npy' of set {} are not contained in {} types!".format( - set_name, self.get_ntypes() - ) - data["real_natoms_vec"] = np.concatenate( - ( - np.tile(np.array([natoms, natoms], dtype=np.int32), (nframes, 1)), - atom_type_nums, - ), - axis=-1, - ) - else: - data["type"] = np.tile(self.atom_type[self.idx_map], (nframes, 1)) - - return data - - def _load_data( - self, - set_name, - key, - nframes, - ndof_, - atomic=False, - must=True, - repeat=1, - high_prec=False, - type_sel=None, - default: float = 0.0, - dtype: Optional[np.dtype] = None, - ): - 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 = 0 - for jj in type_sel: - natoms += np.sum(self.atom_type == jj) - idx_map = self._idx_map_sel(self.atom_type, type_sel) - ndof = ndof_ * natoms - else: - ndof = ndof_ - if dtype is not None: - pass - elif high_prec: - dtype = GLOBAL_ENER_FLOAT_PRECISION - else: - dtype = GLOBAL_NP_FLOAT_PRECISION - path = set_name / (key + ".npy") - if path.is_file(): - data = path.load_numpy().astype(dtype) - try: # YWolfeee: deal with data shape error - if atomic: - data = data.reshape([nframes, natoms, -1]) - data = data[:, idx_map, :] - 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`." - log.error(str(err_message)) - log.error(explanation) - raise ValueError(str(err_message) + ". " + explanation) - if repeat != 1: - data = np.repeat(data, repeat).reshape([nframes, -1]) - return np.float32(1.0), data - elif must: - raise RuntimeError("%s not found!" % path) - else: - data = np.full([nframes, ndof], default, dtype=dtype) - if repeat != 1: - data = np.repeat(data, repeat).reshape([nframes, -1]) - return np.float32(0.0), data - - def _load_type(self, sys_path: DPPath): - atom_type = (sys_path / "type.raw").load_txt(ndmin=1).astype(np.int32) - return atom_type - - def _load_type_mix(self, set_name: DPPath): - 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): - natoms = atom_type.shape[0] - idx = np.arange(natoms) - if self.sort_atoms: - idx_map = np.lexsort((idx, atom_type)) - else: - idx_map = idx - return idx_map - - def _load_type_map(self, sys_path: DPPath): - 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): - pbc = True - if (sys_path / "nopbc").is_file(): - pbc = False - return pbc - def _check_mode(self, set_path: DPPath): - return (set_path / "real_atom_types.npy").is_file() +__all__ = [ + "DeepmdData", +] diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 69a6cbe112..65e87d8ebc 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -1,653 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import collections -import logging -import warnings -from functools import ( - lru_cache, -) -from typing import ( - List, - Optional, -) - -import numpy as np - -from deepmd.common import ( - make_default_mesh, -) -from deepmd.env import ( - GLOBAL_NP_FLOAT_PRECISION, -) -from deepmd.utils import random as dp_random -from deepmd.utils.data import ( - DeepmdData, -) - -log = logging.getLogger(__name__) - - -class DeepmdDataSystem: - """Class for manipulating many data systems. - - It is implemented with the help of DeepmdData - """ - - def __init__( - self, - systems: List[str], - batch_size: int, - test_size: int, - rcut: float, - set_prefix: str = "set", - shuffle_test: bool = True, - type_map: Optional[List[str]] = None, - optional_type_map: bool = True, - modifier=None, - trn_all_set=False, - sys_probs=None, - auto_prob_style="prob_sys_size", - sort_atoms: bool = True, - ): - """Constructor. - - Parameters - ---------- - systems - Specifying the paths to systems - batch_size - The batch size - test_size - The size of test data - rcut - The cut-off radius - set_prefix - Prefix for the directories of different sets - shuffle_test - If the test data are shuffled - type_map - Gives the name of different atom types - optional_type_map - If the type_map.raw in each system is optional - modifier - Data modifier that has the method `modify_data` - trn_all_set - Use all sets as training dataset. Otherwise, if the number of sets is more than 1, the last set is left for test. - sys_probs : list of float - The probabilitis of systems to get the batch. - Summation of positive elements of this list should be no greater than 1. - Element of this list can be negative, the probability of the corresponding system is determined - automatically by the number of batches in the system. - auto_prob_style : str - Determine the probability of systems automatically. The method is assigned by this key and can be - - "prob_uniform" : the probability all the systems are equal, namely 1.0/self.get_nsystems() - - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system - - "prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;..." : - the list of systems is devided into blocks. A block is specified by `stt_idx:end_idx:weight`, - where `stt_idx` is the starting index of the system, `end_idx` is then ending (not including) index of the system, - the probabilities of the systems in this block sums up to `weight`, and the relatively probabilities within this block is proportional - to the number of batches in the system. - sort_atoms : bool - Sort atoms by atom types. Required to enable when the data is directly feeded to - descriptors except mixed types. - """ - # init data - self.rcut = rcut - self.system_dirs = systems - self.nsystems = len(self.system_dirs) - self.data_systems = [] - for ii in self.system_dirs: - self.data_systems.append( - DeepmdData( - ii, - set_prefix=set_prefix, - shuffle_test=shuffle_test, - type_map=type_map, - optional_type_map=optional_type_map, - modifier=modifier, - trn_all_set=trn_all_set, - sort_atoms=sort_atoms, - ) - ) - # check mix_type format - error_format_msg = ( - "if one of the system is of mixed_type format, " - "then all of the systems should be of mixed_type format!" - ) - if self.data_systems[0].mixed_type: - for data_sys in self.data_systems[1:]: - assert data_sys.mixed_type, error_format_msg - self.mixed_type = True - else: - for data_sys in self.data_systems[1:]: - assert not data_sys.mixed_type, error_format_msg - self.mixed_type = False - # batch size - self.batch_size = batch_size - is_auto_bs = False - self.mixed_systems = False - if isinstance(self.batch_size, int): - self.batch_size = self.batch_size * np.ones(self.nsystems, dtype=int) - elif isinstance(self.batch_size, str): - words = self.batch_size.split(":") - if "auto" == words[0]: - is_auto_bs = True - rule = 32 - if len(words) == 2: - rule = int(words[1]) - self.batch_size = self._make_auto_bs(rule) - elif "mixed" == words[0]: - self.mixed_type = True - self.mixed_systems = True - if len(words) == 2: - rule = int(words[1]) - else: - raise RuntimeError("batch size must be specified for mixed systems") - self.batch_size = rule * np.ones(self.nsystems, dtype=int) - else: - raise RuntimeError("unknown batch_size rule " + words[0]) - elif isinstance(self.batch_size, list): - pass - else: - raise RuntimeError("invalid batch_size") - assert isinstance(self.batch_size, (list, np.ndarray)) - assert len(self.batch_size) == self.nsystems - - # natoms, nbatches - ntypes = [] - for ii in self.data_systems: - ntypes.append(ii.get_ntypes()) - self.sys_ntypes = max(ntypes) - self.natoms = [] - self.natoms_vec = [] - self.nbatches = [] - type_map_list = [] - for ii in range(self.nsystems): - self.natoms.append(self.data_systems[ii].get_natoms()) - self.natoms_vec.append( - self.data_systems[ii].get_natoms_vec(self.sys_ntypes).astype(int) - ) - self.nbatches.append( - self.data_systems[ii].get_sys_numb_batch(self.batch_size[ii]) - ) - type_map_list.append(self.data_systems[ii].get_type_map()) - self.type_map = self._check_type_map_consistency(type_map_list) - - # ! altered by Marián Rynik - # test size - # now test size can be set as a percentage of systems data or test size - # can be set for each system individualy in the same manner as batch - # size. This enables one to use systems with diverse number of - # structures and different number of atoms. - self.test_size = test_size - if isinstance(self.test_size, int): - self.test_size = self.test_size * np.ones(self.nsystems, dtype=int) - elif isinstance(self.test_size, str): - words = self.test_size.split("%") - try: - percent = int(words[0]) - except ValueError: - raise RuntimeError("unknown test_size rule " + words[0]) - self.test_size = self._make_auto_ts(percent) - elif isinstance(self.test_size, list): - pass - else: - raise RuntimeError("invalid test_size") - assert isinstance(self.test_size, (list, np.ndarray)) - assert len(self.test_size) == self.nsystems - - # prob of batch, init pick idx - self.prob_nbatches = [float(i) for i in self.nbatches] / np.sum(self.nbatches) - self.pick_idx = 0 - - # derive system probabilities - self.sys_probs = None - self.set_sys_probs(sys_probs, auto_prob_style) - - # check batch and test size - for ii in range(self.nsystems): - chk_ret = self.data_systems[ii].check_batch_size(self.batch_size[ii]) - if chk_ret is not None and not is_auto_bs and not self.mixed_systems: - warnings.warn( - "system %s required batch size is larger than the size of the dataset %s (%d > %d)" - % ( - self.system_dirs[ii], - chk_ret[0], - self.batch_size[ii], - chk_ret[1], - ) - ) - chk_ret = self.data_systems[ii].check_test_size(self.test_size[ii]) - if chk_ret is not None and not is_auto_bs and not self.mixed_systems: - warnings.warn( - "system %s required test size is larger than the size of the dataset %s (%d > %d)" - % (self.system_dirs[ii], chk_ret[0], self.test_size[ii], chk_ret[1]) - ) - - def _load_test(self, ntests=-1): - self.test_data = collections.defaultdict(list) - for ii in range(self.nsystems): - test_system_data = self.data_systems[ii].get_test(ntests=ntests) - for nn in test_system_data: - self.test_data[nn].append(test_system_data[nn]) - - @property - @lru_cache(maxsize=None) - def default_mesh(self) -> List[np.ndarray]: - """Mesh for each system.""" - return [ - make_default_mesh( - self.data_systems[ii].pbc, self.data_systems[ii].mixed_type - ) - for ii in range(self.nsystems) - ] - - def compute_energy_shift(self, rcond=None, key="energy"): - sys_ener = [] - for ss in self.data_systems: - sys_ener.append(ss.avg(key)) - sys_ener = np.concatenate(sys_ener) - sys_tynatom = np.array(self.natoms_vec, dtype=GLOBAL_NP_FLOAT_PRECISION) - sys_tynatom = np.reshape(sys_tynatom, [self.nsystems, -1]) - sys_tynatom = sys_tynatom[:, 2:] - energy_shift, resd, rank, s_value = np.linalg.lstsq( - sys_tynatom, sys_ener, rcond=rcond - ) - return energy_shift - - def add_dict(self, adict: dict) -> None: - """Add items to the data system by a `dict`. - `adict` should have items like - .. code-block:: python. - - adict[key] = { - "ndof": ndof, - "atomic": atomic, - "must": must, - "high_prec": high_prec, - "type_sel": type_sel, - "repeat": repeat, - } - - For the explaination of the keys see `add` - """ - for kk in adict: - self.add( - kk, - adict[kk]["ndof"], - atomic=adict[kk]["atomic"], - must=adict[kk]["must"], - high_prec=adict[kk]["high_prec"], - type_sel=adict[kk]["type_sel"], - repeat=adict[kk]["repeat"], - default=adict[kk]["default"], - ) - - def add( - self, - key: str, - ndof: int, - atomic: bool = False, - must: bool = False, - high_prec: bool = False, - type_sel: Optional[List[int]] = None, - repeat: int = 1, - default: float = 0.0, - ): - """Add a data item that to be loaded. - - Parameters - ---------- - key - The key of the item. The corresponding data is stored in `sys_path/set.*/key.npy` - ndof - The number of dof - atomic - The item is an atomic property. - If False, the size of the data should be nframes x ndof - If True, the size of data should be nframes x natoms x ndof - must - The data file `sys_path/set.*/key.npy` must exist. - If must is False and the data file does not exist, the `data_dict[find_key]` is set to 0.0 - high_prec - Load the data and store in float64, otherwise in float32 - type_sel - Select certain type of atoms - repeat - The data will be repeated `repeat` times. - default, default=0. - Default value of data - """ - for ii in self.data_systems: - ii.add( - key, - ndof, - atomic=atomic, - must=must, - high_prec=high_prec, - repeat=repeat, - type_sel=type_sel, - default=default, - ) - - def reduce(self, key_out, key_in): - """Generate a new item from the reduction of another atom. - - Parameters - ---------- - key_out - The name of the reduced item - key_in - The name of the data item to be reduced - """ - for ii in self.data_systems: - ii.reduce(key_out, key_in) - - 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"): - if sys_probs is None: - if auto_prob_style == "prob_uniform": - prob_v = 1.0 / float(self.nsystems) - probs = [prob_v for ii in range(self.nsystems)] - elif auto_prob_style == "prob_sys_size": - probs = self.prob_nbatches - elif auto_prob_style[:14] == "prob_sys_size;": - probs = prob_sys_size_ext( - auto_prob_style, self.get_nsystems(), self.nbatches - ) - else: - raise RuntimeError("Unknown auto prob style: " + auto_prob_style) - else: - probs = process_sys_probs(sys_probs, self.nbatches) - self.sys_probs = probs - - def get_batch(self, sys_idx: Optional[int] = 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 - # defined as a private variable, self.sys_probs, initialized in __init__(). - # This is to optimize the (vain) efforts in evaluating sys_probs every batch. - """Get a batch of data from the data systems. - - Parameters - ---------- - sys_idx : int - The index of system from which the batch is get. - If sys_idx is not None, `sys_probs` and `auto_prob_style` are ignored - If sys_idx is None, automatically determine the system according to `sys_probs` or `auto_prob_style`, see the following. - This option does not work for mixed systems. - - Returns - ------- - dict - The batch data - """ - if not self.mixed_systems: - b_data = self.get_batch_standard(sys_idx) - else: - b_data = self.get_batch_mixed() - return b_data - - def get_batch_standard(self, sys_idx: Optional[int] = None) -> dict: - """Get a batch of data from the data systems in the standard way. - - Parameters - ---------- - sys_idx : int - The index of system from which the batch is get. - If sys_idx is not None, `sys_probs` and `auto_prob_style` are ignored - If sys_idx is None, automatically determine the system according to `sys_probs` or `auto_prob_style`, see the following. - - Returns - ------- - dict - The batch data - """ - if sys_idx is not None: - self.pick_idx = sys_idx - else: - # prob = self._get_sys_probs(sys_probs, auto_prob_style) - self.pick_idx = dp_random.choice(np.arange(self.nsystems), p=self.sys_probs) - b_data = self.data_systems[self.pick_idx].get_batch( - self.batch_size[self.pick_idx] - ) - b_data["natoms_vec"] = self.natoms_vec[self.pick_idx] - b_data["default_mesh"] = self.default_mesh[self.pick_idx] - return b_data - - def get_batch_mixed(self) -> dict: - """Get a batch of data from the data systems in the mixed way. - - Returns - ------- - dict - The batch data - """ - # mixed systems have a global batch size - batch_size = self.batch_size[0] - batch_data = [] - for _ in range(batch_size): - self.pick_idx = dp_random.choice(np.arange(self.nsystems), p=self.sys_probs) - bb_data = self.data_systems[self.pick_idx].get_batch(1) - bb_data["natoms_vec"] = self.natoms_vec[self.pick_idx] - bb_data["default_mesh"] = self.default_mesh[self.pick_idx] - batch_data.append(bb_data) - b_data = self._merge_batch_data(batch_data) - return b_data - - def _merge_batch_data(self, batch_data: List[dict]) -> dict: - """Merge batch data from different systems. - - Parameters - ---------- - batch_data : list of dict - A list of batch data from different systems. - - Returns - ------- - dict - The merged batch data. - """ - b_data = {} - max_natoms = max(bb["natoms_vec"][0] for bb in batch_data) - # natoms_vec - natoms_vec = np.zeros(2 + self.get_ntypes(), dtype=int) - natoms_vec[0:3] = max_natoms - b_data["natoms_vec"] = natoms_vec - # real_natoms_vec - real_natoms_vec = np.vstack([bb["natoms_vec"] for bb in batch_data]) - b_data["real_natoms_vec"] = real_natoms_vec - # type - type_vec = np.full((len(batch_data), max_natoms), -1, dtype=int) - for ii, bb in enumerate(batch_data): - type_vec[ii, : bb["type"].shape[1]] = bb["type"][0] - b_data["type"] = type_vec - # default_mesh - default_mesh = np.mean([bb["default_mesh"] for bb in batch_data], axis=0) - b_data["default_mesh"] = default_mesh - # other data - data_dict = self.get_data_dict(0) - for kk, vv in data_dict.items(): - if kk not in batch_data[0]: - continue - b_data["find_" + kk] = batch_data[0]["find_" + kk] - if not vv["atomic"]: - b_data[kk] = np.concatenate([bb[kk] for bb in batch_data], axis=0) - else: - b_data[kk] = np.zeros( - (len(batch_data), max_natoms * vv["ndof"] * vv["repeat"]), - dtype=batch_data[0][kk].dtype, - ) - for ii, bb in enumerate(batch_data): - b_data[kk][ii, : bb[kk].shape[1]] = bb[kk][0] - return b_data - - # ! altered by Marián Rynik - def get_test(self, sys_idx: Optional[int] = None, n_test: int = -1): # depreciated - """Get test data from the the data systems. - - Parameters - ---------- - sys_idx - The test dat of system with index `sys_idx` will be returned. - If is None, the currently selected system will be returned. - n_test - Number of test data. If set to -1 all test data will be get. - """ - if not hasattr(self, "test_data"): - self._load_test(ntests=n_test) - if sys_idx is not None: - idx = sys_idx - else: - idx = self.pick_idx - - test_system_data = {} - for nn in self.test_data: - test_system_data[nn] = self.test_data[nn][idx] - test_system_data["natoms_vec"] = self.natoms_vec[idx] - test_system_data["default_mesh"] = self.default_mesh[idx] - return test_system_data - - def get_sys_ntest(self, sys_idx=None): - """Get number of tests for the currently selected system, - or one defined by sys_idx. - """ - if sys_idx is not None: - return self.test_size[sys_idx] - else: - return self.test_size[self.pick_idx] - - def get_type_map(self) -> List[str]: - """Get the type map.""" - return self.type_map - - def get_nbatches(self) -> int: - """Get the total number of batches.""" - return self.nbatches - - def get_ntypes(self) -> int: - """Get the number of types.""" - return self.sys_ntypes - - def get_nsystems(self) -> int: - """Get the number of data systems.""" - return self.nsystems - - def get_sys(self, idx: int) -> DeepmdData: - """Get a certain data system.""" - return self.data_systems[idx] - - def get_batch_size(self) -> int: - """Get the batch size.""" - return self.batch_size - - def _format_name_length(self, name, width): - if len(name) <= width: - return "{: >{}}".format(name, width) - else: - name = name[-(width - 3) :] - name = "-- " + name - return name - - def print_summary(self, name): - # width 65 - sys_width = 42 - log.info( - f"---Summary of DataSystem: {name:13s}-----------------------------------------------" - ) - log.info("found %d system(s):" % self.nsystems) - log.info( - ("%s " % self._format_name_length("system", sys_width)) - + ("%6s %6s %6s %5s %3s" % ("natoms", "bch_sz", "n_bch", "prob", "pbc")) - ) - for ii in range(self.nsystems): - log.info( - "%s %6d %6d %6d %5.3f %3s" - % ( - self._format_name_length(self.system_dirs[ii], sys_width), - self.natoms[ii], - # TODO batch size * nbatches = number of structures - self.batch_size[ii], - self.nbatches[ii], - self.sys_probs[ii], - "T" if self.data_systems[ii].pbc else "F", - ) - ) - log.info( - "--------------------------------------------------------------------------------------" - ) - - def _make_auto_bs(self, rule): - bs = [] - for ii in self.data_systems: - ni = ii.get_natoms() - bsi = rule // ni - if bsi * ni < rule: - bsi += 1 - bs.append(bsi) - return bs - - # ! added by Marián Rynik - def _make_auto_ts(self, percent): - ts = [] - for ii in range(self.nsystems): - ni = self.batch_size[ii] * self.nbatches[ii] - tsi = int(ni * percent / 100) - ts.append(tsi) - - return ts - - def _check_type_map_consistency(self, type_map_list): - ret = [] - for ii in type_map_list: - if ii is not None: - min_len = min([len(ii), len(ret)]) - for idx in range(min_len): - if ii[idx] != ret[idx]: - raise RuntimeError(f"inconsistent type map: {ret!s} {ii!s}") - if len(ii) > len(ret): - ret = ii - return ret - - -def process_sys_probs(sys_probs, nbatch): - sys_probs = np.array(sys_probs) - type_filter = sys_probs >= 0 - assigned_sum_prob = np.sum(type_filter * sys_probs) - # 1e-8 is to handle floating point error; See #1917 - assert ( - assigned_sum_prob <= 1.0 + 1e-8 - ), "the sum of assigned probability should be less than 1" - rest_sum_prob = 1.0 - assigned_sum_prob - if not np.isclose(rest_sum_prob, 0): - rest_nbatch = (1 - type_filter) * nbatch - rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch) - ret_prob = rest_prob + type_filter * sys_probs - else: - ret_prob = sys_probs - assert np.isclose(np.sum(ret_prob), 1), "sum of probs should be 1" - return ret_prob - - -def prob_sys_size_ext(keywords, nsystems, nbatch): - block_str = keywords.split(";")[1:] - block_stt = [] - block_end = [] - block_weights = [] - for ii in block_str: - stt = int(ii.split(":")[0]) - end = int(ii.split(":")[1]) - weight = float(ii.split(":")[2]) - assert weight >= 0, "the weight of a block should be no less than 0" - block_stt.append(stt) - block_end.append(end) - block_weights.append(weight) - nblocks = len(block_str) - block_probs = np.array(block_weights) / np.sum(block_weights) - sys_probs = np.zeros([nsystems]) - for ii in range(nblocks): - nbatch_block = nbatch[block_stt[ii] : block_end[ii]] - tmp_prob = [float(i) for i in nbatch_block] / np.sum(nbatch_block) - sys_probs[block_stt[ii] : block_end[ii]] = tmp_prob * block_probs[ii] - return sys_probs +"""Alias for backward compatibility.""" +from deepmd_utils.utils.data_system import ( + DeepmdDataSystem, + prob_sys_size_ext, + process_sys_probs, +) + +__all__ = [ + "DeepmdDataSystem", + "process_sys_probs", + "prob_sys_size_ext", +] diff --git a/deepmd/utils/errors.py b/deepmd/utils/errors.py index 5d96fa0e6a..683131e48a 100644 --- a/deepmd/utils/errors.py +++ b/deepmd/utils/errors.py @@ -1,4 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd_utils.utils.errors import ( + OutOfMemoryError, +) + + class GraphTooLargeError(Exception): """The graph is too large, exceeding protobuf's hard limit of 2GB.""" @@ -7,5 +12,8 @@ class GraphWithoutTensorError(Exception): pass -class OutOfMemoryError(Exception): - """This error is caused by out-of-memory (OOM).""" +__all__ = [ + "OutOfMemoryError", + "GraphTooLargeError", + "GraphWithoutTensorError", +] diff --git a/deepmd/utils/finetune.py b/deepmd/utils/finetune.py index 4e597b1e05..cc6c0224de 100644 --- a/deepmd/utils/finetune.py +++ b/deepmd/utils/finetune.py @@ -41,12 +41,14 @@ def replace_model_params_with_pretrained_model( pretrained_jdata = json.loads(t_jdata) # Check the model type - assert pretrained_jdata["model"]["descriptor"]["type"] in [ - "se_atten", - "se_atten_v2", - ] and pretrained_jdata["model"]["fitting_net"]["type"] in [ - "ener" - ], "The finetune process only supports models pretrained with 'se_atten' or 'se_atten_v2' descriptor and 'ener' fitting_net!" + assert ( + pretrained_jdata["model"]["descriptor"]["type"] + in [ + "se_atten", + "se_atten_v2", + ] + and pretrained_jdata["model"]["fitting_net"]["type"] in ["ener"] + ), "The finetune process only supports models pretrained with 'se_atten' or 'se_atten_v2' descriptor and 'ener' fitting_net!" # Check the type map pretrained_type_map = pretrained_jdata["model"]["type_map"] diff --git a/deepmd/utils/graph.py b/deepmd/utils/graph.py index 2a795a45a2..ad4ee0224a 100644 --- a/deepmd/utils/graph.py +++ b/deepmd/utils/graph.py @@ -237,6 +237,91 @@ def get_embedding_net_variables_from_graph_def( return embedding_net_variables +def get_extra_embedding_net_suffix(type_one_side: bool): + """Get the extra embedding net suffix according to the value of type_one_side. + + Parameters + ---------- + type_one_side + The value of type_one_side + + Returns + ------- + str + The extra embedding net suffix + """ + if type_one_side: + extra_suffix = "_one_side_ebd" + else: + extra_suffix = "_two_side_ebd" + return extra_suffix + + +def get_variables_from_graph_def_as_numpy_array(graph_def: tf.GraphDef, pattern: str): + """Get variables from the given tf.GraphDef object, with numpy array returns. + + Parameters + ---------- + graph_def + The input tf.GraphDef object + pattern : str + The name of variable + + Returns + ------- + np.ndarray + The numpy array of the variable + """ + node = get_pattern_nodes_from_graph_def(graph_def, pattern)[pattern] + dtype = tf.as_dtype(node.dtype).as_numpy_dtype + tensor_shape = tf.TensorShape(node.tensor_shape).as_list() + if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): + tensor_value = np.frombuffer( + node.tensor_content, + dtype=tf.as_dtype(node.dtype).as_numpy_dtype, + ) + else: + tensor_value = get_tensor_by_type(node, dtype) + return np.reshape(tensor_value, tensor_shape) + + +def get_extra_embedding_net_variables_from_graph_def( + graph_def: tf.GraphDef, suffix: str, extra_suffix: str, layer_size: int +): + """Get extra embedding net variables from the given tf.GraphDef object. + The "extra embedding net" means the embedding net with only type embeddings input, + which occurs in "se_atten_v2" and "se_a_ebd_v2" descriptor. + + Parameters + ---------- + graph_def + The input tf.GraphDef object + suffix : str + The "common" suffix in the descriptor + extra_suffix : str + This value depends on the value of "type_one_side". + It should always be "_one_side_ebd" or "_two_side_ebd" + layer_size : int + The layer size of the embedding net + + Returns + ------- + Dict + The extra embedding net variables within the given tf.GraphDef object + """ + extra_embedding_net_variables = {} + for i in range(1, layer_size + 1): + matrix_pattern = f"filter_type_all{suffix}/matrix_{i}{extra_suffix}" + extra_embedding_net_variables[ + matrix_pattern + ] = get_variables_from_graph_def_as_numpy_array(graph_def, matrix_pattern) + bias_pattern = f"filter_type_all{suffix}/bias_{i}{extra_suffix}" + extra_embedding_net_variables[ + bias_pattern + ] = get_variables_from_graph_def_as_numpy_array(graph_def, bias_pattern) + return extra_embedding_net_variables + + def get_embedding_net_variables(model_file: str, suffix: str = "") -> Dict: """Get the embedding net variables with the given frozen model(model_file). diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 4451f53379..1a526ac5fc 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -1,91 +1,9 @@ -#!/usr/bin/env python3 - # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Tuple, -) - -import numpy as np -from scipy.interpolate import ( - CubicSpline, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.pair_tab import ( + PairTab, ) - -class PairTab: - """Pairwise tabulated potential. - - Parameters - ---------- - filename - File name for the short-range tabulated potential. - The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. - The first colume is the distance between atoms. - The second to the last columes are energies for pairs of certain types. - For example we have two atom types, 0 and 1. - The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. - """ - - def __init__(self, filename: str) -> None: - """Constructor.""" - self.reinit(filename) - - def reinit(self, filename: str) -> None: - """Initialize the tabulated interaction. - - Parameters - ---------- - filename - File name for the short-range tabulated potential. - The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. - The first colume is the distance between atoms. - The second to the last columes are energies for pairs of certain types. - For example we have two atom types, 0 and 1. - The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. - """ - self.vdata = np.loadtxt(filename) - self.rmin = self.vdata[0][0] - self.hh = self.vdata[1][0] - self.vdata[0][0] - self.nspline = self.vdata.shape[0] - 1 - ncol = self.vdata.shape[1] - 1 - n0 = (-1 + np.sqrt(1 + 8 * ncol)) * 0.5 - self.ntypes = int(n0 + 0.1) - assert self.ntypes * (self.ntypes + 1) // 2 == ncol, ( - "number of volumes provided in %s does not match guessed number of types %d" - % (filename, self.ntypes) - ) - self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) - self.tab_data = self._make_data() - - def get(self) -> Tuple[np.array, np.array]: - """Get the serialized table.""" - return self.tab_info, self.tab_data - - def _make_data(self): - data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) - stride = 4 * self.nspline - idx_iter = 0 - xx = self.vdata[:, 0] - for t0 in range(self.ntypes): - for t1 in range(t0, self.ntypes): - vv = self.vdata[:, 1 + idx_iter] - cs = CubicSpline(xx, vv) - dd = cs(xx, 1) - dd *= self.hh - dtmp = np.zeros(stride) - for ii in range(self.nspline): - dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] - dtmp[ii * 4 + 1] = ( - -3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1] - ) - dtmp[ii * 4 + 2] = dd[ii] - dtmp[ii * 4 + 3] = vv[ii] - data[ - (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride - + stride - ] = dtmp - data[ - (t1 * self.ntypes + t0) * stride : (t1 * self.ntypes + t0) * stride - + stride - ] = dtmp - idx_iter += 1 - return data +__all__ = [ + "PairTab", +] diff --git a/deepmd/utils/path.py b/deepmd/utils/path.py index a8e4bc329f..780bc8cabf 100644 --- a/deepmd/utils/path.py +++ b/deepmd/utils/path.py @@ -1,358 +1,13 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import os -from abc import ( - ABC, - abstractmethod, -) -from functools import ( - lru_cache, -) -from pathlib import ( - Path, -) -from typing import ( - List, - Optional, -) - -import h5py -import numpy as np -from wcmatch.glob import ( - globfilter, -) - - -class DPPath(ABC): - """The path class to data system (DeepmdData). - - Parameters - ---------- - path : str - path - """ - - def __new__(cls, path: str): - if cls is DPPath: - if os.path.isdir(path): - return super().__new__(DPOSPath) - elif os.path.isfile(path.split("#")[0]): - # assume h5 if it is not dir - # TODO: check if it is a real h5? or just check suffix? - return super().__new__(DPH5Path) - raise FileNotFoundError("%s not found" % path) - return super().__new__(cls) - - @abstractmethod - def load_numpy(self) -> np.ndarray: - """Load NumPy array. - - Returns - ------- - np.ndarray - loaded NumPy array - """ - - @abstractmethod - def load_txt(self, **kwargs) -> np.ndarray: - """Load NumPy array from text. - - Returns - ------- - np.ndarray - loaded NumPy array - """ - - @abstractmethod - def glob(self, pattern: str) -> List["DPPath"]: - """Search path using the glob pattern. - - Parameters - ---------- - pattern : str - glob pattern - - Returns - ------- - List[DPPath] - list of paths - """ - - @abstractmethod - def rglob(self, pattern: str) -> List["DPPath"]: - """This is like calling :meth:`DPPath.glob()` with `**/` added in front - of the given relative pattern. - - Parameters - ---------- - pattern : str - glob pattern - - Returns - ------- - List[DPPath] - list of paths - """ - - @abstractmethod - def is_file(self) -> bool: - """Check if self is file.""" - - @abstractmethod - def is_dir(self) -> bool: - """Check if self is directory.""" - - @abstractmethod - def __truediv__(self, key: str) -> "DPPath": - """Used for / operator.""" - - @abstractmethod - def __lt__(self, other: "DPPath") -> bool: - """Whether this DPPath is less than other for sorting.""" - - @abstractmethod - def __str__(self) -> str: - """Represent string.""" - - def __repr__(self) -> str: - return f"{type(self)} ({self!s})" - - def __eq__(self, other) -> bool: - return str(self) == str(other) - - def __hash__(self): - return hash(str(self)) - - -class DPOSPath(DPPath): - """The OS path class to data system (DeepmdData) for real directories. - - Parameters - ---------- - path : str - path - """ - - def __init__(self, path: str) -> None: - super().__init__() - if isinstance(path, Path): - self.path = path - else: - self.path = Path(path) - - def load_numpy(self) -> np.ndarray: - """Load NumPy array. - - Returns - ------- - np.ndarray - loaded NumPy array - """ - return np.load(str(self.path)) - - def load_txt(self, **kwargs) -> np.ndarray: - """Load NumPy array from text. - - Returns - ------- - np.ndarray - loaded NumPy array - """ - return np.loadtxt(str(self.path), **kwargs) - - def glob(self, pattern: str) -> List["DPPath"]: - """Search path using the glob pattern. - - Parameters - ---------- - pattern : str - glob pattern - - Returns - ------- - List[DPPath] - list of paths - """ - # currently DPOSPath will only derivative DPOSPath - # TODO: discuss if we want to mix DPOSPath and DPH5Path? - return [type(self)(p) for p in self.path.glob(pattern)] - - def rglob(self, pattern: str) -> List["DPPath"]: - """This is like calling :meth:`DPPath.glob()` with `**/` added in front - of the given relative pattern. - - Parameters - ---------- - pattern : str - glob pattern - - Returns - ------- - List[DPPath] - list of paths - """ - return [type(self)(p) for p in self.path.rglob(pattern)] - - def is_file(self) -> bool: - """Check if self is file.""" - return self.path.is_file() - - def is_dir(self) -> bool: - """Check if self is directory.""" - return self.path.is_dir() - - def __truediv__(self, key: str) -> "DPPath": - """Used for / operator.""" - return type(self)(self.path / key) - - def __lt__(self, other: "DPOSPath") -> bool: - """Whether this DPPath is less than other for sorting.""" - return self.path < other.path - - def __str__(self) -> str: - """Represent string.""" - return str(self.path) - - -class DPH5Path(DPPath): - """The path class to data system (DeepmdData) for HDF5 files. - - Notes - ----- - OS - HDF5 relationship: - directory - Group - file - Dataset - - Parameters - ---------- - path : str - path - """ - - def __init__(self, path: str) -> None: - super().__init__() - # we use "#" to split path - # so we do not support file names containing #... - s = path.split("#") - self.root_path = s[0] - self.root = self._load_h5py(s[0]) - # h5 path: default is the root path - self.name = s[1] if len(s) > 1 else "/" - - @classmethod - @lru_cache(None) - def _load_h5py(cls, path: str) -> h5py.File: - """Load hdf5 file. - - Parameters - ---------- - path : str - path to hdf5 file - """ - # this method has cache to avoid duplicated - # loading from different DPH5Path - # However the file will be never closed? - return h5py.File(path, "r") - - def load_numpy(self) -> np.ndarray: - """Load NumPy array. - - Returns - ------- - np.ndarray - loaded NumPy array - """ - return self.root[self.name][:] - - def load_txt(self, dtype: Optional[np.dtype] = None, **kwargs) -> np.ndarray: - """Load NumPy array from text. - - Returns - ------- - np.ndarray - loaded NumPy array - """ - arr = self.load_numpy() - if dtype: - arr = arr.astype(dtype) - return arr - - def glob(self, pattern: str) -> List["DPPath"]: - """Search path using the glob pattern. - - Parameters - ---------- - pattern : str - glob pattern - - Returns - ------- - List[DPPath] - list of paths - """ - # got paths starts with current path first, which is faster - subpaths = [ii for ii in self._keys if ii.startswith(self.name)] - return [ - type(self)(f"{self.root_path}#{pp}") - for pp in globfilter(subpaths, self._connect_path(pattern)) - ] - - def rglob(self, pattern: str) -> List["DPPath"]: - """This is like calling :meth:`DPPath.glob()` with `**/` added in front - of the given relative pattern. - - Parameters - ---------- - pattern : str - glob pattern - - Returns - ------- - List[DPPath] - list of paths - """ - return self.glob("**" + pattern) - - @property - def _keys(self) -> List[str]: - """Walk all groups and dataset.""" - return self._file_keys(self.root) - - @classmethod - @lru_cache(None) - def _file_keys(cls, file: h5py.File) -> List[str]: - """Walk all groups and dataset.""" - l = [] - file.visit(lambda x: l.append("/" + x)) - return l - - def is_file(self) -> bool: - """Check if self is file.""" - if self.name not in self._keys: - return False - return isinstance(self.root[self.name], h5py.Dataset) - - def is_dir(self) -> bool: - """Check if self is directory.""" - if self.name not in self._keys: - return False - return isinstance(self.root[self.name], h5py.Group) - - def __truediv__(self, key: str) -> "DPPath": - """Used for / operator.""" - return type(self)(f"{self.root_path}#{self._connect_path(key)}") - - def _connect_path(self, path: str) -> str: - """Connect self with path.""" - if self.name.endswith("/"): - return f"{self.name}{path}" - return f"{self.name}/{path}" - - def __lt__(self, other: "DPH5Path") -> bool: - """Whether this DPPath is less than other for sorting.""" - if self.root_path == other.root_path: - return self.name < other.name - return self.root_path < other.root_path - - def __str__(self) -> str: - """Returns path of self.""" - return f"{self.root_path}#{self.name}" +"""Alias for backward compatibility.""" +from deepmd_utils.utils.path import ( + DPH5Path, + DPOSPath, + DPPath, +) + +__all__ = [ + "DPPath", + "DPOSPath", + "DPH5Path", +] diff --git a/deepmd/utils/plugin.py b/deepmd/utils/plugin.py index 2a77b744c5..3b5b297304 100644 --- a/deepmd/utils/plugin.py +++ b/deepmd/utils/plugin.py @@ -1,95 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Base of plugin systems.""" -# copied from https://github.com/deepmodeling/dpdata/blob/a3e76d75de53f6076254de82d18605a010dc3b00/dpdata/plugin.py - -from abc import ( - ABCMeta, -) -from typing import ( - Callable, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.plugin import ( + Plugin, + PluginVariant, + VariantABCMeta, + VariantMeta, ) - -class Plugin: - """A class to register and restore plugins. - - Attributes - ---------- - plugins : Dict[str, object] - plugins - - Examples - -------- - >>> plugin = Plugin() - >>> @plugin.register("xx") - def xxx(): - pass - >>> print(plugin.plugins['xx']) - """ - - def __init__(self): - self.plugins = {} - - def __add__(self, other) -> "Plugin": - self.plugins.update(other.plugins) - return self - - def register(self, key: str) -> Callable[[object], object]: - """Register a plugin. - - Parameters - ---------- - key : str - key of the plugin - - Returns - ------- - Callable[[object], object] - decorator - """ - - def decorator(object: object) -> object: - self.plugins[key] = object - return object - - return decorator - - def get_plugin(self, key) -> object: - """Visit a plugin by key. - - Parameters - ---------- - key : str - key of the plugin - - Returns - ------- - object - the plugin - """ - return self.plugins[key] - - -class VariantMeta: - def __call__(cls, *args, **kwargs): - """Remove `type` and keys that starts with underline.""" - obj = cls.__new__(cls, *args, **kwargs) - kwargs.pop("type", None) - to_pop = [] - for kk in kwargs: - if kk[0] == "_": - to_pop.append(kk) - for kk in to_pop: - kwargs.pop(kk, None) - obj.__init__(*args, **kwargs) - return obj - - -class VariantABCMeta(VariantMeta, ABCMeta): - pass - - -class PluginVariant(metaclass=VariantABCMeta): - """A class to remove `type` from input arguments.""" - - pass +__all__ = [ + "Plugin", + "VariantMeta", + "VariantABCMeta", + "PluginVariant", +] diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 8944419412..09547eeac9 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -1,67 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Optional, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.random import ( + choice, + random, + seed, + shuffle, ) -import numpy as np - -_RANDOM_GENERATOR = np.random.RandomState() - - -def choice(a: np.ndarray, p: Optional[np.ndarray] = None): - """Generates a random sample from a given 1-D array. - - Parameters - ---------- - a : np.ndarray - A random sample is generated from its elements. - p : np.ndarray - The probabilities associated with each entry in a. - - Returns - ------- - np.ndarray - arrays with results and their shapes - """ - return _RANDOM_GENERATOR.choice(a, p=p) - - -def random(size=None): - """Return random floats in the half-open interval [0.0, 1.0). - - Parameters - ---------- - size - Output shape. - - Returns - ------- - np.ndarray - Arrays with results and their shapes. - """ - return _RANDOM_GENERATOR.random_sample(size) - - -def seed(val: Optional[int] = None): - """Seed the generator. - - Parameters - ---------- - val : int - Seed. - """ - _RANDOM_GENERATOR.seed(val) - - -def shuffle(x: np.ndarray): - """Modify a sequence in-place by shuffling its contents. - - Parameters - ---------- - x : np.ndarray - The array or list to be shuffled. - """ - _RANDOM_GENERATOR.shuffle(x) - - -__all__ = ["choice", "random", "seed", "shuffle"] +__all__ = [ + "choice", + "random", + "seed", + "shuffle", +] diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py index d0a167f1dc..2b270b1dbc 100644 --- a/deepmd/utils/tabulate.py +++ b/deepmd/utils/tabulate.py @@ -85,7 +85,10 @@ def __init__( # functype if activation_fn == ACTIVATION_FN_DICT["tanh"]: self.functype = 1 - elif activation_fn == ACTIVATION_FN_DICT["gelu"]: + elif activation_fn in ( + ACTIVATION_FN_DICT["gelu"], + ACTIVATION_FN_DICT["gelu_tf"], + ): self.functype = 2 elif activation_fn == ACTIVATION_FN_DICT["relu"]: self.functype = 3 @@ -330,8 +333,7 @@ def _build_lower( elif isinstance(self.descrpt, deepmd.descriptor.DescrptSeT): tt = np.full((nspline, self.last_layer_size), stride1) tt[ - int((lower - extrapolate * lower) / stride1) - + 1 : ( + int((lower - extrapolate * lower) / stride1) + 1 : ( int((lower - extrapolate * lower) / stride1) + int((upper - lower) / stride0) ), diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py index aadbb3c6e0..c8ab01f7f5 100644 --- a/deepmd/utils/type_embed.py +++ b/deepmd/utils/type_embed.py @@ -16,7 +16,6 @@ nvnmd_cfg, ) from deepmd.utils.graph import ( - get_tensor_by_name_from_graph, get_type_embedding_net_variables_from_graph_def, ) from deepmd.utils.network import ( @@ -109,7 +108,6 @@ def __init__( self.trainable = trainable self.uniform_seed = uniform_seed self.type_embedding_net_variables = None - self.type_embedding_from_graph = None self.padding = padding self.model_type = None @@ -135,8 +133,6 @@ def build( embedded_types The computational graph for embedded types """ - if self.model_type is not None and self.model_type == "compressed_model": - return self.type_embedding_from_graph types = tf.convert_to_tensor(list(range(ntypes)), dtype=tf.int32) ebd_type = tf.cast( tf.one_hot(tf.cast(types, dtype=tf.int32), int(ntypes)), @@ -166,7 +162,7 @@ def build( if self.padding: last_type = tf.cast(tf.zeros([1, self.neuron[-1]]), self.filter_precision) ebd_type = tf.concat([ebd_type, last_type], 0) # (ntypes + 1) * neuron[-1] - self.ebd_type = tf.identity(ebd_type, name="t_typeebd") + self.ebd_type = tf.identity(ebd_type, name="t_typeebd" + suffix) return self.ebd_type def init_variables( @@ -193,5 +189,3 @@ def init_variables( self.type_embedding_net_variables = ( get_type_embedding_net_variables_from_graph_def(graph_def, suffix=suffix) ) - type_embedding = get_tensor_by_name_from_graph(graph, "t_typeebd") - self.type_embedding_from_graph = tf.convert_to_tensor(type_embedding) diff --git a/deepmd/utils/weight_avg.py b/deepmd/utils/weight_avg.py index b344d3bb75..267f89ed28 100644 --- a/deepmd/utils/weight_avg.py +++ b/deepmd/utils/weight_avg.py @@ -1,48 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -from collections import ( - defaultdict, +"""Alias for backward compatibility.""" +from deepmd_utils.utils.weight_avg import ( + weighted_average, ) -from typing import ( - Dict, - List, - Tuple, -) - -import numpy as np - - -def weighted_average(errors: List[Dict[str, Tuple[float, float]]]) -> Dict: - """Compute wighted average of prediction errors (MAE or RMSE) for model. - - Parameters - ---------- - errors : List[Dict[str, Tuple[float, float]]] - List: the error of systems - Dict: the error of quantities, name given by the key - str: the name of the quantity, must starts with 'mae' or 'rmse' - Tuple: (error, weight) - Returns - ------- - Dict - weighted averages - """ - sum_err = defaultdict(float) - sum_siz = defaultdict(int) - for err in errors: - for kk, (ee, ss) in err.items(): - if kk.startswith("mae"): - sum_err[kk] += ee * ss - elif kk.startswith("rmse"): - sum_err[kk] += ee * ee * ss - else: - raise RuntimeError("unknown error type") - sum_siz[kk] += ss - for kk in sum_err.keys(): - if kk.startswith("mae"): - sum_err[kk] = sum_err[kk] / sum_siz[kk] - elif kk.startswith("rmse"): - sum_err[kk] = np.sqrt(sum_err[kk] / sum_siz[kk]) - else: - raise RuntimeError("unknown error type") - return sum_err +__all__ = [ + "weighted_average", +] diff --git a/deepmd_cli/__init__.py b/deepmd_cli/__init__.py deleted file mode 100644 index d295053965..0000000000 --- a/deepmd_cli/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -"""This module contains the entry points for DeePMD-kit. - -If only printing the help message, this module does not call -the main DeePMD-kit module to avoid the slow import of TensorFlow. -""" diff --git a/deepmd_utils/__init__.py b/deepmd_utils/__init__.py new file mode 100644 index 0000000000..1c5314bb7e --- /dev/null +++ b/deepmd_utils/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Untilization methods for DeePMD-kit. + +The __init__ module should not import any modules +for performance. +""" diff --git a/deepmd_utils/common.py b/deepmd_utils/common.py new file mode 100644 index 0000000000..b594c54030 --- /dev/null +++ b/deepmd_utils/common.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import warnings +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Optional, + TypeVar, + Union, +) + +try: + from typing import Literal # python >=3.8 +except ImportError: + from typing_extensions import Literal # type: ignore + +import numpy as np +import yaml + +from deepmd_utils.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd_utils.utils.path import ( + DPPath, +) + +__all__ = [ + "data_requirement", + "add_data_requirement", + "select_idx_map", + "make_default_mesh", + "j_must_have", + "j_loader", + "expand_sys_str", + "get_np_precision", +] + + +if TYPE_CHECKING: + _DICT_VAL = TypeVar("_DICT_VAL") + _PRECISION = Literal["default", "float16", "float32", "float64"] + _ACTIVATION = Literal[ + "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu", "gelu_tf" + ] + __all__.extend( + [ + "_DICT_VAL", + "_PRECISION", + "_ACTIVATION", + ] + ) + + +# TODO this is not a good way to do things. This is some global variable to which +# TODO anyone can write and there is no good way to keep track of the changes +data_requirement = {} + + +def add_data_requirement( + key: str, + ndof: int, + atomic: bool = False, + must: bool = False, + high_prec: bool = False, + type_sel: Optional[bool] = None, + repeat: int = 1, + default: float = 0.0, + dtype: Optional[np.dtype] = None, +): + """Specify data requirements for training. + + Parameters + ---------- + key : str + type of data stored in corresponding `*.npy` file e.g. `forces` or `energy` + ndof : int + number of the degrees of freedom, this is tied to `atomic` parameter e.g. forces + have `atomic=True` and `ndof=3` + atomic : bool, optional + specifies whwther the `ndof` keyworrd applies to per atom quantity or not, + by default False + must : bool, optional + specifi if the `*.npy` data file must exist, by default False + high_prec : bool, optional + if true load data to `np.float64` else `np.float32`, by default False + type_sel : bool, optional + select only certain type of atoms, by default None + repeat : int, optional + if specify repaeat data `repeat` times, by default 1 + default : float, optional, default=0. + default value of data + dtype : np.dtype, optional + the dtype of data, overwrites `high_prec` if provided + """ + data_requirement[key] = { + "ndof": ndof, + "atomic": atomic, + "must": must, + "high_prec": high_prec, + "type_sel": type_sel, + "repeat": repeat, + "default": default, + "dtype": dtype, + } + + +def select_idx_map(atom_types: np.ndarray, select_types: np.ndarray) -> np.ndarray: + """Build map of indices for element supplied element types from all atoms list. + + Parameters + ---------- + atom_types : np.ndarray + array specifing type for each atoms as integer + select_types : np.ndarray + types of atoms you want to find indices for + + Returns + ------- + np.ndarray + indices of types of atoms defined by `select_types` in `atom_types` array + + Warnings + -------- + `select_types` array will be sorted before finding indices in `atom_types` + """ + sort_select_types = np.sort(select_types) + idx_map = [] + for ii in sort_select_types: + idx_map.append(np.where(atom_types == ii)[0]) + return np.concatenate(idx_map) + + +def make_default_mesh(pbc: bool, mixed_type: bool) -> np.ndarray: + """Make mesh. + + Only the size of mesh matters, not the values: + * 6 for PBC, no mixed types + * 0 for no PBC, no mixed types + * 7 for PBC, mixed types + * 1 for no PBC, mixed types + + Parameters + ---------- + pbc : bool + if True, the mesh will be made for periodic boundary conditions + mixed_type : bool + if True, the mesh will be made for mixed types + + Returns + ------- + np.ndarray + mesh + """ + mesh_size = int(pbc) * 6 + int(mixed_type) + default_mesh = np.zeros(mesh_size, dtype=np.int32) + return default_mesh + + +# TODO maybe rename this to j_deprecated and only warn about deprecated keys, +# TODO if the deprecated_key argument is left empty function puppose is only custom +# TODO error since dict[key] already raises KeyError when the key is missing +def j_must_have( + jdata: Dict[str, "_DICT_VAL"], key: str, deprecated_key: List[str] = [] +) -> "_DICT_VAL": + """Assert that supplied dictionary conaines specified key. + + Returns + ------- + _DICT_VAL + value that was store unde supplied key + + Raises + ------ + RuntimeError + if the key is not present + """ + if key not in jdata.keys(): + for ii in deprecated_key: + if ii in jdata.keys(): + warnings.warn(f"the key {ii} is deprecated, please use {key} instead") + return jdata[ii] + else: + raise RuntimeError(f"json database must provide key {key}") + else: + return jdata[key] + + +def j_loader(filename: Union[str, Path]) -> Dict[str, Any]: + """Load yaml or json settings file. + + Parameters + ---------- + filename : Union[str, Path] + path to file + + Returns + ------- + Dict[str, Any] + loaded dictionary + + Raises + ------ + TypeError + if the supplied file is of unsupported type + """ + filepath = Path(filename) + if filepath.suffix.endswith("json"): + with filepath.open() as fp: + return json.load(fp) + elif filepath.suffix.endswith(("yml", "yaml")): + with filepath.open() as fp: + return yaml.safe_load(fp) + else: + raise TypeError("config file must be json, or yaml/yml") + + +# TODO port completely to pathlib when all callers are ported +def expand_sys_str(root_dir: Union[str, Path]) -> List[str]: + """Recursively iterate over directories taking those that contain `type.raw` file. + + Parameters + ---------- + root_dir : Union[str, Path] + starting directory + + Returns + ------- + List[str] + list of string pointing to system directories + """ + root_dir = DPPath(root_dir) + matches = [str(d) for d in root_dir.rglob("*") if (d / "type.raw").is_file()] + if (root_dir / "type.raw").is_file(): + matches.append(str(root_dir)) + return matches + + +def get_np_precision(precision: "_PRECISION") -> np.dtype: + """Get numpy precision constant from string. + + Parameters + ---------- + precision : _PRECISION + string name of numpy constant or default + + Returns + ------- + np.dtype + numpy presicion constant + + Raises + ------ + 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") diff --git a/deepmd_utils/entrypoints/__init__.py b/deepmd_utils/entrypoints/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/deepmd_utils/entrypoints/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/deepmd_utils/entrypoints/doc.py b/deepmd_utils/entrypoints/doc.py new file mode 100644 index 0000000000..9f1fd39095 --- /dev/null +++ b/deepmd_utils/entrypoints/doc.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Module that prints train input arguments docstrings.""" + +from deepmd_utils.utils.argcheck import ( + gen_doc, + gen_json, +) + +__all__ = ["doc_train_input"] + + +def doc_train_input(*, out_type: str = "rst", **kwargs): + """Print out trining input arguments to console.""" + if out_type == "rst": + doc_str = gen_doc(make_anchor=True) + elif out_type == "json": + doc_str = gen_json() + else: + raise RuntimeError("Unsupported out type %s" % out_type) + print(doc_str) diff --git a/deepmd_utils/entrypoints/gui.py b/deepmd_utils/entrypoints/gui.py new file mode 100644 index 0000000000..8b6b9e0a09 --- /dev/null +++ b/deepmd_utils/entrypoints/gui.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""DP-GUI entrypoint.""" + + +def start_dpgui(*, port: int, bind_all: bool, **kwargs): + """Host DP-GUI server. + + Parameters + ---------- + port : int + The port to serve DP-GUI on. + bind_all : bool + Serve on all public interfaces. This will expose your DP-GUI instance + to the network on both IPv4 and IPv6 (where available). + **kwargs + additional arguments + + Raises + ------ + ModuleNotFoundError + The dpgui package is not installed + """ + try: + from dpgui import ( + start_dpgui, + ) + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "To use DP-GUI, please install the dpgui package:\npip install dpgui" + ) from e + start_dpgui(port=port, bind_all=bind_all) diff --git a/deepmd_utils/env.py b/deepmd_utils/env.py new file mode 100644 index 0000000000..b1d4958ed8 --- /dev/null +++ b/deepmd_utils/env.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os + +import numpy as np + +__all__ = [ + "GLOBAL_NP_FLOAT_PRECISION", + "GLOBAL_ENER_FLOAT_PRECISION", + "global_float_prec", +] + +# FLOAT_PREC +dp_float_prec = os.environ.get("DP_INTERFACE_PREC", "high").lower() +if dp_float_prec in ("high", ""): + # default is high + GLOBAL_NP_FLOAT_PRECISION = np.float64 + GLOBAL_ENER_FLOAT_PRECISION = np.float64 + global_float_prec = "double" +elif dp_float_prec == "low": + GLOBAL_NP_FLOAT_PRECISION = np.float32 + GLOBAL_ENER_FLOAT_PRECISION = np.float64 + global_float_prec = "float" +else: + raise RuntimeError( + "Unsupported float precision option: %s. Supported: high," + "low. Please set precision with environmental variable " + "DP_INTERFACE_PREC." % dp_float_prec + ) diff --git a/deepmd_utils/loggers/__init__.py b/deepmd_utils/loggers/__init__.py new file mode 100644 index 0000000000..39aa76139d --- /dev/null +++ b/deepmd_utils/loggers/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Module taking care of logging duties.""" + +from .loggers import ( + set_log_handles, +) + +__all__ = ["set_log_handles"] diff --git a/deepmd_utils/loggers/loggers.py b/deepmd_utils/loggers/loggers.py new file mode 100644 index 0000000000..015581f6bd --- /dev/null +++ b/deepmd_utils/loggers/loggers.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Logger initialization for package.""" + +import logging +import os +from typing import ( + TYPE_CHECKING, + Optional, +) + +if TYPE_CHECKING: + from pathlib import ( + Path, + ) + + from mpi4py import ( + MPI, + ) + + _MPI_APPEND_MODE = MPI.MODE_CREATE | MPI.MODE_APPEND + +logging.getLogger(__name__) + +__all__ = ["set_log_handles"] + +# logger formater +FFORMATTER = logging.Formatter( + "[%(asctime)s] %(app_name)s %(levelname)-7s %(name)-45s %(message)s" +) +CFORMATTER = logging.Formatter( + # "%(app_name)s %(levelname)-7s |-> %(name)-45s %(message)s" + "%(app_name)s %(levelname)-7s %(message)s" +) +FFORMATTER_MPI = logging.Formatter( + "[%(asctime)s] %(app_name)s rank:%(rank)-2s %(levelname)-7s %(name)-45s %(message)s" +) +CFORMATTER_MPI = logging.Formatter( + # "%(app_name)s rank:%(rank)-2s %(levelname)-7s |-> %(name)-45s %(message)s" + "%(app_name)s rank:%(rank)-2s %(levelname)-7s %(message)s" +) + + +class _AppFilter(logging.Filter): + """Add field `app_name` to log messages.""" + + def filter(self, record): + record.app_name = "DEEPMD" + return True + + +class _MPIRankFilter(logging.Filter): + """Add MPI rank number to log messages, adds field `rank`.""" + + def __init__(self, rank: int) -> None: + super().__init__(name="MPI_rank_id") + self.mpi_rank = str(rank) + + def filter(self, record): + record.rank = self.mpi_rank + return True + + +class _MPIMasterFilter(logging.Filter): + """Filter that lets through only messages emited from rank==0.""" + + def __init__(self, rank: int) -> None: + super().__init__(name="MPI_master_log") + self.mpi_rank = rank + + def filter(self, record): + if self.mpi_rank == 0: + return True + else: + return False + + +class _MPIFileStream: + """Wrap MPI.File` so it has the same API as python file streams. + + Parameters + ---------- + filename : Path + disk location of the file stream + MPI : MPI + MPI communicator object + mode : str, optional + file write mode, by default _MPI_APPEND_MODE + """ + + def __init__( + self, filename: "Path", MPI: "MPI", mode: str = "_MPI_APPEND_MODE" + ) -> None: + self.stream = MPI.File.Open(MPI.COMM_WORLD, filename, mode) + self.stream.Set_atomicity(True) + self.name = "MPIfilestream" + + def write(self, msg: str): + """Write to MPI shared file stream. + + Parameters + ---------- + msg : str + message to write + """ + b = bytearray() + b.extend(map(ord, msg)) + self.stream.Write_shared(b) + + def close(self): + """Synchronize and close MPI file stream.""" + self.stream.Sync() + self.stream.Close() + + +class _MPIHandler(logging.FileHandler): + """Emulate `logging.FileHandler` with MPI shared File that all ranks can write to. + + Parameters + ---------- + filename : Path + file path + MPI : MPI + MPI communicator object + mode : str, optional + file access mode, by default "_MPI_APPEND_MODE" + """ + + def __init__( + self, + filename: "Path", + MPI: "MPI", + mode: str = "_MPI_APPEND_MODE", + ) -> None: + self.MPI = MPI + super().__init__(filename, mode=mode, encoding=None, delay=False) + + def _open(self): + return _MPIFileStream(self.baseFilename, self.MPI, self.mode) + + def setStream(self, stream): + """Stream canot 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 +): + """Set desired level for package loggers and add file handlers. + + Parameters + ---------- + level : int + logging level + log_path : Optional[str] + path to log file, if None logs will be send only to console. If the parent + directory does not exist it will be automatically created, by default None + mpi_log : Optional[str], optional + mpi log type. Has three options. `master` will output logs to file and console + only from rank==0. `collect` will write messages from all ranks to one file + opened under rank==0 and to console. `workers` will open one log file for each + worker designated by its rank, console behaviour is the same as for `collect`. + If this argument is specified, package 'mpi4py' must be already installed. + by default None + + Raises + ------ + RuntimeError + If the argument `mpi_log` is specified, package `mpi4py` is not installed. + + References + ---------- + https://groups.google.com/g/mpi4py/c/SaNzc8bdj6U + https://stackoverflow.com/questions/35869137/avoid-tensorflow-print-on-standard-error + https://stackoverflow.com/questions/56085015/suppress-openmp-debug-messages-when-running-tensorflow-on-cpu + + Notes + ----- + Logging levels: + + +---------+--------------+----------------+----------------+----------------+ + | | our notation | python logging | tensorflow cpp | OpenMP | + +=========+==============+================+================+================+ + | debug | 10 | 10 | 0 | 1/on/true/yes | + +---------+--------------+----------------+----------------+----------------+ + | info | 20 | 20 | 1 | 0/off/false/no | + +---------+--------------+----------------+----------------+----------------+ + | warning | 30 | 30 | 2 | 0/off/false/no | + +---------+--------------+----------------+----------------+----------------+ + | error | 40 | 40 | 3 | 0/off/false/no | + +---------+--------------+----------------+----------------+----------------+ + + """ + # silence logging for OpenMP when running on CPU if level is any other than debug + if level <= 10: + os.environ["KMP_WARNINGS"] = "FALSE" + + # set TF cpp internal logging level + os.environ["TF_CPP_MIN_LOG_LEVEL"] = str(int((level / 10) - 1)) + + # get root logger + root_log = logging.getLogger("deepmd") + root_log.propagate = False + + root_log.setLevel(level) + + # check if arguments are present + MPI = None + if mpi_log: + try: + from mpi4py import ( + MPI, + ) + except ImportError as e: + raise RuntimeError( + "You cannot specify 'mpi_log' when mpi4py not installed" + ) from e + + # * add console handler ************************************************************ + ch = logging.StreamHandler() + if MPI: + rank = MPI.COMM_WORLD.Get_rank() + if mpi_log == "master": + ch.setFormatter(CFORMATTER) + ch.addFilter(_MPIMasterFilter(rank)) + else: + ch.setFormatter(CFORMATTER_MPI) + ch.addFilter(_MPIRankFilter(rank)) + else: + ch.setFormatter(CFORMATTER) + + ch.setLevel(level) + ch.addFilter(_AppFilter()) + # clean old handlers before adding new one + root_log.handlers.clear() + root_log.addHandler(ch) + + # * add file handler *************************************************************** + if log_path: + # create directory + log_path.parent.mkdir(exist_ok=True, parents=True) + + fh = None + + if mpi_log == "master": + rank = MPI.COMM_WORLD.Get_rank() + if rank == 0: + fh = logging.FileHandler(log_path, mode="w") + fh.addFilter(_MPIMasterFilter(rank)) + fh.setFormatter(FFORMATTER) + elif mpi_log == "collect": + rank = MPI.COMM_WORLD.Get_rank() + fh = _MPIHandler(log_path, MPI, mode=MPI.MODE_WRONLY | MPI.MODE_CREATE) + fh.addFilter(_MPIRankFilter(rank)) + fh.setFormatter(FFORMATTER_MPI) + elif mpi_log == "workers": + rank = MPI.COMM_WORLD.Get_rank() + # if file has suffix than inser rank number before suffix + # e.g deepmd.log -> deepmd_.log + # if no suffix is present, insert rank as suffix + # e.g. deepmdlog -> deepmdlog. + if log_path.suffix: + worker_log = (log_path.parent / f"{log_path.stem}_{rank}").with_suffix( + log_path.suffix + ) + else: + worker_log = log_path.with_suffix(f".{rank}") + + fh = logging.FileHandler(worker_log, mode="w") + fh.setFormatter(FFORMATTER) + else: + fh = logging.FileHandler(log_path, mode="w") + fh.setFormatter(FFORMATTER) + + if fh: + fh.setLevel(level) + fh.addFilter(_AppFilter()) + root_log.addHandler(fh) diff --git a/deepmd_cli/main.py b/deepmd_utils/main.py similarity index 96% rename from deepmd_cli/main.py rename to deepmd_utils/main.py index bffc1c6911..19afaeee1f 100644 --- a/deepmd_cli/main.py +++ b/deepmd_utils/main.py @@ -1,4 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +"""The entry points for DeePMD-kit. + +If only printing the help message, this module does not call +the main DeePMD-kit module to avoid the slow import of TensorFlow. +""" import argparse import logging import textwrap @@ -8,7 +13,7 @@ ) try: - from deepmd_cli._version import version as __version__ + from deepmd_utils._version import version as __version__ except ImportError: __version__ = "unknown" @@ -547,10 +552,26 @@ def main_parser() -> argparse.ArgumentParser: parents=[parser_log], help="train nvnmd model", formatter_class=argparse.ArgumentDefaultsHelpFormatter, + epilog=textwrap.dedent( + """\ + examples: + dp train-nvnmd input_cnn.json -s s1 + dp train-nvnmd input_qnn.json -s s2 + dp train-nvnmd input_cnn.json -s s1 --restart model.ckpt + dp train-nvnmd input_cnn.json -s s2 --init-model model.ckpt + """ + ), ) parser_train_nvnmd.add_argument( "INPUT", help="the input parameter file in json format" ) + parser_train_nvnmd.add_argument( + "-i", + "--init-model", + type=str, + default=None, + help="Initialize the model by the provided path prefix of checkpoint files.", + ) parser_train_nvnmd.add_argument( "-r", "--restart", diff --git a/deepmd_utils/model_format/__init__.py b/deepmd_utils/model_format/__init__.py new file mode 100644 index 0000000000..253bca3507 --- /dev/null +++ b/deepmd_utils/model_format/__init__.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .common import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) +from .env_mat import ( + EnvMat, +) +from .network import ( + EmbeddingNet, + FittingNet, + NativeLayer, + NativeNet, + NetworkCollection, + load_dp_model, + make_embedding_network, + make_fitting_network, + make_multilayer_network, + save_dp_model, + traverse_model_dict, +) +from .output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, + fitting_check_output, + get_deriv_name, + get_reduce_name, + model_check_output, +) +from .se_e2_a import ( + DescrptSeA, +) + +__all__ = [ + "DescrptSeA", + "EnvMat", + "make_multilayer_network", + "make_embedding_network", + "make_fitting_network", + "EmbeddingNet", + "FittingNet", + "NativeLayer", + "NativeNet", + "NetworkCollection", + "NativeOP", + "load_dp_model", + "save_dp_model", + "traverse_model_dict", + "PRECISION_DICT", + "DEFAULT_PRECISION", + "ModelOutputDef", + "FittingOutputDef", + "OutputVariableDef", + "model_check_output", + "fitting_check_output", + "get_reduce_name", + "get_deriv_name", +] diff --git a/deepmd_utils/model_format/common.py b/deepmd_utils/model_format/common.py new file mode 100644 index 0000000000..d032e5d5df --- /dev/null +++ b/deepmd_utils/model_format/common.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from abc import ( + ABC, +) + +import numpy as np + +PRECISION_DICT = { + "float16": np.float16, + "float32": np.float32, + "float64": np.float64, + "half": np.float16, + "single": np.float32, + "double": np.float64, +} +DEFAULT_PRECISION = "float64" + + +class NativeOP(ABC): + """The unit operation of a native model.""" + + def call(self, *args, **kwargs): + """Forward pass in NumPy implementation.""" + raise NotImplementedError + + def __call__(self, *args, **kwargs): + """Forward pass in NumPy implementation.""" + return self.call(*args, **kwargs) diff --git a/deepmd_utils/model_format/env_mat.py b/deepmd_utils/model_format/env_mat.py new file mode 100644 index 0000000000..7822bd7d0c --- /dev/null +++ b/deepmd_utils/model_format/env_mat.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Optional, + Union, +) + +import numpy as np + +from .common import ( + NativeOP, +) + + +def compute_smooth_weight( + distance: np.ndarray, + rmin: float, + rmax: float, +): + """Compute smooth weight for descriptor elements.""" + min_mask = distance <= rmin + max_mask = distance >= rmax + mid_mask = np.logical_not(np.logical_or(min_mask, max_mask)) + uu = (distance - rmin) / (rmax - rmin) + vv = uu * uu * uu * (-6.0 * uu * uu + 15.0 * uu - 10.0) + 1.0 + return vv * mid_mask + min_mask + + +def _make_env_mat( + nlist, + coord, + rcut: float, + ruct_smth: float, +): + """Make smooth environment matrix.""" + nf, nloc, nnei = nlist.shape + # nf x nall x 3 + coord = coord.reshape(nf, -1, 3) + mask = nlist >= 0 + nlist = nlist * mask + # nf x (nloc x nnei) x 3 + index = np.tile(nlist.reshape(nf, -1, 1), (1, 1, 3)) + coord_r = np.take_along_axis(coord, index, 1) + # nf x nloc x nnei x 3 + coord_r = coord_r.reshape(nf, nloc, nnei, 3) + # nf x nloc x 1 x 3 + coord_l = coord[:, :nloc].reshape(nf, -1, 1, 3) + # nf x nloc x nnei x 3 + diff = coord_r - coord_l + # nf x nloc x nnei + length = np.linalg.norm(diff, axis=-1, keepdims=True) + # for index 0 nloc atom + length = length + ~np.expand_dims(mask, -1) + t0 = 1 / length + t1 = diff / length**2 + weight = compute_smooth_weight(length, ruct_smth, rcut) + env_mat_se_a = np.concatenate([t0, t1], axis=-1) * weight * np.expand_dims(mask, -1) + return env_mat_se_a, diff * np.expand_dims(mask, -1), weight + + +class EnvMat(NativeOP): + def __init__( + self, + rcut, + rcut_smth, + ): + self.rcut = rcut + self.rcut_smth = rcut_smth + + def call( + self, + coord_ext: np.ndarray, + atype_ext: np.ndarray, + nlist: np.ndarray, + davg: Optional[np.ndarray] = None, + dstd: Optional[np.ndarray] = None, + ) -> Union[np.ndarray, np.ndarray]: + """Compute the environment matrix. + + Parameters + ---------- + nlist + The neighbor list. shape: nf x nloc x nnei + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + davg + The data avg. shape: nt x nnei x 4 + dstd + The inverse of data std. shape: nt x nnei x 4 + + Returns + ------- + env_mat + The environment matrix. shape: nf x nloc x nnei x 4 + switch + The value of switch function. shape: nf x nloc x nnei + """ + em, sw = self._call(nlist, coord_ext) + nf, nloc, nnei = nlist.shape + atype = atype_ext[:, :nloc] + if davg is not None: + em -= davg[atype] + if dstd is not None: + em /= dstd[atype] + return em, sw + + def _call( + self, + nlist, + coord_ext, + ): + em, diff, ww = _make_env_mat(nlist, coord_ext, self.rcut, self.rcut_smth) + return em, ww + + def serialize( + self, + ) -> dict: + return { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + } + + @classmethod + def deserialize( + cls, + data: dict, + ) -> "EnvMat": + return cls(**data) diff --git a/deepmd_utils/model_format/network.py b/deepmd_utils/model_format/network.py new file mode 100644 index 0000000000..71ed659787 --- /dev/null +++ b/deepmd_utils/model_format/network.py @@ -0,0 +1,692 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Native DP model format for multiple backends. + +See issue #2982 for more information. +""" +import copy +import itertools +import json +from typing import ( + ClassVar, + Dict, + List, + Optional, + Union, +) + +import h5py +import numpy as np + +try: + from deepmd_utils._version import version as __version__ +except ImportError: + __version__ = "unknown" + +from .common import ( + DEFAULT_PRECISION, + PRECISION_DICT, + NativeOP, +) + + +def traverse_model_dict(model_obj, callback: callable, is_variable: bool = False): + """Traverse a model dict and call callback on each variable. + + Parameters + ---------- + model_obj : object + The model object to traverse. + callback : callable + The callback function to call on each variable. + is_variable : bool, optional + Whether the current node is a variable. + + Returns + ------- + object + The model object after traversing. + """ + if isinstance(model_obj, dict): + for kk, vv in model_obj.items(): + model_obj[kk] = traverse_model_dict( + vv, callback, is_variable=is_variable or kk == "@variables" + ) + elif isinstance(model_obj, list): + for ii, vv in enumerate(model_obj): + model_obj[ii] = traverse_model_dict(vv, callback, is_variable=is_variable) + elif is_variable: + model_obj = callback(model_obj) + return model_obj + + +class Counter: + """A callable counter. + + Examples + -------- + >>> counter = Counter() + >>> counter() + 0 + >>> counter() + 1 + """ + + def __init__(self): + self.count = -1 + + def __call__(self): + self.count += 1 + return self.count + + +def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] = None): + """Save a DP model to a file in the native format. + + Parameters + ---------- + filename : str + The filename to save to. + model_dict : dict + The model dict to save. + extra_info : dict, optional + Extra meta information to save. + """ + model_dict = model_dict.copy() + variable_counter = Counter() + if extra_info is not None: + extra_info = extra_info.copy() + else: + extra_info = {} + with h5py.File(filename, "w") as f: + model_dict = traverse_model_dict( + model_dict, + lambda x: f.create_dataset( + f"variable_{variable_counter():04d}", data=x + ).name, + ) + save_dict = { + "model": model_dict, + "software": "deepmd-kit", + "version": __version__, + **extra_info, + } + f.attrs["json"] = json.dumps(save_dict, separators=(",", ":")) + + +def load_dp_model(filename: str) -> dict: + """Load a DP model from a file in the native format. + + Parameters + ---------- + filename : str + The filename to load from. + + Returns + ------- + dict + The loaded model dict, including meta information. + """ + with h5py.File(filename, "r") as f: + model_dict = json.loads(f.attrs["json"]) + model_dict = traverse_model_dict(model_dict, lambda x: f[x][()].copy()) + return model_dict + + +class NativeLayer(NativeOP): + """Native representation of a layer. + + Parameters + ---------- + w : np.ndarray, optional + The weights of the layer. + b : np.ndarray, optional + The biases of the layer. + idt : np.ndarray, optional + The identity matrix of the layer. + activation_function : str, optional + The activation function of the layer. + resnet : bool, optional + Whether the layer is a residual layer. + """ + + def __init__( + self, + num_in, + num_out, + bias: bool = True, + use_timestep: bool = False, + activation_function: Optional[str] = None, + resnet: bool = False, + precision: str = DEFAULT_PRECISION, + ) -> None: + prec = PRECISION_DICT[precision.lower()] + self.precision = precision + rng = np.random.default_rng() + self.w = rng.normal(size=(num_in, num_out)).astype(prec) + self.b = rng.normal(size=(num_out,)).astype(prec) if bias else None + self.idt = rng.normal(size=(num_out,)).astype(prec) if use_timestep else None + self.activation_function = ( + activation_function if activation_function is not None else "none" + ) + self.resnet = resnet + self.check_type_consistency() + self.check_shape_consistency() + + def serialize(self) -> dict: + """Serialize the layer to a dict. + + Returns + ------- + dict + The serialized layer. + """ + data = { + "w": self.w, + "b": self.b, + "idt": self.idt, + } + return { + "bias": self.b is not None, + "use_timestep": self.idt is not None, + "activation_function": self.activation_function, + "resnet": self.resnet, + "precision": self.precision, + "@variables": data, + } + + @classmethod + def deserialize(cls, data: dict) -> "NativeLayer": + """Deserialize the layer from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + data = copy.deepcopy(data) + variables = data.pop("@variables") + assert variables["w"] is not None and len(variables["w"].shape) == 2 + num_in, num_out = variables["w"].shape + obj = cls( + num_in, + num_out, + **data, + ) + obj.w, obj.b, obj.idt = ( + variables["w"], + variables.get("b", None), + variables.get("idt", None), + ) + obj.check_shape_consistency() + return obj + + def check_shape_consistency(self): + if self.b is not None and self.w.shape[1] != self.b.shape[0]: + raise ValueError( + f"dim 1 of w {self.w.shape[1]} is not equal to shape " + f"of b {self.b.shape[0]}", + ) + if self.idt is not None and self.w.shape[1] != self.idt.shape[0]: + raise ValueError( + f"dim 1 of w {self.w.shape[1]} is not equal to shape " + f"of idt {self.idt.shape[0]}", + ) + + def check_type_consistency(self): + precision = self.precision + + def check_var(var): + if var is not None: + # assertion "float64" == "double" would fail + assert PRECISION_DICT[var.dtype.name] is PRECISION_DICT[precision] + + check_var(self.w) + check_var(self.b) + check_var(self.idt) + + def __setitem__(self, key, value): + if key in ("w", "matrix"): + self.w = value + elif key in ("b", "bias"): + self.b = value + elif key == "idt": + self.idt = value + elif key == "activation_function": + self.activation_function = value + elif key == "resnet": + self.resnet = value + elif key == "precision": + self.precision = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("w", "matrix"): + return self.w + elif key in ("b", "bias"): + return self.b + elif key == "idt": + return self.idt + elif key == "activation_function": + return self.activation_function + elif key == "resnet": + return self.resnet + elif key == "precision": + return self.precision + else: + raise KeyError(key) + + def dim_in(self) -> int: + return self.w.shape[0] + + def dim_out(self) -> int: + return self.w.shape[1] + + def call(self, x: np.ndarray) -> np.ndarray: + """Forward pass. + + Parameters + ---------- + x : np.ndarray + The input. + + Returns + ------- + np.ndarray + The output. + """ + if self.w is None or self.activation_function is None: + raise ValueError("w, b, and activation_function must be set") + if self.activation_function == "tanh": + fn = np.tanh + elif self.activation_function.lower() == "none": + + def fn(x): + return x + else: + raise NotImplementedError(self.activation_function) + y = ( + np.matmul(x, self.w) + self.b + if self.b is not None + else np.matmul(x, self.w) + ) + y = fn(y) + if self.idt is not None: + y *= self.idt + if self.resnet and self.w.shape[1] == self.w.shape[0]: + y += x + elif self.resnet and self.w.shape[1] == 2 * self.w.shape[0]: + y += np.concatenate([x, x], axis=-1) + return y + + +def make_multilayer_network(T_NetworkLayer, ModuleBase): + class NN(ModuleBase): + """Native representation of a neural network. + + Parameters + ---------- + layers : list[NativeLayer], optional + The layers of the network. + """ + + def __init__(self, layers: Optional[List[dict]] = None) -> None: + super().__init__() + if layers is None: + layers = [] + self.layers = [T_NetworkLayer.deserialize(layer) for layer in layers] + self.check_shape_consistency() + + def serialize(self) -> dict: + """Serialize the network to a dict. + + Returns + ------- + dict + The serialized network. + """ + return {"layers": [layer.serialize() for layer in self.layers]} + + @classmethod + def deserialize(cls, data: dict) -> "NN": + """Deserialize the network from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + return cls(data["layers"]) + + def __getitem__(self, key): + assert isinstance(key, int) + return self.layers[key] + + def __setitem__(self, key, value): + assert isinstance(key, int) + self.layers[key] = value + + def check_shape_consistency(self): + for ii in range(len(self.layers) - 1): + if self.layers[ii].dim_out() != self.layers[ii + 1].dim_in(): + raise ValueError( + f"the dim of layer {ii} output {self.layers[ii].dim_out} ", + f"does not match the dim of layer {ii+1} ", + f"output {self.layers[ii].dim_out}", + ) + + def call(self, x): + """Forward pass. + + Parameters + ---------- + x : np.ndarray + The input. + + Returns + ------- + np.ndarray + The output. + """ + for layer in self.layers: + x = layer(x) + return x + + return NN + + +NativeNet = make_multilayer_network(NativeLayer, NativeOP) + + +def make_embedding_network(T_Network, T_NetworkLayer): + class EN(T_Network): + """The embedding network. + + Parameters + ---------- + in_dim + Input dimension. + neuron + The number of neurons in each layer. The output dimension + is the same as the dimension of the last layer. + activation_function + The activation function. + resnet_dt + Use time step at the resnet architecture. + precision + Floating point precision for the model paramters. + + """ + + def __init__( + self, + in_dim, + neuron: List[int] = [24, 48, 96], + activation_function: str = "tanh", + resnet_dt: bool = False, + precision: str = DEFAULT_PRECISION, + ): + layers = [] + i_in = in_dim + for idx, ii in enumerate(neuron): + i_ot = ii + layers.append( + T_NetworkLayer( + i_in, + i_ot, + bias=True, + use_timestep=resnet_dt, + activation_function=activation_function, + resnet=True, + precision=precision, + ).serialize() + ) + i_in = i_ot + super().__init__(layers) + self.in_dim = in_dim + self.neuron = neuron + self.activation_function = activation_function + self.resnet_dt = resnet_dt + self.precision = precision + + def serialize(self) -> dict: + """Serialize the network to a dict. + + Returns + ------- + dict + The serialized network. + """ + return { + "in_dim": self.in_dim, + "neuron": self.neuron.copy(), + "activation_function": self.activation_function, + "resnet_dt": self.resnet_dt, + "precision": self.precision, + "layers": [layer.serialize() for layer in self.layers], + } + + @classmethod + def deserialize(cls, data: dict) -> "EmbeddingNet": + """Deserialize the network from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + data = copy.deepcopy(data) + layers = data.pop("layers") + obj = cls(**data) + super(EN, obj).__init__(layers) + return obj + + return EN + + +EmbeddingNet = make_embedding_network(NativeNet, NativeLayer) + + +def make_fitting_network(T_EmbeddingNet, T_Network, T_NetworkLayer): + class FN(T_EmbeddingNet): + """The fitting network. It may be implemented as an embedding + net connected with a linear output layer. + + Parameters + ---------- + in_dim + Input dimension. + out_dim + Output dimension + neuron + The number of neurons in each hidden layer. + activation_function + The activation function. + resnet_dt + Use time step at the resnet architecture. + precision + Floating point precision for the model paramters. + bias_out + The last linear layer has bias. + + """ + + def __init__( + self, + in_dim, + out_dim, + neuron: List[int] = [24, 48, 96], + activation_function: str = "tanh", + resnet_dt: bool = False, + precision: str = DEFAULT_PRECISION, + bias_out: bool = True, + ): + super().__init__( + in_dim, + neuron=neuron, + activation_function=activation_function, + resnet_dt=resnet_dt, + precision=precision, + ) + i_in, i_ot = neuron[-1], out_dim + self.layers.append( + T_NetworkLayer( + i_in, + i_ot, + bias=bias_out, + use_timestep=False, + activation_function=None, + resnet=False, + precision=precision, + ) + ) + self.out_dim = out_dim + self.bias_out = bias_out + + def serialize(self) -> dict: + """Serialize the network to a dict. + + Returns + ------- + dict + The serialized network. + """ + return { + "in_dim": self.in_dim, + "out_dim": self.out_dim, + "neuron": self.neuron.copy(), + "activation_function": self.activation_function, + "resnet_dt": self.resnet_dt, + "precision": self.precision, + "bias_out": self.bias_out, + "layers": [layer.serialize() for layer in self.layers], + } + + @classmethod + def deserialize(cls, data: dict) -> "FittingNet": + """Deserialize the network from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + data = copy.deepcopy(data) + layers = data.pop("layers") + obj = cls(**data) + T_Network.__init__(obj, layers) + return obj + + return FN + + +FittingNet = make_fitting_network(EmbeddingNet, NativeNet, NativeLayer) + + +class NetworkCollection: + """A collection of networks for multiple elements. + + The number of dimesions for types might be 0, 1, or 2. + - 0: embedding or fitting with type embedding, in () + - 1: embedding with type_one_side, or fitting, in (type_i) + - 2: embedding without type_one_side, in (type_i, type_j) + + Parameters + ---------- + ndim : int + The number of dimensions. + network_type : str, optional + The type of the network. + networks : dict, optional + The networks to initialize with. + """ + + # subclass may override this + NETWORK_TYPE_MAP: ClassVar[Dict[str, type]] = { + "network": NativeNet, + "embedding_network": EmbeddingNet, + "fitting_network": FittingNet, + } + + def __init__( + self, + ndim: int, + ntypes: int, + network_type: str = "network", + networks: List[Union[NativeNet, dict]] = [], + ): + self.ndim = ndim + self.ntypes = ntypes + self.network_type = self.NETWORK_TYPE_MAP[network_type] + self._networks = [None for ii in range(ntypes**ndim)] + for ii, network in enumerate(networks): + self[ii] = network + if len(networks): + self.check_completeness() + + def check_completeness(self): + """Check whether the collection is complete. + + Raises + ------ + RuntimeError + If the collection is incomplete. + """ + for tt in itertools.product(range(self.ntypes), repeat=self.ndim): + if self[tuple(tt)] is None: + raise RuntimeError(f"network for {tt} not found") + + def _convert_key(self, key): + if isinstance(key, int): + idx = key + else: + if isinstance(key, tuple): + pass + elif isinstance(key, str): + key = tuple([int(tt) for tt in key.split("_")[1:]]) + else: + raise TypeError(key) + assert isinstance(key, tuple) + assert len(key) == self.ndim + idx = sum([tt * self.ntypes**ii for ii, tt in enumerate(key)]) + return idx + + def __getitem__(self, key): + return self._networks[self._convert_key(key)] + + def __setitem__(self, key, value): + if isinstance(value, self.network_type): + pass + elif isinstance(value, dict): + value = self.network_type.deserialize(value) + else: + raise TypeError(value) + self._networks[self._convert_key(key)] = value + + def serialize(self) -> dict: + """Serialize the networks to a dict. + + Returns + ------- + dict + The serialized networks. + """ + network_type_map_inv = {v: k for k, v in self.NETWORK_TYPE_MAP.items()} + network_type_name = network_type_map_inv[self.network_type] + return { + "ndim": self.ndim, + "ntypes": self.ntypes, + "network_type": network_type_name, + "networks": [nn.serialize() for nn in self._networks], + } + + @classmethod + def deserialize(cls, data: dict) -> "NetworkCollection": + """Deserialize the networks from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + return cls(**data) diff --git a/deepmd_utils/model_format/output_def.py b/deepmd_utils/model_format/output_def.py new file mode 100644 index 0000000000..268dc21ea6 --- /dev/null +++ b/deepmd_utils/model_format/output_def.py @@ -0,0 +1,281 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, + List, + Tuple, +) + + +def check_shape( + shape: List[int], + def_shape: List[int], +): + """Check if the shape satisfies the defined shape.""" + assert len(shape) == len(def_shape) + if def_shape[-1] == -1: + if list(shape[:-1]) != def_shape[:-1]: + raise ValueError(f"{shape[:-1]} shape not matching def {def_shape[:-1]}") + else: + if list(shape) != def_shape: + raise ValueError(f"{shape} shape not matching def {def_shape}") + + +def check_var(var, var_def): + if var_def.atomic: + # var.shape == [nf, nloc, *var_def.shape] + if len(var.shape) != len(var_def.shape) + 2: + raise ValueError(f"{var.shape[2:]} length not matching def {var_def.shape}") + check_shape(list(var.shape[2:]), var_def.shape) + else: + # var.shape == [nf, *var_def.shape] + if len(var.shape) != len(var_def.shape) + 1: + raise ValueError(f"{var.shape[1:]} length not matching def {var_def.shape}") + check_shape(list(var.shape[1:]), var_def.shape) + + +def model_check_output(cls): + """Check if the output of the Model is consistent with the definition. + + Two methods are assumed to be provided by the Model: + 1. Model.output_def that gives the output definition. + 2. Model.__call__ that defines the forward path of the model. + + """ + + class wrapper(cls): + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.md = self.output_def() + + def __call__( + self, + *args, + **kwargs, + ): + ret = cls.__call__(self, *args, **kwargs) + for kk in self.md.keys_outp(): + dd = self.md[kk] + check_var(ret[kk], dd) + if dd.reduciable: + rk = get_reduce_name(kk) + check_var(ret[rk], self.md[rk]) + if dd.differentiable: + dnr, dnc = get_deriv_name(kk) + check_var(ret[dnr], self.md[dnr]) + check_var(ret[dnc], self.md[dnc]) + return ret + + return wrapper + + +def fitting_check_output(cls): + """Check if the output of the Fitting is consistent with the definition. + + Two methods are assumed to be provided by the Fitting: + 1. Fitting.output_def that gives the output definition. + 2. Fitting.__call__ defines the forward path of the fitting. + + """ + + class wrapper(cls): + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.md = self.output_def() + + def __call__( + self, + *args, + **kwargs, + ): + ret = cls.__call__(self, *args, **kwargs) + for kk in self.md.keys(): + dd = self.md[kk] + check_var(ret[kk], dd) + return ret + + return wrapper + + +class OutputVariableDef: + """Defines the shape and other properties of the one output variable. + + It is assume that the fitting network output variables for each + local atom. This class defines one output variable, including its + name, shape, reducibility and differentiability. + + Parameters + ---------- + name + Name of the output variable. Notice that the xxxx_redu, + xxxx_derv_c, xxxx_derv_r are reserved names that should + not be used to define variables. + shape + The shape of the variable. e.g. energy should be [1], + dipole should be [3], polarizabilty should be [3,3]. + reduciable + If the variable is reduced. + differentiable + If the variable is differentiated with respect to coordinates + of atoms and cell tensor (pbc case). Only reduciable variable + are differentiable. + + """ + + def __init__( + self, + name: str, + shape: List[int], + reduciable: bool = False, + differentiable: bool = False, + atomic: bool = True, + ): + self.name = name + self.shape = list(shape) + self.atomic = atomic + self.reduciable = reduciable + self.differentiable = differentiable + if not self.reduciable and self.differentiable: + raise ValueError("only reduciable variable are differentiable") + + +class FittingOutputDef: + """Defines the shapes and other properties of the fitting network outputs. + + It is assume that the fitting network output variables for each + local atom. This class defines all the outputs. + + Parameters + ---------- + var_defs + List of output variable definitions. + + """ + + def __init__( + self, + var_defs: List[OutputVariableDef], + ): + self.var_defs = {vv.name: vv for vv in var_defs} + + def __getitem__( + self, + key: str, + ) -> OutputVariableDef: + return self.var_defs[key] + + def get_data(self) -> Dict[str, OutputVariableDef]: + return self.var_defs + + def keys(self): + return self.var_defs.keys() + + +class ModelOutputDef: + """Defines the shapes and other properties of the model outputs. + + The model reduce and differentiate fitting outputs if applicable. + If a variable is named by foo, then the reduced variable is called + foo_redu, the derivative w.r.t. coordinates is called foo_derv_r + and the derivative w.r.t. cell is called foo_derv_c. + + Parameters + ---------- + fit_defs + Definition for the fitting net output + + """ + + def __init__( + self, + fit_defs: FittingOutputDef, + ): + self.def_outp = fit_defs + self.def_redu = do_reduce(self.def_outp) + self.def_derv_r, self.def_derv_c = do_derivative(self.def_outp) + self.var_defs: Dict[str, OutputVariableDef] = {} + for ii in [ + self.def_outp.get_data(), + self.def_redu, + self.def_derv_c, + self.def_derv_r, + ]: + self.var_defs.update(ii) + + def __getitem__( + self, + key: str, + ) -> OutputVariableDef: + return self.var_defs[key] + + def get_data( + self, + key: str, + ) -> Dict[str, OutputVariableDef]: + return self.var_defs + + def keys(self): + return self.var_defs.keys() + + def keys_outp(self): + return self.def_outp.keys() + + def keys_redu(self): + return self.def_redu.keys() + + def keys_derv_r(self): + return self.def_derv_r.keys() + + def keys_derv_c(self): + return self.def_derv_c.keys() + + +def get_reduce_name(name: str) -> str: + return name + "_redu" + + +def get_deriv_name(name: str) -> Tuple[str, str]: + return name + "_derv_r", name + "_derv_c" + + +def do_reduce( + def_outp: FittingOutputDef, +) -> Dict[str, OutputVariableDef]: + def_redu: Dict[str, OutputVariableDef] = {} + for kk, vv in def_outp.get_data().items(): + if vv.reduciable: + rk = get_reduce_name(kk) + def_redu[rk] = OutputVariableDef( + rk, vv.shape, reduciable=False, differentiable=False, atomic=False + ) + return def_redu + + +def do_derivative( + def_outp: FittingOutputDef, +) -> Tuple[Dict[str, OutputVariableDef], Dict[str, OutputVariableDef]]: + def_derv_r: Dict[str, OutputVariableDef] = {} + def_derv_c: Dict[str, OutputVariableDef] = {} + for kk, vv in def_outp.get_data().items(): + if vv.differentiable: + rkr, rkc = get_deriv_name(kk) + def_derv_r[rkr] = OutputVariableDef( + rkr, + vv.shape + [3], # noqa: RUF005 + reduciable=False, + differentiable=False, + ) + def_derv_c[rkc] = OutputVariableDef( + rkc, + vv.shape + [3, 3], # noqa: RUF005 + reduciable=True, + differentiable=False, + ) + return def_derv_r, def_derv_c diff --git a/deepmd_utils/model_format/se_e2_a.py b/deepmd_utils/model_format/se_e2_a.py new file mode 100644 index 0000000000..b9143ee360 --- /dev/null +++ b/deepmd_utils/model_format/se_e2_a.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np + +try: + from deepmd_utils._version import version as __version__ +except ImportError: + __version__ = "unknown" + +import copy +from typing import ( + Any, + List, + Optional, +) + +from .common import ( + DEFAULT_PRECISION, + NativeOP, +) +from .env_mat import ( + EnvMat, +) +from .network import ( + EmbeddingNet, + NetworkCollection, +) + + +class DescrptSeA(NativeOP): + r"""DeepPot-SE constructed from all information (both angular and radial) of + atomic configurations. The embedding takes the distance between atoms as input. + + The descriptor :math:`\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}` is given by [1]_ + + .. math:: + \mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_< + + where :math:`\mathcal{R}^i \in \mathbb{R}^{N \times 4}` is the coordinate + matrix, and each row of :math:`\mathcal{R}^i` can be constructed as follows + + .. math:: + (\mathcal{R}^i)_j = [ + \begin{array}{c} + s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} + \end{array} + ] + + where :math:`\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})` is + the relative coordinate and :math:`r_{ji}=\lVert \mathbf{R}_{ji} \lVert` is its norm. + The switching function :math:`s(r)` is defined as: + + .. math:: + s(r)= + \begin{cases} + \frac{1}{r}, & r