diff --git a/.github/workflows/build_cc.yml b/.github/workflows/build_cc.yml new file mode 100644 index 0000000000..8029263061 --- /dev/null +++ b/.github/workflows/build_cc.yml @@ -0,0 +1,36 @@ +on: + push: + pull_request: +name: Build C++ +jobs: + testpython: + name: Build C++ + runs-on: ubuntu-20.04 + strategy: + matrix: + include: + - float_prec: high + variant: cpu + - float_prec: low + variant: cpu + - float_prec: high + variant: gpu + - float_prec: low + variant: gpu + steps: + - uses: actions/checkout@master + with: + submodules: true + - run: sudo apt update && sudo apt install g++-7 + - run: sudo apt install nvidia-cuda-toolkit + if: matrix.variant == 'gpu' + - run: source/install/build_cc.sh + env: + FLOAT_PREC: ${{ matrix.float_prec }} + CC: gcc-7 + CXX: g++-7 + - run: source/install/build_lammps.sh + env: + FLOAT_PREC: ${{ matrix.float_prec }} + CC: gcc-7 + CXX: g++-7 diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml new file mode 100644 index 0000000000..2d522e5cf3 --- /dev/null +++ b/.github/workflows/build_wheel.yml @@ -0,0 +1,67 @@ +name: Build and upload to PyPI + +on: + push: + pull_request: + +jobs: + build_wheels: + name: Build wheels on ${{ matrix.os }} + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-18.04] #, windows-latest, macos-latest] + + steps: + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + name: Install Python + with: + python-version: '3.8' + + - name: Install cibuildwheel + run: | + python -m pip install cibuildwheel + + - name: Build wheels + env: + CIBW_BUILD: "cp36-* cp37-* cp38-*" + CIBW_BEFORE_BUILD: pip install tensorflow && sed -i 's/libresolv.so.2"/libresolv.so.2", "libtensorflow_framework.so.2"/g' /opt/_internal/tools/lib/python*/site-packages/auditwheel/policy/policy.json + CIBW_SKIP: "*-win32 *-manylinux_i686" + run: | + python -m cibuildwheel --output-dir wheelhouse + - uses: actions/upload-artifact@v2 + with: + path: ./wheelhouse/*.whl + + build_sdist: + name: Build source distribution + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + name: Install Python + with: + python-version: '3.8' + - run: pip install -U scikit-build tensorflow setuptools_scm + - name: Build sdist + run: python setup.py sdist + + - uses: actions/upload-artifact@v2 + with: + path: dist/*.tar.gz + + upload_pypi: + needs: [build_wheels, build_sdist] + runs-on: ubuntu-latest + if: startsWith(github.event.ref, 'refs/tags/v') + steps: + - uses: actions/download-artifact@v2 + with: + name: artifact + path: dist + + - uses: pypa/gh-action-pypi-publish@master + with: + user: __token__ + password: ${{ secrets.pypi_password }} diff --git a/.github/workflows/lint_python.yml b/.github/workflows/lint_python.yml new file mode 100644 index 0000000000..92a609edc1 --- /dev/null +++ b/.github/workflows/lint_python.yml @@ -0,0 +1,33 @@ +on: + push: + pull_request: +name: Lint Python +jobs: + lintpython: + name: Lint Python + runs-on: ubuntu-20.04 + strategy: + matrix: + python-version: [3.8] + + steps: + - uses: actions/checkout@v1 + - uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python }} + - name: Install requirements + run: pip install -r requirements.txt + - uses: marian-code/python-lint-annotate@v2.5.0 + with: + python-root-list: "./deepmd/*.py ./deepmd/*/*.py ./source/train/*.py ./source/tests/*.py ./source/op/*.py" + use-black: true + use-isort: true + use-mypy: true + use-pycodestyle: true + use-pydocstyle: true + extra-pycodestyle-options: "--max-line-length=88" + use-pylint: false + use-flake8: false + use-vulture: true + conda-python-version: "3.8" + diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml new file mode 100644 index 0000000000..5192eda6e9 --- /dev/null +++ b/.github/workflows/test_cc.yml @@ -0,0 +1,11 @@ +on: + push: + pull_request: +name: Test C++ +jobs: + testpython: + name: Test C++ + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@master + - run: source/install/test_cc.sh diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml new file mode 100644 index 0000000000..dadd749b36 --- /dev/null +++ b/.github/workflows/test_python.yml @@ -0,0 +1,73 @@ +on: + push: + pull_request: +name: Test Python +jobs: + testpython: + name: Test Python + runs-on: ubuntu-18.04 + strategy: + matrix: + include: + - python: 3.6 + gcc: 4.8 + tf: 1.8 + - python: 3.6 + gcc: 4.8 + tf: 1.12 + - python: 3.6 + gcc: 4.8 + tf: 1.14 + - python: 3.6 + gcc: 5 + tf: 1.14 + - python: 3.6 + gcc: 8 + tf: 1.14 + - python: 3.7 + gcc: 5 + tf: 1.14 + - python: 3.7 + gcc: 6 + tf: 1.14 + - python: 3.7 + gcc: 7 + tf: 1.14 + - python: 3.7 + gcc: 8 + tf: 1.14 + - python: 3.7 + gcc: 5 + tf: 2.3 + - python: 3.7 + gcc: 8 + tf: 2.3 + - python: 3.8 + gcc: 5 + tf: 2.3 + - python: 3.8 + gcc: 8 + tf: 2.3 + + steps: + - uses: actions/checkout@master + - uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python }} + - name: pip cache + uses: actions/cache@v2 + with: + path: ~/.cache/pip + key: + ${{ runner.os }}-pip-${{ hashFiles('**/setup.py') }}-${{ hashFiles('**/requirements.txt') }} + restore-keys: | + ${{ runner.os }}-pip- + - run: | + sudo apt update + sudo apt install gcc-${{ matrix.gcc }} g++-${{ matrix.gcc }} + - run: pip install -e .[cpu,test] codecov + env: + CC: gcc-${{ matrix.gcc }} + CXX: g++-${{ matrix.gcc }} + TENSORFLOW_VERSION: ${{ matrix.tf }} + - run: pytest --cov=deepmd source/tests && codecov diff --git a/.gitignore b/.gitignore index b392cdbca5..41eb111e4e 100644 --- a/.gitignore +++ b/.gitignore @@ -26,3 +26,6 @@ venv* .vscode/** _build _templates +API_CC +dp/** +build_lammps/** \ No newline at end of file diff --git a/.gitmodules b/.gitmodules index c2225c5f76..5373ec05b3 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ -[submodule "source/op/cuda/cub"] - path = source/op/cuda/cub - url = git://github.com/NVlabs/cub.git +[submodule "source/lib/src/cuda/cub"] + path = source/lib/src/cuda/cub + url = git://github.com/NVIDIA/cub.git diff --git a/.travis.yml b/.travis.yml deleted file mode 100644 index b434af6819..0000000000 --- a/.travis.yml +++ /dev/null @@ -1,95 +0,0 @@ -language: python -addons: - apt: - sources: - - ubuntu-toolchain-r-test - packages: - - gcc-4.8 - - g++-4.8 - - gcc-5 - - g++-5 - - gcc-6 - - g++-6 - - gcc-7 - - g++-7 - - gcc-8 - - g++-8 -jobs: - include: - - stage: unit tests - python: 3.6 - env: - - CC=gcc-4.8 - - CXX=g++-4.8 - - TENSORFLOW_VERSION=1.8 - - python: 3.6 - env: - - CC=gcc-4.8 - - CXX=g++-4.8 - - TENSORFLOW_VERSION=1.12 - - python: 3.6 - env: - - CC=gcc-4.8 - - CXX=g++-4.8 - - TENSORFLOW_VERSION=1.14 - - python: 3.6 - env: - - CC=gcc-5 - - CXX=g++-5 - - TENSORFLOW_VERSION=1.14 - - python: 3.6 - env: - - CC=gcc-8 - - CXX=g++-8 - - TENSORFLOW_VERSION=1.14 - - python: 3.7 - env: - - CC=gcc-5 - - CXX=g++-5 - - TENSORFLOW_VERSION=1.14 - - python: 3.7 - env: - - CC=gcc-6 - - CXX=g++-6 - - TENSORFLOW_VERSION=1.14 - - python: 3.7 - env: - - CC=gcc-7 - - CXX=g++-7 - - TENSORFLOW_VERSION=1.14 - - python: 3.7 - env: - - CC=gcc-8 - - CXX=g++-8 - - TENSORFLOW_VERSION=1.14 - - python: 3.7 - env: - - CC=gcc-5 - - CXX=g++-5 - - TENSORFLOW_VERSION=2.3 - - python: 3.7 - env: - - CC=gcc-8 - - CXX=g++-8 - - TENSORFLOW_VERSION=2.3 - - stage: build whls - services: docker - env: - - TWINE_USERNAME=__token__ - - CIBW_BUILD="cp36-* cp37-*" - - CIBW_BEFORE_BUILD="sed -i 's/libresolv.so.2\"/libresolv.so.2\", \"libtensorflow_framework.so.2\"/g' \$(find / -name policy.json)" - - CIBW_SKIP="*-win32 *-manylinux_i686" - - CC=gcc-7 - - CXX=g++-7 - - TENSORFLOW_VERSION=2.3 - install: - - python -m pip install twine cibuildwheel==1.6.3 scikit-build setuptools_scm - script: - - python -m cibuildwheel --output-dir wheelhouse - - python setup.py sdist - after_success: - - if [[ $TRAVIS_TAG ]]; then python -m twine upload wheelhouse/*; python -m twine upload dist/*.tar.gz; fi -install: - - pip install .[cpu,test] -script: - - cd source/tests && python -m unittest diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000000..86199588d7 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,136 @@ +# DeePMD-kit Contributing Guide + +Welcome to [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit)! + +## What you can contribute + +You can either make a code contribution, help improve our document or offer help to other users. Your help is always appreciated. Come and have fun! + +### Code contribution +You can start from any one of the following items to help improve deepmd-kit + +- Smash a bug +- Implement a feature or add a patch, whatever you think deepmd-kit is missing +- Browse [issues](https://github.com/deepmodeling/deepmd-kit/issues), find an issue labeled enhancement or bug, and help to solve it. + +See [here](#before-you-contribute) for some before-hand heads-up. + +See [here](#how-to-contribute) to learn how to contribute. + +### Document improvement +You can start from any one of the following items to help improve [DeePMD-kit Docs](https://deepmd.readthedocs.io/en/latest/?badge=latest): + +- Fix typos or format (punctuation, space, indentation, code block, etc.) +- Fix or update inappropriate or outdated descriptions +- Add missing content (sentence, paragraph, or a new document) +- Translate docs changes from English to Chinese + +### Offer help +You can help other users of deepmd-kit in the following way + +- Submit, reply to, and resolve [issues](https://github.com/deepmodeling/deepmd-kit/issues) +- (Advanced) Review Pull Requests created by others + +## Before you contribute +### Overview of DeePMD-kit +Currently, we maintain two main branch: +- master: stable branch with version tag +- devel : branch for developers + +### Developer guide +See [here](doc/development/index.md) for coding conventions, API and other needs-to-know of the code. + +## How to contribute +Please perform the following steps to create your Pull Request to this repository. If don't like to use commands, you can also use [GitHub Desktop](https://desktop.github.com/), which is easier to get started. Go to [git documentation](https://git-scm.com/doc) if you want to really master git. + +### Step 1: Fork the repository + +1. Visit the project: +2. Click the **Fork** button on the top right and wait it to finish. + +### Step 2: Clone the forked repository to local storage and set configurations + +1. Clone your own repo, not the public repo (from deepmodeling) ! And change the branch to devel. + ```bash + git clone https://github.com/$username/deepmd-kit.git + # Replace `$username` with your GitHub ID + + git checkout devel + ``` + +2. Add deepmodeling's repo as your remote repo, we can name it "upstream". And fetch upstream's latest codes to your workstation. + ```bash + git remote add upstream https://github.com/deepmodeling/deepmd-kit.git + # After you add a remote repo, your local repo will be automatically named "origin". + + git fetch upstream + + # If your current codes are behind the latest codes, you should merge latest codes first. + # Notice you should merge from "devel"! + git merge upstream/devel + ``` + +3. Modify your codes and design unit tests. + +4. Commit your changes + ```bash + git status # Checks the local status + git add ... # Adds the file(s) you want to commit. If you want to commit all changes, you can directly use `git add.` + git commit -m "commit-message: update the xx" + ``` + +5. Push the changed codes to your original repo on github. + ```bash + git push origin devel + ``` + +### Alternatively: Create a new branch + +1. Get your local master up-to-date with upstream/master. + + ```bash + cd $working_dir/deepmd-kit + git fetch upstream + git checkout master + git rebase upstream/master + ``` + +2. Create a new branch based on the master branch. + + ```bash + git checkout -b new-branch-name + ``` + +3. Modify your codes and design unit tests. + +4. Commit your changes + + ```bash + git status # Checks the local status + git add ... # Adds the file(s) you want to commit. If you want to commit all changes, you can directly use `git add.` + git commit -m "commit-message: update the xx" + ``` + +5. Keep your branch in sync with upstream/master + + ```bash + # While on your new branch + git fetch upstream + git rebase upstream/master + ``` + +6. Push your changes to the remote + + ```bash + git push -u origin new-branch-name # "-u" is used to track the remote branch from origin + ``` + +### Step 3: Create a pull request + +1. Visit your fork at (replace `$username` with your GitHub ID) +2. Click `pull requests`, followed by `New pull request` and `Compare & pull request` to create your PR. + +Now, your PR is successfully submitted! After this PR is merged, you will automatically become a contributor to DeePMD-kit. + +## Contact us +E-mail: contact@deepmodeling.org diff --git a/README.md b/README.md index 778277ffe9..bac29b4ddd 100644 --- a/README.md +++ b/README.md @@ -5,12 +5,13 @@ # Table of contents - [About DeePMD-kit](#about-deepmd-kit) + - [Highlights in v2.0](#highlights-in-deepmd-kit-v2.0) - [Highlighted features](#highlighted-features) - - [Code structure](#code-structure) - [License and credits](#license-and-credits) - [Deep Potential in a nutshell](#deep-potential-in-a-nutshell) - [Download and install](#download-and-install) - [Use DeePMD-kit](#use-deepmd-kit) +- [Code structure](#code-structure) - [Troubleshooting](#troubleshooting) # About DeePMD-kit @@ -18,37 +19,22 @@ DeePMD-kit is a package written in Python/C++, designed to minimize the effort r For more information, check the [documentation](https://deepmd.readthedocs.io/). +# Highlights in DeePMD-kit v2.0 +* [Model compression](doc/getting-started.md#compress-a-model). Accelerate the efficiency of model inference for 4-15 times. +* [New descriptors](doc/getting-started.md#write-the-input-script). Including [`se_e2_r`](doc/train-se-e2-r.md) and [`se_e3`](doc/train-se-e3.md). +* [Hybridization of descriptors](doc/train-hybrid.md). Hybrid descriptor constructed from concatenation of several descriptors. +* [Atom type embedding](doc/train-se-e2-a-tebd.md). Enable atom type embedding to decline training complexity and refine performance. +* Training and inference the dipole (vector) and polarizability (matrix). +* Split of training and validation dataset. +* Optimized training on GPUs. + ## Highlighted features -* **interfaced with TensorFlow**, one of the most popular deep learning frameworks, making the training process highly automatic and efficient. +* **interfaced with TensorFlow**, one of the most popular deep learning frameworks, making the training process highly automatic and efficient, in addition Tensorboard can be used to visualize training procedure. * **interfaced with high-performance classical MD and quantum (path-integral) MD packages**, i.e., LAMMPS and i-PI, respectively. * **implements the Deep Potential series models**, which have been successfully applied to finite and extended systems including organic molecules, metals, semiconductors, and insulators, etc. * **implements MPI and GPU supports**, makes it highly efficient for high performance parallel and distributed computing. * **highly modularized**, easy to adapt to different descriptors for deep learning based potential energy models. -## Code structure -The code is organized as follows: - -* `data/raw`: tools manipulating the raw data files. - -* `examples`: example json parameter files. - -* `source/3rdparty`: third-party packages used by DeePMD-kit. - -* `source/cmake`: cmake scripts for building. - -* `source/ipi`: source code of i-PI client. - -* `source/lib`: source code of DeePMD-kit library. - -* `source/lmp`: source code of Lammps module. - -* `source/op`: tensorflow op implementation. working with library. - -* `source/scripts`: Python script for model freezing. - -* `source/train`: Python modules and scripts for training and testing. - - ## License and credits The project DeePMD-kit is licensed under [GNU LGPLv3.0](./LICENSE). If you use this code in any future publications, please cite this using @@ -74,59 +60,52 @@ One may manually install DeePMD-kit by following the instuctions on [installing # Use DeePMD-kit -The typical procedure of using DeePMD-kit includes 5 steps +The typical procedure of using DeePMD-kit includes the following steps -1. [Prepare data](doc/use-deepmd-kit.md#prepare-data) -2. [Train a model](doc/use-deepmd-kit.md#train-a-model) -3. [Freeze the model](doc/use-deepmd-kit.md#freeze-a-model) -4. [Test the model](doc/use-deepmd-kit.md#test-a-model) -5. [Inference the model in python](doc/use-deepmd-kit.md#model-inference) or using the model in other molecular simulation packages like [LAMMPS](doc/use-deepmd-kit.md#run-md-with-lammps), [i-PI](doc/use-deepmd-kit.md#run-path-integral-md-with-i-pi) or [ASE](doc/use-deepmd-kit.md#use-deep-potential-with-ase). +1. [Prepare data](doc/getting-started.md#prepare-data) +2. [Train a model](doc/getting-started.md#train-a-model) +3. [Analyze training with Tensorboard](doc/tensorboard.md) +4. [Freeze the model](doc/getting-started.md#freeze-a-model) +5. [Test the model](doc/getting-started.md#test-a-model) +6. [Compress the model](doc/getting-started.md#compress-a-model) +7. [Inference the model in python](doc/getting-started.md#model-inference) or using the model in other molecular simulation packages like [LAMMPS](doc/getting-started.md#run-md-with-lammps), [i-PI](doc/getting-started.md#run-path-integral-md-with-i-pi) or [ASE](doc/getting-started.md#use-deep-potential-with-ase). -A quick-start on using DeePMD-kit can be found [here](doc/use-deepmd-kit.md). +A quick-start on using DeePMD-kit can be found [here](doc/getting-started.md). -A full [document](doc/train-input.rst) on options in the training input script is available. +A full [document](doc/train-input-auto.rst) on options in the training input script is available. -# Troubleshooting -In consequence of various differences of computers or systems, problems may occur. Some common circumstances are listed as follows. -If other unexpected problems occur, you're welcome to contact us for help. +# Code structure +The code is organized as follows: + +* `data/raw`: tools manipulating the raw data files. + +* `examples`: examples. -## Model compatability +* `deepmd`: DeePMD-kit python modules. -When the version of DeePMD-kit used to training model is different from the that of DeePMD-kit running MDs, one has the problem of model compatability. +* `source/api_cc`: source code of DeePMD-kit C++ API. -DeePMD-kit guarantees that the codes with the same major and minor revisions are compatible. That is to say v0.12.5 is compatible to v0.12.0, but is not compatible to v0.11.0 nor v1.0.0. +* `source/ipi`: source code of i-PI client. + +* `source/lib`: source code of DeePMD-kit library. + +* `source/lmp`: source code of Lammps module. + +* `source/op`: tensorflow op implementation. working with library. -## Installation: inadequate versions of gcc/g++ -Sometimes you may use a gcc/g++ of version <4.9. If you have a gcc/g++ of version > 4.9, say, 7.2.0, you may choose to use it by doing -```bash -export CC=/path/to/gcc-7.2.0/bin/gcc -export CXX=/path/to/gcc-7.2.0/bin/g++ -``` -If, for any reason, for example, you only have a gcc/g++ of version 4.8.5, you can still compile all the parts of TensorFlow and most of the parts of DeePMD-kit. i-Pi will be disabled automatically. +# Troubleshooting -## Installation: build files left in DeePMD-kit -When you try to build a second time when installing DeePMD-kit, files produced before may contribute to failure. Thus, you may clear them by -```bash -cd build -rm -r * -``` -and redo the `cmake` process. +See the [troubleshooting page](doc/troubleshooting/index.md). -## MD: cannot run LAMMPS after installing a new version of DeePMD-kit -This typically happens when you install a new version of DeePMD-kit and copy directly the generated `USER-DEEPMD` to a LAMMPS source code folder and re-install LAMMPS. -To solve this problem, it suffices to first remove `USER-DEEPMD` from LAMMPS source code by -```bash -make no-user-deepmd -``` -and then install the new `USER-DEEPMD`. +# Contributing -If this does not solve your problem, try to decompress the LAMMPS source tarball and install LAMMPS from scratch again, which typically should be very fast. +See [DeePMD-kit Contributing Guide](CONTRIBUTING.md) to become a contributor! 🤓 -[1]: http://www.global-sci.com/galley/CiCP-2017-0213.pdf +[1]: https://arxiv.org/abs/1707.01478 [2]: https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.120.143001 -[3]:https://arxiv.org/abs/1805.09003 -[4]:https://aip.scitation.org/doi/full/10.1063/1.5027645 +[3]: https://arxiv.org/abs/1805.09003 +[4]: https://aip.scitation.org/doi/full/10.1063/1.5027645 diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 0000000000..b882b3989c --- /dev/null +++ b/codecov.yml @@ -0,0 +1,2 @@ +ignore: + - "source/**/tests" diff --git a/data/raw/raw_to_set.sh b/data/raw/raw_to_set.sh index 8e4b917adb..a89ef3a872 100755 --- a/data/raw/raw_to_set.sh +++ b/data/raw/raw_to_set.sh @@ -17,6 +17,10 @@ test -f force.raw && split force.raw -l $nline_per_set -d -a 3 force.raw test -f virial.raw && split virial.raw -l $nline_per_set -d -a 3 virial.raw test -f atom_ener.raw && split atom_ener.raw -l $nline_per_set -d -a 3 atom_ener.raw test -f fparam.raw && split fparam.raw -l $nline_per_set -d -a 3 fparam.raw +test -f dipole.raw && split dipole.raw -l $nline_per_set -d -a 3 dipole.raw +test -f polarizability.raw && split polarizability.raw -l $nline_per_set -d -a 3 polarizability.raw +test -f atomic_dipole.raw && split atomic_dipole.raw -l $nline_per_set -d -a 3 atomic_dipole.raw +test -f atomic_polarizability.raw && split atomic_polarizability.raw -l $nline_per_set -d -a 3 atomic_polarizability.raw nset=`ls | grep box.raw[0-9] | wc -l` nset_1=$(($nset-1)) @@ -34,6 +38,8 @@ do test -f virial.raw$pi && mv virial.raw$pi set.$pi/virial.raw test -f atom_ener.raw$pi && mv atom_ener.raw$pi set.$pi/atom_ener.raw test -f fparam.raw$pi && mv fparam.raw$pi set.$pi/fparam.raw + test -f atomic_dipole.raw$pi && mv atomic_dipole.raw$pi set.$pi/atomic_dipole.raw + test -f atomic_polarizability.raw$pi && mv atomic_polarizability.raw$pi set.$pi/atomic_polarizability.raw cd set.$pi python -c 'import numpy as np; data = np.loadtxt("box.raw" , ndmin = 2); data = data.astype (np.float32); np.save ("box", data)' @@ -72,6 +78,34 @@ if os.path.isfile("fparam.raw"): data = np.loadtxt("fparam.raw", ndmin = 2); data = data.astype (np.float32); np.save ("fparam", data) +' + python -c \ +'import numpy as np; import os.path; +if os.path.isfile("dipole.raw"): + data = np.loadtxt("dipole.raw", ndmin = 2); + data = data.astype (np.float32); + np.save ("dipole", data) +' + python -c \ +'import numpy as np; import os.path; +if os.path.isfile("polarizability.raw"): + data = np.loadtxt("polarizability.raw", ndmin = 2); + data = data.astype (np.float32); + np.save ("polarizability", data) +' + python -c \ +'import numpy as np; import os.path; +if os.path.isfile("atomic_dipole.raw"): + data = np.loadtxt("atomic_dipole.raw", ndmin = 2); + data = data.astype (np.float32); + np.save ("atomic_dipole", data) +' + python -c \ +'import numpy as np; import os.path; +if os.path.isfile("atomic_polarizability.raw"): + data = np.loadtxt("atomic_polarizability.raw", ndmin = 2); + data = data.astype (np.float32); + np.save ("atomic_polarizability", data) ' rm *.raw cd ../ diff --git a/deepmd/__init__.py b/deepmd/__init__.py index 231145b989..64386b754c 100644 --- a/deepmd/__init__.py +++ b/deepmd/__init__.py @@ -1,10 +1,11 @@ +"""Root of the deepmd package, exposes all public classes and submodules.""" + +import deepmd.utils.network as network + +from . import cluster, descriptor, fit, loss, utils from .env import set_mkl -from .DeepEval import DeepEval -from .DeepPot import DeepPot -from .DeepDipole import DeepDipole -from .DeepPolar import DeepPolar -from .DeepPolar import DeepGlobalPolar -from .DeepWFC import DeepWFC +from .infer import DeepPotential +from .infer.data_modifier import DipoleChargeModifier set_mkl() @@ -13,3 +14,14 @@ except ImportError: from .__about__ import __version__ +__all__ = [ + "descriptor", + "fit", + "loss", + "utils", + "cluster", + "network", + "DeepEval", + "DeepPotential", + "DipoleChargeModifier", +] diff --git a/deepmd/__main__.py b/deepmd/__main__.py index 04b1c9aca5..2dea15ee78 100644 --- a/deepmd/__main__.py +++ b/deepmd/__main__.py @@ -1,5 +1,6 @@ -from .main import main +"""Package dp entry point.""" + +from .entrypoints.main import main if __name__ == '__main__': main() - diff --git a/deepmd/calculator.py b/deepmd/calculator.py new file mode 100644 index 0000000000..7ca7943d7a --- /dev/null +++ b/deepmd/calculator.py @@ -0,0 +1,103 @@ +"""ASE calculator interface module.""" + +from typing import TYPE_CHECKING, Dict, List, Optional, Union +from pathlib import Path + +from deepmd import DeepPotential +from ase.calculators.calculator import Calculator, all_changes + +if TYPE_CHECKING: + from ase import Atoms + +__all__ = ["DP"] + + +class DP(Calculator): + """Implementation of ASE deepmd calculator. + + Implemented propertie are `energy`, `forces` and `stress` + + Parameters + ---------- + model : Union[str, Path] + path to the model + label : str, optional + calculator label, by default "DP" + 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 + + Examples + -------- + Compute potential energy + + >>> from ase import Atoms + >>> from deepmd.calculator import DP + >>> water = Atoms('H2O', + >>> positions=[(0.7601, 1.9270, 1), + >>> (1.9575, 1, 1), + >>> (1., 1., 1.)], + >>> cell=[100, 100, 100], + >>> calculator=DP(model="frozen_model.pb")) + >>> print(water.get_potential_energy()) + >>> print(water.get_forces()) + + Run BFGS structure optimization + + >>> from ase.optimize import BFGS + >>> dyn = BFGS(water) + >>> dyn.run(fmax=1e-6) + >>> print(water.get_positions()) + """ + + name = "DP" + implemented_properties = ["energy", "forces", "stress"] + + def __init__( + self, + model: Union[str, "Path"], + label: str = "DP", + type_dict: Dict[str, int] = None, + **kwargs + ) -> None: + Calculator.__init__(self, label=label, **kwargs) + self.dp = DeepPotential(str(Path(model).resolve())) + if type_dict: + self.type_dict = type_dict + else: + self.type_dict = dict( + zip(self.dp.get_type_map(), range(self.dp.get_ntypes())) + ) + + def calculate( + self, + atoms: Optional["Atoms"] = None, + properties: List[str] = ["energy", "forces", "stress"], + system_changes: List[str] = all_changes, + ): + """Run calculation with deepmd model. + + Parameters + ---------- + atoms : Optional[Atoms], optional + atoms object to run the calculation on, by default None + properties : List[str], optional + unused, only for function signature compatibility, + by default ["energy", "forces", "stress"] + system_changes : List[str], optional + unused, only for function signature compatibility, by default all_changes + """ + if atoms is not None: + self.atoms = atoms.copy() + + coord = self.atoms.get_positions().reshape([1, -1]) + if sum(self.atoms.get_pbc()) > 0: + cell = self.atoms.get_cell().reshape([1, -1]) + else: + cell = None + symbols = self.atoms.get_chemical_symbols() + atype = [self.type_dict[k] for k in symbols] + e, f, v = self.dp.eval(coords=coord, cells=cell, atom_types=atype) + self.results["energy"] = e[0] + self.results["forces"] = f[0] + self.results["stress"] = v[0] diff --git a/deepmd/cluster/__init__.py b/deepmd/cluster/__init__.py new file mode 100644 index 0000000000..1875b21f9b --- /dev/null +++ b/deepmd/cluster/__init__.py @@ -0,0 +1,22 @@ +"""Module that reads node resources, auto detects if running local or on SLURM.""" + +from .local import get_resource as get_local_res +from .slurm import get_resource as get_slurm_res +import os +from typing import List, Tuple, Optional + +__all__ = ["get_resource"] + + +def get_resource() -> Tuple[str, List[str], Optional[List[int]]]: + """Get local or slurm resources: nodename, nodelist, and gpus. + + Returns + ------- + Tuple[str, List[str], Optional[List[int]]] + nodename, nodelist, and gpus + """ + if "SLURM_JOB_NODELIST" in os.environ: + return get_slurm_res() + else: + return get_local_res() diff --git a/deepmd/cluster/local.py b/deepmd/cluster/local.py new file mode 100644 index 0000000000..37a2e99b6e --- /dev/null +++ b/deepmd/cluster/local.py @@ -0,0 +1,26 @@ +"""Get local GPU resources from `CUDA_VISIBLE_DEVICES` enviroment variable.""" + +import os +import socket +from typing import List, Tuple, Optional + +__all__ = ["get_resource"] + + +def get_resource() -> Tuple[str, List[str], Optional[List[int]]]: + """Get local resources: nodename, nodelist, and gpus. + + Returns + ------- + Tuple[str, List[str], Optional[List[int]]] + nodename, nodelist, and gpus + """ + nodename = socket.gethostname() + nodelist = [nodename] + gpus_env = os.getenv("CUDA_VISIBLE_DEVICES", None) + if not gpus_env: + gpus = None + else: + gpus = [gpu for gpu in gpus_env.split(",")] + + return nodename, nodelist, gpus diff --git a/deepmd/cluster/slurm.py b/deepmd/cluster/slurm.py new file mode 100644 index 0000000000..df4ac3dbf9 --- /dev/null +++ b/deepmd/cluster/slurm.py @@ -0,0 +1,84 @@ +"""MOdule to get resources on SLURM cluster. + +References +---------- +https://github.com/deepsense-ai/tensorflow_on_slurm #### +""" + +import re +import os +from typing import List, Tuple, Optional, Iterable + +__all__ = ["get_resource"] + + +def get_resource() -> Tuple[str, List[str], Optional[List[int]]]: + """Get SLURM resources: nodename, nodelist, and gpus. + + Returns + ------- + Tuple[str, List[str], Optional[List[int]]] + nodename, nodelist, and gpus + + Raises + ------ + RuntimeError + if number of nodes could not be retrieved + ValueError + list of nodes is not of the same length sa number of nodes + ValueError + if current nodename is not found in node list + """ + nodelist = _expand_nodelist(os.environ["SLURM_JOB_NODELIST"]) + nodename = os.environ["SLURMD_NODENAME"] + num_nodes_env = os.getenv("SLURM_JOB_NUM_NODES") + if num_nodes_env: + num_nodes = int(num_nodes_env) + else: + raise RuntimeError("Could not get SLURM number of nodes") + + if len(nodelist) != num_nodes: + raise ValueError( + f"Number of slurm nodes {len(nodelist)} not equal to {num_nodes}" + ) + if nodename not in nodelist: + raise ValueError( + f"Nodename({nodename}) not in nodelist({nodelist}). This should not happen!" + ) + gpus_env = os.getenv("CUDA_VISIBLE_DEVICES") + if not gpus_env: + gpus = None + else: + gpus = [int(gpu) for gpu in gpus_env.split(",")] + return nodename, nodelist, gpus + + +def _pad_zeros(iterable: Iterable, length: int): + return (str(t).rjust(length, "0") for t in iterable) + + +def _expand_ids(ids: str) -> List[str]: + result = [] + for _id in ids.split(","): + if "-" in _id: + str_end = _id.split("-")[1] + begin, end = [int(token) for token in _id.split("-")] + result.extend(_pad_zeros(range(begin, end + 1), len(str_end))) + else: + result.append(_id) + return result + + +def _expand_nodelist(nodelist: str) -> List[str]: + result = [] + interval_list = nodelist.split(",") + for interval in interval_list: + match = re.search(r"(.*)\[(.*)\]", interval) + if match: + prefix = match.group(1) + ids = match.group(2) + ids_list = _expand_ids(ids) + result.extend([f"{prefix}{_id}" for _id in ids_list]) + else: + result.append(interval) + return result diff --git a/deepmd/common.py b/deepmd/common.py new file mode 100644 index 0000000000..e4613aeabb --- /dev/null +++ b/deepmd/common.py @@ -0,0 +1,485 @@ +"""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, + Tuple, + TypeVar, + Union, +) + +import numpy as np +import yaml + +from deepmd.env import op_module, tf +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION, GLOBAL_NP_FLOAT_PRECISION + +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"] + _PRECISION = Literal["default", "float16", "float32", "float64"] + +# define constants +PRECISION_DICT = { + "default": GLOBAL_TF_FLOAT_PRECISION, + "float16": tf.float16, + "float32": tf.float32, + "float64": tf.float64, +} + + +def gelu(x: tf.Tensor) -> tf.Tensor: + """Gaussian Error Linear Unit. + + This is a smoother version of the RELU. + + Parameters + ---------- + x : tf.Tensor + float Tensor to perform activation + + Returns + ------- + `x` with the GELU activation applied + + References + ---------- + Original paper + https://arxiv.org/abs/1606.08415 + """ + return op_module.gelu(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, + "softplus": tf.nn.softplus, + "sigmoid": tf.sigmoid, + "tanh": tf.nn.tanh, + "gelu": gelu, +} + + +def add_data_requirement( + key: str, + ndof: int, + atomic: bool = False, + must: bool = False, + high_prec: bool = False, + type_sel: bool = None, + repeat: int = 1, +): + """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 tru 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 + """ + data_requirement[key] = { + "ndof": ndof, + "atomic": atomic, + "must": must, + "high_prec": high_prec, + "type_sel": type_sel, + "repeat": repeat, + } + + +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 = np.array([], dtype=int) + for ii in sort_select_types: + idx_map = np.append(idx_map, np.where(atom_types == ii)) + return idx_map + + +# TODO not really sure if the docstring is right the purpose of this is a bit unclear +def make_default_mesh( + test_box: np.ndarray, cell_size: float = 3.0 +) -> np.ndarray: + """Get number of cells of size=`cell_size` fit into average box. + + Parameters + ---------- + test_box : np.ndarray + numpy array with cells of shape Nx9 + cell_size : float, optional + length of one cell, by default 3.0 + + Returns + ------- + np.ndarray + mesh for supplied boxes, how many cells fit in each direction + """ + cell_lengths = np.linalg.norm(test_box.reshape([-1, 3, 3]), axis=2) + avg_cell_lengths = np.average(cell_lengths, axis=0) + ncell = (avg_cell_lengths / cell_size).astype(np.int32) + ncell[ncell < 2] = 2 + default_mesh = np.zeros(6, dtype=np.int32) + default_mesh[3:6] = ncell + return default_mesh + + +# TODO not an ideal approach, every class uses this to parse arguments on its own, json +# TODO should be parsed once and the parsed result passed to all objects that need it +class ClassArg: + """Class that take care of input json/yaml parsing. + + The rules for parsing are defined by the `add` method, than `parse` is called to + process the supplied dict + + Attributes + ---------- + arg_dict: Dict[str, Any] + dictionary containing parsing rules + alias_map: Dict[str, Any] + dictionary with keyword aliases + """ + + def __init__(self) -> None: + self.arg_dict = {} + self.alias_map = {} + + def add( + self, + key: str, + types_: Union[type, List[type]], + alias: Optional[Union[str, List[str]]] = None, + default: Any = None, + must: bool = False, + ) -> "ClassArg": + """Add key to be parsed. + + Parameters + ---------- + key : str + key name + types_ : Union[type, List[type]] + list of allowed key types + alias : Optional[Union[str, List[str]]], optional + alias for the key, by default None + default : Any, optional + default value for the key, by default None + must : bool, optional + if the key is mandatory, by default False + + Returns + ------- + ClassArg + instance with added key + """ + if not isinstance(types_, list): + types = [types_] + else: + types = types_ + if alias is not None: + if not isinstance(alias, list): + alias_ = [alias] + else: + alias_ = alias + else: + alias_ = [] + + self.arg_dict[key] = { + "types": types, + "alias": alias_, + "value": default, + "must": must, + } + for ii in alias_: + self.alias_map[ii] = key + + return self + + def _add_single(self, key: str, data: Any): + vtype = type(data) + if data is None: + return data + if not (vtype in self.arg_dict[key]["types"]): + for tp in self.arg_dict[key]["types"]: + try: + vv = tp(data) + except TypeError: + pass + else: + break + else: + raise TypeError( + f"cannot convert provided key {key} to type(s) " + f'{self.arg_dict[key]["types"]} ' + ) + else: + vv = data + self.arg_dict[key]["value"] = vv + + def _check_must(self): + for kk in self.arg_dict: + if self.arg_dict[kk]["must"] and self.arg_dict[kk]["value"] is None: + raise RuntimeError(f"key {kk} must be provided") + + def parse(self, jdata: Dict[str, Any]) -> Dict[str, Any]: + """Parse input dictionary, use the rules defined by add method. + + Parameters + ---------- + jdata : Dict[str, Any] + loaded json/yaml data + + Returns + ------- + Dict[str, Any] + parsed dictionary + """ + for kk in jdata.keys(): + if kk in self.arg_dict: + key = kk + self._add_single(key, jdata[kk]) + else: + if kk in self.alias_map: + key = self.alias_map[kk] + self._add_single(key, jdata[kk]) + self._check_must() + return self.get_dict() + + def get_dict(self) -> Dict[str, Any]: + """Get dictionary built from rules defined by add method. + + Returns + ------- + Dict[str, Any] + settings dictionary with default values + """ + ret = {} + for kk in self.arg_dict.keys(): + ret[kk] = self.arg_dict[kk]["value"] + return ret + + +# 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: "_ACTIVATION", +) -> Callable[[tf.Tensor], tf.Tensor]: + """Get activation function callable based on string name. + + Parameters + ---------- + activation_fn : _ACTIVATION + one of the defined activation functions + + Returns + ------- + Callable[[tf.Tensor], tf.Tensor] + correspondingg TF callable + + Raises + ------ + RuntimeError + if unknown activation function is specified + """ + if activation_fn not in ACTIVATION_FN_DICT: + raise RuntimeError(f"{activation_fn} is not a valid activation function") + return ACTIVATION_FN_DICT[activation_fn] + + +def get_precision(precision: "_PRECISION") -> Any: + """Convert str to TF DType constant. + + Parameters + ---------- + precision : _PRECISION + one of the allowed precisions + + Returns + ------- + tf.python.framework.dtypes.DType + appropriate TF constant + + Raises + ------ + RuntimeError + if supplied precision string does not have acorresponding TF constant + """ + if precision not in PRECISION_DICT: + raise RuntimeError(f"{precision} is not a valid precision") + 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 + """ + matches = [str(d) for d in Path(root_dir).rglob("*") if (d / "type.raw").is_file()] + if (Path(root_dir) / "type.raw").is_file(): + matches += [root_dir] + return matches + + +def docstring_parameter(*sub: Tuple[str, ...]): + """Add parameters to object docstring. + + Parameters + ---------- + sub: Tuple[str, ...] + list of strings that will be inserted into prepared locations in docstring. + + Note + ---- + Can be used on both object and classes. + """ + + @wraps + def dec(obj: "_OBJ") -> "_OBJ": + if obj.__doc__ is not None: + obj.__doc__ = obj.__doc__.format(*sub) + return obj + + return dec + + +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/descriptor/__init__.py b/deepmd/descriptor/__init__.py new file mode 100644 index 0000000000..7c3d910091 --- /dev/null +++ b/deepmd/descriptor/__init__.py @@ -0,0 +1,9 @@ +from .hybrid import DescrptHybrid +from .se_a import DescrptSeA +from .se_r import DescrptSeR +from .se_ar import DescrptSeAR +from .se_t import DescrptSeT +from .se_a_ebd import DescrptSeAEbd +from .se_a_ef import DescrptSeAEf +from .se_a_ef import DescrptSeAEfLower +from .loc_frame import DescrptLocFrame diff --git a/deepmd/descriptor/hybrid.py b/deepmd/descriptor/hybrid.py new file mode 100644 index 0000000000..a0748591e5 --- /dev/null +++ b/deepmd/descriptor/hybrid.py @@ -0,0 +1,216 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import ClassArg +from deepmd.env import op_module +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +# from deepmd.descriptor import DescrptLocFrame +# from deepmd.descriptor import DescrptSeA +# from deepmd.descriptor import DescrptSeT +# from deepmd.descriptor import DescrptSeAEbd +# from deepmd.descriptor import DescrptSeAEf +# from deepmd.descriptor import DescrptSeR +from .se_a import DescrptSeA +from .se_r import DescrptSeR +from .se_ar import DescrptSeAR +from .se_t import DescrptSeT +from .se_a_ebd import DescrptSeAEbd +from .se_a_ef import DescrptSeAEf +from .loc_frame import DescrptLocFrame + +class DescrptHybrid (): + def __init__ (self, + descrpt_list : list + ) -> None : + """ + Constructor + + Parameters + ---------- + descrpt_list : list + Build a descriptor from the concatenation of the list of descriptors. + """ + if descrpt_list == [] or descrpt_list is None: + raise RuntimeError('cannot build descriptor from an empty list of descriptors.') + # args = ClassArg()\ + # .add('list', list, must = True) + # class_data = args.parse(jdata) + # dict_list = class_data['list'] + self.descrpt_list = descrpt_list + self.numb_descrpt = len(self.descrpt_list) + for ii in range(1, self.numb_descrpt): + assert(self.descrpt_list[ii].get_ntypes() == + self.descrpt_list[ 0].get_ntypes()), \ + f'number of atom types in {ii}th descrptor does not match others' + + + def get_rcut (self) -> float: + """ + Returns the cut-off radius + """ + all_rcut = [ii.get_rcut() for ii in self.descrpt_list] + return np.max(all_rcut) + + + def get_ntypes (self) -> int: + """ + Returns the number of atom types + """ + return self.descrpt_list[0].get_ntypes() + + + def get_dim_out (self) -> int: + """ + Returns the output dimension of this descriptor + """ + all_dim_out = [ii.get_dim_out() for ii in self.descrpt_list] + return sum(all_dim_out) + + + def get_nlist_i(self, + ii : int + ) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Parameters + ---------- + ii : int + Get the neighbor information of the ii-th descriptor + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ + return self.descrpt_list[ii].nlist, self.descrpt_list[ii].rij, self.descrpt_list[ii].sel_a, self.descrpt_list[ii].sel_r + + + def compute_input_stats (self, + data_coord : list, + data_box : list, + data_atype : list, + natoms_vec : list, + mesh : list, + input_dict : dict + ) -> None : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ + for ii in self.descrpt_list: + ii.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict) + + + def build (self, + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ + with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : + t_rcut = tf.constant(self.get_rcut(), + name = 'rcut', + dtype = GLOBAL_TF_FLOAT_PRECISION) + t_ntypes = tf.constant(self.get_ntypes(), + name = 'ntypes', + dtype = tf.int32) + all_dout = [] + for idx,ii in enumerate(self.descrpt_list): + dout = ii.build(coord_, atype_, natoms, box_, mesh, input_dict, suffix=suffix+f'_{idx}', reuse=reuse) + dout = tf.reshape(dout, [-1, ii.get_dim_out()]) + all_dout.append(dout) + dout = tf.concat(all_dout, axis = 1) + dout = tf.reshape(dout, [-1, natoms[0] * self.get_dim_out()]) + return dout + + + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ + for idx,ii in enumerate(self.descrpt_list): + ff, vv, av = ii.prod_force_virial(atom_ener, natoms) + if idx == 0: + force = ff + virial = vv + atom_virial = av + else: + force += ff + virial += vv + atom_virial += av + return force, virial, atom_virial diff --git a/source/train/DescrptLocFrame.py b/deepmd/descriptor/loc_frame.py similarity index 52% rename from source/train/DescrptLocFrame.py rename to deepmd/descriptor/loc_frame.py index 69c1473db0..d8cd4f6366 100644 --- a/source/train/DescrptLocFrame.py +++ b/deepmd/descriptor/loc_frame.py @@ -1,23 +1,53 @@ import numpy as np +from typing import Tuple, List + from deepmd.env import tf -from deepmd.common import ClassArg -from deepmd.RunOptions import global_tf_float_precision -from deepmd.RunOptions import global_np_float_precision +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION from deepmd.env import op_module from deepmd.env import default_tf_session_config class DescrptLocFrame () : - def __init__(self, jdata): - args = ClassArg()\ - .add('sel_a', list, must = True) \ - .add('sel_r', list, must = True) \ - .add('rcut', float, default = 6.0) \ - .add('axis_rule',list, must = True) - class_data = args.parse(jdata) - self.sel_a = class_data['sel_a'] - self.sel_r = class_data['sel_r'] - self.axis_rule = class_data['axis_rule'] - self.rcut_r = class_data['rcut'] + def __init__(self, + rcut: float, + sel_a : List[int], + sel_r : List[int], + axis_rule : List[int] + ) -> None: + """ + Constructor + + Parameters + rcut + The cut-off radius + sel_a : list[str] + 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. + sel_r : list[str] + 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. + axis_rule: list[int] + The length should be 6 times of the number of types. + - 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 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+4]: type of the atom defining the second axis of type-i atom.\n\n\ + - axis_rule[i*6+5]: 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. + """ + # args = ClassArg()\ + # .add('sel_a', list, must = True) \ + # .add('sel_r', list, must = True) \ + # .add('rcut', float, default = 6.0) \ + # .add('axis_rule',list, must = True) + # class_data = args.parse(jdata) + self.sel_a = sel_a + self.sel_r = sel_r + self.axis_rule = axis_rule + self.rcut_r = rcut # ntypes and rcut_a === -1 self.ntypes = len(self.sel_a) assert(self.ntypes == len(self.sel_r)) @@ -33,13 +63,13 @@ def __init__(self, jdata): self.dstd = None self.place_holders = {} - avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) - std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) sub_graph = tf.Graph() with sub_graph.as_default(): name_pfx = 'd_lf_' for ii in ['coord', 'box']: - self.place_holders[ii] = tf.placeholder(global_np_float_precision, [None, None], name = name_pfx+'t_'+ii) + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') @@ -59,24 +89,65 @@ def __init__(self, jdata): self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config) - def get_rcut (self) : + def get_rcut (self) -> float: + """ + Returns the cut-off radisu + """ return self.rcut_r - def get_ntypes (self) : + def get_ntypes (self) -> int: + """ + Returns the number of atom types + """ return self.ntypes - def get_dim_out (self) : + def get_dim_out (self) -> int: + """ + Returns the output dimension of this descriptor + """ return self.ndescrpt - def get_nlist (self) : + def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ return self.nlist, self.rij, self.sel_a, self.sel_r def compute_input_stats (self, - data_coord, - data_box, - data_atype, - natoms_vec, - mesh) : + data_coord : list, + data_box : list, + data_atype : list, + natoms_vec : list, + mesh : list, + input_dict : dict + ) -> None : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ all_davg = [] all_dstd = [] if True: @@ -105,13 +176,45 @@ def compute_input_stats (self, def build (self, - coord_, - atype_, - natoms, - box_, - mesh, - suffix = '', - reuse = None): + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ davg = self.davg dstd = self.dstd with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : @@ -121,18 +224,18 @@ def build (self, dstd = np.ones ([self.ntypes, self.ndescrpt]) t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), name = 'rcut', - dtype = global_tf_float_precision) + dtype = GLOBAL_TF_FLOAT_PRECISION) t_ntypes = tf.constant(self.ntypes, name = 'ntypes', dtype = tf.int32) self.t_avg = tf.get_variable('t_avg', davg.shape, - dtype = global_tf_float_precision, + dtype = GLOBAL_TF_FLOAT_PRECISION, trainable = False, initializer = tf.constant_initializer(davg)) self.t_std = tf.get_variable('t_std', dstd.shape, - dtype = global_tf_float_precision, + dtype = GLOBAL_TF_FLOAT_PRECISION, trainable = False, initializer = tf.constant_initializer(dstd)) @@ -154,13 +257,45 @@ def build (self, sel_r = self.sel_r, axis_rule = self.axis_rule) self.descrpt = tf.reshape(self.descrpt, [-1, self.ndescrpt]) + tf.summary.histogram('descrpt', self.descrpt) + tf.summary.histogram('rij', self.rij) + tf.summary.histogram('nlist', self.nlist) + return self.descrpt - def get_rot_mat(self) : + def get_rot_mat(self) -> tf.Tensor: + """ + Get rotational matrix + """ return self.rot_mat - def prod_force_virial(self, atom_ener, natoms) : + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ [net_deriv] = tf.gradients (atom_ener, self.descrpt) + tf.summary.histogram('net_derivative', net_deriv) net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) force = op_module.prod_force (net_deriv_reshape, self.descrpt_deriv, @@ -178,6 +313,9 @@ def prod_force_virial(self, atom_ener, natoms) : natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) + tf.summary.histogram('force', force) + tf.summary.histogram('virial', virial) + tf.summary.histogram('atom_virial', atom_virial) return force, virial, atom_virial diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py new file mode 100644 index 0000000000..b4f1292d48 --- /dev/null +++ b/deepmd/descriptor/se_a.py @@ -0,0 +1,702 @@ +import math +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, get_np_precision +from deepmd.utils.argcheck import list_to_doc +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift +from deepmd.utils.tabulate import DeepTabulate +from deepmd.utils.type_embed import embed_atom_type + +class DescrptSeA (): + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: int = None, + type_one_side: bool = True, + exclude_types: List[List[int]] = [], + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + uniform_seed: bool = False + ) -> None: + """ + Constructor + + Parameters + ---------- + rcut + The cut-off radius + rcut_smth + From where the environment matrix should be smoothed + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + axis_neuron + Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix) + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + type_one_side + Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets + exclude_types : list[int] + The Excluded types + set_davg_zero + Set the shift of embedding net input to zero. + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + self.sel_a = sel + self.rcut_r = rcut + self.rcut_r_smth = rcut_smth + self.filter_neuron = neuron + self.n_axis_neuron = axis_neuron + self.filter_resnet_dt = resnet_dt + self.seed = seed + self.uniform_seed = uniform_seed + self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron) + self.trainable = trainable + self.filter_activation_fn = get_activation_func(activation_function) + self.filter_precision = get_precision(precision) + self.filter_np_precision = get_np_precision(precision) + self.exclude_types = set() + for tt in exclude_types: + assert(len(tt) == 2) + self.exclude_types.add((tt[0], tt[1])) + self.exclude_types.add((tt[1], tt[0])) + self.set_davg_zero = set_davg_zero + self.type_one_side = type_one_side + if self.type_one_side and len(exclude_types) != 0: + raise RuntimeError('"type_one_side" is not compatible with "exclude_types"') + + # descrpt config + self.sel_r = [ 0 for ii in range(len(self.sel_a)) ] + self.ntypes = len(self.sel_a) + assert(self.ntypes == len(self.sel_r)) + self.rcut_a = -1 + # numb of neighbors and numb of descrptors + self.nnei_a = np.cumsum(self.sel_a)[-1] + self.nnei_r = np.cumsum(self.sel_r)[-1] + self.nnei = self.nnei_a + self.nnei_r + self.ndescrpt_a = self.nnei_a * 4 + self.ndescrpt_r = self.nnei_r * 1 + self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r + self.useBN = False + self.dstd = None + self.davg = None + self.compress = False + self.place_holders = {} + nei_type = np.array([]) + for ii in range(self.ntypes): + nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii])) # like a mask + self.nei_type = tf.constant(nei_type, dtype = tf.int32) + + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + sub_graph = tf.Graph() + with sub_graph.as_default(): + name_pfx = 'd_sea_' + for ii in ['coord', 'box']: + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) + self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') + self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') + self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') + self.stat_descrpt, descrpt_deriv, rij, nlist \ + = op_module.prod_env_mat_a(self.place_holders['coord'], + self.place_holders['type'], + self.place_holders['natoms_vec'], + self.place_holders['box'], + self.place_holders['default_mesh'], + tf.constant(avg_zero), + tf.constant(std_ones), + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config) + + + def get_rcut (self) -> float: + """ + Returns the cut-off radius + """ + return self.rcut_r + + def get_ntypes (self) -> int: + """ + Returns the number of atom types + """ + return self.ntypes + + def get_dim_out (self) -> int: + """ + Returns the output dimension of this descriptor + """ + return self.filter_neuron[-1] * self.n_axis_neuron + + def get_dim_rot_mat_1 (self) -> int: + """ + Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3 + """ + return self.filter_neuron[-1] + + def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ + return self.nlist, self.rij, self.sel_a, self.sel_r + + def compute_input_stats (self, + data_coord : list, + data_box : list, + data_atype : list, + natoms_vec : list, + mesh : list, + input_dict : dict + ) -> None : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ + all_davg = [] + all_dstd = [] + if True: + sumr = [] + suma = [] + sumn = [] + sumr2 = [] + suma2 = [] + for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) : + sysr,sysr2,sysa,sysa2,sysn \ + = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm) + sumr.append(sysr) + suma.append(sysa) + sumn.append(sysn) + sumr2.append(sysr2) + suma2.append(sysa2) + sumr = np.sum(sumr, axis = 0) + suma = np.sum(suma, axis = 0) + sumn = np.sum(sumn, axis = 0) + sumr2 = np.sum(sumr2, axis = 0) + suma2 = np.sum(suma2, axis = 0) + for type_i in range(self.ntypes) : + davgunit = [sumr[type_i]/(sumn[type_i]+1e-15), 0, 0, 0] + dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]) + ] + davg = np.tile(davgunit, self.ndescrpt // 4) + dstd = np.tile(dstdunit, self.ndescrpt // 4) + all_davg.append(davg) + all_dstd.append(dstd) + + if not self.set_davg_zero: + self.davg = np.array(all_davg) + self.dstd = np.array(all_dstd) + + def enable_compression(self, + min_nbor_dist : float, + model_file : str = 'frozon_model.pb', + table_extrapolate : float = 5, + table_stride_1 : float = 0.01, + table_stride_2 : float = 0.1, + check_frequency : int = -1 + ) -> None: + """ + Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + model_file + The original frozen model, which will be compressed by the program + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + self.compress = True + self.model_file = model_file + self.table_config = [table_extrapolate, table_stride_1, table_stride_2, check_frequency] + self.table = DeepTabulate(self.model_file, self.type_one_side) + self.lower, self.upper \ + = self.table.build(min_nbor_dist, + table_extrapolate, + table_stride_1, + table_stride_2) + + def build (self, + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ + davg = self.davg + dstd = self.dstd + with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : + if davg is None: + davg = np.zeros([self.ntypes, self.ndescrpt]) + if dstd is None: + dstd = np.ones ([self.ntypes, self.ndescrpt]) + t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), + name = 'rcut', + dtype = GLOBAL_TF_FLOAT_PRECISION) + t_ntypes = tf.constant(self.ntypes, + name = 'ntypes', + dtype = tf.int32) + t_ndescrpt = tf.constant(self.ndescrpt, + name = 'ndescrpt', + dtype = tf.int32) + t_sel = tf.constant(self.sel_a, + name = 'sel', + dtype = tf.int32) + self.t_avg = tf.get_variable('t_avg', + davg.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(davg)) + self.t_std = tf.get_variable('t_std', + dstd.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(dstd)) + + coord = tf.reshape (coord_, [-1, natoms[1] * 3]) + box = tf.reshape (box_, [-1, 9]) + atype = tf.reshape (atype_, [-1, natoms[1]]) + + self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ + = op_module.prod_env_mat_a (coord, + atype, + natoms, + box, + mesh, + self.t_avg, + self.t_std, + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + # only used when tensorboard was set as true + tf.summary.histogram('descrpt', self.descrpt) + tf.summary.histogram('rij', self.rij) + tf.summary.histogram('nlist', self.nlist) + + self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') + self.rij = tf.identity(self.rij, name = 'o_rij') + self.nlist = tf.identity(self.nlist, name = 'o_nlist') + + self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, + atype, + natoms, + input_dict, + suffix = suffix, + reuse = reuse, + trainable = self.trainable) + + # only used when tensorboard was set as true + tf.summary.histogram('embedding_net_output', self.dout) + return self.dout + + def get_rot_mat(self) -> tf.Tensor: + """ + Get rotational matrix + """ + return self.qmat + + + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ + [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape) + tf.summary.histogram('net_derivative', net_deriv) + net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) + force \ + = op_module.prod_force_se_a (net_deriv_reshape, + self.descrpt_deriv, + self.nlist, + natoms, + n_a_sel = self.nnei_a, + n_r_sel = self.nnei_r) + virial, atom_virial \ + = op_module.prod_virial_se_a (net_deriv_reshape, + self.descrpt_deriv, + self.rij, + self.nlist, + natoms, + n_a_sel = self.nnei_a, + n_r_sel = self.nnei_r) + tf.summary.histogram('force', force) + tf.summary.histogram('virial', virial) + tf.summary.histogram('atom_virial', atom_virial) + + return force, virial, atom_virial + + + def _pass_filter(self, + inputs, + atype, + natoms, + input_dict, + reuse = None, + suffix = '', + trainable = True) : + if input_dict is not None: + type_embedding = input_dict.get('type_embedding', None) + else: + type_embedding = None + start_index = 0 + inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) + output = [] + output_qmat = [] + if not self.type_one_side and type_embedding is None: + for type_i in range(self.ntypes): + inputs_i = tf.slice (inputs, + [ 0, start_index* self.ndescrpt], + [-1, natoms[2+type_i]* self.ndescrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) + layer, qmat = self._filter(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) + layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_out()]) + qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_rot_mat_1() * 3]) + output.append(layer) + output_qmat.append(qmat) + start_index += natoms[2+type_i] + else : + inputs_i = inputs + inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) + type_i = -1 + layer, qmat = self._filter(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn, type_embedding=type_embedding) + layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) + qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3]) + output.append(layer) + output_qmat.append(qmat) + output = tf.concat(output, axis = 1) + output_qmat = tf.concat(output_qmat, axis = 1) + return output, output_qmat + + + def _compute_dstats_sys_smth (self, + data_coord, + data_box, + data_atype, + natoms_vec, + mesh) : + dd_all \ + = self.sub_sess.run(self.stat_descrpt, + feed_dict = { + self.place_holders['coord']: data_coord, + self.place_holders['type']: data_atype, + self.place_holders['natoms_vec']: natoms_vec, + self.place_holders['box']: data_box, + self.place_holders['default_mesh']: mesh, + }) + natoms = natoms_vec + dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]]) + start_index = 0 + sysr = [] + sysa = [] + sysn = [] + sysr2 = [] + sysa2 = [] + for type_i in range(self.ntypes): + end_index = start_index + self.ndescrpt * natoms[2+type_i] + dd = dd_all[:, start_index:end_index] + dd = np.reshape(dd, [-1, self.ndescrpt]) + start_index = end_index + # compute + dd = np.reshape (dd, [-1, 4]) + ddr = dd[:,:1] + dda = dd[:,1:] + sumr = np.sum(ddr) + suma = np.sum(dda) / 3. + sumn = dd.shape[0] + sumr2 = np.sum(np.multiply(ddr, ddr)) + suma2 = np.sum(np.multiply(dda, dda)) / 3. + sysr.append(sumr) + sysa.append(suma) + sysn.append(sumn) + sysr2.append(sumr2) + sysa2.append(suma2) + return sysr, sysr2, sysa, sysa2, sysn + + + def _compute_std (self,sumv2, sumv, sumn) : + if sumn == 0: + return 1e-2 + val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn)) + if np.abs(val) < 1e-2: + val = 1e-2 + return val + + + def _concat_type_embedding( + self, + xyz_scatter, + nframes, + natoms, + type_embedding, + ): + te_out_dim = type_embedding.get_shape().as_list()[-1] + nei_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(self.nei_type,dtype=tf.int32)) #nnei*nchnl + nei_embed = tf.tile(nei_embed,(nframes*natoms[0],1)) + nei_embed = tf.reshape(nei_embed,[-1,te_out_dim]) + embedding_input = tf.concat([xyz_scatter,nei_embed],1) + if not self.type_one_side: + atm_embed = embed_atom_type(self.ntypes, natoms, type_embedding) + atm_embed = tf.tile(atm_embed,(1,self.nnei)) + atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) + embedding_input = tf.concat([embedding_input,atm_embed],1) + return embedding_input + + + def _filter_lower( + self, + type_i, + type_input, + start_index, + incrs_index, + inputs, + nframes, + natoms, + type_embedding=None, + is_exclude = False, + activation_fn = None, + bavg = 0.0, + stddev = 1.0, + trainable = True, + suffix = '', + ): + """ + input env matrix, returns R.G + """ + outputs_size = [1] + self.filter_neuron + # cut-out inputs + # with natom x (nei_type_i x 4) + inputs_i = tf.slice (inputs, + [ 0, start_index* 4], + [-1, incrs_index* 4] ) + shape_i = inputs_i.get_shape().as_list() + # with (natom x nei_type_i) x 4 + inputs_reshape = tf.reshape(inputs_i, [-1, 4]) + # with (natom x nei_type_i) x 1 + xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1]) + if type_embedding is not None: + type_embedding = tf.cast(type_embedding, self.filter_precision) + xyz_scatter = self._concat_type_embedding( + xyz_scatter, nframes, natoms, type_embedding) + if self.compress: + raise RuntimeError('compression of type embedded descriptor is not supported at the moment') + # with (natom x nei_type_i) x out_size + if self.compress and (not is_exclude): + info = [self.lower, self.upper, self.upper * self.table_config[0], self.table_config[1], self.table_config[2], self.table_config[3]] + if self.type_one_side: + net = 'filter_-1_net_' + str(type_i) + else: + net = 'filter_' + str(type_input) + '_net_' + str(type_i) + return op_module.tabulate_fusion(self.table.data[net].astype(self.filter_np_precision), info, xyz_scatter, tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), last_layer_size = outputs_size[-1]) + else: + if (not is_exclude): + xyz_scatter = embedding_net( + xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + name_suffix = suffix, + stddev = stddev, + bavg = bavg, + seed = self.seed, + trainable = trainable, + uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + else: + w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=GLOBAL_TF_FLOAT_PRECISION) + xyz_scatter = tf.matmul(xyz_scatter, w) + # natom x nei_type_i x out_size + xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1])) + return tf.matmul(tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]), xyz_scatter, transpose_a = True) + + + def _filter( + self, + inputs, + type_input, + natoms, + type_embedding = None, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + trainable = True): + nframes = tf.shape(tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]))[0] + # natom x (nei x 4) + shape = inputs.get_shape().as_list() + outputs_size = [1] + self.filter_neuron + outputs_size_2 = self.n_axis_neuron + with tf.variable_scope(name, reuse=reuse): + start_index = 0 + type_i = 0 + if type_embedding is None: + for type_i in range(self.ntypes): + ret = self._filter_lower( + type_i, type_input, + start_index, self.sel_a[type_i], + inputs, + nframes, + natoms, + type_embedding = type_embedding, + is_exclude = (type_input, type_i) in self.exclude_types, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + trainable = trainable, + suffix = "_"+str(type_i)) + if type_i == 0: + xyz_scatter_1 = ret + else: + xyz_scatter_1+= ret + start_index += self.sel_a[type_i] + else : + xyz_scatter_1 = self._filter_lower( + type_i, type_input, + start_index, np.cumsum(self.sel_a)[-1], + inputs, + nframes, + natoms, + type_embedding = type_embedding, + is_exclude = False, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + trainable = trainable) + # natom x nei x outputs_size + # xyz_scatter = tf.concat(xyz_scatter_total, axis=1) + # natom x nei x 4 + # inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4]) + # natom x 4 x outputs_size + # xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True) + xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) + # natom x 4 x outputs_size_2 + xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2]) + # # natom x 3 x outputs_size_2 + # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) + # natom x 3 x outputs_size_1 + qmat = tf.slice(xyz_scatter_1, [0,1,0], [-1, 3, -1]) + # natom x outputs_size_1 x 3 + qmat = tf.transpose(qmat, perm = [0, 2, 1]) + # natom x outputs_size x outputs_size_2 + result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a = True) + # natom x (outputs_size x outputs_size_2) + result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]]) + + return result, qmat diff --git a/deepmd/descriptor/se_a_ebd.py b/deepmd/descriptor/se_a_ebd.py new file mode 100644 index 0000000000..7630cabfc9 --- /dev/null +++ b/deepmd/descriptor/se_a_ebd.py @@ -0,0 +1,516 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import ClassArg, get_activation_func, get_precision, add_data_requirement +from deepmd.utils.network import one_layer +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net +from .se_a import DescrptSeA + +class DescrptSeAEbd (DescrptSeA): + def __init__ (self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: int = None, + type_one_side: bool = True, + type_nchanl : int = 2, + type_nlayer : int = 1, + numb_aparam : int = 0, + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + exclude_types: List[int] = [], + ) -> None: + """ + Constructor + + Parameters + ---------- + rcut + The cut-off radius + rcut_smth + From where the environment matrix should be smoothed + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + axis_neuron + Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix) + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + type_one_side + Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets + type_nchanl + Number of channels for type representation + type_nlayer + Number of hidden layers for the type embedding net (skip connected). + numb_aparam + Number of atomic parameters. If >0 it will be embedded with atom types. + set_davg_zero + Set the shift of embedding net input to zero. + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + """ + # args = ClassArg()\ + # .add('type_nchanl', int, default = 4) \ + # .add('type_nlayer', int, default = 2) \ + # .add('type_one_side', bool, default = True) \ + # .add('numb_aparam', int, default = 0) + # class_data = args.parse(jdata) + DescrptSeA.__init__(self, + rcut, + rcut_smth, + sel, + neuron = neuron, + axis_neuron = axis_neuron, + resnet_dt = resnet_dt, + trainable = trainable, + seed = seed, + type_one_side = type_one_side, + set_davg_zero = set_davg_zero, + activation_function = activation_function, + precision = precision + ) + self.type_nchanl = type_nchanl + self.type_nlayer = type_nlayer + self.type_one_side = type_one_side + self.numb_aparam = numb_aparam + if self.numb_aparam > 0: + add_data_requirement('aparam', 3, atomic=True, must=True, high_prec=False) + + + def build (self, + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ + nei_type = np.array([]) + for ii in range(self.ntypes): + nei_type = np.append(nei_type, ii * np.ones(self.sel_a[ii])) + self.nei_type = tf.get_variable('t_nei_type', + [self.nnei], + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(nei_type)) + self.dout = DescrptSeA.build(self, coord_, atype_, natoms, box_, mesh, input_dict, suffix = suffix, reuse = reuse) + tf.summary.histogram('embedding_net_output', self.dout) + + return self.dout + + + def _type_embed(self, + atype, + ndim = 1, + reuse = None, + suffix = '', + trainable = True): + ebd_type = tf.cast(atype, self.filter_precision) + ebd_type = ebd_type / float(self.ntypes) + ebd_type = tf.reshape(ebd_type, [-1, ndim]) + for ii in range(self.type_nlayer): + name = 'type_embed_layer_' + str(ii) + ebd_type = one_layer(ebd_type, + self.type_nchanl, + activation_fn = self.filter_activation_fn, + precision = self.filter_precision, + name = name, + reuse = reuse, + seed = self.seed + ii, + trainable = trainable) + name = 'type_embed_layer_' + str(self.type_nlayer) + ebd_type = one_layer(ebd_type, + self.type_nchanl, + activation_fn = None, + precision = self.filter_precision, + name = name, + reuse = reuse, + seed = self.seed + ii, + trainable = trainable) + ebd_type = tf.reshape(ebd_type, [tf.shape(atype)[0], self.type_nchanl]) + return ebd_type + + + def _embedding_net(self, + inputs, + natoms, + filter_neuron, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + ''' + inputs: nf x na x (nei x 4) + outputs: nf x na x nei x output_size + ''' + # natom x (nei x 4) + inputs = tf.reshape(inputs, [-1, self.ndescrpt]) + shape = inputs.get_shape().as_list() + outputs_size = [1] + filter_neuron + with tf.variable_scope(name, reuse=reuse): + xyz_scatter_total = [] + # with natom x (nei x 4) + inputs_i = inputs + shape_i = inputs_i.get_shape().as_list() + # with (natom x nei) x 4 + inputs_reshape = tf.reshape(inputs_i, [-1, 4]) + # with (natom x nei) x 1 + xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0,0],[-1,1]),[-1,1]) + # with (natom x nei) x out_size + xyz_scatter = embedding_net(xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + stddev = stddev, + bavg = bavg, + seed = seed, + trainable = trainable) + # natom x nei x out_size + xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1]//4, outputs_size[-1])) + xyz_scatter_total.append(xyz_scatter) + # natom x nei x outputs_size + xyz_scatter = tf.concat(xyz_scatter_total, axis=1) + # nf x natom x nei x outputs_size + xyz_scatter = tf.reshape(xyz_scatter, [tf.shape(inputs)[0], natoms[0], self.nnei, outputs_size[-1]]) + return xyz_scatter + + + def _type_embedding_net_two_sides(self, + mat_g, + atype, + natoms, + name = '', + reuse = None, + seed = None, + trainable = True): + outputs_size = self.filter_neuron[-1] + nframes = tf.shape(mat_g)[0] + # (nf x natom x nei) x (outputs_size x chnl x chnl) + mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) + mat_g = one_layer(mat_g, + outputs_size * self.type_nchanl * self.type_nchanl, + activation_fn = None, + precision = self.filter_precision, + name = name+'_amplify', + reuse = reuse, + seed = self.seed, + trainable = trainable) + # nf x natom x nei x outputs_size x chnl x chnl + mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl, self.type_nchanl]) + # nf x natom x outputs_size x chnl x nei x chnl + mat_g = tf.transpose(mat_g, perm = [0, 1, 3, 4, 2, 5]) + # nf x natom x outputs_size x chnl x (nei x chnl) + mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.type_nchanl, self.nnei * self.type_nchanl]) + + # nei x nchnl + ebd_nei_type = self._type_embed(self.nei_type, + reuse = reuse, + trainable = True, + suffix = '') + # (nei x nchnl) + ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl]) + # (nframes x natom) x nchnl + ebd_atm_type = self._type_embed(atype, + reuse = True, + trainable = True, + suffix = '') + # (nframes x natom x nchnl) + ebd_atm_type = tf.reshape(ebd_atm_type, [nframes * natoms[0] * self.type_nchanl]) + + # nf x natom x outputs_size x chnl x (nei x chnl) + mat_g = tf.multiply(mat_g, ebd_nei_type) + # nf x natom x outputs_size x chnl x nei x chnl + mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.type_nchanl, self.nnei, self.type_nchanl]) + # nf x natom x outputs_size x chnl x nei + mat_g = tf.reduce_mean(mat_g, axis = 5) + # outputs_size x nei x nf x natom x chnl + mat_g = tf.transpose(mat_g, perm = [2, 4, 0, 1, 3]) + # outputs_size x nei x (nf x natom x chnl) + mat_g = tf.reshape(mat_g, [outputs_size, self.nnei, nframes * natoms[0] * self.type_nchanl]) + # outputs_size x nei x (nf x natom x chnl) + mat_g = tf.multiply(mat_g, ebd_atm_type) + # outputs_size x nei x nf x natom x chnl + mat_g = tf.reshape(mat_g, [outputs_size, self.nnei, nframes, natoms[0], self.type_nchanl]) + # outputs_size x nei x nf x natom + mat_g = tf.reduce_mean(mat_g, axis = 4) + # nf x natom x nei x outputs_size + mat_g = tf.transpose(mat_g, perm = [2, 3, 1, 0]) + # (nf x natom) x nei x outputs_size + mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) + return mat_g + + + def _type_embedding_net_one_side(self, + mat_g, + atype, + natoms, + name = '', + reuse = None, + seed = None, + trainable = True): + outputs_size = self.filter_neuron[-1] + nframes = tf.shape(mat_g)[0] + # (nf x natom x nei) x (outputs_size x chnl x chnl) + mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) + mat_g = one_layer(mat_g, + outputs_size * self.type_nchanl, + activation_fn = None, + precision = self.filter_precision, + name = name+'_amplify', + reuse = reuse, + seed = self.seed, + trainable = trainable) + # nf x natom x nei x outputs_size x chnl + mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl]) + # nf x natom x outputs_size x nei x chnl + mat_g = tf.transpose(mat_g, perm = [0, 1, 3, 2, 4]) + # nf x natom x outputs_size x (nei x chnl) + mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.nnei * self.type_nchanl]) + + # nei x nchnl + ebd_nei_type = self._type_embed(self.nei_type, + reuse = reuse, + trainable = True, + suffix = '') + # (nei x nchnl) + ebd_nei_type = tf.reshape(ebd_nei_type, [self.nnei * self.type_nchanl]) + + # nf x natom x outputs_size x (nei x chnl) + mat_g = tf.multiply(mat_g, ebd_nei_type) + # nf x natom x outputs_size x nei x chnl + mat_g = tf.reshape(mat_g, [nframes, natoms[0], outputs_size, self.nnei, self.type_nchanl]) + # nf x natom x outputs_size x nei + mat_g = tf.reduce_mean(mat_g, axis = 4) + # nf x natom x nei x outputs_size + mat_g = tf.transpose(mat_g, perm = [0, 1, 3, 2]) + # (nf x natom) x nei x outputs_size + mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) + return mat_g + + + def _type_embedding_net_one_side_aparam(self, + mat_g, + atype, + natoms, + aparam, + name = '', + reuse = None, + seed = None, + trainable = True): + outputs_size = self.filter_neuron[-1] + nframes = tf.shape(mat_g)[0] + # (nf x natom x nei) x (outputs_size x chnl x chnl) + mat_g = tf.reshape(mat_g, [nframes * natoms[0] * self.nnei, outputs_size]) + mat_g = one_layer(mat_g, + outputs_size * self.type_nchanl, + activation_fn = None, + precision = self.filter_precision, + name = name+'_amplify', + reuse = reuse, + seed = self.seed, + trainable = trainable) + # nf x natom x nei x outputs_size x chnl + mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size, self.type_nchanl]) + # outputs_size x nf x natom x nei x chnl + mat_g = tf.transpose(mat_g, perm = [3, 0, 1, 2, 4]) + # outputs_size x (nf x natom x nei x chnl) + mat_g = tf.reshape(mat_g, [outputs_size, nframes * natoms[0] * self.nnei * self.type_nchanl]) + # nf x natom x nnei + embed_type = tf.tile(tf.reshape(self.nei_type, [1, self.nnei]), + [nframes * natoms[0], 1]) + # (nf x natom x nnei) x 1 + embed_type = tf.reshape(embed_type, [nframes * natoms[0] * self.nnei, 1]) + # nf x (natom x naparam) + aparam = tf.reshape(aparam, [nframes, -1]) + # nf x natom x nnei x naparam + embed_aparam = op_module.map_aparam(aparam, self.nlist, natoms, n_a_sel = self.nnei_a, n_r_sel = self.nnei_r) + # (nf x natom x nnei) x naparam + embed_aparam = tf.reshape(embed_aparam, [nframes * natoms[0] * self.nnei, self.numb_aparam]) + # (nf x natom x nnei) x (naparam+1) + embed_input = tf.concat((embed_type, embed_aparam), axis = 1) + + # (nf x natom x nnei) x nchnl + ebd_nei_type = self._type_embed(embed_input, + ndim = self.numb_aparam + 1, + reuse = reuse, + trainable = True, + suffix = '') + # (nf x natom x nei x nchnl) + ebd_nei_type = tf.reshape(ebd_nei_type, [nframes * natoms[0] * self.nnei * self.type_nchanl]) + + # outputs_size x (nf x natom x nei x chnl) + mat_g = tf.multiply(mat_g, ebd_nei_type) + # outputs_size x nf x natom x nei x chnl + mat_g = tf.reshape(mat_g, [outputs_size, nframes, natoms[0], self.nnei, self.type_nchanl]) + # outputs_size x nf x natom x nei + mat_g = tf.reduce_mean(mat_g, axis = 4) + # nf x natom x nei x outputs_size + mat_g = tf.transpose(mat_g, perm = [1, 2, 3, 0]) + # (nf x natom) x nei x outputs_size + mat_g = tf.reshape(mat_g, [nframes * natoms[0], self.nnei, outputs_size]) + return mat_g + + + def _pass_filter(self, + inputs, + atype, + natoms, + input_dict, + reuse = None, + suffix = '', + trainable = True) : + # nf x na x ndescrpt + # nf x na x (nnei x 4) + inputs = tf.reshape(inputs, [-1, natoms[0], self.ndescrpt]) + layer, qmat = self._ebd_filter(tf.cast(inputs, self.filter_precision), + atype, + natoms, + input_dict, + name='filter_type_all'+suffix, + reuse=reuse, + seed = self.seed, + trainable = trainable, + activation_fn = self.filter_activation_fn) + output = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) + output_qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3]) + return output, output_qmat + + + def _ebd_filter(self, + inputs, + atype, + natoms, + input_dict, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + trainable = True): + outputs_size = self.filter_neuron[-1] + outputs_size_2 = self.n_axis_neuron + # nf x natom x (nei x 4) + nframes = tf.shape(inputs)[0] + shape = tf.reshape(inputs, [-1, self.ndescrpt]).get_shape().as_list() + + # nf x natom x nei x outputs_size + mat_g = self._embedding_net(inputs, + natoms, + self.filter_neuron, + activation_fn = activation_fn, + stddev = stddev, + bavg = bavg, + name = name, + reuse = reuse, + seed = seed, + trainable = trainable) + # nf x natom x nei x outputs_size + mat_g = tf.reshape(mat_g, [nframes, natoms[0], self.nnei, outputs_size]) + + # (nf x natom) x nei x outputs_size + if self.type_one_side: + if self.numb_aparam > 0: + aparam = input_dict['aparam'] + xyz_scatter \ + = self._type_embedding_net_one_side_aparam(mat_g, + atype, + natoms, + aparam, + name = name, + reuse = reuse, + seed = seed, + trainable = trainable) + else: + xyz_scatter \ + = self._type_embedding_net_one_side(mat_g, + atype, + natoms, + name = name, + reuse = reuse, + seed = seed, + trainable = trainable) + else: + xyz_scatter \ + = self._type_embedding_net_two_sides(mat_g, + atype, + natoms, + name = name, + reuse = reuse, + seed = seed, + trainable = trainable) + + # natom x nei x 4 + inputs_reshape = tf.reshape(inputs, [-1, shape[1]//4, 4]) + # natom x 4 x outputs_size + xyz_scatter_1 = tf.matmul(inputs_reshape, xyz_scatter, transpose_a = True) + xyz_scatter_1 = xyz_scatter_1 * (4.0 / shape[1]) + # natom x 4 x outputs_size_2 + xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2]) + # # natom x 3 x outputs_size_2 + # qmat = tf.slice(xyz_scatter_2, [0,1,0], [-1, 3, -1]) + # natom x 3 x outputs_size_1 + qmat = tf.slice(xyz_scatter_1, [0,1,0], [-1, 3, -1]) + # natom x outputs_size_2 x 3 + qmat = tf.transpose(qmat, perm = [0, 2, 1]) + # natom x outputs_size x outputs_size_2 + result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a = True) + # natom x (outputs_size x outputs_size_2) + result = tf.reshape(result, [-1, outputs_size_2 * outputs_size]) + + return result, qmat + diff --git a/deepmd/descriptor/se_a_ef.py b/deepmd/descriptor/se_a_ef.py new file mode 100644 index 0000000000..bdba470d39 --- /dev/null +++ b/deepmd/descriptor/se_a_ef.py @@ -0,0 +1,559 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import add_data_requirement,get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from .se_a import DescrptSeA + +class DescrptSeAEf (): + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__(self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: int = None, + type_one_side: bool = True, + exclude_types: List[int] = [], + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + uniform_seed = False + ) -> None: + """ + Constructor + + Parameters + ---------- + rcut + The cut-off radius + rcut_smth + From where the environment matrix should be smoothed + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + axis_neuron + Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix) + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + type_one_side + Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets + exclude_types : list[int] + The Excluded types + set_davg_zero + Set the shift of embedding net input to zero. + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + self.descrpt_para = DescrptSeAEfLower( + op_module.descrpt_se_a_ef_para, + rcut, + rcut_smth, + sel, + neuron, + axis_neuron, + resnet_dt, + trainable, + seed, + type_one_side, + exclude_types, + set_davg_zero, + activation_function, + precision, + uniform_seed, + ) + self.descrpt_vert = DescrptSeAEfLower( + op_module.descrpt_se_a_ef_vert, + rcut, + rcut_smth, + sel, + neuron, + axis_neuron, + resnet_dt, + trainable, + seed, + type_one_side, + exclude_types, + set_davg_zero, + activation_function, + precision, + uniform_seed, + ) + + def get_rcut (self) -> float: + """ + Returns the cut-off radisu + """ + return self.descrpt_vert.rcut_r + + def get_ntypes (self) -> int: + """ + Returns the number of atom types + """ + return self.descrpt_vert.ntypes + + def get_dim_out (self) -> int: + """ + Returns the output dimension of this descriptor + """ + return self.descrpt_vert.get_dim_out() + self.descrpt_para.get_dim_out() + + def get_dim_rot_mat_1 (self) -> int: + """ + Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3 + """ + return self.descrpt_vert.filter_neuron[-1] + + def get_rot_mat(self) -> tf.Tensor: + """ + Get rotational matrix + """ + return self.qmat + + + def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ + return \ + self.descrpt_vert.nlist, \ + self.descrpt_vert.rij, \ + self.descrpt_vert.sel_a, \ + self.descrpt_vert.sel_r + + def compute_input_stats (self, + data_coord : list, + data_box : list, + data_atype : list, + natoms_vec : list, + mesh : list, + input_dict : dict + ) -> None : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ + self.descrpt_vert.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict) + self.descrpt_para.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict) + + def build (self, + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs. Should have 'efield'. + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ + self.dout_vert = self.descrpt_vert.build(coord_, atype_, natoms, box_, mesh, input_dict) + self.dout_para = self.descrpt_para.build(coord_, atype_, natoms, box_, mesh, input_dict, reuse = True) + coord = tf.reshape(coord_, [-1, natoms[1] * 3]) + nframes = tf.shape(coord)[0] + self.dout_vert = tf.reshape(self.dout_vert, [nframes * natoms[0], self.descrpt_vert.get_dim_out()]) + self.dout_para = tf.reshape(self.dout_para, [nframes * natoms[0], self.descrpt_para.get_dim_out()]) + self.dout = tf.concat([self.dout_vert, self.dout_para], axis = 1) + self.dout = tf.reshape(self.dout, [nframes, natoms[0] * self.get_dim_out()]) + self.qmat = self.descrpt_vert.qmat + self.descrpt_para.qmat + + tf.summary.histogram('embedding_net_output', self.dout) + + return self.dout + + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ + f_vert, v_vert, av_vert \ + = self.descrpt_vert.prod_force_virial(atom_ener, natoms) + f_para, v_para, av_para \ + = self.descrpt_para.prod_force_virial(atom_ener, natoms) + force = f_vert + f_para + virial = v_vert + v_para + atom_vir = av_vert + av_para + return force, virial, atom_vir + + +class DescrptSeAEfLower (DescrptSeA): + """ + Helper class for implementing DescrptSeAEf + """ + def __init__ (self, + op, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + axis_neuron: int = 8, + resnet_dt: bool = False, + trainable: bool = True, + seed: int = None, + type_one_side: bool = True, + exclude_types: List[int] = [], + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + uniform_seed : bool = False, + ) -> None: + DescrptSeA.__init__( + self, + rcut, + rcut_smth, + sel, + neuron, + axis_neuron, + resnet_dt, + trainable, + seed, + type_one_side, + exclude_types, + set_davg_zero, + activation_function, + precision, + uniform_seed + ) + # DescrptSeA.__init__(self, **jdata) + # args = ClassArg()\ + # .add('sel', list, must = True) \ + # .add('rcut', float, default = 6.0) \ + # .add('rcut_smth',float, default = 5.5) \ + # .add('neuron', list, default = [10, 20, 40]) \ + # .add('axis_neuron', int, default = 4, alias = 'n_axis_neuron') \ + # .add('resnet_dt',bool, default = False) \ + # .add('trainable',bool, default = True) \ + # .add('seed', int) + # class_data = args.parse(jdata) + # self.sel_a = class_data['sel'] + # self.rcut_r = class_data['rcut'] + # self.rcut_r_smth = class_data['rcut_smth'] + # self.filter_neuron = class_data['neuron'] + # self.n_axis_neuron = class_data['axis_neuron'] + # self.filter_resnet_dt = class_data['resnet_dt'] + # self.seed = class_data['seed'] + # self.trainable = class_data['trainable'] + self.sel_a = sel + self.rcut_r = rcut + self.rcut_r_smth = rcut_smth + self.filter_neuron = neuron + self.n_axis_neuron = axis_neuron + self.filter_resnet_dt = resnet_dt + self.seed = seed + self.trainable = trainable + self.op = op + + # descrpt config + self.sel_r = [ 0 for ii in range(len(self.sel_a)) ] + self.ntypes = len(self.sel_a) + assert(self.ntypes == len(self.sel_r)) + self.rcut_a = -1 + # numb of neighbors and numb of descrptors + self.nnei_a = np.cumsum(self.sel_a)[-1] + self.nnei_r = np.cumsum(self.sel_r)[-1] + self.nnei = self.nnei_a + self.nnei_r + self.ndescrpt_a = self.nnei_a * 4 + self.ndescrpt_r = self.nnei_r * 1 + self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r + self.useBN = False + self.dstd = None + self.davg = None + + add_data_requirement('efield', 3, atomic=True, must=True, high_prec=False) + + self.place_holders = {} + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + sub_graph = tf.Graph() + with sub_graph.as_default(): + name_pfx = 'd_sea_ef_' + for ii in ['coord', 'box']: + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) + self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') + self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') + self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') + self.place_holders['efield'] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name=name_pfx+'t_efield') + self.stat_descrpt, descrpt_deriv, rij, nlist \ + = self.op(self.place_holders['coord'], + self.place_holders['type'], + self.place_holders['natoms_vec'], + self.place_holders['box'], + self.place_holders['default_mesh'], + self.place_holders['efield'], + tf.constant(avg_zero), + tf.constant(std_ones), + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config) + + + + def compute_input_stats (self, + data_coord, + data_box, + data_atype, + natoms_vec, + mesh, + input_dict) : + data_efield = input_dict['efield'] + all_davg = [] + all_dstd = [] + if True: + sumr = [] + suma = [] + sumn = [] + sumr2 = [] + suma2 = [] + for cc,bb,tt,nn,mm,ee in zip(data_coord,data_box,data_atype,natoms_vec,mesh,data_efield) : + sysr,sysr2,sysa,sysa2,sysn \ + = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm,ee) + sumr.append(sysr) + suma.append(sysa) + sumn.append(sysn) + sumr2.append(sysr2) + suma2.append(sysa2) + sumr = np.sum(sumr, axis = 0) + suma = np.sum(suma, axis = 0) + sumn = np.sum(sumn, axis = 0) + sumr2 = np.sum(sumr2, axis = 0) + suma2 = np.sum(suma2, axis = 0) + for type_i in range(self.ntypes) : + davgunit = [sumr[type_i]/sumn[type_i], 0, 0, 0] + dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]) + ] + davg = np.tile(davgunit, self.ndescrpt // 4) + dstd = np.tile(dstdunit, self.ndescrpt // 4) + all_davg.append(davg) + all_dstd.append(dstd) + + self.davg = np.array(all_davg) + self.dstd = np.array(all_dstd) + + def _normalize_3d(self, a): + na = tf.norm(a, axis = 1) + na = tf.tile(tf.reshape(na, [-1,1]), tf.constant([1, 3])) + return tf.divide(a, na) + + def build (self, + coord_, + atype_, + natoms, + box_, + mesh, + input_dict, + suffix = '', + reuse = None): + efield = input_dict['efield'] + davg = self.davg + dstd = self.dstd + with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : + if davg is None: + davg = np.zeros([self.ntypes, self.ndescrpt]) + if dstd is None: + dstd = np.ones ([self.ntypes, self.ndescrpt]) + t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), + name = 'rcut', + dtype = GLOBAL_TF_FLOAT_PRECISION) + t_ntypes = tf.constant(self.ntypes, + name = 'ntypes', + dtype = tf.int32) + t_ndescrpt = tf.constant(self.ndescrpt, + name = 'ndescrpt', + dtype = tf.int32) + t_sel = tf.constant(self.sel_a, + name = 'sel', + dtype = tf.int32) + self.t_avg = tf.get_variable('t_avg', + davg.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(davg)) + self.t_std = tf.get_variable('t_std', + dstd.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(dstd)) + + coord = tf.reshape (coord_, [-1, natoms[1] * 3]) + box = tf.reshape (box_, [-1, 9]) + atype = tf.reshape (atype_, [-1, natoms[1]]) + efield = tf.reshape(efield, [-1, 3]) + efield = self._normalize_3d(efield) + efield = tf.reshape(efield, [-1, natoms[0] * 3]) + + self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ + = self.op (coord, + atype, + natoms, + box, + mesh, + efield, + self.t_avg, + self.t_std, + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + + self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') + self.rij = tf.identity(self.rij, name = 'o_rij') + self.nlist = tf.identity(self.nlist, name = 'o_nlist') + + # only used when tensorboard was set as true + tf.summary.histogram('descrpt', self.descrpt) + tf.summary.histogram('rij', self.rij) + tf.summary.histogram('nlist', self.nlist) + + self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, atype, natoms, input_dict, suffix = suffix, reuse = reuse, trainable = self.trainable) + tf.summary.histogram('embedding_net_output', self.dout) + + return self.dout + + + + def _compute_dstats_sys_smth (self, + data_coord, + data_box, + data_atype, + natoms_vec, + mesh, + data_efield) : + dd_all \ + = self.sub_sess.run(self.stat_descrpt, + feed_dict = { + self.place_holders['coord']: data_coord, + self.place_holders['type']: data_atype, + self.place_holders['natoms_vec']: natoms_vec, + self.place_holders['box']: data_box, + self.place_holders['default_mesh']: mesh, + self.place_holders['efield']: data_efield, + }) + natoms = natoms_vec + dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]]) + start_index = 0 + sysr = [] + sysa = [] + sysn = [] + sysr2 = [] + sysa2 = [] + for type_i in range(self.ntypes): + end_index = start_index + self.ndescrpt * natoms[2+type_i] + dd = dd_all[:, start_index:end_index] + dd = np.reshape(dd, [-1, self.ndescrpt]) + start_index = end_index + # compute + dd = np.reshape (dd, [-1, 4]) + ddr = dd[:,:1] + dda = dd[:,1:] + sumr = np.sum(ddr) + suma = np.sum(dda) / 3. + sumn = dd.shape[0] + sumr2 = np.sum(np.multiply(ddr, ddr)) + suma2 = np.sum(np.multiply(dda, dda)) / 3. + sysr.append(sumr) + sysa.append(suma) + sysn.append(sumn) + sysr2.append(sumr2) + sysa2.append(suma2) + return sysr, sysr2, sysa, sysa2, sysn + + diff --git a/source/train/DescrptSeAR.py b/deepmd/descriptor/se_ar.py similarity index 80% rename from source/train/DescrptSeAR.py rename to deepmd/descriptor/se_ar.py index dadc2f3d95..8ded2c5849 100644 --- a/source/train/DescrptSeAR.py +++ b/deepmd/descriptor/se_ar.py @@ -2,8 +2,8 @@ from deepmd.env import tf from deepmd.common import ClassArg -from deepmd.DescrptSeA import DescrptSeA -from deepmd.DescrptSeR import DescrptSeR +from .se_a import DescrptSeA +from .se_r import DescrptSeR from deepmd.env import op_module class DescrptSeAR (): @@ -14,8 +14,8 @@ def __init__ (self, jdata): class_data = args.parse(jdata) self.param_a = class_data['a'] self.param_r = class_data['r'] - self.descrpt_a = DescrptSeA(self.param_a) - self.descrpt_r = DescrptSeR(self.param_r) + self.descrpt_a = DescrptSeA(**self.param_a) + self.descrpt_r = DescrptSeR(**self.param_r) assert(self.descrpt_a.get_ntypes() == self.descrpt_r.get_ntypes()) self.davg = None self.dstd = None @@ -40,9 +40,10 @@ def compute_input_stats (self, data_box, data_atype, natoms_vec, - mesh) : - self.descrpt_a.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh) - self.descrpt_r.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh) + mesh, + input_dict) : + self.descrpt_a.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict) + self.descrpt_r.compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict) self.davg = [self.descrpt_a.davg, self.descrpt_r.davg] self.dstd = [self.descrpt_a.dstd, self.descrpt_r.dstd] @@ -53,6 +54,7 @@ def build (self, natoms, box, mesh, + input_dict, suffix = '', reuse = None): davg = self.davg @@ -64,12 +66,14 @@ def build (self, dstd = [np.ones ([self.descrpt_a.ntypes, self.descrpt_a.ndescrpt]), np.ones ([self.descrpt_r.ntypes, self.descrpt_r.ndescrpt])] # dout - self.dout_a = self.descrpt_a.build(coord_, atype_, natoms, box, mesh, suffix=suffix+'_a', reuse=reuse) - self.dout_r = self.descrpt_r.build(coord_, atype_, natoms, box, mesh, suffix=suffix , reuse=reuse) + self.dout_a = self.descrpt_a.build(coord_, atype_, natoms, box, mesh, input_dict, suffix=suffix+'_a', reuse=reuse) + self.dout_r = self.descrpt_r.build(coord_, atype_, natoms, box, mesh, input_dict, suffix=suffix , reuse=reuse) self.dout_a = tf.reshape(self.dout_a, [-1, self.descrpt_a.get_dim_out()]) self.dout_r = tf.reshape(self.dout_r, [-1, self.descrpt_r.get_dim_out()]) self.dout = tf.concat([self.dout_a, self.dout_r], axis = 1) self.dout = tf.reshape(self.dout, [-1, natoms[0] * self.get_dim_out()]) + + tf.summary.histogram('embedding_net_output', self.dout) return self.dout @@ -79,6 +83,9 @@ def prod_force_virial(self, atom_ener, natoms) : force = f_a + f_r virial = v_a + v_r atom_virial = av_a + av_r + tf.summary.histogram('force', force) + tf.summary.histogram('virial', virial) + tf.summary.histogram('atom_virial', atom_virial) return force, virial, atom_virial diff --git a/source/train/DescrptSeR.py b/deepmd/descriptor/se_r.py similarity index 55% rename from source/train/DescrptSeR.py rename to deepmd/descriptor/se_r.py index 5d4bc4e4de..6dad7947fa 100644 --- a/source/train/DescrptSeR.py +++ b/deepmd/descriptor/se_r.py @@ -1,44 +1,96 @@ import numpy as np +from typing import Tuple, List + from deepmd.env import tf -from deepmd.common import ClassArg, get_activation_func, get_precision -from deepmd.RunOptions import global_tf_float_precision -from deepmd.RunOptions import global_np_float_precision +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION from deepmd.env import op_module from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift class DescrptSeR (): - def __init__ (self, jdata): - args = ClassArg()\ - .add('sel', list, must = True) \ - .add('rcut', float, default = 6.0) \ - .add('rcut_smth',float, default = 0.5) \ - .add('neuron', list, default = [10, 20, 40]) \ - .add('resnet_dt',bool, default = False) \ - .add('trainable',bool, default = True) \ - .add('seed', int) \ - .add('type_one_side', bool, default = False) \ - .add('exclude_types', list, default = []) \ - .add('set_davg_zero', bool, default = False) \ - .add("activation_function", str, default = "tanh") \ - .add("precision", str, default = "default") - class_data = args.parse(jdata) - self.sel_r = class_data['sel'] - self.rcut = class_data['rcut'] - self.rcut_smth = class_data['rcut_smth'] - self.filter_neuron = class_data['neuron'] - self.filter_resnet_dt = class_data['resnet_dt'] - self.seed = class_data['seed'] - self.trainable = class_data['trainable'] - self.filter_activation_fn = get_activation_func(class_data["activation_function"]) - self.filter_precision = get_precision(class_data['precision']) - exclude_types = class_data['exclude_types'] + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + resnet_dt: bool = False, + trainable: bool = True, + seed: int = None, + type_one_side: bool = True, + exclude_types: List[int] = [], + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + uniform_seed: bool = False + ) -> None: + """ + Constructor + + Parameters + ---------- + rcut + The cut-off radius + rcut_smth + From where the environment matrix should be smoothed + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + type_one_side + Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets + exclude_types : list[int] + The Excluded types + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + # args = ClassArg()\ + # .add('sel', list, must = True) \ + # .add('rcut', float, default = 6.0) \ + # .add('rcut_smth',float, default = 0.5) \ + # .add('neuron', list, default = [10, 20, 40]) \ + # .add('resnet_dt',bool, default = False) \ + # .add('trainable',bool, default = True) \ + # .add('seed', int) \ + # .add('type_one_side', bool, default = False) \ + # .add('exclude_types', list, default = []) \ + # .add('set_davg_zero', bool, default = False) \ + # .add("activation_function", str, default = "tanh") \ + # .add("precision", str, default = "default") + # class_data = args.parse(jdata) + self.sel_r = sel + self.rcut = rcut + self.rcut_smth = rcut_smth + self.filter_neuron = neuron + self.filter_resnet_dt = resnet_dt + self.seed = seed + self.uniform_seed = uniform_seed + self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron) + self.trainable = trainable + self.filter_activation_fn = get_activation_func(activation_function) + self.filter_precision = get_precision(precision) + exclude_types = exclude_types self.exclude_types = set() for tt in exclude_types: assert(len(tt) == 2) self.exclude_types.add((tt[0], tt[1])) self.exclude_types.add((tt[1], tt[0])) - self.set_davg_zero = class_data['set_davg_zero'] - self.type_one_side = class_data['type_one_side'] + self.set_davg_zero = set_davg_zero + self.type_one_side = type_one_side # descrpt config self.sel_a = [ 0 for ii in range(len(self.sel_r)) ] @@ -55,18 +107,18 @@ def __init__ (self, jdata): self.dstd = None self.place_holders = {} - avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) - std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(global_np_float_precision) + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) sub_graph = tf.Graph() with sub_graph.as_default(): name_pfx = 'd_ser_' for ii in ['coord', 'box']: - self.place_holders[ii] = tf.placeholder(global_np_float_precision, [None, None], name = name_pfx+'t_'+ii) + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') self.stat_descrpt, descrpt_deriv, rij, nlist \ - = op_module.descrpt_se_r(self.place_holders['coord'], + = op_module.prod_env_mat_r(self.place_holders['coord'], self.place_holders['type'], self.place_holders['natoms_vec'], self.place_holders['box'], @@ -80,23 +132,63 @@ def __init__ (self, jdata): def get_rcut (self) : + """ + Returns the cut-off radisu + """ return self.rcut def get_ntypes (self) : + """ + Returns the number of atom types + """ return self.ntypes def get_dim_out (self) : + """ + Returns the output dimension of this descriptor + """ return self.filter_neuron[-1] def get_nlist (self) : + """ + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ return self.nlist, self.rij, self.sel_a, self.sel_r def compute_input_stats (self, - data_coord, - data_box, - data_atype, - natoms_vec, - mesh) : + data_coord, + data_box, + data_atype, + natoms_vec, + mesh, + input_dict) : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ all_davg = [] all_dstd = [] sumr = [] @@ -125,13 +217,45 @@ def compute_input_stats (self, def build (self, - coord_, - atype_, - natoms, - box_, - mesh, - suffix = '', - reuse = None): + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ davg = self.davg dstd = self.dstd with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : @@ -141,7 +265,7 @@ def build (self, dstd = np.ones ([self.ntypes, self.ndescrpt]) t_rcut = tf.constant(self.rcut, name = 'rcut', - dtype = global_tf_float_precision) + dtype = GLOBAL_TF_FLOAT_PRECISION) t_ntypes = tf.constant(self.ntypes, name = 'ntypes', dtype = tf.int32) @@ -153,12 +277,12 @@ def build (self, dtype = tf.int32) self.t_avg = tf.get_variable('t_avg', davg.shape, - dtype = global_tf_float_precision, + dtype = GLOBAL_TF_FLOAT_PRECISION, trainable = False, initializer = tf.constant_initializer(davg)) self.t_std = tf.get_variable('t_std', dstd.shape, - dtype = global_tf_float_precision, + dtype = GLOBAL_TF_FLOAT_PRECISION, trainable = False, initializer = tf.constant_initializer(dstd)) @@ -167,7 +291,7 @@ def build (self, atype = tf.reshape (atype_, [-1, natoms[1]]) self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ - = op_module.descrpt_se_r (coord, + = op_module.prod_env_mat_r(coord, atype, natoms, box, @@ -184,13 +308,44 @@ def build (self, self.rij = tf.identity(self.rij, name = 'o_rij') self.nlist = tf.identity(self.nlist, name = 'o_nlist') + # only used when tensorboard was set as true + tf.summary.histogram('descrpt', self.descrpt) + tf.summary.histogram('rij', self.rij) + tf.summary.histogram('nlist', self.nlist) + self.dout = self._pass_filter(self.descrpt_reshape, natoms, suffix = suffix, reuse = reuse, trainable = self.trainable) + tf.summary.histogram('embedding_net_output', self.dout) return self.dout - def prod_force_virial(self, atom_ener, natoms) : + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape) + tf.summary.histogram('net_derivative', net_deriv) net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) force \ = op_module.prod_force_se_r (net_deriv_reshape, @@ -203,6 +358,10 @@ def prod_force_virial(self, atom_ener, natoms) : self.rij, self.nlist, natoms) + tf.summary.histogram('force', force) + tf.summary.histogram('virial', virial) + tf.summary.histogram('atom_virial', atom_virial) + return force, virial, atom_virial @@ -221,7 +380,7 @@ def _pass_filter(self, [ 0, start_index* self.ndescrpt], [-1, natoms[2+type_i]* self.ndescrpt] ) inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) - layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, seed = self.seed, trainable = trainable, activation_fn = self.filter_activation_fn) + layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_'+str(type_i)+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[2+type_i] * self.get_dim_out()]) output.append(layer) start_index += natoms[2+type_i] @@ -229,7 +388,7 @@ def _pass_filter(self, inputs_i = inputs inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) type_i = -1 - layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, seed = self.seed, trainable = trainable, activation_fn = self.filter_activation_fn) + layer = self._filter_r(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) output.append(layer) output = tf.concat(output, axis = 1) @@ -289,7 +448,6 @@ def _filter_r(self, bavg=0.0, name='linear', reuse=None, - seed=None, trainable = True): # natom x nei outputs_size = [1] + self.filter_neuron @@ -307,38 +465,20 @@ def _filter_r(self, # with (natom x nei_type_i) x 1 xyz_scatter = tf.reshape(inputs_i, [-1, 1]) if (type_input, type_i) not in self.exclude_types: - for ii in range(1, len(outputs_size)): - w = tf.get_variable('matrix_'+str(ii)+'_'+str(type_i), - [outputs_size[ii - 1], outputs_size[ii]], - self.filter_precision, - tf.random_normal_initializer(stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), seed = seed), - trainable = trainable) - b = tf.get_variable('bias_'+str(ii)+'_'+str(type_i), - [1, outputs_size[ii]], - self.filter_precision, - tf.random_normal_initializer(stddev=stddev, mean = bavg, seed = seed), - trainable = trainable) - hidden = tf.reshape(activation_fn(tf.matmul(xyz_scatter, w) + b), [-1, outputs_size[ii]]) - if self.filter_resnet_dt : - idt = tf.get_variable('idt_'+str(ii)+'_'+str(type_i), - [1, outputs_size[ii]], - self.filter_precision, - tf.random_normal_initializer(stddev=0.001, mean = 1.0, seed = seed), - trainable = trainable) - if outputs_size[ii] == outputs_size[ii-1]: - if self.filter_resnet_dt : - xyz_scatter += hidden * idt - else : - xyz_scatter += hidden - elif outputs_size[ii] == outputs_size[ii-1] * 2: - if self.filter_resnet_dt : - xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden * idt - else : - xyz_scatter = tf.concat([xyz_scatter,xyz_scatter], 1) + hidden - else: - xyz_scatter = hidden + xyz_scatter = embedding_net(xyz_scatter, + self.filter_neuron, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + name_suffix = "_"+str(type_i), + stddev = stddev, + bavg = bavg, + seed = self.seed, + trainable = trainable, + uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift else: - w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=global_tf_float_precision) + w = tf.zeros((outputs_size[0], outputs_size[-1]), dtype=GLOBAL_TF_FLOAT_PRECISION) xyz_scatter = tf.matmul(xyz_scatter, w) # natom x nei_type_i x out_size xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1], outputs_size[-1])) diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py new file mode 100644 index 0000000000..3206990402 --- /dev/null +++ b/deepmd/descriptor/se_t.py @@ -0,0 +1,513 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net, embedding_net_rand_seed_shift + +class DescrptSeT (): + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + rcut: float, + rcut_smth: float, + sel: List[str], + neuron: List[int] = [24,48,96], + resnet_dt: bool = False, + trainable: bool = True, + seed: int = None, + set_davg_zero: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + uniform_seed: bool = False + ) -> None: + """ + Constructor + + Parameters + ---------- + rcut + The cut-off radius + rcut_smth + From where the environment matrix should be smoothed + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + set_davg_zero + Set the shift of embedding net input to zero. + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + self.sel_a = sel + self.rcut_r = rcut + self.rcut_r_smth = rcut_smth + self.filter_neuron = neuron + self.filter_resnet_dt = resnet_dt + self.seed = seed + self.uniform_seed = uniform_seed + self.seed_shift = embedding_net_rand_seed_shift(self.filter_neuron) + self.trainable = trainable + self.filter_activation_fn = get_activation_func(activation_function) + self.filter_precision = get_precision(precision) + # self.exclude_types = set() + # for tt in exclude_types: + # assert(len(tt) == 2) + # self.exclude_types.add((tt[0], tt[1])) + # self.exclude_types.add((tt[1], tt[0])) + self.set_davg_zero = set_davg_zero + + # descrpt config + self.sel_r = [ 0 for ii in range(len(self.sel_a)) ] + self.ntypes = len(self.sel_a) + assert(self.ntypes == len(self.sel_r)) + self.rcut_a = -1 + # numb of neighbors and numb of descrptors + self.nnei_a = np.cumsum(self.sel_a)[-1] + self.nnei_r = np.cumsum(self.sel_r)[-1] + self.nnei = self.nnei_a + self.nnei_r + self.ndescrpt_a = self.nnei_a * 4 + self.ndescrpt_r = self.nnei_r * 1 + self.ndescrpt = self.ndescrpt_a + self.ndescrpt_r + self.useBN = False + self.dstd = None + self.davg = None + + self.place_holders = {} + avg_zero = np.zeros([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + std_ones = np.ones ([self.ntypes,self.ndescrpt]).astype(GLOBAL_NP_FLOAT_PRECISION) + sub_graph = tf.Graph() + with sub_graph.as_default(): + name_pfx = 'd_sea_' + for ii in ['coord', 'box']: + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name = name_pfx+'t_'+ii) + self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name=name_pfx+'t_type') + self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name=name_pfx+'t_natoms') + self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name=name_pfx+'t_mesh') + self.stat_descrpt, descrpt_deriv, rij, nlist \ + = op_module.prod_env_mat_a(self.place_holders['coord'], + self.place_holders['type'], + self.place_holders['natoms_vec'], + self.place_holders['box'], + self.place_holders['default_mesh'], + tf.constant(avg_zero), + tf.constant(std_ones), + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config) + + + def get_rcut (self) -> float: + """ + Returns the cut-off radisu + """ + return self.rcut_r + + def get_ntypes (self) -> int: + """ + Returns the number of atom types + """ + return self.ntypes + + def get_dim_out (self) -> int: + """ + Returns the output dimension of this descriptor + """ + return self.filter_neuron[-1] + + def get_nlist (self) -> Tuple[tf.Tensor, tf.Tensor, List[int], List[int]]: + """ + Returns + ------- + nlist + Neighbor list + rij + The relative distance between the neighbor and the center atom. + sel_a + The number of neighbors with full information + sel_r + The number of neighbors with only radial information + """ + return self.nlist, self.rij, self.sel_a, self.sel_r + + def compute_input_stats (self, + data_coord : list, + data_box : list, + data_atype : list, + natoms_vec : list, + mesh : list, + input_dict : dict + ) -> None : + """ + Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics. + + Parameters + ---------- + data_coord + The coordinates. Can be generated by deepmd.model.make_stat_input + data_box + The box. Can be generated by deepmd.model.make_stat_input + data_atype + The atom types. Can be generated by deepmd.model.make_stat_input + natoms_vec + The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input + mesh + The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input + input_dict + Dictionary for additional input + """ + all_davg = [] + all_dstd = [] + if True: + sumr = [] + suma = [] + sumn = [] + sumr2 = [] + suma2 = [] + for cc,bb,tt,nn,mm in zip(data_coord,data_box,data_atype,natoms_vec,mesh) : + sysr,sysr2,sysa,sysa2,sysn \ + = self._compute_dstats_sys_smth(cc,bb,tt,nn,mm) + sumr.append(sysr) + suma.append(sysa) + sumn.append(sysn) + sumr2.append(sysr2) + suma2.append(sysa2) + sumr = np.sum(sumr, axis = 0) + suma = np.sum(suma, axis = 0) + sumn = np.sum(sumn, axis = 0) + sumr2 = np.sum(sumr2, axis = 0) + suma2 = np.sum(suma2, axis = 0) + for type_i in range(self.ntypes) : + davgunit = [sumr[type_i]/sumn[type_i], 0, 0, 0] + dstdunit = [self._compute_std(sumr2[type_i], sumr[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]), + self._compute_std(suma2[type_i], suma[type_i], sumn[type_i]) + ] + davg = np.tile(davgunit, self.ndescrpt // 4) + dstd = np.tile(dstdunit, self.ndescrpt // 4) + all_davg.append(davg) + all_dstd.append(dstd) + + if not self.set_davg_zero: + self.davg = np.array(all_davg) + self.dstd = np.array(all_dstd) + + + def build (self, + coord_ : tf.Tensor, + atype_ : tf.Tensor, + natoms : tf.Tensor, + box_ : tf.Tensor, + mesh : tf.Tensor, + input_dict : dict, + reuse : bool = None, + suffix : str = '' + ) -> tf.Tensor: + """ + Build the computational graph for the descriptor + + Parameters + ---------- + coord_ + The coordinate of atoms + atype_ + The type of atoms + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + mesh + For historical reasons, only the length of the Tensor matters. + if size of mesh == 6, pbc is assumed. + if size of mesh == 0, no-pbc is assumed. + input_dict + Dictionary for additional inputs + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + descriptor + The output descriptor + """ + davg = self.davg + dstd = self.dstd + with tf.variable_scope('descrpt_attr' + suffix, reuse = reuse) : + if davg is None: + davg = np.zeros([self.ntypes, self.ndescrpt]) + if dstd is None: + dstd = np.ones ([self.ntypes, self.ndescrpt]) + t_rcut = tf.constant(np.max([self.rcut_r, self.rcut_a]), + name = 'rcut', + dtype = GLOBAL_TF_FLOAT_PRECISION) + t_ntypes = tf.constant(self.ntypes, + name = 'ntypes', + dtype = tf.int32) + t_ndescrpt = tf.constant(self.ndescrpt, + name = 'ndescrpt', + dtype = tf.int32) + t_sel = tf.constant(self.sel_a, + name = 'sel', + dtype = tf.int32) + self.t_avg = tf.get_variable('t_avg', + davg.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(davg)) + self.t_std = tf.get_variable('t_std', + dstd.shape, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(dstd)) + + coord = tf.reshape (coord_, [-1, natoms[1] * 3]) + box = tf.reshape (box_, [-1, 9]) + atype = tf.reshape (atype_, [-1, natoms[1]]) + + self.descrpt, self.descrpt_deriv, self.rij, self.nlist \ + = op_module.prod_env_mat_a (coord, + atype, + natoms, + box, + mesh, + self.t_avg, + self.t_std, + rcut_a = self.rcut_a, + rcut_r = self.rcut_r, + rcut_r_smth = self.rcut_r_smth, + sel_a = self.sel_a, + sel_r = self.sel_r) + + self.descrpt_reshape = tf.reshape(self.descrpt, [-1, self.ndescrpt]) + self.descrpt_reshape = tf.identity(self.descrpt_reshape, name = 'o_rmat') + self.descrpt_deriv = tf.identity(self.descrpt_deriv, name = 'o_rmat_deriv') + self.rij = tf.identity(self.rij, name = 'o_rij') + self.nlist = tf.identity(self.nlist, name = 'o_nlist') + + self.dout, self.qmat = self._pass_filter(self.descrpt_reshape, + atype, + natoms, + input_dict, + suffix = suffix, + reuse = reuse, + trainable = self.trainable) + + return self.dout + + + def prod_force_virial(self, + atom_ener : tf.Tensor, + natoms : tf.Tensor + ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + """ + Compute force and virial + + Parameters + ---------- + atom_ener + The atomic energy + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + Return + ------ + force + The force on atoms + virial + The total virial + atom_virial + The atomic virial + """ + [net_deriv] = tf.gradients (atom_ener, self.descrpt_reshape) + net_deriv_reshape = tf.reshape (net_deriv, [-1, natoms[0] * self.ndescrpt]) + force \ + = op_module.prod_force_se_a (net_deriv_reshape, + self.descrpt_deriv, + self.nlist, + natoms, + n_a_sel = self.nnei_a, + n_r_sel = self.nnei_r) + virial, atom_virial \ + = op_module.prod_virial_se_a (net_deriv_reshape, + self.descrpt_deriv, + self.rij, + self.nlist, + natoms, + n_a_sel = self.nnei_a, + n_r_sel = self.nnei_r) + return force, virial, atom_virial + + + def _pass_filter(self, + inputs, + atype, + natoms, + input_dict, + reuse = None, + suffix = '', + trainable = True) : + start_index = 0 + inputs = tf.reshape(inputs, [-1, self.ndescrpt * natoms[0]]) + output = [] + output_qmat = [] + inputs_i = inputs + inputs_i = tf.reshape(inputs_i, [-1, self.ndescrpt]) + type_i = -1 + layer, qmat = self._filter(tf.cast(inputs_i, self.filter_precision), type_i, name='filter_type_all'+suffix, natoms=natoms, reuse=reuse, trainable = trainable, activation_fn = self.filter_activation_fn) + layer = tf.reshape(layer, [tf.shape(inputs)[0], natoms[0] * self.get_dim_out()]) + # qmat = tf.reshape(qmat, [tf.shape(inputs)[0], natoms[0] * self.get_dim_rot_mat_1() * 3]) + output.append(layer) + # output_qmat.append(qmat) + output = tf.concat(output, axis = 1) + # output_qmat = tf.concat(output_qmat, axis = 1) + return output, None + + + def _compute_dstats_sys_smth (self, + data_coord, + data_box, + data_atype, + natoms_vec, + mesh) : + dd_all \ + = self.sub_sess.run(self.stat_descrpt, + feed_dict = { + self.place_holders['coord']: data_coord, + self.place_holders['type']: data_atype, + self.place_holders['natoms_vec']: natoms_vec, + self.place_holders['box']: data_box, + self.place_holders['default_mesh']: mesh, + }) + natoms = natoms_vec + dd_all = np.reshape(dd_all, [-1, self.ndescrpt * natoms[0]]) + start_index = 0 + sysr = [] + sysa = [] + sysn = [] + sysr2 = [] + sysa2 = [] + for type_i in range(self.ntypes): + end_index = start_index + self.ndescrpt * natoms[2+type_i] + dd = dd_all[:, start_index:end_index] + dd = np.reshape(dd, [-1, self.ndescrpt]) + start_index = end_index + # compute + dd = np.reshape (dd, [-1, 4]) + ddr = dd[:,:1] + dda = dd[:,1:] + sumr = np.sum(ddr) + suma = np.sum(dda) / 3. + sumn = dd.shape[0] + sumr2 = np.sum(np.multiply(ddr, ddr)) + suma2 = np.sum(np.multiply(dda, dda)) / 3. + sysr.append(sumr) + sysa.append(suma) + sysn.append(sumn) + sysr2.append(sumr2) + sysa2.append(suma2) + return sysr, sysr2, sysa, sysa2, sysn + + + def _compute_std (self,sumv2, sumv, sumn) : + val = np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn)) + if np.abs(val) < 1e-2: + val = 1e-2 + return val + + + def _filter(self, + inputs, + type_input, + natoms, + activation_fn=tf.nn.tanh, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + trainable = True): + # natom x (nei x 4) + shape = inputs.get_shape().as_list() + outputs_size = [1] + self.filter_neuron + with tf.variable_scope(name, reuse=reuse): + start_index_i = 0 + result = None + for type_i in range(self.ntypes): + # cut-out inputs + # with natom x (nei_type_i x 4) + inputs_i = tf.slice (inputs, + [ 0, start_index_i *4], + [-1, self.sel_a[type_i] *4] ) + start_index_i += self.sel_a[type_i] + nei_type_i = self.sel_a[type_i] + shape_i = inputs_i.get_shape().as_list() + assert(shape_i[1] == nei_type_i * 4) + # with natom x nei_type_i x 4 + env_i = tf.reshape(inputs_i, [-1, nei_type_i, 4]) + # with natom x nei_type_i x 3 + env_i = tf.slice(env_i, [0, 0, 1], [-1, -1, -1]) + start_index_j = 0 + for type_j in range(type_i, self.ntypes): + # with natom x (nei_type_j x 4) + inputs_j = tf.slice (inputs, + [ 0, start_index_j *4], + [-1, self.sel_a[type_j] *4] ) + start_index_j += self.sel_a[type_j] + nei_type_j = self.sel_a[type_j] + shape_j = inputs_j.get_shape().as_list() + assert(shape_j[1] == nei_type_j * 4) + # with natom x nei_type_j x 4 + env_j = tf.reshape(inputs_j, [-1, nei_type_j, 4]) + # with natom x nei_type_i x 3 + env_j = tf.slice(env_j, [0, 0, 1], [-1, -1, -1]) + # with natom x nei_type_i x nei_type_j + env_ij = tf.einsum('ijm,ikm->ijk', env_i, env_j) + # with (natom x nei_type_i x nei_type_j) + ebd_env_ij = tf.reshape(env_ij, [-1, 1]) + # with (natom x nei_type_i x nei_type_j) x out_size + ebd_env_ij = embedding_net(ebd_env_ij, + self.filter_neuron, + self.filter_precision, + activation_fn = activation_fn, + resnet_dt = self.filter_resnet_dt, + name_suffix = f"_{type_i}_{type_j}", + stddev = stddev, + bavg = bavg, + seed = self.seed, + trainable = trainable, + uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # with natom x nei_type_i x nei_type_j x out_size + ebd_env_ij = tf.reshape(ebd_env_ij, [-1, nei_type_i, nei_type_j, outputs_size[-1]]) + # with natom x out_size + res_ij = tf.einsum('ijk,ijkm->im', env_ij, ebd_env_ij) + res_ij = res_ij * (1.0 / float(nei_type_i) / float(nei_type_j)) + if result is None: + result = res_ij + else: + result += res_ij + return result, None + diff --git a/deepmd/entrypoints/__init__.py b/deepmd/entrypoints/__init__.py new file mode 100644 index 0000000000..2f8bae9bbd --- /dev/null +++ b/deepmd/entrypoints/__init__.py @@ -0,0 +1,20 @@ +"""Submodule that contains all the DeePMD-Kit entry point scripts.""" + +from .compress import compress +from .config import config +from .doc import doc_train_input +from .freeze import freeze +from .test import test +from .train import train +from .transfer import transfer + +__all__ = [ + "config", + "doc_train_input", + "freeze", + "test", + "train", + "transfer", + "compress", + "doc_train_input", +] diff --git a/deepmd/entrypoints/compress.py b/deepmd/entrypoints/compress.py new file mode 100644 index 0000000000..2193f0c1f9 --- /dev/null +++ b/deepmd/entrypoints/compress.py @@ -0,0 +1,115 @@ +"""Compress a model, which including tabulating the embedding-net.""" + +import json +import logging +from typing import Optional + +from deepmd.common import j_loader +from deepmd.utils.argcheck import normalize +from deepmd.utils.compat import convert_input_v0_v1 + +from .freeze import freeze +from .train import train +from .transfer import transfer + +__all__ = ["compress"] + +log = logging.getLogger(__name__) + + +def compress( + *, + INPUT: str, + input: str, + output: str, + extrapolate: int, + stride: float, + frequency: str, + checkpoint_folder: str, + mpi_log: str, + log_path: Optional[str], + log_level: int, + **kwargs +): + """Compress model. + + The table is composed of fifth-order polynomial coefficients and is assembled from + two sub-tables. The first table takes the stride(parameter) as it's uniform stride, + while the second table takes 10 * stride as it's uniform stride. The range of the + first table is automatically detected by deepmd-kit, while the second table ranges + from the first table's upper boundary(upper) to the extrapolate(parameter) * upper. + + Parameters + ---------- + INPUT : str + input json/yaml control file + input : str + frozen model file to compress + output : str + compressed model filename + extrapolate : int + scale of model extrapolation + stride : float + uniform stride of tabulation's first table + frequency : str + frequency of tabulation overflow check + checkpoint_folder : str + trining checkpoint folder for freezing + mpi_log : str + mpi logging mode for training + log_path : Optional[str] + if speccified log will be written to this file + log_level : int + logging level + """ + jdata = j_loader(INPUT) + if "model" not in jdata.keys(): + jdata = convert_input_v0_v1(jdata, warning=True, dump="input_v1_compat.json") + jdata["model"]["compress"] = {} + jdata["model"]["compress"]["type"] = 'se_e2_a' + jdata["model"]["compress"]["compress"] = True + jdata["model"]["compress"]["model_file"] = input + jdata["model"]["compress"]["table_config"] = [ + extrapolate, + stride, + 10 * stride, + int(frequency), + ] + # be careful here, if one want to refine the model + jdata["training"]["numb_steps"] = jdata["training"]["save_freq"] + jdata = normalize(jdata) + + + # check the descriptor info of the input file + assert ( + jdata["model"]["descriptor"]["type"] == "se_a" or jdata["model"]["descriptor"]["type"] == "se_e2_a" + ), "Model compression error: descriptor type must be se_a or se_e2_a!" + assert ( + jdata["model"]["descriptor"]["resnet_dt"] is False + ), "Model compression error: descriptor resnet_dt must be false!" + + # stage 1: training or refining the model with tabulation + log.info("\n\n") + log.info("stage 1: train or refine the model with tabulation") + control_file = "compress.json" + with open(control_file, "w") as fp: + json.dump(jdata, fp, indent=4) + train( + INPUT=control_file, + init_model=None, + restart=None, + output=control_file, + mpi_log=mpi_log, + log_level=log_level, + log_path=log_path, + ) + + # stage 2: freeze the model + log.info("\n\n") + log.info("stage 2: freeze the model") + freeze(checkpoint_folder=checkpoint_folder, output=output, node_names=None) + + # stage 3: transfer the model + log.info("\n\n") + log.info("stage 3: transfer the model") + transfer(old_model=input, raw_model=output, output=output) diff --git a/deepmd/entrypoints/config.py b/deepmd/entrypoints/config.py new file mode 100644 index 0000000000..0c98bb5f82 --- /dev/null +++ b/deepmd/entrypoints/config.py @@ -0,0 +1,354 @@ +#!/usr/bin/env python3 +"""Quickly create a configuration file for smooth model.""" + +import json +import yaml +from pathlib import Path +from typing import Any, Dict, List, Tuple + +import numpy as np + +__all__ = ["config"] + + +DEFAULT_DATA: Dict[str, Any] = { + "use_smooth": True, + "sel_a": [], + "rcut_smth": -1, + "rcut": -1, + "filter_neuron": [20, 40, 80], + "filter_resnet_dt": False, + "axis_neuron": 8, + "fitting_neuron": [240, 240, 240], + "fitting_resnet_dt": True, + "coord_norm": True, + "type_fitting_net": False, + "systems": [], + "set_prefix": "set", + "stop_batch": -1, + "batch_size": -1, + "start_lr": 0.001, + "decay_steps": -1, + "decay_rate": 0.95, + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "seed": 1, + "disp_file": "lcurve.out", + "disp_freq": 1000, + "numb_test": 10, + "save_freq": 10000, + "save_ckpt": "model.ckpt", + "disp_training": True, + "time_training": True, +} + + +def valid_dir(path: Path): + """Check if directory is a valid deepmd system directory. + + Parameters + ---------- + path : Path + path to directory + + Raises + ------ + OSError + if `type.raw` is missing on dir or `box.npy` or `coord.npy` are missing in one + of the sets subdirs + """ + if not (path / "type.raw").is_file(): + raise OSError + for ii in path.glob("set.*"): + if not (ii / "box.npy").is_file(): + raise OSError + if not (ii / "coord.npy").is_file(): + raise OSError + + +def load_systems(dirs: List[Path]) -> Tuple[List[np.ndarray], List[np.ndarray]]: + """Load systems to memory for disk. + + Parameters + ---------- + dirs : List[Path] + list of system directories paths + + Returns + ------- + Tuple[List[np.ndarray], List[np.ndarray]] + atoms types and structure cells formated as Nx9 array + """ + all_type = [] + all_box = [] + for d in dirs: + sys_type = np.loadtxt(d / "type.raw", dtype=int) + sys_box = np.vstack([np.load(s / "box.npy") for s in d.glob("set.*")]) + all_type.append(sys_type) + all_box.append(sys_box) + return all_type, all_box + + +def get_system_names() -> List[Path]: + """Get system directory paths from stdin. + + Returns + ------- + List[Path] + list of system directories paths + """ + dirs = input("Enter system path(s) (seperated by space, wild card supported): \n") + system_dirs = [] + for dir_str in dirs.split(): + found_dirs = Path.cwd().glob(dir_str) + for d in found_dirs: + valid_dir(d) + system_dirs.append(d) + + return system_dirs + + +def get_rcut() -> float: + """Get rcut from stdin from user. + + Returns + ------- + float + input rcut lenght converted to float + + Raises + ------ + ValueError + if rcut is smaller than 0.0 + """ + dv = 6.0 + rcut_input = input(f"Enter rcut (default {dv:.1f} A): \n") + try: + rcut = float(rcut_input) + except ValueError as e: + print(f"invalid rcut: {e} setting to default: {dv:.1f}") + rcut = dv + if rcut <= 0: + raise ValueError("rcut should be > 0") + return rcut + + +def get_batch_size_rule() -> int: + """Get minimal batch size from user from stdin. + + Returns + ------- + int + size of the batch + + Raises + ------ + ValueError + if batch size is <= 0 + """ + dv = 32 + matom_input = input( + f"Enter the minimal number of atoms in a batch (default {dv:d}: \n" + ) + try: + matom = int(matom_input) + except ValueError as e: + print(f"invalid batch size: {e} setting to default: {dv:d}") + matom = dv + if matom <= 0: + raise ValueError("the number should be > 0") + return matom + + +def get_stop_batch() -> int: + """Get stop batch from user from stdin. + + Returns + ------- + int + size of the batch + + Raises + ------ + ValueError + if stop batch is <= 0 + """ + dv = 1000000 + sb_input = input(f"Enter the stop batch (default {dv:d}): \n") + try: + sb = int(sb_input) + except ValueError as e: + print(f"invalid stop batch: {e} setting to default: {dv:d}") + sb = dv + if sb <= 0: + raise ValueError("the number should be > 0") + return sb + + +def get_ntypes(all_type: List[np.ndarray]) -> int: + """Count number of unique elements. + + Parameters + ---------- + all_type : List[np.ndarray] + list with arrays specifying elements of structures + + Returns + ------- + int + number of unique elements + """ + return len(np.unique(all_type)) + + +def get_max_density( + all_type: List[np.ndarray], all_box: List[np.ndarray] +) -> np.ndarray: + """Compute maximum density in suppliedd cells. + + Parameters + ---------- + all_type : List[np.ndarray] + list with arrays specifying elements of structures + all_box : List[np.ndarray] + list with arrays specifying cells for all structures + + Returns + ------- + float + maximum atom density in all supplies structures for each element individually + """ + ntypes = get_ntypes(all_type) + all_max = [] + for tt, bb in zip(all_type, all_box): + vv = np.reshape(bb, [-1, 3, 3]) + vv = np.linalg.det(vv) + min_v = np.min(vv) + type_count = [] + for ii in range(ntypes): + type_count.append(sum(tt == ii)) + max_den = type_count / min_v + all_max.append(max_den) + all_max = np.max(all_max, axis=0) + return all_max + + +def suggest_sel( + all_type: List[np.ndarray], + all_box: List[np.ndarray], + rcut: float, + ratio: float = 1.5, +) -> List[int]: + """Suggest selection parameter. + + Parameters + ---------- + all_type : List[np.ndarray] + list with arrays specifying elements of structures + all_box : List[np.ndarray] + list with arrays specifying cells for all structures + rcut : float + cutoff radius + ratio : float, optional + safety margin to add to estimated value, by default 1.5 + + Returns + ------- + List[int] + [description] + """ + max_den = get_max_density(all_type, all_box) + return [int(ii) for ii in max_den * 4.0 / 3.0 * np.pi * rcut ** 3 * ratio] + + +def suggest_batch_size(all_type: List[np.ndarray], min_atom: int) -> List[int]: + """Get suggestion for batch size. + + Parameters + ---------- + all_type : List[np.ndarray] + list with arrays specifying elements of structures + min_atom : int + minimal number of atoms in batch + + Returns + ------- + List[int] + suggested batch sizes for each system + """ + bs = [] + for ii in all_type: + natoms = len(ii) + tbs = min_atom // natoms + if (min_atom // natoms) * natoms != min_atom: + tbs += 1 + bs.append(tbs) + return bs + + +def suggest_decay(stop_batch: int) -> Tuple[int, float]: + """Suggest number of decay steps and decay rate. + + Parameters + ---------- + stop_batch : int + stop batch number + + Returns + ------- + Tuple[int, float] + number of decay steps and decay rate + """ + decay_steps = int(stop_batch // 200) + decay_rate = 0.95 + return decay_steps, decay_rate + + +def config(*, output: str, **kwargs): + """Auto config file generator. + + Parameters + ---------- + output: str + file to write config file + + Raises + ------ + RuntimeError + if user does not input any systems + ValueError + if output file is of wrong type + """ + all_sys = get_system_names() + if len(all_sys) == 0: + raise RuntimeError("no system specified") + rcut = get_rcut() + matom = get_batch_size_rule() + stop_batch = get_stop_batch() + + all_type, all_box = load_systems(all_sys) + sel = suggest_sel(all_type, all_box, rcut, ratio=1.5) + bs = suggest_batch_size(all_type, matom) + decay_steps, decay_rate = suggest_decay(stop_batch) + + jdata = DEFAULT_DATA.copy() + jdata["systems"] = [str(ii) for ii in all_sys] + jdata["sel_a"] = sel + jdata["rcut"] = rcut + jdata["rcut_smth"] = rcut - 0.2 + jdata["stop_batch"] = stop_batch + jdata["batch_size"] = bs + jdata["decay_steps"] = decay_steps + jdata["decay_rate"] = decay_rate + + with open(output, "w") as fp: + if output.endswith("json"): + json.dump(jdata, fp, indent=4) + elif output.endswith(("yml", "yaml")): + yaml.safe_dump(jdata, fp, default_flow_style=False) + else: + raise ValueError("output file must be of type json or yaml") diff --git a/deepmd/entrypoints/doc.py b/deepmd/entrypoints/doc.py new file mode 100644 index 0000000000..941efd61c2 --- /dev/null +++ b/deepmd/entrypoints/doc.py @@ -0,0 +1,11 @@ +"""Module that prints train input arguments docstrings.""" + +from deepmd.utils.argcheck import gen_doc + +__all__ = ["doc_train_input"] + + +def doc_train_input(): + """Print out trining input arguments to console.""" + doc_str = gen_doc(make_anchor=True) + print(doc_str) diff --git a/deepmd/entrypoints/freeze.py b/deepmd/entrypoints/freeze.py new file mode 100755 index 0000000000..e11ac0c906 --- /dev/null +++ b/deepmd/entrypoints/freeze.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Script for freezing TF trained graph so it can be used with LAMMPS and i-PI. + +References +---------- +https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc +""" + +from deepmd.env import tf +from deepmd.env import op_module +from os.path import abspath + +# load grad of force module +import deepmd.op + +from typing import List, Optional + +__all__ = ["freeze"] + + +def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> List[str]: + """Get node names based on model type. + + Parameters + ---------- + model_type : str + str type of model + modifier_type : Optional[str], optional + modifier type if any, by default None + + Returns + ------- + List[str] + list with all node names to freeze + + Raises + ------ + RuntimeError + if unknown model type + """ + nodes = [ + "descrpt_attr/rcut", + "descrpt_attr/ntypes", + "model_attr/tmap", + "model_attr/model_type", + "model_attr/model_version", + ] + + if model_type == "ener": + nodes += [ + "o_energy", + "o_force", + "o_virial", + "o_atom_energy", + "o_atom_virial", + "fitting_attr/dfparam", + "fitting_attr/daparam", + ] + elif model_type == "wfc": + nodes += [ + "o_wfc", + "model_attr/sel_type", + "model_attr/output_dim", + ] + elif model_type == "dipole": + nodes += [ + "o_dipole", + "o_rmat", + "o_rmat_deriv", + "o_nlist", + "o_rij", + "descrpt_attr/sel", + "descrpt_attr/ndescrpt", + "model_attr/sel_type", + "model_attr/output_dim", + ] + elif model_type == "polar": + nodes += [ + "o_polar", + "model_attr/sel_type", + "model_attr/output_dim", + ] + elif model_type == "global_polar": + nodes += [ + "o_global_polar", + "model_attr/sel_type", + "model_attr/output_dim", + ] + else: + raise RuntimeError(f"unknow model type {model_type}") + if modifier_type == "dipole_charge": + nodes += [ + "modifier_attr/type", + "modifier_attr/mdl_name", + "modifier_attr/mdl_charge_map", + "modifier_attr/sys_charge_map", + "modifier_attr/ewald_h", + "modifier_attr/ewald_beta", + "dipole_charge/descrpt_attr/rcut", + "dipole_charge/descrpt_attr/ntypes", + "dipole_charge/model_attr/tmap", + "dipole_charge/model_attr/model_type", + "o_dm_force", + "dipole_charge/model_attr/sel_type", + "dipole_charge/o_dipole", + "dipole_charge/model_attr/output_dim", + "o_dm_virial", + "o_dm_av", + ] + return nodes + + +def freeze( + *, checkpoint_folder: str, output: str, node_names: Optional[str] = None, **kwargs +): + """Freeze the graph in supplied folder. + + Parameters + ---------- + checkpoint_folder : str + location of the folder with model + output : str + output file name + node_names : Optional[str], optional + names of nodes to output, by default None + """ + # We retrieve our checkpoint fullpath + checkpoint = tf.train.get_checkpoint_state(checkpoint_folder) + input_checkpoint = checkpoint.model_checkpoint_path + + # expand the output file to full path + output_graph = abspath(output) + + # Before exporting our graph, we need to precise what is our output node + # This is how TF decides what part of the Graph he has to keep + # and what part it can dump + # NOTE: this variable is plural, because you can have multiple output nodes + # node_names = "energy_test,force_test,virial_test,t_rcut" + + # We clear devices to allow TensorFlow to control + # on which device it will load operations + clear_devices = True + + # We import the meta graph and retrieve a Saver + saver = tf.train.import_meta_graph( + f"{input_checkpoint}.meta", clear_devices=clear_devices + ) + + # We retrieve the protobuf graph definition + graph = tf.get_default_graph() + input_graph_def = graph.as_graph_def() + nodes = [n.name for n in input_graph_def.node] + + # We start a session and restore the graph weights + with tf.Session() as sess: + saver.restore(sess, input_checkpoint) + model_type = sess.run("model_attr/model_type:0", feed_dict={}).decode("utf-8") + if "modifier_attr/type" in nodes: + modifier_type = sess.run("modifier_attr/type:0", feed_dict={}).decode( + "utf-8" + ) + else: + modifier_type = None + if node_names is None: + output_node_list = _make_node_names(model_type, modifier_type) + else: + output_node_list = node_names.split(",") + print(f"The following nodes will be frozen: {output_node_list}") + + # We use a built-in TF helper to export variables to constants + output_graph_def = tf.graph_util.convert_variables_to_constants( + sess, # The session is used to retrieve the weights + input_graph_def, # The graph_def is used to retrieve the nodes + output_node_list, # The output node names are used to select the usefull nodes + ) + + # Finally we serialize and dump the output graph to the filesystem + with tf.gfile.GFile(output_graph, "wb") as f: + f.write(output_graph_def.SerializeToString()) + print(f"{len(output_graph_def.node):d} ops in the final graph.") diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py new file mode 100644 index 0000000000..23bb30ccde --- /dev/null +++ b/deepmd/entrypoints/main.py @@ -0,0 +1,358 @@ +"""DeePMD-Kit entry point module.""" + +import argparse +import logging +from pathlib import Path +from typing import List, Optional + +from deepmd.entrypoints import ( + compress, + config, + doc_train_input, + freeze, + test, + train, + transfer, +) +from deepmd.loggers import set_log_handles + +__all__ = ["main", "parse_args", "get_ll"] + + +def get_ll(log_level: str) -> int: + """Convert string to python logging level. + + Parameters + ---------- + log_level : str + allowed input values are: DEBUG, INFO, WARNING, ERROR, 3, 2, 1, 0 + + Returns + ------- + int + one of python logging module log levels - 10, 20, 30 or 40 + """ + if log_level.isdigit(): + int_level = (4 - int(log_level)) * 10 + else: + int_level = getattr(logging, log_level) + + return int_level + + +def parse_args(args: Optional[List[str]] = None): + """DeePMD-Kit commandline options argument parser. + + Parameters + ---------- + args: List[str] + list of command line arguments, main purpose is testing default option None + takes arguments from sys.argv + """ + parser = argparse.ArgumentParser( + description="DeePMD-kit: A deep learning package for many-body potential energy" + " representation and molecular dynamics", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + subparsers = parser.add_subparsers(title="Valid subcommands", dest="command") + + # * logging options parser ********************************************************* + # with use of the parent argument this options will be added to every parser + parser_log = argparse.ArgumentParser( + add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser_log.add_argument( + "-v", + "--log-level", + choices=["DEBUG", "3", "INFO", "2", "WARNING", "1", "ERROR", "0"], + default="INFO", + help="set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO " + "and 3=DEBUG", + ) + parser_log.add_argument( + "-l", + "--log-path", + type=str, + default=None, + help="set log file to log messages to disk, if not specified, the logs will " + "only be output to console", + ) + # * mpi logging parser ************************************************************* + parser_mpi_log = argparse.ArgumentParser( + add_help=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser_mpi_log.add_argument( + "-m", + "--mpi-log", + type=str, + default="master", + choices=("master", "collect", "workers"), + help="Set the manner of logging when running with MPI. 'master' logs only on " + "main process, 'collect' broadcasts logs from workers to master and 'workers' " + "means each process will output its own log", + ) + + # * config script ****************************************************************** + parser_cfig = subparsers.add_parser( + "config", + parents=[parser_log], + help="fast configuration of parameter file for smooth model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser_cfig.add_argument( + "-o", "--output", type=str, default="input.json", help="the output json file" + ) + + # * transfer script **************************************************************** + parser_transfer = subparsers.add_parser( + "transfer", parents=[parser_log], help="pass parameters to another model" + ) + parser_transfer.add_argument( + "-r", + "--raw-model", + default="raw_frozen_model.pb", + type=str, + help="the model receiving parameters", + ) + parser_transfer.add_argument( + "-O", + "--old-model", + default="old_frozen_model.pb", + type=str, + help="the model providing parameters", + ) + parser_transfer.add_argument( + "-o", + "--output", + default="frozen_model.pb", + type=str, + help="the model after passing parameters", + ) + + # * config parser ****************************************************************** + parser_train = subparsers.add_parser( + "train", + parents=[parser_log, parser_mpi_log], + help="train a model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser_train.add_argument( + "INPUT", help="the input parameter file in json or yaml format" + ) + parser_train.add_argument( + "-i", + "--init-model", + type=str, + default=None, + help="Initialize the model by the provided checkpoint.", + ) + parser_train.add_argument( + "-r", + "--restart", + type=str, + default=None, + help="Restart the training from the provided checkpoint.", + ) + parser_train.add_argument( + "-o", + "--output", + type=str, + default="out.json", + help="The output file of the parameters used in training.", + ) + + # * freeze script ****************************************************************** + parser_frz = subparsers.add_parser( + "freeze", + parents=[parser_log], + help="freeze the model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser_frz.add_argument( + "-c", + "--checkpoint-folder", + type=str, + default=".", + help="path to checkpoint folder", + ) + parser_frz.add_argument( + "-o", + "--output", + type=str, + default="frozen_model.pb", + help="name of graph, will output to the checkpoint folder", + ) + parser_frz.add_argument( + "-n", + "--node-names", + type=str, + default=None, + help="the frozen nodes, if not set, determined from the model type", + ) + + # * test script ******************************************************************** + parser_tst = subparsers.add_parser( + "test", + parents=[parser_log], + help="test the model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser_tst.add_argument( + "-m", + "--model", + default="frozen_model.pb", + type=str, + help="Frozen model file to import", + ) + parser_tst.add_argument( + "-s", + "--system", + default=".", + type=str, + help="The system dir. Recursively detect systems in this directory", + ) + parser_tst.add_argument( + "-S", "--set-prefix", default="set", type=str, help="The set prefix" + ) + parser_tst.add_argument( + "-n", "--numb-test", default=100, type=int, help="The number of data for test" + ) + parser_tst.add_argument( + "-r", "--rand-seed", type=int, default=None, help="The random seed" + ) + parser_tst.add_argument( + "--shuffle-test", action="store_true", default=False, help="Shuffle test data" + ) + parser_tst.add_argument( + "-d", + "--detail-file", + type=str, + default=None, + help="File where details of energy force and virial accuracy will be written", + ) + parser_tst.add_argument( + "-a", + "--atomic", + action="store_true", + default=False, + help="Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)", + ) + + # * compress model ***************************************************************** + # Compress a model, which including tabulating the embedding-net. + # The table is composed of fifth-order polynomial coefficients and is assembled + # from two sub-tables. The first table takes the stride(parameter) as it's uniform + # stride, while the second table takes 10 * stride as it\s uniform stride + #  The range of the first table is automatically detected by deepmd-kit, while the + # second table ranges from the first table's upper boundary(upper) to the + # extrapolate(parameter) * upper. + parser_compress = subparsers.add_parser( + "compress", + parents=[parser_log, parser_mpi_log], + help="compress a model", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser_compress.add_argument( + "INPUT", + help="The input parameter file in json or yaml format, which should be " + "consistent with the original model parameter file", + ) + parser_compress.add_argument( + "-i", + "--input", + default="frozen_model.pb", + type=str, + help="The original frozen model, which will be compressed by the deepmd-kit", + ) + parser_compress.add_argument( + "-o", + "--output", + default="frozen_model_compress.pb", + type=str, + help="The compressed model", + ) + parser_compress.add_argument( + "-e", + "--extrapolate", + default=5, + type=int, + help="The scale of model extrapolation", + ) + parser_compress.add_argument( + "-s", + "--stride", + default=0.01, + type=float, + help="The uniform stride of tabulation's first table, the second table will " + "use 10 * stride as it's uniform stride", + ) + parser_compress.add_argument( + "-f", + "--frequency", + default=-1, + type=int, + help="The frequency of tabulation overflow check(If the input environment " + "matrix overflow the first or second table range). " + "By default do not check the overflow", + ) + parser_compress.add_argument( + "-c", + "--checkpoint-folder", + type=str, + default=".", + help="path to checkpoint folder", + ) + + # * print docs script ************************************************************** + subparsers.add_parser( + "doc-train-input", + parents=[parser_log], + help="print the documentation (in rst format) of input training parameters.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + parsed_args = parser.parse_args(args=args) + if parsed_args.command is None: + parser.print_help() + else: + parsed_args.log_level = get_ll(parsed_args.log_level) + + return parsed_args + + +def main(): + """DeePMD-Kit entry point. + + Raises + ------ + RuntimeError + if no command was input + """ + args = parse_args() + + # do not set log handles for None, it is useless + # log handles for train will be set separatelly + # when the use of MPI will be determined in `RunOptions` + if args.command not in (None, "train"): + set_log_handles(args.log_level, Path(args.log_path) if args.log_path else None) + + dict_args = vars(args) + + if args.command == "train": + train(**dict_args) + elif args.command == "freeze": + freeze(**dict_args) + elif args.command == "config": + config(**dict_args) + elif args.command == "test": + test(**dict_args) + elif args.command == "transfer": + transfer(**dict_args) + elif args.command == "compress": + compress(**dict_args) + elif args.command == "doc-train-input": + doc_train_input() + elif args.command is None: + pass + else: + raise RuntimeError(f"unknown command {args.command}") diff --git a/deepmd/entrypoints/test.py b/deepmd/entrypoints/test.py new file mode 100644 index 0000000000..d1284641a9 --- /dev/null +++ b/deepmd/entrypoints/test.py @@ -0,0 +1,606 @@ +"""Test trained DeePMD model.""" +import logging +from pathlib import Path +from typing import TYPE_CHECKING, List, Dict, Optional, Tuple + +import numpy as np +from deepmd import DeepPotential +from deepmd.common import expand_sys_str +from deepmd.utils.data import DeepmdData +from deepmd.utils.weight_avg import weighted_average + +if TYPE_CHECKING: + from deepmd.infer import DeepDipole, DeepPolar, DeepPot, DeepWFC + from deepmd.infer.deep_eval import DeepTensor + +__all__ = ["test"] + +log = logging.getLogger(__name__) + + +def test( + *, + model: str, + system: str, + set_prefix: str, + numb_test: int, + rand_seed: Optional[int], + shuffle_test: bool, + detail_file: str, + atomic: bool, + **kwargs, +): + """Test model predictions. + + Parameters + ---------- + model : str + path where model is stored + system : str + system directory + set_prefix : str + string prefix of set + numb_test : int + munber of tests to do + rand_seed : Optional[int] + seed for random generator + shuffle_test : bool + whether to shuffle tests + detail_file : Optional[str] + file where test details will be output + atomic : bool + whether per atom quantities should be computed + + Raises + ------ + RuntimeError + if no valid system was found + """ + all_sys = expand_sys_str(system) + if len(all_sys) == 0: + raise RuntimeError("Did not find valid system") + err_coll = [] + siz_coll = [] + + # init random seed + if rand_seed is not None: + np.random.seed(rand_seed % (2 ** 32)) + + # init model + dp = DeepPotential(model) + + for cc, system in enumerate(all_sys): + log.info("# ---------------output of dp test--------------- ") + log.info(f"# testing system : {system}") + + # create data class + tmap = dp.get_type_map() if dp.model_type == "ener" else None + data = DeepmdData(system, set_prefix, shuffle_test=shuffle_test, type_map=tmap) + + if dp.model_type == "ener": + err = test_ener( + dp, + data, + system, + numb_test, + detail_file, + atomic, + append_detail=(cc != 0), + ) + elif dp.model_type == "dipole": + err = test_dipole(dp, data, numb_test, detail_file, atomic) + elif dp.model_type == "polar": + err = test_polar(dp, data, numb_test, detail_file, atomic=atomic) + elif dp.model_type == "global_polar": # should not appear in this new version + log.warning("Global polar model is not currently supported. Please directly use the polar mode and change loss parameters.") + err = test_polar(dp, data, numb_test, detail_file, atomic=False) # YWolfeee: downward compatibility + log.info("# ----------------------------------------------- ") + err_coll.append(err) + + avg_err = weighted_average(err_coll) + + if len(all_sys) != len(err_coll): + log.warning("Not all systems are tested! Check if the systems are valid") + + if len(all_sys) > 1: + log.info("# ----------weighted average of errors----------- ") + log.info(f"# number of systems : {len(all_sys)}") + if dp.model_type == "ener": + print_ener_sys_avg(avg_err) + elif dp.model_type == "dipole": + print_dipole_sys_avg(avg_err) + elif dp.model_type == "polar": + print_polar_sys_avg(avg_err) + elif dp.model_type == "global_polar": + print_polar_sys_avg(avg_err) + elif dp.model_type == "wfc": + print_wfc_sys_avg(avg_err) + log.info("# ----------------------------------------------- ") + + +def rmse(diff: np.ndarray) -> np.ndarray: + """Calculate average root mean square error. + + Parameters + ---------- + diff: np.ndarray + difference + + Returns + ------- + np.ndarray + array with normalized difference + """ + return np.sqrt(np.average(diff * diff)) + + +def save_txt_file( + fname: Path, data: np.ndarray, header: str = "", append: bool = False +): + """Save numpy array to test file. + + Parameters + ---------- + fname : str + filename + data : np.ndarray + data to save to disk + header : str, optional + header string to use in file, by default "" + append : bool, optional + if true file will be appended insted of overwriting, by default False + """ + flags = "ab" if append else "w" + with fname.open(flags) as fp: + np.savetxt(fp, data, header=header) + + +def test_ener( + dp: "DeepPot", + data: DeepmdData, + system: str, + numb_test: int, + detail_file: Optional[str], + has_atom_ener: bool, + append_detail: bool = False, +) -> Tuple[List[np.ndarray], List[int]]: + """Test energy type model. + + Parameters + ---------- + dp : DeepPot + instance of deep potential + data: DeepmdData + data container object + system : str + system directory + numb_test : int + munber of tests to do + detail_file : Optional[str] + file where test details will be output + has_atom_ener : bool + whether per atom quantities should be computed + append_detail : bool, optional + if true append output detail file, by default False + + Returns + ------- + Tuple[List[np.ndarray], List[int]] + arrays with results and their shapes + """ + data.add("energy", 1, atomic=False, must=False, high_prec=True) + data.add("force", 3, atomic=True, must=False, high_prec=False) + data.add("virial", 9, atomic=False, must=False, high_prec=False) + if dp.has_efield: + data.add("efield", 3, atomic=True, must=True, high_prec=False) + if has_atom_ener: + data.add("atom_ener", 1, atomic=True, must=True, high_prec=False) + if dp.get_dim_fparam() > 0: + data.add( + "fparam", dp.get_dim_fparam(), atomic=False, must=True, high_prec=False + ) + if dp.get_dim_aparam() > 0: + data.add("aparam", dp.get_dim_aparam(), atomic=True, must=True, high_prec=False) + + test_data = data.get_test() + natoms = len(test_data["type"][0]) + nframes = test_data["box"].shape[0] + numb_test = min(nframes, numb_test) + + coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) + box = test_data["box"][:numb_test] + if dp.has_efield: + efield = test_data["efield"][:numb_test].reshape([numb_test, -1]) + else: + efield = None + if not data.pbc: + box = None + atype = test_data["type"][0] + if dp.get_dim_fparam() > 0: + fparam = test_data["fparam"][:numb_test] + else: + fparam = None + if dp.get_dim_aparam() > 0: + aparam = test_data["aparam"][:numb_test] + else: + aparam = None + + ret = dp.eval( + coord, + box, + atype, + fparam=fparam, + aparam=aparam, + atomic=has_atom_ener, + efield=efield, + ) + energy = ret[0] + force = ret[1] + virial = ret[2] + energy = energy.reshape([numb_test, 1]) + force = force.reshape([numb_test, -1]) + virial = virial.reshape([numb_test, 9]) + if has_atom_ener: + ae = ret[3] + av = ret[4] + ae = ae.reshape([numb_test, -1]) + av = av.reshape([numb_test, -1]) + + rmse_e = rmse(energy - test_data["energy"][:numb_test].reshape([-1, 1])) + rmse_f = rmse(force - test_data["force"][:numb_test]) + rmse_v = rmse(virial - test_data["virial"][:numb_test]) + rmse_ea = rmse_e / natoms + rmse_va = rmse_v / natoms + if has_atom_ener: + rmse_ae = rmse( + test_data["atom_ener"][:numb_test].reshape([-1]) - ae.reshape([-1]) + ) + + # print ("# energies: %s" % energy) + log.info(f"# number of test data : {numb_test:d} ") + log.info(f"Energy RMSE : {rmse_e:e} eV") + log.info(f"Energy RMSE/Natoms : {rmse_ea:e} eV") + log.info(f"Force RMSE : {rmse_f:e} eV/A") + log.info(f"Virial RMSE : {rmse_v:e} eV") + log.info(f"Virial RMSE/Natoms : {rmse_va:e} eV") + if has_atom_ener: + log.info(f"Atomic ener RMSE : {rmse_ae:e} eV") + + if detail_file is not None: + detail_path = Path(detail_file) + + pe = np.concatenate( + ( + np.reshape(test_data["energy"][:numb_test], [-1, 1]), + np.reshape(energy, [-1, 1]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".e.out"), + pe, + header="%s: data_e pred_e" % system, + append=append_detail, + ) + pf = np.concatenate( + ( + np.reshape(test_data["force"][:numb_test], [-1, 3]), + np.reshape(force, [-1, 3]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".f.out"), + pf, + header="%s: data_fx data_fy data_fz pred_fx pred_fy pred_fz" % system, + append=append_detail, + ) + pv = np.concatenate( + ( + np.reshape(test_data["virial"][:numb_test], [-1, 9]), + np.reshape(virial, [-1, 9]), + ), + axis=1, + ) + save_txt_file( + detail_path.with_suffix(".v.out"), + pv, + header=f"{system}: data_vxx data_vxy data_vxz data_vyx data_vyy " + "data_vyz data_vzx data_vzy data_vzz pred_vxx pred_vxy pred_vxz pred_vyx " + "pred_vyy pred_vyz pred_vzx pred_vzy pred_vzz", + append=append_detail, + ) + return { + "rmse_ea" : (rmse_ea, energy.size), + "rmse_f" : (rmse_f, force.size), + "rmse_va" : (rmse_va, virial.size), + } + + +def print_ener_sys_avg(avg: Dict[str,float]): + """Print errors summary for energy type potential. + + Parameters + ---------- + avg : np.ndarray + array with summaries + """ + log.info(f"Energy RMSE/Natoms : {avg['rmse_ea']:e} eV") + log.info(f"Force RMSE : {avg['rmse_f']:e} eV/A") + log.info(f"Virial RMSE/Natoms : {avg['rmse_va']:e} eV") + + +def run_test(dp: "DeepTensor", test_data: dict, numb_test: int): + """Run tests. + + Parameters + ---------- + dp : DeepTensor + instance of deep potential + test_data : dict + dictionary with test data + numb_test : int + munber of tests to do + + Returns + ------- + [type] + [description] + """ + nframes = test_data["box"].shape[0] + numb_test = min(nframes, numb_test) + + coord = test_data["coord"][:numb_test].reshape([numb_test, -1]) + box = test_data["box"][:numb_test] + atype = test_data["type"][0] + prediction = dp.eval(coord, box, atype) + + return prediction.reshape([numb_test, -1]), numb_test, atype + + +def test_wfc( + dp: "DeepWFC", + data: DeepmdData, + numb_test: int, + detail_file: Optional[str], +) -> Tuple[List[np.ndarray], List[int]]: + """Test energy type model. + + Parameters + ---------- + dp : DeepPot + instance of deep potential + data: DeepmdData + data container object + numb_test : int + munber of tests to do + detail_file : Optional[str] + file where test details will be output + + Returns + ------- + Tuple[List[np.ndarray], List[int]] + arrays with results and their shapes + """ + data.add( + "wfc", 12, atomic=True, must=True, high_prec=False, type_sel=dp.get_sel_type() + ) + test_data = data.get_test() + wfc, numb_test, _ = run_test(dp, test_data, numb_test) + rmse_f = rmse(wfc - test_data["wfc"][:numb_test]) + + log.info("# number of test data : {numb_test:d} ") + log.info("WFC RMSE : {rmse_f:e} eV/A") + + if detail_file is not None: + detail_path = Path(detail_file) + pe = np.concatenate( + ( + np.reshape(test_data["wfc"][:numb_test], [-1, 12]), + np.reshape(wfc, [-1, 12]), + ), + axis=1, + ) + np.savetxt( + detail_path.with_suffix(".out"), + pe, + header="ref_wfc(12 dofs) predicted_wfc(12 dofs)", + ) + return { + 'rmse' : (rmse_f, wfc.size) + } + + +def print_wfc_sys_avg(avg): + """Print errors summary for wfc type potential. + + Parameters + ---------- + avg : np.ndarray + array with summaries + """ + log.info(f"WFC RMSE : {avg['rmse']:e} eV/A") + + +def test_polar( + dp: "DeepPolar", + data: DeepmdData, + numb_test: int, + detail_file: Optional[str], + *, + atomic: bool, +) -> Tuple[List[np.ndarray], List[int]]: + """Test energy type model. + + Parameters + ---------- + dp : DeepPot + instance of deep potential + data: DeepmdData + data container object + numb_test : int + munber of tests to do + detail_file : Optional[str] + file where test details will be output + global_polar : bool + wheter to use glovbal version of polar potential + + Returns + ------- + Tuple[List[np.ndarray], List[int]] + arrays with results and their shapes + """ + data.add( + "polarizability" if not atomic else "atomic_polarizability", + 9, + atomic=atomic, + must=True, + high_prec=False, + type_sel=dp.get_sel_type(), + ) + + test_data = data.get_test() + polar, numb_test, atype = run_test(dp, test_data, numb_test) + + sel_type = dp.get_sel_type() + sel_natoms = 0 + for ii in sel_type: + sel_natoms += sum(atype == ii) + + # YWolfeee: do summation in global polar mode + if not atomic: + polar = np.sum(polar.reshape((polar.shape[0],-1,9)),axis=1) + rmse_f = rmse(polar - test_data["polarizability"][:numb_test]) + rmse_fs = rmse_f / np.sqrt(sel_natoms) + rmse_fa = rmse_f / sel_natoms + else: + rmse_f = rmse(polar - test_data["atomic_polarizability"][:numb_test]) + + log.info(f"# number of test data : {numb_test:d} ") + log.info(f"Polarizability RMSE : {rmse_f:e}") + if not atomic: + log.info(f"Polarizability RMSE/sqrtN : {rmse_fs:e}") + log.info(f"Polarizability RMSE/N : {rmse_fa:e}") + log.info(f"The unit of error is the same as the unit of provided label.") + + if detail_file is not None: + detail_path = Path(detail_file) + + pe = np.concatenate( + ( + np.reshape(test_data["polarizability"][:numb_test], [-1, 9]), + np.reshape(polar, [-1, 9]), + ), + axis=1, + ) + np.savetxt( + detail_path.with_suffix(".out"), + pe, + header="data_pxx data_pxy data_pxz data_pyx data_pyy data_pyz data_pzx " + "data_pzy data_pzz pred_pxx pred_pxy pred_pxz pred_pyx pred_pyy pred_pyz " + "pred_pzx pred_pzy pred_pzz", + ) + return { + "rmse" : (rmse_f, polar.size) + } + + +def print_polar_sys_avg(avg): + """Print errors summary for polar type potential. + + Parameters + ---------- + avg : np.ndarray + array with summaries + """ + log.info(f"Polarizability RMSE : {avg['rmse']:e} eV/A") + + +def test_dipole( + dp: "DeepDipole", + data: DeepmdData, + numb_test: int, + detail_file: Optional[str], + atomic: bool, +) -> Tuple[List[np.ndarray], List[int]]: + """Test energy type model. + + Parameters + ---------- + dp : DeepPot + instance of deep potential + data: DeepmdData + data container object + numb_test : int + munber of tests to do + detail_file : Optional[str] + file where test details will be output + atomic : bool + whether atomic dipole is provided + + Returns + ------- + Tuple[List[np.ndarray], List[int]] + arrays with results and their shapes + """ + data.add( + "dipole" if not atomic else "atomic_dipole", + 3, + atomic=atomic, + must=True, + high_prec=False, + type_sel=dp.get_sel_type() + ) + test_data = data.get_test() + dipole, numb_test, atype = run_test(dp, test_data, numb_test) + + sel_type = dp.get_sel_type() + sel_natoms = 0 + for ii in sel_type: + sel_natoms += sum(atype == ii) + + # do summation in atom dimension + if not atomic: + dipole = np.sum(dipole.reshape((dipole.shape[0], -1, 3)),axis=1) + rmse_f = rmse(dipole - test_data["dipole"][:numb_test]) + rmse_fs = rmse_f / np.sqrt(sel_natoms) + rmse_fa = rmse_f / sel_natoms + else: + rmse_f = rmse(dipole - test_data["atomic_dipole"][:numb_test]) + + log.info(f"# number of test data : {numb_test:d}") + log.info(f"Dipole RMSE : {rmse_f:e}") + if not atomic: + log.info(f"Dipole RMSE/sqrtN : {rmse_fs:e}") + log.info(f"Dipole RMSE/N : {rmse_fa:e}") + log.info(f"The unit of error is the same as the unit of provided label.") + + if detail_file is not None: + detail_path = Path(detail_file) + + pe = np.concatenate( + ( + np.reshape(test_data["dipole"][:numb_test], [-1, 3]), + np.reshape(dipole, [-1, 3]), + ), + axis=1, + ) + np.savetxt( + detail_path.with_suffix(".out"), + pe, + header="data_x data_y data_z pred_x pred_y pred_z", + ) + return { + 'rmse' : (rmse_f, dipole.size) + } + + +def print_dipole_sys_avg(avg): + """Print errors summary for dipole type potential. + + Parameters + ---------- + avg : np.ndarray + array with summaries + """ + log.info(f"Dipole RMSE : {avg['rmse']:e} eV/A") diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py new file mode 100755 index 0000000000..1ff0cff02e --- /dev/null +++ b/deepmd/entrypoints/train.py @@ -0,0 +1,328 @@ +"""DeePMD training entrypoint script. + +Can handle local or distributed training. +""" + +import json +import logging +import time +import os +from typing import Dict, TYPE_CHECKING, List, Optional, Any + +import numpy as np +from deepmd.common import data_requirement, expand_sys_str, j_loader, j_must_have +from deepmd.env import tf +from deepmd.infer.data_modifier import DipoleChargeModifier +from deepmd.train.run_options import BUILD, CITATION, WELCOME, RunOptions +from deepmd.train.trainer import DPTrainer +from deepmd.utils.argcheck import normalize +from deepmd.utils.compat import updata_deepmd_input +from deepmd.utils.data_system import DeepmdDataSystem + +if TYPE_CHECKING: + from deepmd.run_options import TFServerV1 + +__all__ = ["train"] + +log = logging.getLogger(__name__) + + +def create_done_queue( + cluster_spec: tf.train.ClusterSpec, task_index: int +) -> tf.FIFOQueue: + """Create FIFO queue for distributed tasks. + + Parameters + ---------- + cluster_spec : tf.train.ClusterSpec + tf cluster specification object + task_index : int + identifying index of a task + + Returns + ------- + tf.FIFOQueue + tf distributed FIFI queue + """ + with tf.device(f"/job:ps/task:{task_index:d}"): + queue = tf.FIFOQueue( + cluster_spec.num_tasks("worker"), + tf.int32, + shared_name=f"done_queue{task_index}", + ) + return queue + + +def wait_done_queue( + cluster_spec: tf.train.ClusterSpec, + server: "TFServerV1", + queue: tf.FIFOQueue, + task_index: int, +): + """Wait until all enqued operation in tf distributed queue are finished. + + Parameters + ---------- + cluster_spec : tf.train.ClusterSpec + tf cluster specification object + server : TFServerV1 + tf server specification object + queue : tf.FIFOQueue + tf distributed queue + task_index : int + identifying index of a task + """ + with tf.Session(server.target) as sess: + for i in range(cluster_spec.num_tasks("worker")): + sess.run(queue.dequeue()) + log.debug(f"ps:{task_index:d} received done from worker:{i:d}") + log.debug(f"ps:{task_index:f} quitting") + + +def connect_done_queue( + cluster_spec: tf.train.ClusterSpec, task_index: int +) -> List[tf.Operation]: + """Create tf FIFO queue filling operations. + + Parameters + ---------- + cluster_spec : tf.train.ClusterSpec + tf cluster specification object + task_index : int + identifying index of a task + + Returns + ------- + List[tf.Operation] + list of tf operations that will populate the queue + """ + done_ops = [] + for i in range(cluster_spec.num_tasks("ps")): + with tf.device(f"/job:ps/task:{i:d}"): + queue = tf.FIFOQueue( + cluster_spec.num_tasks("worker"), tf.int32, shared_name=f"done_queue{i}" + ) + done_ops.append(queue.enqueue(task_index)) + return done_ops + + +def fill_done_queue( + cluster_spec: tf.train.ClusterSpec, + server: "TFServerV1", + done_ops: List[tf.Operation], + task_index: int, +): + """Run specified operations that will fill the tf distributed FIFO queue. + + Parameters + ---------- + cluster_spec : tf.train.ClusterSpec + tf cluster specification object + server : TFServerV1 + tf server specification object + done_ops : List[tf.Operation] + a list of tf operations that will fill the queue + task_index : int + identifying index of a task + """ + with tf.Session(server.target) as sess: + for i in range(cluster_spec.num_tasks("ps")): + sess.run(done_ops[i]) + log.debug(f"worker:{task_index:d} sending done to ps:{i:d}") + + +def train( + *, + INPUT: str, + init_model: Optional[str], + restart: Optional[str], + output: str, + mpi_log: str, + log_level: int, + log_path: Optional[str], + **kwargs, +): + """Run DeePMD model training. + + Parameters + ---------- + INPUT : str + json/yaml control file + init_model : Optional[str] + path to checkpoint folder or None + restart : Optional[str] + path to checkpoint folder or None + output : str + path for dump file with arguments + mpi_log : str + mpi logging mode + log_level : int + logging level defined by int 0-3 + log_path : Optional[str] + logging file path or None if logs are to be output only to stdout + + Raises + ------ + RuntimeError + if distributed training job nem is wrong + """ + # load json database + jdata = j_loader(INPUT) + + jdata = updata_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") + + jdata = normalize(jdata) + with open(output, "w") as fp: + json.dump(jdata, fp, indent=4) + + # run options + run_opt = RunOptions( + init_model=init_model, + restart=restart, + log_path=log_path, + log_level=log_level, + mpi_log=mpi_log, + try_distrib=jdata.get("with_distrib", False), + ) + + for message in WELCOME + CITATION + BUILD: + log.info(message) + + run_opt.print_resource_summary() + + if run_opt.is_distrib: + # distributed training + if run_opt.my_job_name == "ps": + queue = create_done_queue(run_opt.cluster_spec, run_opt.my_task_index) + wait_done_queue( + run_opt.cluster_spec, run_opt.server, queue, run_opt.my_task_index + ) + # server.join() + elif run_opt.my_job_name == "worker": + done_ops = connect_done_queue(run_opt.cluster_spec, run_opt.my_task_index) + _do_work(jdata, run_opt) + fill_done_queue( + run_opt.cluster_spec, run_opt.server, done_ops, run_opt.my_task_index + ) + else: + raise RuntimeError("unknown job name") + else: + # serial training + _do_work(jdata, run_opt) + + +def _do_work(jdata: Dict[str, Any], run_opt: RunOptions): + """Run serial model training. + + Parameters + ---------- + jdata : Dict[str, Any] + arguments read form json/yaml control file + run_opt : RunOptions + object with run configuration + + Raises + ------ + RuntimeError + If unsupported modifier type is selected for model + """ + # make necessary checks + assert "training" in jdata + + # init the model + model = DPTrainer(jdata, run_opt=run_opt) + rcut = model.model.get_rcut() + type_map = model.model.get_type_map() + if len(type_map) == 0: + ipt_type_map = None + else: + ipt_type_map = type_map + + #  init random seed + seed = jdata["training"].get("seed", None) + if seed is not None: + seed = seed % (2 ** 32) + np.random.seed(seed) + + # setup data modifier + modifier = get_modifier(jdata["model"].get("modifier", None)) + + # init data + train_data = get_data(jdata["training"]["training_data"], rcut, ipt_type_map, modifier) + train_data.print_summary("training") + if jdata["training"].get("validation_data", None) is not None: + valid_data = get_data(jdata["training"]["validation_data"], rcut, ipt_type_map, modifier) + valid_data.print_summary("validation") + else: + valid_data = None + + # get training info + stop_batch = j_must_have(jdata["training"], "numb_steps") + model.build(train_data, stop_batch) + + # train the model with the provided systems in a cyclic way + start_time = time.time() + model.train(train_data, valid_data) + end_time = time.time() + log.info("finished training") + log.info(f"wall time: {(end_time - start_time):.3f} s") + + +def get_data(jdata: Dict[str, Any], rcut, type_map, modifier): + systems = j_must_have(jdata, "systems") + if isinstance(systems, str): + systems = expand_sys_str(systems) + help_msg = 'Please check your setting for data systems' + # check length of systems + if len(systems) == 0: + msg = 'cannot find valid a data system' + log.fatal(msg) + raise IOError(msg, help_msg) + # rougly check all items in systems are valid + for ii in systems: + if (not os.path.isdir(ii)): + msg = f'dir {ii} is not a valid dir' + log.fatal(msg) + raise IOError(msg, help_msg) + if (not os.path.isfile(os.path.join(ii, 'type.raw'))): + msg = f'dir {ii} is not a valid data system dir' + log.fatal(msg) + raise IOError(msg, help_msg) + + batch_size = j_must_have(jdata, "batch_size") + sys_probs = jdata.get("sys_probs", None) + auto_prob = jdata.get("auto_prob", "prob_sys_size") + + data = DeepmdDataSystem( + systems=systems, + batch_size=batch_size, + test_size=1, # to satisfy the old api + shuffle_test=True, # to satisfy the old api + rcut=rcut, + type_map=type_map, + modifier=modifier, + trn_all_set=True, # sample from all sets + sys_probs=sys_probs, + auto_prob_style=auto_prob + ) + data.add_dict(data_requirement) + + return data + + +def get_modifier(modi_data=None): + modifier: Optional[DipoleChargeModifier] + if modi_data is not None: + if modi_data["type"] == "dipole_charge": + modifier = DipoleChargeModifier( + modi_data["model_name"], + modi_data["model_charge_map"], + modi_data["sys_charge_map"], + modi_data["ewald_h"], + modi_data["ewald_beta"], + ) + else: + raise RuntimeError("unknown modifier type " + str(modi_data["type"])) + else: + modifier = None + return modifier diff --git a/deepmd/entrypoints/transfer.py b/deepmd/entrypoints/transfer.py new file mode 100644 index 0000000000..0af45a4244 --- /dev/null +++ b/deepmd/entrypoints/transfer.py @@ -0,0 +1,257 @@ +"""Module used for transfering parameters between models.""" + +from typing import Dict, Optional, Sequence, Tuple +from deepmd.env import tf +import re +import numpy as np +import logging + +__all__ = ["transfer"] + +log = logging.getLogger(__name__) + +PRECISION_MAPPING: Dict[int, type] = { + 1: np.float32, + 2: np.float64, + 19: np.float16, +} + + +@np.vectorize +def convert_number(number: int) -> float: + binary = bin(number).replace("0b", "").zfill(16) + sign = int(binary[0]) * -2 + 1 + exp = int(binary[1:6], 2) + frac = (int(binary[6:], 2) + 2 ** 10) * (2 ** -10) + return sign * (2 ** (exp - 15)) * frac + + +def convert_matrix( + matrix: np.ndarray, shape: Sequence[int], dtype: Optional[type] = None +) -> np.ndarray: + """Convert matrix of integers to self defined binary format. + + Parameters + ---------- + matrix : np.ndarray + array of ints + shape : Sequence[int] + shape to cast resulting array to + dtype : Optional[type] + type that finall array will be cast to, If None no casting will take place + + Returns + ------- + np.ndarray + array cast to required format + """ + conv = convert_number(matrix.flatten()).reshape(shape) + if dtype: + conv = conv.astype(dtype) + + return conv + + +def transfer(*, old_model: str, raw_model: str, output: str, **kwargs): + """Transfer operation from old fron graph to new prepared raw graph. + + Parameters + ---------- + old_model : str + frozen old graph model + raw_model : str + new model that will accept ops from old model + output : str + new model with transfered parameters will be saved to this location + """ + raw_graph = load_graph(raw_model) + old_graph = load_graph(old_model) + log.info(f"{len(raw_graph.as_graph_def().node)} ops in the raw graph") + log.info(f"{len(old_graph.as_graph_def().node)} ops in the old graph") + + new_graph_def = transform_graph(raw_graph, old_graph) + with tf.gfile.GFile(output, mode="wb") as f: + f.write(new_graph_def.SerializeToString()) + log.info("the output model is saved in " + output) + + +def load_graph(graph_name: str) -> tf.Graph: + """Load graph from passed in path. + + Parameters + ---------- + graph_name : str + path to frozen graph on disk + + Returns + ------- + tf.Graph + tf graph object + """ + graph_def = tf.GraphDef() + with open(graph_name, "rb") as f: + graph_def.ParseFromString(f.read()) + with tf.Graph().as_default() as graph: + tf.import_graph_def(graph_def, name="") + return graph + + +def transform_graph(raw_graph: tf.Graph, old_graph: tf.Graph) -> tf.Graph: + """Trasform old graph into new. + + Parameters + ---------- + raw_graph : tf.Graph + graph receiving parameters from the old one + old_graph : tf.Graph + graph providing parameters + + Returns + ------- + tf.Graph + new graph with parameters transfered form the old one + """ + old_graph_def = old_graph.as_graph_def() + raw_graph_def = raw_graph.as_graph_def() + raw_graph_node = load_transform_node(raw_graph_def) + old_graph_node = load_transform_node(old_graph_def) + + for node in raw_graph_def.node: + if node.name not in raw_graph_node.keys(): + continue + + old_node = old_graph_node[node.name] + raw_node = raw_graph_node[node.name] + cp_attr = CopyNodeAttr(node) + + check_dim(raw_graph_node, old_graph_node, node.name) + tensor_shape = [dim.size for dim in raw_node.tensor_shape.dim] + old_graph_dtype = PRECISION_MAPPING[old_node.dtype] + raw_graph_dtype = PRECISION_MAPPING[raw_node.dtype] + log.info( + f"{node.name} is passed from old graph({old_graph_dtype}) " + f"to raw graph({raw_graph_dtype})" + ) + + if raw_graph_dtype == np.float16: + if old_graph_dtype == np.float64 or old_graph_dtype == np.float32: + if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): + tensor = np.frombuffer(old_node.tensor_content).astype(raw_graph_dtype) + cp_attr.from_array(tensor, tf.float16, shape = tensor_shape) + else: + tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype) + cp_attr.from_array(tensor, tf.float16, [1]) + + elif old_graph_dtype[1] == "float16": + tensor = convertMatrix(np.array(old_node.half_val), tensor_shape) + cp_attr.from_array(tensor, raw_graph_dtype) + + elif raw_graph_dtype == np.float64 or raw_graph_dtype == np.float32: + if old_graph_dtype == np.float64 or old_graph_dtype == np.float32: + if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): + tensor = np.frombuffer(old_node.tensor_content).astype(raw_graph_dtype) + cp_attr.from_str(tensor) + else: + tensor = load_tensor(old_node, old_graph_dtype, raw_graph_dtype) + cp_attr.from_array(tensor, raw_graph_dtype, shape=[1]) + + elif old_graph_dtype == np.float16: + if (len(tensor_shape) != 1) or (tensor_shape[0] != 1): + tensor = convertMatrix(np.array(old_node.half_val), tensor_shape).astype(raw_graph_dtype) + cp_attr.from_str(tensor) + else: + tensor = convertMatrix(np.array(old_node.half_val), tensor_shape).astype(raw_graph_dtype) + cp_attr.from_array(tensor, raw_graph_dtype) + + return raw_graph_def + + +class CopyNodeAttr: + def __init__(self, node) -> None: + self.node = node + + def from_array( + self, tensor: np.ndarray, dtype: type, shape: Optional[Sequence[int]] = None + ): + if shape is None: + shape = tensor.shape + self.node.attr["value"].CopyFrom( + tf.AttrValue(tensor=tf.make_tensor_proto(tensor, dtype, shape)) + ) + + def from_str(self, tensor: np.ndarray): + self.node.attr["value"].tensor.tensor_content = tensor.tostring() + + +def load_tensor(node: tf.Tensor, dtype_old: type, dtype_new: type) -> np.ndarray: + if dtype_old == np.float64: + tensor = np.array(node.double_val).astype(dtype_new) + elif dtype_old == np.float32: + tensor = np.array(node.float_val).astype(dtype_new) + + return tensor + + +def check_dim(raw_graph_node: tf.Tensor, old_graph_node: tf.Tensor, node_name: str): + """Check if dimensions of tensor in old and new graph is equal. + + Parameters + ---------- + raw_graph_node : tf.Tensor + node of the receiving graph + old_graph_node : tf.Tensor + node of the graph from which will node be extracted + node_name : str + name of the node + + Raises + ------ + RuntimeError + if node dimension do not match + """ + raw_graph_dim = raw_graph_node[node_name].tensor_shape + old_graph_dim = old_graph_node[node_name].tensor_shape + if raw_graph_dim != old_graph_dim: + raise RuntimeError( + f"old graph {old_graph_dim} and raw graph {raw_graph_dim} " + f"has different {node_name} dim" + ) + + +def load_transform_node(graph: tf.Graph) -> Dict[str, tf.Tensor]: + """Load nodes and their names from graph to dict. + + Parameters + ---------- + graph : tf.Graph + tensforflow graph + + Returns + ------- + Dict[str, tf.Tensor] + mapping on graph node names and corresponding tensors + """ + transform_node_pattern = re.compile( + r"filter_type_\d+/matrix_\d+_\d+|" + r"filter_type_\d+/bias_\d+_\d+|" + r"filter_type_\d+/idt_\d+_\d+|" + r"layer_\d+_type_\d+/matrix|" + r"layer_\d+_type_\d+/bias|" + r"layer_\d+_type_\d+/idt|" + r"final_layer_type_\d+/matrix|" + r"descrpt_attr/t_avg|" + r"descrpt_attr/t_std|" + r"final_layer_type_\d+/bias|" + r"fitting_attr/t_fparam_avg|" + r"fitting_attr/t_fparam_istd|" + r"fitting_attr/t_aparam_avg|" + r"fitting_attr/t_aparam_istd|" + r"model_attr/t_tab_info|" + r"model_attr/t_tab_data|" + ) + + transform_node = {} + for node in graph.node: + if transform_node_pattern.fullmatch(node.name) is not None: + transform_node[node.name] = node.attr["value"].tensor + return transform_node diff --git a/deepmd/env.py b/deepmd/env.py new file mode 100644 index 0000000000..8c6937b7f7 --- /dev/null +++ b/deepmd/env.py @@ -0,0 +1,225 @@ +"""Module that sets tensorflow working environment and exports inportant constants.""" + +import os +from pathlib import Path +import logging +import platform +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Any +import numpy as np +from imp import reload +from configparser import ConfigParser + +if TYPE_CHECKING: + from types import ModuleType + +# import tensorflow v1 compatability +try: + import tensorflow.compat.v1 as tf + + tf.disable_v2_behavior() +except ImportError: + import tensorflow as tf + +__all__ = [ + "GLOBAL_CONFIG", + "GLOBAL_TF_FLOAT_PRECISION", + "GLOBAL_NP_FLOAT_PRECISION", + "GLOBAL_ENER_FLOAT_PRECISION", + "global_float_prec", + "global_cvt_2_tf_float", + "global_cvt_2_ener_float", + "MODEL_VERSION", + "SHARED_LIB_MODULE", + "default_tf_session_config", + "op_module", + "op_grads_module", +] + +SHARED_LIB_MODULE = "op" + +def set_env_if_empty(key: str, value: str, verbose: bool = True): + """Set environment variable only if it is empty. + + Parameters + ---------- + key : str + env variable name + value : str + env variable value + verbose : bool, optional + if True action will be logged, by default True + """ + if os.environ.get(key) is None: + os.environ[key] = value + if verbose: + logging.warn( + f"Environment variable {key} is empty. Use the default value {value}" + ) + + +def set_mkl(): + """Tuning MKL for the best performance. + + References + ---------- + TF overview + https://www.tensorflow.org/guide/performance/overview + + Fixing an issue in numpy built by MKL + https://github.com/ContinuumIO/anaconda-issues/issues/11367 + https://github.com/numpy/numpy/issues/12374 + + check whether the numpy is built by mkl, see + https://github.com/numpy/numpy/issues/14751 + """ + if "mkl_rt" in np.__config__.get_info("blas_mkl_info").get("libraries", []): + set_env_if_empty("KMP_BLOCKTIME", "0") + set_env_if_empty("KMP_AFFINITY", "granularity=fine,verbose,compact,1,0") + reload(np) + + +def set_tf_default_nthreads(): + """Set TF internal number of threads to default=automatic selection. + + Notes + ----- + `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS` + control TF configuration of multithreading. + """ + set_env_if_empty("TF_INTRA_OP_PARALLELISM_THREADS", "0", verbose=False) + set_env_if_empty("TF_INTER_OP_PARALLELISM_THREADS", "0", verbose=False) + + +def get_tf_default_nthreads() -> Tuple[int, int]: + """Get TF paralellism settings. + + Returns + ------- + Tuple[int, int] + number of `TF_INTRA_OP_PARALLELISM_THREADS` and + `TF_INTER_OP_PARALLELISM_THREADS` + """ + return int(os.environ.get("TF_INTRA_OP_PARALLELISM_THREADS", "0")), int( + os.environ.get("TF_INTER_OP_PARALLELISM_THREADS", "0") + ) + + +def get_tf_session_config() -> Any: + """Configure tensorflow session. + + Returns + ------- + Any + session configure object + """ + set_tf_default_nthreads() + intra, inter = get_tf_default_nthreads() + return tf.ConfigProto( + intra_op_parallelism_threads=intra, inter_op_parallelism_threads=inter + ) + +default_tf_session_config = get_tf_session_config() + +def get_module(module_name: str) -> "ModuleType": + """Load force module. + + Returns + ------- + ModuleType + loaded force module + + Raises + ------ + FileNotFoundError + if module is not found in directory + """ + if platform.system() == "Windows": + ext = ".dll" + elif platform.system() == "Darwin": + ext = ".dylib" + else: + ext = ".so" + + module_file = ( + (Path(__file__).parent / SHARED_LIB_MODULE / module_name) + .with_suffix(ext) + .resolve() + ) + + if not module_file.is_file(): + raise FileNotFoundError(f"module {module_name} does not exist") + else: + module = tf.load_op_library(str(module_file)) + return module + + +op_module = get_module("libop_abi") +op_grads_module = get_module("libop_grads") + + +def _get_package_constants( + config_file: Path = Path(__file__).parent / "pkg_config/run_config.ini", +) -> Dict[str, str]: + """Read package constants set at compile time by CMake to dictionary. + + Parameters + ---------- + config_file : str, optional + path to CONFIG file, by default "config/run_config.ini" + + Returns + ------- + Dict[str, str] + dictionary with package constants + """ + config = ConfigParser() + config.read(config_file) + return dict(config.items("CONFIG")) + +GLOBAL_CONFIG = _get_package_constants() +MODEL_VERSION = GLOBAL_CONFIG["model_version"] + +if GLOBAL_CONFIG["precision"] == "-DHIGH_PREC": + GLOBAL_TF_FLOAT_PRECISION = tf.float64 + GLOBAL_NP_FLOAT_PRECISION = np.float64 + GLOBAL_ENER_FLOAT_PRECISION = np.float64 + global_float_prec = "double" +else: + GLOBAL_TF_FLOAT_PRECISION = tf.float32 + GLOBAL_NP_FLOAT_PRECISION = np.float32 + GLOBAL_ENER_FLOAT_PRECISION = np.float64 + global_float_prec = "float" + + +def global_cvt_2_tf_float(xx: tf.Tensor) -> tf.Tensor: + """Cast tensor to globally set TF precision. + + Parameters + ---------- + xx : tf.Tensor + input tensor + + Returns + ------- + tf.Tensor + output tensor cast to `GLOBAL_TF_FLOAT_PRECISION` + """ + return tf.cast(xx, GLOBAL_TF_FLOAT_PRECISION) + + +def global_cvt_2_ener_float(xx: tf.Tensor) -> tf.Tensor: + """Cast tensor to globally set energy precision. + + Parameters + ---------- + xx : tf.Tensor + input tensor + + Returns + ------- + tf.Tensor + output tensor cast to `GLOBAL_ENER_FLOAT_PRECISION` + """ + return tf.cast(xx, GLOBAL_ENER_FLOAT_PRECISION) + + diff --git a/deepmd/fit/__init__.py b/deepmd/fit/__init__.py new file mode 100644 index 0000000000..71d582e8ca --- /dev/null +++ b/deepmd/fit/__init__.py @@ -0,0 +1,6 @@ +from .ener import EnerFitting +from .wfc import WFCFitting +from .dipole import DipoleFittingSeA +from .polar import PolarFittingSeA +from .polar import GlobalPolarFittingSeA +from .polar import PolarFittingLocFrame diff --git a/deepmd/fit/dipole.py b/deepmd/fit/dipole.py new file mode 100644 index 0000000000..73562951dc --- /dev/null +++ b/deepmd/fit/dipole.py @@ -0,0 +1,165 @@ +import warnings +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.network import one_layer, one_layer_rand_seed_shift +from deepmd.descriptor import DescrptSeA + +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION + +class DipoleFittingSeA () : + """ + Fit the atomic dipole with descriptor se_a + """ + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + descrpt : tf.Tensor, + neuron : List[int] = [120,120,120], + resnet_dt : bool = True, + sel_type : List[int] = None, + seed : int = None, + activation_function : str = 'tanh', + precision : str = 'default', + uniform_seed: bool = False + ) -> None: + """ + Constructor + + Parameters + ---------- + descrpt : tf.Tensor + The descrptor + neuron : List[int] + Number of neurons in each hidden layer of the fitting net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + sel_type : List[int] + The atom types selected to have an atomic dipole prediction. If is None, all atoms are selected. + seed : int + Random seed for initializing the network parameters. + activation_function : str + The activation function in the embedding net. Supported options are {0} + precision : str + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + if not isinstance(descrpt, DescrptSeA) : + raise RuntimeError('DipoleFittingSeA only supports DescrptSeA') + self.ntypes = descrpt.get_ntypes() + self.dim_descrpt = descrpt.get_dim_out() + # args = ClassArg()\ + # .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\ + # .add('resnet_dt', bool, default = True)\ + # .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'dipole_type')\ + # .add('seed', int)\ + # .add("activation_function", str, default = "tanh")\ + # .add('precision', str, default = "default") + # class_data = args.parse(jdata) + self.n_neuron = neuron + self.resnet_dt = resnet_dt + self.sel_type = sel_type + if self.sel_type is None: + self.sel_type = [ii for ii in range(self.ntypes)] + self.sel_type = sel_type + self.seed = seed + self.uniform_seed = uniform_seed + self.seed_shift = one_layer_rand_seed_shift() + self.fitting_activation_fn = get_activation_func(activation_function) + self.fitting_precision = get_precision(precision) + self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1() + self.dim_rot_mat = self.dim_rot_mat_1 * 3 + self.useBN = False + + def get_sel_type(self) -> int: + """ + Get selected type + """ + return self.sel_type + + def get_out_size(self) -> int: + """ + Get the output size. Should be 3 + """ + return 3 + + def build (self, + input_d : tf.Tensor, + rot_mat : tf.Tensor, + natoms : tf.Tensor, + reuse : bool = None, + suffix : str = '') -> tf.Tensor: + """ + Build the computational graph for fitting net + + Parameters + ---------- + input_d + The input descriptor + rot_mat + The rotation matrix from the descriptor. + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Return + ------ + dipole + The atomic dipole. + """ + start_index = 0 + inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) + + count = 0 + for type_i in range(self.ntypes): + # cut-out inputs + inputs_i = tf.slice (inputs, + [ 0, start_index* self.dim_descrpt], + [-1, natoms[2+type_i]* self.dim_descrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) + rot_mat_i = tf.slice (rot_mat, + [ 0, start_index* self.dim_rot_mat], + [-1, natoms[2+type_i]* self.dim_rot_mat] ) + rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) + start_index += natoms[2+type_i] + if not type_i in self.sel_type : + continue + layer = inputs_i + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + else : + layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # (nframes x natoms) x naxis + final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # (nframes x natoms) x 1 * naxis + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 1, self.dim_rot_mat_1]) + # (nframes x natoms) x 1 x 3(coord) + final_layer = tf.matmul(final_layer, rot_mat_i) + # nframes x natoms x 3 + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3]) + + # concat the results + if count == 0: + outs = final_layer + else: + outs = tf.concat([outs, final_layer], axis = 1) + count += 1 + + tf.summary.histogram('fitting_net_output', outs) + return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) + # return tf.reshape(outs, [tf.shape(inputs)[0] * natoms[0] * 3 // 3]) diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py new file mode 100644 index 0000000000..0d2932a76b --- /dev/null +++ b/deepmd/fit/ener.py @@ -0,0 +1,445 @@ +import warnings +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.network import one_layer, one_layer_rand_seed_shift +from deepmd.descriptor import DescrptLocFrame +from deepmd.descriptor import DescrptSeA +from deepmd.utils.type_embed import embed_atom_type + +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION + +class EnerFitting (): + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + descrpt : tf.Tensor, + neuron : List[int] = [120,120,120], + resnet_dt : bool = True, + numb_fparam : int = 0, + numb_aparam : int = 0, + rcond : float = 1e-3, + tot_ener_zero : bool = False, + trainable : List[bool] = None, + seed : int = None, + atom_ener : List[float] = [], + activation_function : str = 'tanh', + precision : str = 'default', + uniform_seed: bool = False + ) -> None: + """ + Constructor + + Parameters + ---------- + descrpt + The descrptor + neuron + Number of neurons in each hidden layer of the fitting net + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + numb_fparam + Number of frame parameter + numb_aparam + Number of atomic parameter + rcond + The condition number for the regression of atomic energy. + tot_ener_zero + Force the total energy to zero. Useful for the charge fitting. + trainable + If the weights of fitting net are trainable. + Suppose that we have N_l hidden layers in the fitting net, + this list is of length N_l + 1, specifying if the hidden layers and the output layer are trainable. + seed + Random seed for initializing the network parameters. + atom_ener + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + # model param + self.ntypes = descrpt.get_ntypes() + self.dim_descrpt = descrpt.get_dim_out() + # args = ()\ + # .add('numb_fparam', int, default = 0)\ + # .add('numb_aparam', int, default = 0)\ + # .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\ + # .add('resnet_dt', bool, default = True)\ + # .add('rcond', float, default = 1e-3) \ + # .add('tot_ener_zero', bool, default = False) \ + # .add('seed', int) \ + # .add('atom_ener', list, default = [])\ + # .add("activation_function", str, default = "tanh")\ + # .add("precision", str, default = "default")\ + # .add("trainable", [list, bool], default = True) + self.numb_fparam = numb_fparam + self.numb_aparam = numb_aparam + self.n_neuron = neuron + self.resnet_dt = resnet_dt + self.rcond = rcond + self.seed = seed + self.uniform_seed = uniform_seed + self.seed_shift = one_layer_rand_seed_shift() + self.tot_ener_zero = tot_ener_zero + self.fitting_activation_fn = get_activation_func(activation_function) + self.fitting_precision = get_precision(precision) + self.trainable = trainable + if self.trainable is None: + self.trainable = [True for ii in range(len(self.n_neuron) + 1)] + if type(self.trainable) is bool: + self.trainable = [self.trainable] * (len(self.n_neuron)+1) + assert(len(self.trainable) == len(self.n_neuron) + 1), 'length of trainable should be that of n_neuron + 1' + self.atom_ener = [] + for at, ae in enumerate(atom_ener): + if ae is not None: + self.atom_ener.append(tf.constant(ae, GLOBAL_TF_FLOAT_PRECISION, name = "atom_%d_ener" % at)) + else: + self.atom_ener.append(None) + self.useBN = False + self.bias_atom_e = None + # data requirement + if self.numb_fparam > 0 : + add_data_requirement('fparam', self.numb_fparam, atomic=False, must=True, high_prec=False) + self.fparam_avg = None + self.fparam_std = None + self.fparam_inv_std = None + if self.numb_aparam > 0: + add_data_requirement('aparam', self.numb_aparam, atomic=True, must=True, high_prec=False) + self.aparam_avg = None + self.aparam_std = None + self.aparam_inv_std = None + + def get_numb_fparam(self) -> int: + """ + Get the number of frame parameters + """ + return self.numb_fparam + + def get_numb_aparam(self) -> int: + """ + Get the number of atomic parameters + """ + return self.numb_fparam + + def compute_output_stats(self, + all_stat: dict + ) -> None: + """ + Compute the ouput statistics + + Parameters + ---------- + all_stat + must have the following components: + all_stat['energy'] of shape n_sys x n_batch x n_frame + can be prepared by model.make_stat_input + """ + self.bias_atom_e = self._compute_output_stats(all_stat, rcond = self.rcond) + + @classmethod + def _compute_output_stats(self, all_stat, rcond = 1e-3): + data = all_stat['energy'] + # data[sys_idx][batch_idx][frame_idx] + sys_ener = np.array([]) + for ss in range(len(data)): + sys_data = [] + for ii in range(len(data[ss])): + for jj in range(len(data[ss][ii])): + sys_data.append(data[ss][ii][jj]) + sys_data = np.concatenate(sys_data) + sys_ener = np.append(sys_ener, np.average(sys_data)) + data = all_stat['natoms_vec'] + sys_tynatom = np.array([]) + nsys = len(data) + for ss in range(len(data)): + sys_tynatom = np.append(sys_tynatom, data[ss][0].astype(np.float64)) + sys_tynatom = np.reshape(sys_tynatom, [nsys,-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 compute_input_stats(self, + all_stat : dict, + protection : float = 1e-2) -> None: + """ + Compute the input statistics + + Parameters: + all_stat + if numb_fparam > 0 must have all_stat['fparam'] + if numb_aparam > 0 must have all_stat['aparam'] + can be prepared by model.make_stat_input + protection + Divided-by-zero protection + """ + # stat fparam + if self.numb_fparam > 0: + cat_data = np.concatenate(all_stat['fparam'], axis = 0) + cat_data = np.reshape(cat_data, [-1, self.numb_fparam]) + self.fparam_avg = np.average(cat_data, axis = 0) + self.fparam_std = np.std(cat_data, axis = 0) + for ii in range(self.fparam_std.size): + if self.fparam_std[ii] < protection: + self.fparam_std[ii] = protection + self.fparam_inv_std = 1./self.fparam_std + # stat aparam + if self.numb_aparam > 0: + sys_sumv = [] + sys_sumv2 = [] + sys_sumn = [] + for ss_ in all_stat['aparam'] : + ss = np.reshape(ss_, [-1, self.numb_aparam]) + sys_sumv.append(np.sum(ss, axis = 0)) + sys_sumv2.append(np.sum(np.multiply(ss, ss), axis = 0)) + sys_sumn.append(ss.shape[0]) + sumv = np.sum(sys_sumv, axis = 0) + sumv2 = np.sum(sys_sumv2, axis = 0) + sumn = np.sum(sys_sumn) + self.aparam_avg = (sumv)/sumn + self.aparam_std = self._compute_std(sumv2, sumv, sumn) + for ii in range(self.aparam_std.size): + if self.aparam_std[ii] < protection: + self.aparam_std[ii] = protection + self.aparam_inv_std = 1./self.aparam_std + + + def _compute_std (self, sumv2, sumv, sumn) : + return np.sqrt(sumv2/sumn - np.multiply(sumv/sumn, sumv/sumn)) + + def _build_lower( + self, + start_index, + natoms, + inputs, + fparam = None, + aparam = None, + bias_atom_e = 0.0, + suffix = '', + reuse = None + ): + # cut-out inputs + inputs_i = tf.slice (inputs, + [ 0, start_index* self.dim_descrpt], + [-1, natoms* self.dim_descrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) + layer = inputs_i + if fparam is not None: + ext_fparam = tf.tile(fparam, [1, natoms]) + ext_fparam = tf.reshape(ext_fparam, [-1, self.numb_fparam]) + ext_fparam = tf.cast(ext_fparam,self.fitting_precision) + layer = tf.concat([layer, ext_fparam], axis = 1) + if aparam is not None: + ext_aparam = tf.slice(aparam, + [ 0, start_index * self.numb_aparam], + [-1, natoms * self.numb_aparam]) + ext_aparam = tf.reshape(ext_aparam, [-1, self.numb_aparam]) + ext_aparam = tf.cast(ext_aparam,self.fitting_precision) + layer = tf.concat([layer, ext_aparam], axis = 1) + + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer( + layer, + self.n_neuron[ii], + name='layer_'+str(ii)+suffix, + reuse=reuse, + seed = self.seed, + use_timestep = self.resnet_dt, + activation_fn = self.fitting_activation_fn, + precision = self.fitting_precision, + trainable = self.trainable[ii], + uniform_seed = self.uniform_seed) + else : + layer = one_layer( + layer, + self.n_neuron[ii], + name='layer_'+str(ii)+suffix, + reuse=reuse, + seed = self.seed, + activation_fn = self.fitting_activation_fn, + precision = self.fitting_precision, + trainable = self.trainable[ii], + uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + final_layer = one_layer( + layer, + 1, + activation_fn = None, + bavg = bias_atom_e, + name='final_layer'+suffix, + reuse=reuse, + seed = self.seed, + precision = self.fitting_precision, + trainable = self.trainable[-1], + uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + + return final_layer + + + + def build (self, + inputs : tf.Tensor, + natoms : tf.Tensor, + input_dict : dict = {}, + reuse : bool = None, + suffix : str = '', + ) -> tf.Tensor: + """ + Build the computational graph for fitting net + + Parameters + ---------- + inputs + The input descriptor + input_dict + Additional dict for inputs. + if numb_fparam > 0, should have input_dict['fparam'] + if numb_aparam > 0, should have input_dict['aparam'] + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Return + ------ + ener + The system energy + """ + bias_atom_e = self.bias_atom_e + if self.numb_fparam > 0 and ( self.fparam_avg is None or self.fparam_inv_std is None ): + raise RuntimeError('No data stat result. one should do data statisitic, before build') + if self.numb_aparam > 0 and ( self.aparam_avg is None or self.aparam_inv_std is None ): + raise RuntimeError('No data stat result. one should do data statisitic, before build') + + 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) + if self.numb_fparam > 0: + t_fparam_avg = tf.get_variable('t_fparam_avg', + self.numb_fparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.fparam_avg)) + t_fparam_istd = tf.get_variable('t_fparam_istd', + self.numb_fparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.fparam_inv_std)) + if self.numb_aparam > 0: + t_aparam_avg = tf.get_variable('t_aparam_avg', + self.numb_aparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.aparam_avg)) + t_aparam_istd = tf.get_variable('t_aparam_istd', + self.numb_aparam, + dtype = GLOBAL_TF_FLOAT_PRECISION, + trainable = False, + initializer = tf.constant_initializer(self.aparam_inv_std)) + + inputs = tf.cast(tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + inputs_zero = tf.zeros_like(inputs, dtype=GLOBAL_TF_FLOAT_PRECISION) + + + if bias_atom_e is not None : + assert(len(bias_atom_e) == self.ntypes) + + fparam = None + aparam = None + if self.numb_fparam > 0 : + fparam = input_dict['fparam'] + fparam = tf.reshape(fparam, [-1, self.numb_fparam]) + fparam = (fparam - t_fparam_avg) * t_fparam_istd + if self.numb_aparam > 0 : + aparam = input_dict['aparam'] + aparam = tf.reshape(aparam, [-1, self.numb_aparam]) + aparam = (aparam - t_aparam_avg) * t_aparam_istd + aparam = tf.reshape(aparam, [-1, self.numb_aparam * natoms[0]]) + + if input_dict is not None: + type_embedding = input_dict.get('type_embedding', None) + else: + type_embedding = None + if type_embedding is not None: + atype_embed = embed_atom_type(self.ntypes, natoms, type_embedding) + atype_embed = tf.tile(atype_embed,[tf.shape(inputs)[0],1]) + else: + atype_embed = None + + if atype_embed is None: + start_index = 0 + for type_i in range(self.ntypes): + if bias_atom_e is None : + type_bias_ae = 0.0 + else : + type_bias_ae = bias_atom_e[type_i] + final_layer = self._build_lower( + start_index, natoms[2+type_i], + inputs, fparam, aparam, + bias_atom_e=type_bias_ae, suffix='_type_'+str(type_i)+suffix, reuse=reuse + ) + # concat the results + if type_i < len(self.atom_ener) and self.atom_ener[type_i] is not None: + zero_layer = self._build_lower( + start_index, natoms[2+type_i], + inputs_zero, fparam, aparam, + bias_atom_e=type_bias_ae, suffix='_zero_type_'+str(type_i)+suffix, reuse=reuse + ) + final_layer += self.atom_ener[type_i] - zero_layer + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i]]) + # concat the results + if type_i == 0: + outs = final_layer + else: + outs = tf.concat([outs, final_layer], axis = 1) + start_index += natoms[2+type_i] + # with type embedding + else: + if len(self.atom_ener) > 0: + raise RuntimeError("setting atom_ener is not supported by type embedding") + atype_embed = tf.cast(atype_embed, self.fitting_precision) + type_shape = atype_embed.get_shape().as_list() + inputs = tf.concat( + [tf.reshape(inputs,[-1,self.dim_descrpt]),atype_embed], + axis=1 + ) + self.dim_descrpt = self.dim_descrpt + type_shape[1] + inputs = tf.cast(tf.reshape(inputs, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + final_layer = self._build_lower( + 0, natoms[0], + inputs, fparam, aparam, + bias_atom_e=0.0, suffix=suffix, reuse=reuse + ) + outs = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[0]]) + + if self.tot_ener_zero: + force_tot_ener = 0.0 + outs = tf.reshape(outs, [-1, natoms[0]]) + outs_mean = tf.reshape(tf.reduce_mean(outs, axis = 1), [-1, 1]) + outs_mean = outs_mean - tf.ones_like(outs_mean, dtype = GLOBAL_TF_FLOAT_PRECISION) * (force_tot_ener/global_cvt_2_tf_float(natoms[0])) + outs = outs - outs_mean + outs = tf.reshape(outs, [-1]) + + tf.summary.histogram('fitting_net_output', outs) + return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) + + diff --git a/deepmd/fit/polar.py b/deepmd/fit/polar.py new file mode 100644 index 0000000000..33c5be378a --- /dev/null +++ b/deepmd/fit/polar.py @@ -0,0 +1,479 @@ +import warnings +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.network import one_layer, one_layer_rand_seed_shift +from deepmd.descriptor import DescrptLocFrame +from deepmd.descriptor import DescrptSeA + +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION + + +class PolarFittingLocFrame () : + """ + Fitting polarizability with local frame descriptor. not supported anymore. + """ + def __init__ (self, jdata, descrpt) : + if not isinstance(descrpt, DescrptLocFrame) : + raise RuntimeError('PolarFittingLocFrame only supports DescrptLocFrame') + self.ntypes = descrpt.get_ntypes() + self.dim_descrpt = descrpt.get_dim_out() + args = ClassArg()\ + .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\ + .add('resnet_dt', bool, default = True)\ + .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\ + .add('seed', int)\ + .add("activation_function", str, default = "tanh")\ + .add('precision', str, default = "default") + class_data = args.parse(jdata) + self.n_neuron = class_data['neuron'] + self.resnet_dt = class_data['resnet_dt'] + self.sel_type = class_data['sel_type'] + self.seed = class_data['seed'] + self.fitting_activation_fn = get_activation_func(class_data["activation_function"]) + self.fitting_precision = get_precision(class_data['precision']) + self.useBN = False + + def get_sel_type(self): + return self.sel_type + + def get_out_size(self): + return 9 + + def build (self, + input_d, + rot_mat, + natoms, + reuse = None, + suffix = '') : + start_index = 0 + inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]]) + + count = 0 + for type_i in range(self.ntypes): + # cut-out inputs + inputs_i = tf.slice (inputs, + [ 0, start_index* self.dim_descrpt], + [-1, natoms[2+type_i]* self.dim_descrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) + rot_mat_i = tf.slice (rot_mat, + [ 0, start_index* 9], + [-1, natoms[2+type_i]* 9] ) + rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3]) + start_index += natoms[2+type_i] + if not type_i in self.sel_type : + continue + layer = inputs_i + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision) + else : + layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision) + # (nframes x natoms) x 9 + final_layer = one_layer(layer, 9, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision) + # (nframes x natoms) x 3 x 3 + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], 3, 3]) + # (nframes x natoms) x 3 x 3 + final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1]) + # (nframes x natoms) x 3 x 3(coord) + final_layer = tf.matmul(final_layer, rot_mat_i) + # (nframes x natoms) x 3(coord) x 3(coord) + final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True) + # nframes x natoms x 3 x 3 + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3]) + + # concat the results + if count == 0: + outs = final_layer + else: + outs = tf.concat([outs, final_layer], axis = 1) + count += 1 + + tf.summary.histogram('fitting_net_output', outs) + return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) + + +class PolarFittingSeA () : + """ + Fit the atomic polarizability with descriptor se_a + """ + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + descrpt : tf.Tensor, + neuron : List[int] = [120,120,120], + resnet_dt : bool = True, + sel_type : List[int] = None, + fit_diag : bool = True, + scale : List[float] = None, + shift_diag : bool = True, # YWolfeee: will support the user to decide whether to use this function + #diag_shift : List[float] = None, YWolfeee: will not support the user to assign a shift + seed : int = None, + activation_function : str = 'tanh', + precision : str = 'default', + uniform_seed: bool = False + ) -> None: + """ + Constructor + + Parameters + ---------- + descrpt : tf.Tensor + The descrptor + neuron : List[int] + Number of neurons in each hidden layer of the fitting net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + sel_type : List[int] + The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected. + fit_diag : bool + Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix. + scale : List[float] + The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i] + diag_shift : List[float] + The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale. + seed : int + Random seed for initializing the network parameters. + activation_function : str + The activation function in the embedding net. Supported options are {0} + precision : str + The precision of the embedding net parameters. Supported options are {1} + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + if not isinstance(descrpt, DescrptSeA) : + raise RuntimeError('PolarFittingSeA only supports DescrptSeA') + self.ntypes = descrpt.get_ntypes() + self.dim_descrpt = descrpt.get_dim_out() + # args = ClassArg()\ + # .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\ + # .add('resnet_dt', bool, default = True)\ + # .add('fit_diag', bool, default = True)\ + # .add('diag_shift', [list,float], default = [0.0 for ii in range(self.ntypes)])\ + # .add('scale', [list,float], default = [1.0 for ii in range(self.ntypes)])\ + # .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\ + # .add('seed', int)\ + # .add("activation_function", str , default = "tanh")\ + # .add('precision', str, default = "default") + # class_data = args.parse(jdata) + self.n_neuron = neuron + self.resnet_dt = resnet_dt + self.sel_type = sel_type + self.fit_diag = fit_diag + self.seed = seed + self.uniform_seed = uniform_seed + self.seed_shift = one_layer_rand_seed_shift() + #self.diag_shift = diag_shift + self.shift_diag = shift_diag + self.scale = scale + self.fitting_activation_fn = get_activation_func(activation_function) + self.fitting_precision = get_precision(precision) + if self.sel_type is None: + self.sel_type = [ii for ii in range(self.ntypes)] + if self.scale is None: + self.scale = [1.0 for ii in range(self.ntypes)] + #if self.diag_shift is None: + # self.diag_shift = [0.0 for ii in range(self.ntypes)] + if type(self.sel_type) is not list: + self.sel_type = [self.sel_type] + self.constant_matrix = np.zeros(len(self.sel_type)) # len(sel_type) x 1, store the average diagonal value + #if type(self.diag_shift) is not list: + # self.diag_shift = [self.diag_shift] + if type(self.scale) is not list: + self.scale = [self.scale] + self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1() + self.dim_rot_mat = self.dim_rot_mat_1 * 3 + self.useBN = False + + def get_sel_type(self) -> List[int]: + """ + Get selected atom types + """ + return self.sel_type + + def get_out_size(self) -> int: + """ + Get the output size. Should be 9 + """ + return 9 + + def compute_input_stats(self, + all_stat, + protection = 1e-2): + """ + Compute the input statistics + + Parameters: + all_stat + Dictionary of inputs. + can be prepared by model.make_stat_input + protection + Divided-by-zero protection + """ + if not ('polarizability' in all_stat.keys()): + self.avgeig = np.zeros([9]) + warnings.warn('no polarizability data, cannot do data stat. use zeros as guess') + return + data = all_stat['polarizability'] + all_tmp = [] + for ss in range(len(data)): + tmp = np.concatenate(data[ss], axis = 0) + tmp = np.reshape(tmp, [-1, 3, 3]) + tmp,_ = np.linalg.eig(tmp) + tmp = np.absolute(tmp) + tmp = np.sort(tmp, axis = 1) + all_tmp.append(tmp) + all_tmp = np.concatenate(all_tmp, axis = 1) + self.avgeig = np.average(all_tmp, axis = 0) + + # YWolfeee: support polar normalization, initialize to a more appropriate point + if self.shift_diag: + mean_polar = np.zeros([len(self.sel_type), 9]) + sys_matrix, polar_bias = [], [] + for ss in range(len(all_stat['type'])): + atom_has_polar = [w for w in all_stat['type'][ss][0] if (w in self.sel_type)] # select atom with polar + if all_stat['find_atomic_polarizability'][ss] > 0.0: + for itype in range(len(self.sel_type)): # Atomic polar mode, should specify the atoms + index_lis = [index for index, w in enumerate(atom_has_polar) \ + if atom_has_polar[index] == self.sel_type[itype]] # select index in this type + + sys_matrix.append(np.zeros((1,len(self.sel_type)))) + sys_matrix[-1][0,itype] = len(index_lis) + + polar_bias.append(np.sum( + all_stat['atomic_polarizability'][ss].reshape((-1,9))[index_lis],axis=0).reshape((1,9))) + else: # No atomic polar in this system, so it should have global polar + if not all_stat['find_polarizability'][ss] > 0.0: # This system is jsut a joke? + continue + # Till here, we have global polar + sys_matrix.append(np.zeros((1,len(self.sel_type)))) # add a line in the equations + for itype in range(len(self.sel_type)): # Atomic polar mode, should specify the atoms + index_lis = [index for index, w in enumerate(atom_has_polar) \ + if atom_has_polar[index] == self.sel_type[itype]] # select index in this type + + sys_matrix[-1][0,itype] = len(index_lis) + + # 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) + atom_polar,_,_,_ \ + = np.linalg.lstsq(matrix, bias, rcond = 1e-3) + for itype in range(len(self.sel_type)): + self.constant_matrix[itype] = np.mean(np.diagonal(atom_polar[itype].reshape((3,3)))) + + def build (self, + input_d : tf.Tensor, + rot_mat : tf.Tensor, + natoms : tf.Tensor, + reuse : bool = None, + suffix : str = '') : + """ + Build the computational graph for fitting net + + Parameters + ---------- + input_d + The input descriptor + rot_mat + The rotation matrix from the descriptor. + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Return + ------ + atomic_polar + The atomic polarizability + """ + start_index = 0 + inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + rot_mat = tf.reshape(rot_mat, [-1, self.dim_rot_mat * natoms[0]]) + + count = 0 + for type_i in range(self.ntypes): + # cut-out inputs + inputs_i = tf.slice (inputs, + [ 0, start_index* self.dim_descrpt], + [-1, natoms[2+type_i]* self.dim_descrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) + rot_mat_i = tf.slice (rot_mat, + [ 0, start_index* self.dim_rot_mat], + [-1, natoms[2+type_i]* self.dim_rot_mat] ) + rot_mat_i = tf.reshape(rot_mat_i, [-1, self.dim_rot_mat_1, 3]) + start_index += natoms[2+type_i] + if not type_i in self.sel_type : + continue + layer = inputs_i + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + else : + layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + if self.fit_diag : + bavg = np.zeros(self.dim_rot_mat_1) + # bavg[0] = self.avgeig[0] + # bavg[1] = self.avgeig[1] + # bavg[2] = self.avgeig[2] + # (nframes x natoms) x naxis + final_layer = one_layer(layer, self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # (nframes x natoms) x naxis + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1]) + # (nframes x natoms) x naxis x naxis + final_layer = tf.matrix_diag(final_layer) + else : + bavg = np.zeros(self.dim_rot_mat_1*self.dim_rot_mat_1) + # bavg[0*self.dim_rot_mat_1+0] = self.avgeig[0] + # bavg[1*self.dim_rot_mat_1+1] = self.avgeig[1] + # bavg[2*self.dim_rot_mat_1+2] = self.avgeig[2] + # (nframes x natoms) x (naxis x naxis) + final_layer = one_layer(layer, self.dim_rot_mat_1*self.dim_rot_mat_1, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, bavg = bavg, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # (nframes x natoms) x naxis x naxis + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.dim_rot_mat_1, self.dim_rot_mat_1]) + # (nframes x natoms) x naxis x naxis + final_layer = final_layer + tf.transpose(final_layer, perm = [0,2,1]) + # (nframes x natoms) x naxis x 3(coord) + final_layer = tf.matmul(final_layer, rot_mat_i) + # (nframes x natoms) x 3(coord) x 3(coord) + final_layer = tf.matmul(rot_mat_i, final_layer, transpose_a = True) + # nframes x natoms x 3 x 3 + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], 3, 3]) + # shift and scale + sel_type_idx = self.sel_type.index(type_i) + final_layer = final_layer * self.scale[sel_type_idx] + final_layer = final_layer + self.constant_matrix[sel_type_idx] * tf.eye(3, batch_shape=[tf.shape(inputs)[0], natoms[2+type_i]], dtype = GLOBAL_TF_FLOAT_PRECISION) + + # concat the results + if count == 0: + outs = final_layer + else: + outs = tf.concat([outs, final_layer], axis = 1) + count += 1 + + tf.summary.histogram('fitting_net_output', outs) + return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) + + +class GlobalPolarFittingSeA () : + """ + Fit the system polarizability with descriptor se_a + """ + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__ (self, + descrpt : tf.Tensor, + neuron : List[int] = [120,120,120], + resnet_dt : bool = True, + sel_type : List[int] = None, + fit_diag : bool = True, + scale : List[float] = None, + diag_shift : List[float] = None, + seed : int = None, + activation_function : str = 'tanh', + precision : str = 'default' + ) -> None: + """ + Constructor + + Parameters + ---------- + descrpt : tf.Tensor + The descrptor + neuron : List[int] + Number of neurons in each hidden layer of the fitting net + resnet_dt : bool + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + sel_type : List[int] + The atom types selected to have an atomic polarizability prediction + fit_diag : bool + Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix. + scale : List[float] + The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i] + diag_shift : List[float] + The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale. + seed : int + Random seed for initializing the network parameters. + activation_function : str + The activation function in the embedding net. Supported options are {0} + precision : str + The precision of the embedding net parameters. Supported options are {1} + """ + if not isinstance(descrpt, DescrptSeA) : + raise RuntimeError('GlobalPolarFittingSeA only supports DescrptSeA') + self.ntypes = descrpt.get_ntypes() + self.dim_descrpt = descrpt.get_dim_out() + self.polar_fitting = PolarFittingSeA(descrpt, + neuron, + resnet_dt, + sel_type, + fit_diag, + scale, + diag_shift, + seed, + activation_function, + precision) + + def get_sel_type(self) -> int: + """ + Get selected atom types + """ + return self.polar_fitting.get_sel_type() + + def get_out_size(self) -> int: + """ + Get the output size. Should be 9 + """ + return self.polar_fitting.get_out_size() + + def build (self, + input_d, + rot_mat, + natoms, + reuse = None, + suffix = '') -> tf.Tensor: + """ + Build the computational graph for fitting net + + Parameters + ---------- + input_d + The input descriptor + rot_mat + The rotation matrix from the descriptor. + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Return + ------ + polar + The system polarizability + """ + inputs = tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]) + outs = self.polar_fitting.build(input_d, rot_mat, natoms, reuse, suffix) + # nframes x natoms x 9 + outs = tf.reshape(outs, [tf.shape(inputs)[0], -1, 9]) + outs = tf.reduce_sum(outs, axis = 1) + tf.summary.histogram('fitting_net_output', outs) + return tf.reshape(outs, [-1]) + diff --git a/deepmd/fit/wfc.py b/deepmd/fit/wfc.py new file mode 100644 index 0000000000..fce82bd04b --- /dev/null +++ b/deepmd/fit/wfc.py @@ -0,0 +1,103 @@ +import warnings +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import ClassArg, add_data_requirement, get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.network import one_layer, one_layer_rand_seed_shift +from deepmd.descriptor import DescrptLocFrame + +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION + +class WFCFitting () : + """ + Fitting Wannier function centers (WFCs) with local frame descriptor. Not supported anymore. + """ + def __init__ (self, jdata, descrpt): + if not isinstance(descrpt, DescrptLocFrame) : + raise RuntimeError('WFC only supports DescrptLocFrame') + self.ntypes = descrpt.get_ntypes() + self.dim_descrpt = descrpt.get_dim_out() + args = ClassArg()\ + .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\ + .add('resnet_dt', bool, default = True)\ + .add('wfc_numb', int, must = True)\ + .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'wfc_type')\ + .add('seed', int)\ + .add("activation_function", str, default = "tanh")\ + .add('precision', str, default = "default")\ + .add('uniform_seed', bool, default = False) + class_data = args.parse(jdata) + self.n_neuron = class_data['neuron'] + self.resnet_dt = class_data['resnet_dt'] + self.wfc_numb = class_data['wfc_numb'] + self.sel_type = class_data['sel_type'] + self.seed = class_data['seed'] + self.uniform_seed = class_data['uniform_seed'] + self.seed_shift = one_layer_rand_seed_shift() + self.fitting_activation_fn = get_activation_func(class_data["activation_function"]) + self.fitting_precision = get_precision(class_data['precision']) + self.useBN = False + + + def get_sel_type(self): + return self.sel_type + + def get_wfc_numb(self): + return self.wfc_numb + + def get_out_size(self): + return self.wfc_numb * 3 + + def build (self, + input_d, + rot_mat, + natoms, + reuse = None, + suffix = '') : + start_index = 0 + inputs = tf.cast(tf.reshape(input_d, [-1, self.dim_descrpt * natoms[0]]), self.fitting_precision) + rot_mat = tf.reshape(rot_mat, [-1, 9 * natoms[0]]) + + count = 0 + for type_i in range(self.ntypes): + # cut-out inputs + inputs_i = tf.slice (inputs, + [ 0, start_index* self.dim_descrpt], + [-1, natoms[2+type_i]* self.dim_descrpt] ) + inputs_i = tf.reshape(inputs_i, [-1, self.dim_descrpt]) + rot_mat_i = tf.slice (rot_mat, + [ 0, start_index* 9], + [-1, natoms[2+type_i]* 9] ) + rot_mat_i = tf.reshape(rot_mat_i, [-1, 3, 3]) + start_index += natoms[2+type_i] + if not type_i in self.sel_type : + continue + layer = inputs_i + for ii in range(0,len(self.n_neuron)) : + if ii >= 1 and self.n_neuron[ii] == self.n_neuron[ii-1] : + layer+= one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, use_timestep = self.resnet_dt, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + else : + layer = one_layer(layer, self.n_neuron[ii], name='layer_'+str(ii)+'_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, activation_fn = self.fitting_activation_fn, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # (nframes x natoms) x (nwfc x 3) + final_layer = one_layer(layer, self.wfc_numb * 3, activation_fn = None, name='final_layer_type_'+str(type_i)+suffix, reuse=reuse, seed = self.seed, precision = self.fitting_precision, uniform_seed = self.uniform_seed) + if (not self.uniform_seed) and (self.seed is not None): self.seed += self.seed_shift + # (nframes x natoms) x nwfc(wc) x 3(coord_local) + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0] * natoms[2+type_i], self.wfc_numb, 3]) + # (nframes x natoms) x nwfc(wc) x 3(coord) + final_layer = tf.matmul(final_layer, rot_mat_i) + # nframes x natoms x nwfc(wc) x 3(coord_local) + final_layer = tf.reshape(final_layer, [tf.shape(inputs)[0], natoms[2+type_i], self.wfc_numb, 3]) + + # concat the results + if count == 0: + outs = final_layer + else: + outs = tf.concat([outs, final_layer], axis = 1) + count += 1 + + tf.summary.histogram('fitting_net_output', outs) + return tf.cast(tf.reshape(outs, [-1]), GLOBAL_TF_FLOAT_PRECISION) diff --git a/deepmd/infer/__init__.py b/deepmd/infer/__init__.py new file mode 100644 index 0000000000..c0c7c73182 --- /dev/null +++ b/deepmd/infer/__init__.py @@ -0,0 +1,74 @@ +"""Submodule containing all the implemented potentials.""" + +from pathlib import Path +from typing import Union + +from .data_modifier import DipoleChargeModifier +from .deep_dipole import DeepDipole +from .deep_eval import DeepEval +from .deep_polar import DeepGlobalPolar, DeepPolar +from .deep_pot import DeepPot +from .deep_wfc import DeepWFC +from .ewald_recp import EwaldRecp + +__all__ = [ + "DeepPotential", + "DeepDipole", + "DeepEval", + "DeepGlobalPolar", + "DeepPolar", + "DeepPot", + "DeepWFC", + "DipoleChargeModifier", + "EwaldRecp", +] + + +def DeepPotential( + model_file: Union[str, Path], + load_prefix: str = "load", + default_tf_graph: bool = False, +) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]: + """Factory function that will inialize appropriate potential read from `model_file`. + + Parameters + ---------- + model_file: str + The name of the frozen model file. + load_prefix: str + 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 + + Returns + ------- + Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC] + one of the available potentials + + Raises + ------ + RuntimeError + if model file does not correspond to any implementd potential + """ + mf = Path(model_file) + + model_type = DeepEval( + mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph + ).model_type + + if model_type == "ener": + dp = DeepPot(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + elif model_type == "dipole": + dp = DeepDipole(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + elif model_type == "polar": + dp = DeepPolar(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + elif model_type == "global_polar": + dp = DeepGlobalPolar( + mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph + ) + elif model_type == "wfc": + dp = DeepWFC(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + else: + raise RuntimeError(f"unknow model type {model_type}") + + return dp diff --git a/source/train/DataModifier.py b/deepmd/infer/data_modifier.py similarity index 83% rename from source/train/DataModifier.py rename to deepmd/infer/data_modifier.py index 0afa5f1c4d..c6fec62564 100644 --- a/source/train/DataModifier.py +++ b/deepmd/infer/data_modifier.py @@ -1,24 +1,43 @@ import os import numpy as np -from deepmd.DeepDipole import DeepDipole +from typing import Tuple, List + +from deepmd.infer.deep_dipole import DeepDipole +from deepmd.infer.ewald_recp import EwaldRecp from deepmd.env import tf from deepmd.common import select_idx_map, make_default_mesh -from deepmd.EwaldRecp import EwaldRecp -from deepmd.RunOptions import global_tf_float_precision -from deepmd.RunOptions import global_np_float_precision -from deepmd.RunOptions import global_ener_float_precision -from deepmd.RunOptions import global_cvt_2_tf_float -from deepmd.RunOptions import global_cvt_2_ener_float +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import global_cvt_2_ener_float from deepmd.env import op_module class DipoleChargeModifier(DeepDipole): def __init__(self, - model_name, - model_charge_map, - sys_charge_map, - ewald_h = 1, - ewald_beta = 1): + model_name : str, + model_charge_map : List[float], + sys_charge_map : List[float], + ewald_h : float = 1, + ewald_beta : float = 1 + ) -> None: + """ + Constructor + + Parameters + ---------- + model_name + The model file for the DeepDipole model + model_charge_map + Gives the amount of charge for the wfcc + sys_charge_map + Gives the amount of charge for the real atoms + ewald_h + Grid spacing of the reciprocal part of Ewald sum. Unit: A + ewald_beta + Splitting parameter of the Ewald sum. Unit: A^{-1} + """ # the dipole model is loaded with prefix 'dipole_charge' self.modifier_prefix = 'dipole_charge' # init dipole model @@ -49,7 +68,11 @@ def __init__(self, self.force = None self.ntypes = len(self.sel_a) - def build_fv_graph(self): + + def build_fv_graph(self) -> tf.Tensor: + """ + Build the computational graph for the force and virial inference. + """ with tf.variable_scope('modifier_attr') : t_mdl_name = tf.constant(self.model_name, name = 'mdl_name', @@ -72,8 +95,9 @@ def build_fv_graph(self): with self.graph.as_default(): return self._build_fv_graph_inner() + def _build_fv_graph_inner(self): - self.t_ef = tf.placeholder(global_tf_float_precision, [None], name = 't_ef') + self.t_ef = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name = 't_ef') nf = 10 nfxnas = 64*nf nfxna = 192*nf @@ -165,7 +189,7 @@ def _enrich(self, dipole, dof = 3): sel_start_idx += self.t_natoms[2+type_i] else: di = tf.zeros([tf.shape(dipole)[0], self.t_natoms[2+type_i] * dof], - dtype = global_tf_float_precision) + dtype = GLOBAL_TF_FLOAT_PRECISION) coll.append(di) return tf.concat(coll, axis = 1) @@ -182,9 +206,39 @@ def _slice_descrpt_deriv(self, deriv): return tf.concat(coll, axis = 1) - def eval(self, coord, box, atype, eval_fv = True): + def eval(self, + coord : np.array, + box : np.array, + atype : np.array, + eval_fv : bool = True + ) -> Tuple[np.array, np.array, np.array]: + """ + Evaluate the modification + + Parameters + ---------- + coord + The coordinates of atoms + box + The simulation region. PBC is assumed + atype + The atom types + eval_fv + Evaluate force and virial + + Returns + ------- + tot_e + The energy modification + tot_f + The force modification + tot_v + The virial modification + """ + atype = np.array(atype, dtype=int) coord, atype, imap = self.sort_input(coord, atype) - natoms = coord.shape[1] // 3 + # natoms = coord.shape[1] // 3 + natoms = atype.size nframes = coord.shape[0] box = np.reshape(box, [nframes, 9]) atype = np.reshape(atype, [natoms]) @@ -227,7 +281,7 @@ def eval(self, coord, box, atype, eval_fv = True): corr_v = [] corr_av = [] for ii in range(0,nframes,batch_size): - f, v, av = self.eval_fv(coord[ii:ii+batch_size], box[ii:ii+batch_size], atype, ext_f[ii:ii+batch_size]) + f, v, av = self._eval_fv(coord[ii:ii+batch_size], box[ii:ii+batch_size], atype, ext_f[ii:ii+batch_size]) corr_f.append(f) corr_v.append(v) corr_av.append(av) @@ -257,7 +311,7 @@ def eval(self, coord, box, atype, eval_fv = True): return tot_e, tot_f, tot_v - def eval_fv(self, coords, cells, atom_types, ext_f) : + def _eval_fv(self, coords, cells, atom_types, ext_f) : # reshape the inputs cells = np.reshape(cells, [-1, 9]) nframes = cells.shape[0] @@ -321,7 +375,26 @@ def _extend_system(self, coord, box, atype, charge): return all_coord, all_charge, dipole - def modify_data(self, data): + def modify_data(self, + data : dict) -> None: + """ + Modify data. + + Parameters + ---------- + data + Internal data of DeepmdData. + Be a dict, has the following keys + - coord coordinates + - box simulation box + - type atom types + - find_energy tells if data has energy + - find_force tells if data has force + - find_virial tells if data has virial + - energy energy + - force force + - virial virial + """ if 'find_energy' not in data and 'find_force' not in data and 'find_virial' not in data: return diff --git a/deepmd/infer/deep_dipole.py b/deepmd/infer/deep_dipole.py new file mode 100644 index 0000000000..d158cbafac --- /dev/null +++ b/deepmd/infer/deep_dipole.py @@ -0,0 +1,55 @@ +from deepmd.infer.deep_tensor import DeepTensor + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + + +class DeepDipole(DeepTensor): + """Constructor. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + load_prefix: str + 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 + + Warnings + -------- + For developers: `DeepTensor` initializer must be called at the end after + `self.tensors` are modified because it uses the data in `self.tensors` dict. + Do not chanage the order! + """ + + def __init__( + self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False + ) -> None: + + # use this in favor of dict update to move attribute from class to + # instance namespace + self.tensors = dict( + { + # output tensor + "t_tensor": "o_dipole:0", + }, + **self.tensors + ) + + DeepTensor.__init__( + self, + model_file, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + ) + + def get_dim_fparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") + + def get_dim_aparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py new file mode 100644 index 0000000000..41fa9ebd2f --- /dev/null +++ b/deepmd/infer/deep_eval.py @@ -0,0 +1,249 @@ +import os +from typing import List, Optional, TYPE_CHECKING + +import numpy as np +from deepmd.common import make_default_mesh +from deepmd.env import default_tf_session_config, tf, MODEL_VERSION + +if TYPE_CHECKING: + from pathlib import Path + + +class DeepEval: + """Common methods for DeepPot, DeepWFC, DeepPolar, ...""" + + _model_type: Optional[str] = None + _model_version: Optional[str] = None + load_prefix: str # set by subclass + + def __init__( + self, + model_file: "Path", + load_prefix: str = "load", + default_tf_graph: bool = False + ): + self.graph = self._load_graph( + model_file, prefix=load_prefix, default_tf_graph=default_tf_graph + ) + self.load_prefix = load_prefix + + # graph_compatable should be called after graph and prefix are set + if not self._graph_compatable(): + raise RuntimeError( + f"model in graph (version {self.model_version}) is incompatible" + f"with the model (version {MODEL_VERSION}) supported by the current code." + ) + + @property + def model_type(self) -> str: + """Get type of model. + + :type:str + """ + if not self._model_type: + t_mt = self._get_tensor("model_attr/model_type:0") + sess = tf.Session(graph=self.graph, config=default_tf_session_config) + [mt] = sess.run([t_mt], feed_dict={}) + self._model_type = mt.decode("utf-8") + return self._model_type + + @property + def model_version(self) -> str: + """Get type of model. + + :type:str + """ + if not self._model_version: + try: + t_mt = self._get_tensor("model_attr/model_version:0") + sess = tf.Session(graph=self.graph, config=default_tf_session_config) + [mt] = sess.run([t_mt], feed_dict={}) + self._model_version = mt.decode("utf-8") + except KeyError: + # For deepmd-kit version 0.x - 1.x, set model version to 0.0 + self._model_version = "0.0" + return self._model_version + + def _graph_compatable( + self + ) -> bool : + """ Check the model compatability + + Return + bool + If the model stored in the graph file is compatable with the current code + """ + model_version_major = int(self.model_version.split('.')[0]) + model_version_minor = int(self.model_version.split('.')[1]) + MODEL_VERSION_MAJOR = int(MODEL_VERSION.split('.')[0]) + MODEL_VERSION_MINOR = int(MODEL_VERSION.split('.')[1]) + if model_version_major == 0: + # We plan to support model generated from v1.3 + # We have no way to distinguish versions earlier than v1.3 + return True + if (model_version_major != MODEL_VERSION_MAJOR) or \ + (model_version_minor > MODEL_VERSION_MINOR) : + return False + else: + return True + + def _get_tensor( + self, tensor_name: str, attr_name: Optional[str] = None + ) -> tf.Tensor: + """Get TF graph tensor and assign it to class namespace. + + Parameters + ---------- + tensor_name : str + name of tensor to get + attr_name : Optional[str], optional + if specified, class attribute with this name will be created and tensor will + be assigned to it, by default None + + Returns + ------- + tf.Tensor + loaded tensor + """ + tensor_path = os.path.join(self.load_prefix, tensor_name) + tensor = self.graph.get_tensor_by_name(tensor_path) + if attr_name: + setattr(self, attr_name, tensor) + return tensor + else: + return tensor + + @staticmethod + def _load_graph( + frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False + ): + # We load the protobuf file from the disk and parse it to retrieve the + # unserialized graph_def + with tf.gfile.GFile(str(frozen_graph_filename), "rb") as f: + graph_def = tf.GraphDef() + graph_def.ParseFromString(f.read()) + + if default_tf_graph: + tf.import_graph_def( + graph_def, + input_map=None, + return_elements=None, + name=prefix, + producer_op_list=None + ) + graph = tf.get_default_graph() + else : + # Then, we can use again a convenient built-in function to import + # a graph_def into the current default Graph + with tf.Graph().as_default() as graph: + tf.import_graph_def( + graph_def, + input_map=None, + return_elements=None, + name=prefix, + producer_op_list=None + ) + + return graph + + @staticmethod + def sort_input( + coord : np.array, atom_type : np.array, sel_atoms : List[int] = None + ): + """ + Sort atoms in the system according their types. + + Parameters + ---------- + coord + The coordinates of atoms. + Should be of shape [nframes, natoms, 3] + atom_type + The type of atoms + Should be of shape [natoms] + sel_atom + The selected atoms by type + + Returns + ------- + coord_out + The coordinates after sorting + atom_type_out + The atom types after sorting + idx_map + The index mapping from the input to the output. + For example coord_out = coord[:,idx_map,:] + sel_atom_type + Only output if sel_atoms is not None + The sorted selected atom types + sel_idx_map + Only output if sel_atoms is not None + The index mapping from the selected atoms to sorted selected atoms. + """ + if sel_atoms is not None: + selection = [False] * np.size(atom_type) + for ii in sel_atoms: + selection += (atom_type == ii) + sel_atom_type = atom_type[selection] + natoms = atom_type.size + idx = np.arange (natoms) + idx_map = np.lexsort ((idx, atom_type)) + nframes = coord.shape[0] + coord = coord.reshape([nframes, -1, 3]) + coord = np.reshape(coord[:,idx_map,:], [nframes, -1]) + atom_type = atom_type[idx_map] + if sel_atoms is not None: + sel_natoms = np.size(sel_atom_type) + sel_idx = np.arange(sel_natoms) + sel_idx_map = np.lexsort((sel_idx, sel_atom_type)) + sel_atom_type = sel_atom_type[sel_idx_map] + return coord, atom_type, idx_map, sel_atom_type, sel_idx_map + else: + return coord, atom_type, idx_map + + @staticmethod + def reverse_map(vec : np.ndarray, imap : List[int]) -> np.ndarray: + """Reverse mapping of a vector according to the index map + + Parameters + ---------- + vec + Input vector. Be of shape [nframes, natoms, -1] + imap + Index map. Be of shape [natoms] + + Returns + ------- + vec_out + Reverse mapped vector. + """ + ret = np.zeros(vec.shape) + for idx,ii in enumerate(imap) : + ret[:,ii,:] = vec[:,idx,:] + return ret + + + def make_natoms_vec(self, atom_types : np.ndarray) -> np.ndarray : + """Make the natom vector used by deepmd-kit. + + Parameters + ---------- + atom_types + The type of atoms + + Returns + ------- + natoms + The number of atoms. This tensor has the length of Ntypes + 2 + 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_vec = np.zeros (self.ntypes+2).astype(int) + natoms = atom_types.size + natoms_vec[0] = natoms + natoms_vec[1] = natoms + for ii in range (self.ntypes) : + natoms_vec[ii+2] = np.count_nonzero(atom_types == ii) + return natoms_vec diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py new file mode 100644 index 0000000000..7ee02cf1c8 --- /dev/null +++ b/deepmd/infer/deep_polar.py @@ -0,0 +1,140 @@ +from deepmd.infer.deep_tensor import DeepTensor +import numpy as np + +from typing import TYPE_CHECKING, List, Optional + +if TYPE_CHECKING: + from pathlib import Path + + +class DeepPolar(DeepTensor): + """Constructor. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + load_prefix: str + 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 + + Warnings + -------- + For developers: `DeepTensor` initializer must be called at the end after + `self.tensors` are modified because it uses the data in `self.tensors` dict. + Do not chanage the order! + """ + + def __init__( + self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False + ) -> None: + + # use this in favor of dict update to move attribute from class to + # instance namespace + self.tensors = dict( + { + # output tensor + "t_tensor": "o_polar:0", + }, + **self.tensors + ) + + DeepTensor.__init__( + self, + model_file, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + ) + + def get_dim_fparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") + + def get_dim_aparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") + + +class DeepGlobalPolar(DeepTensor): + """Constructor. + + Parameters + ---------- + model_file : str + The name of the frozen model file. + load_prefix: str + 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 + """ + + def __init__( + self, model_file: str, load_prefix: str = "load", default_tf_graph: bool = False + ) -> None: + + self.tensors.update( + { + "t_sel_type": "model_attr/sel_type:0", + # output tensor + "t_tensor": "o_global_polar:0", + } + ) + + DeepTensor.__init__( + self, + model_file, + 9, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + ) + + def eval( + self, + coords: np.array, + cells: np.array, + atom_types: List[int], + atomic: bool = True, + fparam: Optional[np.array] = None, + aparam: Optional[np.array] = None, + efield: Optional[np.array] = None, + ) -> np.array: + """Evaluate the model. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Calculate the atomic energy and virial + fparam + Not used in this model + aparam + Not used in this model + efield + Not used in this model + + Returns + ------- + tensor + The returned tensor + If atomic == False then of size nframes x variable_dof + else of size nframes x natoms x variable_dof + """ + return DeepTensor.eval(self, coords, cells, atom_types, atomic=False) + + def get_dim_fparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") + + def get_dim_aparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py new file mode 100644 index 0000000000..a8e70d5a72 --- /dev/null +++ b/deepmd/infer/deep_pot.py @@ -0,0 +1,335 @@ +import logging +from typing import TYPE_CHECKING, List, Optional, Tuple + +import numpy as np +from deepmd.common import make_default_mesh +from deepmd.env import default_tf_session_config, tf +from deepmd.infer.data_modifier import DipoleChargeModifier +from deepmd.infer.deep_eval import DeepEval + +if TYPE_CHECKING: + from pathlib import Path + +log = logging.getLogger(__name__) + + +class DeepPot(DeepEval): + """Constructor. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + load_prefix: str + 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 + + Warnings + -------- + For developers: `DeepTensor` initializer must be called at the end after + `self.tensors` are modified because it uses the data in `self.tensors` dict. + Do not chanage the order! + """ + + def __init__( + self, + model_file: "Path", + load_prefix: str = "load", + default_tf_graph: bool = False + ) -> 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 + # instance namespace + self.tensors = dict( + { + # descrpt attrs + "t_ntypes": "descrpt_attr/ntypes:0", + "t_rcut": "descrpt_attr/rcut:0", + # fitting attrs + "t_dfparam": "fitting_attr/dfparam:0", + "t_daparam": "fitting_attr/daparam:0", + # model attrs + "t_tmap": "model_attr/tmap:0", + # inputs + "t_coord": "t_coord:0", + "t_type": "t_type:0", + "t_natoms": "t_natoms:0", + "t_box": "t_box:0", + "t_mesh": "t_mesh:0", + # add output tensors + "t_energy": "o_energy:0", + "t_force": "o_force:0", + "t_virial": "o_virial:0", + "t_ae": "o_atom_energy:0", + "t_av": "o_atom_virial:0" + }, + ) + DeepEval.__init__( + self, + model_file, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph + ) + + # load optional tensors + operations = [op.name for op in self.graph.get_operations()] + # check if the graph has these operations: + # if yes add them + if 't_efield' in operations: + self._get_tensor("t_efield:0", "t_efield") + self.has_efield = True + else: + log.debug(f"Could not get tensor 't_efield:0'") + self.t_efield = None + self.has_efield = False + + if 'load/t_fparam' in operations: + self.tensors.update({"t_fparam": "t_fparam:0"}) + self.has_fparam = True + else: + log.debug(f"Could not get tensor 't_fparam:0'") + self.t_fparam = None + self.has_fparam = False + + if 'load/t_aparam' in operations: + self.tensors.update({"t_aparam": "t_aparam:0"}) + self.has_aparam = True + else: + log.debug(f"Could not get tensor 't_aparam:0'") + self.t_aparam = None + self.has_aparam = False + + # now load tensors to object attributes + for attr_name, tensor_name in self.tensors.items(): + self._get_tensor(tensor_name, attr_name) + + # start a tf session associated to the graph + self.sess = tf.Session(graph=self.graph, config=default_tf_session_config) + self._run_default_sess() + self.tmap = self.tmap.decode('UTF-8').split() + + # setup modifier + try: + t_modifier_type = self._get_tensor("modifier_attr/type:0") + self.modifier_type = self.sess.run(t_modifier_type).decode("UTF-8") + except (ValueError, KeyError): + self.modifier_type = None + + if self.modifier_type == "dipole_charge": + t_mdl_name = self._get_tensor("modifier_attr/mdl_name:0") + t_mdl_charge_map = self._get_tensor("modifier_attr/mdl_charge_map:0") + t_sys_charge_map = self._get_tensor("modifier_attr/sys_charge_map:0") + t_ewald_h = self._get_tensor("modifier_attr/ewald_h:0") + t_ewald_beta = self._get_tensor("modifier_attr/ewald_beta:0") + [mdl_name, mdl_charge_map, sys_charge_map, ewald_h, ewald_beta] = self.sess.run([t_mdl_name, t_mdl_charge_map, t_sys_charge_map, t_ewald_h, t_ewald_beta]) + mdl_charge_map = [int(ii) for ii in mdl_charge_map.decode("UTF-8").split()] + sys_charge_map = [int(ii) for ii in sys_charge_map.decode("UTF-8").split()] + self.dm = DipoleChargeModifier(mdl_name, mdl_charge_map, sys_charge_map, ewald_h = ewald_h, ewald_beta = ewald_beta) + + def _run_default_sess(self): + [self.ntypes, self.rcut, self.dfparam, self.daparam, self.tmap] = self.sess.run( + [self.t_ntypes, self.t_rcut, self.t_dfparam, self.t_daparam, self.t_tmap] + ) + + def get_ntypes(self) -> int: + """Get the number of atom types of this model.""" + return self.ntypes + + def get_rcut(self) -> float: + """Get the cut-off radius of this model.""" + return self.rcut + + def get_type_map(self) -> List[int]: + """Get the type map (element name of the atom types) of this model.""" + return self.tmap + + def get_sel_type(self) -> List[int]: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this DP.""" + return self.dfparam + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this DP.""" + return self.daparam + + def eval( + self, + coords: np.array, + cells: np.array, + atom_types: List[int], + atomic: bool = False, + fparam: Optional[np.array] = None, + aparam: Optional[np.array] = None, + efield: Optional[np.array] = None + ) -> Tuple[np.ndarray, ...]: + """Evaluate the energy, force and virial by using this DP. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Calculate the atomic energy and virial + fparam + The frame parameter. + The array can be of size : + - nframes x dim_fparam. + - dim_fparam. Then all frames are assumed to be provided with the same fparam. + aparam + The atomic parameter + The array can be of size : + - nframes x natoms x dim_aparam. + - natoms x dim_aparam. Then all frames are assumed to be provided with the same aparam. + - dim_aparam. Then all frames and atoms are provided with the same aparam. + efield + The external field on atoms. + The array should be of size nframes x natoms x 3 + + Returns + ------- + energy + The system energy. + force + The force on each atom + virial + The virial + atom_energy + The atomic energy. Only returned when atomic == True + atom_virial + The atomic virial. Only returned when atomic == True + """ + if atomic: + if self.modifier_type is not None: + raise RuntimeError('modifier does not support atomic modification') + return self._eval_inner(coords, cells, atom_types, fparam = fparam, aparam = aparam, atomic = atomic, efield = efield) + else : + e, f, v = self._eval_inner(coords, cells, atom_types, fparam = fparam, aparam = aparam, atomic = atomic, efield = efield) + if self.modifier_type is not None: + me, mf, mv = self.dm.eval(coords, cells, atom_types) + e += me.reshape(e.shape) + f += mf.reshape(f.shape) + v += mv.reshape(v.shape) + return e, f, v + + def _eval_inner( + self, + coords, + cells, + atom_types, + fparam=None, + aparam=None, + atomic=False, + efield=None + ): + # standarize the shape of inputs + atom_types = np.array(atom_types, dtype = int).reshape([-1]) + natoms = atom_types.size + coords = np.reshape(np.array(coords), [-1, natoms * 3]) + nframes = coords.shape[0] + if cells is None: + pbc = False + # make cells to work around the requirement of pbc + cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9]) + else: + pbc = True + cells = np.array(cells).reshape([nframes, 9]) + + if self.has_fparam : + assert(fparam is not None) + fparam = np.array(fparam) + if self.has_aparam : + assert(aparam is not None) + aparam = np.array(aparam) + if self.has_efield : + assert(efield is not None), "you are using a model with external field, parameter efield should be provided" + efield = np.array(efield) + + # reshape the inputs + if self.has_fparam : + fdim = self.get_dim_fparam() + if fparam.size == nframes * fdim : + fparam = np.reshape(fparam, [nframes, fdim]) + elif fparam.size == fdim : + fparam = np.tile(fparam.reshape([-1]), [nframes, 1]) + else : + raise RuntimeError('got wrong size of frame param, should be either %d x %d or %d' % (nframes, fdim, fdim)) + if self.has_aparam : + fdim = self.get_dim_aparam() + if aparam.size == nframes * natoms * fdim: + aparam = np.reshape(aparam, [nframes, natoms * fdim]) + elif aparam.size == natoms * fdim : + aparam = np.tile(aparam.reshape([-1]), [nframes, 1]) + elif aparam.size == fdim : + aparam = np.tile(aparam.reshape([-1]), [nframes, natoms]) + else : + raise RuntimeError('got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d' % (nframes, natoms, fdim, natoms, fdim, fdim)) + + # sort inputs + coords, atom_types, imap = self.sort_input(coords, atom_types) + if self.has_efield: + efield = np.reshape(efield, [nframes, natoms, 3]) + efield = efield[:,imap,:] + efield = np.reshape(efield, [nframes, natoms*3]) + + # make natoms_vec and default_mesh + natoms_vec = self.make_natoms_vec(atom_types) + assert(natoms_vec[0] == natoms) + + # evaluate + feed_dict_test = {} + feed_dict_test[self.t_natoms] = natoms_vec + feed_dict_test[self.t_type ] = np.tile(atom_types, [nframes, 1]).reshape([-1]) + t_out = [self.t_energy, + self.t_force, + self.t_virial] + if atomic : + t_out += [self.t_ae, + self.t_av] + + feed_dict_test[self.t_coord] = np.reshape(coords, [-1]) + feed_dict_test[self.t_box ] = np.reshape(cells , [-1]) + if self.has_efield: + feed_dict_test[self.t_efield]= np.reshape(efield, [-1]) + if pbc: + feed_dict_test[self.t_mesh ] = make_default_mesh(cells) + else: + feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32) + 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]) + v_out = self.sess.run (t_out, feed_dict = feed_dict_test) + energy = v_out[0] + force = v_out[1] + virial = v_out[2] + if atomic: + ae = v_out[3] + av = v_out[4] + + # reverse map of the outputs + force = self.reverse_map(np.reshape(force, [nframes,-1,3]), imap) + if atomic : + ae = self.reverse_map(np.reshape(ae, [nframes,-1,1]), imap) + av = self.reverse_map(np.reshape(av, [nframes,-1,9]), imap) + + energy = np.reshape(energy, [nframes, 1]) + force = np.reshape(force, [nframes, natoms, 3]) + virial = np.reshape(virial, [nframes, 9]) + if atomic: + ae = np.reshape(ae, [nframes, natoms, 1]) + av = np.reshape(av, [nframes, natoms, 9]) + return energy, force, virial, ae, av + else : + return energy, force, virial diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py new file mode 100644 index 0000000000..24a7832a32 --- /dev/null +++ b/deepmd/infer/deep_tensor.py @@ -0,0 +1,172 @@ +import os +from typing import List, Optional, TYPE_CHECKING + +import numpy as np +from deepmd.common import make_default_mesh +from deepmd.env import default_tf_session_config, tf +from deepmd.infer.deep_eval import DeepEval + +if TYPE_CHECKING: + from pathlib import Path + +class DeepTensor(DeepEval): + """Evaluates a tensor model. + + Constructor + + Parameters + ---------- + model_file: str + The name of the frozen model file. + load_prefix: str + 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 + """ + + tensors = { + # descriptor attrs + "t_ntypes": "descrpt_attr/ntypes:0", + "t_rcut": "descrpt_attr/rcut:0", + # model attrs + "t_tmap": "model_attr/tmap:0", + "t_sel_type": "model_attr/sel_type:0", + "t_ouput_dim": "model_attr/output_dim:0", + # inputs + "t_coord": "t_coord:0", + "t_type": "t_type:0", + "t_natoms": "t_natoms:0", + "t_box": "t_box:0", + "t_mesh": "t_mesh:0", + } + + def __init__( + self, + model_file: "Path", + load_prefix: str = 'load', + default_tf_graph: bool = False + ) -> None: + DeepEval.__init__( + self, + model_file, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph + ) + # now load tensors to object attributes + for attr_name, tensor_name in self.tensors.items(): + self._get_tensor(tensor_name, attr_name) + + # start a tf session associated to the graph + self.sess = tf.Session(graph=self.graph, config=default_tf_session_config) + self._run_default_sess() + self.tmap = self.tmap.decode('UTF-8').split() + + def _run_default_sess(self): + [self.ntypes, self.rcut, self.tmap, self.tselt, self.output_dim] \ + = self.sess.run( + [self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type, self.t_ouput_dim] + ) + + def get_ntypes(self) -> int: + """Get the number of atom types of this model.""" + return self.ntypes + + def get_rcut(self) -> float: + """Get the cut-off radius of this model.""" + return self.rcut + + def get_type_map(self) -> List[int]: + """Get the type map (element name of the atom types) of this model.""" + return self.tmap + + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model.""" + return self.tselt + + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this DP.""" + return self.dfparam + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this DP.""" + return self.daparam + + def eval( + self, + coords: np.array, + cells: np.array, + atom_types: List[int], + atomic: bool = True, + fparam: Optional[np.array] = None, + aparam: Optional[np.array] = None, + efield: Optional[np.array] = None + ) -> np.array: + """Evaluate the model. + + Parameters + ---------- + coords + The coordinates of atoms. + The array should be of size nframes x natoms x 3 + cells + The cell of the region. + If None then non-PBC is assumed, otherwise using PBC. + The array should be of size nframes x 9 + atom_types + The atom types + The list should contain natoms ints + atomic + Calculate the atomic energy and virial + fparam + Not used in this model + aparam + Not used in this model + efield + Not used in this model + + Returns + ------- + tensor + The returned tensor + If atomic == False then of size nframes x output_dim + else of size nframes x natoms x output_dim + """ + # standarize the shape of inputs + coords = np.array(coords) + cells = np.array(cells) + atom_types = np.array(atom_types, dtype = int) + + # reshape the inputs + cells = np.reshape(cells, [-1, 9]) + nframes = cells.shape[0] + coords = np.reshape(coords, [nframes, -1]) + natoms = coords.shape[1] // 3 + + # sort inputs + coords, atom_types, imap, sel_at, sel_imap = self.sort_input(coords, atom_types, sel_atoms = self.get_sel_type()) + + # make natoms_vec and default_mesh + natoms_vec = self.make_natoms_vec(atom_types) + assert(natoms_vec[0] == natoms) + + # evaluate + tensor = [] + feed_dict_test = {} + feed_dict_test[self.t_natoms] = natoms_vec + feed_dict_test[self.t_type ] = np.tile(atom_types, [nframes,1]).reshape([-1]) + t_out = [self.t_tensor] + 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(cells) + v_out = self.sess.run (t_out, feed_dict = feed_dict_test) + tensor = v_out[0] + + # reverse map of the outputs + if atomic: + tensor = np.array(tensor) + tensor = self.reverse_map(np.reshape(tensor, [nframes,-1,self.output_dim]), sel_imap) + tensor = np.reshape(tensor, [nframes, len(sel_at), self.output_dim]) + else: + tensor = np.reshape(tensor, [nframes, self.output_dim]) + + return tensor diff --git a/deepmd/infer/deep_wfc.py b/deepmd/infer/deep_wfc.py new file mode 100644 index 0000000000..40d3cd6a5c --- /dev/null +++ b/deepmd/infer/deep_wfc.py @@ -0,0 +1,54 @@ +from deepmd.infer.deep_tensor import DeepTensor + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pathlib import Path + + +class DeepWFC(DeepTensor): + """Constructor. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + load_prefix: str + 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 + + Warnings + -------- + For developers: `DeepTensor` initializer must be called at the end after + `self.tensors` are modified because it uses the data in `self.tensors` dict. + Do not chanage the order! + """ + + def __init__( + self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False + ) -> None: + + # use this in favor of dict update to move attribute from class to + # instance namespace + self.tensors = dict( + { + # output tensor + "t_tensor": "o_wfc:0", + }, + **self.tensors + ) + DeepTensor.__init__( + self, + model_file, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + ) + + def get_dim_fparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") + + def get_dim_aparam(self) -> int: + """Unsupported in this model.""" + raise NotImplementedError("This model type does not support this attribute") diff --git a/source/train/EwaldRecp.py b/deepmd/infer/ewald_recp.py similarity index 53% rename from source/train/EwaldRecp.py rename to deepmd/infer/ewald_recp.py index 2517669454..68ee06c552 100644 --- a/source/train/EwaldRecp.py +++ b/deepmd/infer/ewald_recp.py @@ -1,26 +1,41 @@ import numpy as np +from typing import Tuple, List + from deepmd.env import tf from deepmd.common import ClassArg -from deepmd.RunOptions import global_tf_float_precision -from deepmd.RunOptions import global_np_float_precision -from deepmd.RunOptions import global_ener_float_precision -from deepmd.RunOptions import global_cvt_2_tf_float -from deepmd.RunOptions import global_cvt_2_ener_float +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import global_cvt_2_ener_float from deepmd.env import op_module from deepmd.env import default_tf_session_config class EwaldRecp () : + """ + Evaluate the reciprocal part of the Ewald sum + """ def __init__(self, hh, beta): + """ + Constructor + + Parameters + ---------- + hh + Grid spacing of the reciprocal part of Ewald sum. Unit: A + beta + Splitting parameter of the Ewald sum. Unit: A^{-1} + """ self.hh = hh self.beta = beta with tf.Graph().as_default() as graph: # place holders self.t_nloc = tf.placeholder(tf.int32, [1], name = "t_nloc") - self.t_coord = tf.placeholder(global_tf_float_precision, [None], name='t_coord') - self.t_charge = tf.placeholder(global_tf_float_precision, [None], name='t_charge') - self.t_box = tf.placeholder(global_tf_float_precision, [None], name='t_box') + self.t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_coord') + self.t_charge = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_charge') + self.t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_box') # output self.t_energy, self.t_force, self.t_virial \ = op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box, @@ -29,9 +44,31 @@ def __init__(self, self.sess = tf.Session(graph=graph, config=default_tf_session_config) def eval(self, - coord, - charge, - box) : + coord : np.array, + charge : np.array, + box : np.array + ) -> Tuple[np.array, np.array, np.array] : + """ + Evaluate + + Parameters + ---------- + coord + The coordinates of atoms + charge + The atomic charge + box + The simulation region. PBC is assumed + + Returns + ------- + e + The energy + f + The force + v + The virial + """ coord = np.array(coord) charge = np.array(charge) box = np.array(box) diff --git a/deepmd/loggers/__init__.py b/deepmd/loggers/__init__.py new file mode 100644 index 0000000000..8eb1868ae3 --- /dev/null +++ b/deepmd/loggers/__init__.py @@ -0,0 +1,5 @@ +"""Module taking care of logging duties.""" + +from .loggers import set_log_handles + +__all__ = ["set_log_handles"] diff --git a/deepmd/loggers/loggers.py b/deepmd/loggers/loggers.py new file mode 100644 index 0000000000..f787ff1e1a --- /dev/null +++ b/deepmd/loggers/loggers.py @@ -0,0 +1,268 @@ +"""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, + MPI: Optional["MPI"] = 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 than also `MPI` object must be passed in. + by default None + MPI : Optional[MPI, optional] + `MPI` communicator object, must be specified if `mpi_log` is specified, + by default None + + Raises + ------ + RuntimeError + if only one of the arguments `mpi_log`, `MPI` is specified + + 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() + + # remove all old handlers + root_log.setLevel(level) + for hdlr in root_log.handlers[:]: + root_log.removeHandler(hdlr) + + # check if arguments are present + if (mpi_log and not MPI) or (not mpi_log and MPI): + raise RuntimeError("You cannot specify only one of 'mpi_log', 'MPI' arguments") + + # * 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()) + 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/__init__.py b/deepmd/loss/__init__.py new file mode 100644 index 0000000000..94655dd734 --- /dev/null +++ b/deepmd/loss/__init__.py @@ -0,0 +1,4 @@ +from .ener import EnerStdLoss +from .ener import EnerDipoleLoss +from .tensor import TensorLoss + diff --git a/source/train/Loss.py b/deepmd/loss/ener.py similarity index 54% rename from source/train/Loss.py rename to deepmd/loss/ener.py index 1f336325a3..f4979e9b1a 100644 --- a/source/train/Loss.py +++ b/deepmd/loss/ener.py @@ -2,47 +2,50 @@ from deepmd.env import tf from deepmd.common import ClassArg, add_data_requirement -from deepmd.RunOptions import global_cvt_2_tf_float -from deepmd.RunOptions import global_cvt_2_ener_float +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import global_cvt_2_ener_float class EnerStdLoss () : - def __init__ (self, jdata, **kwarg) : - self.starter_learning_rate = kwarg['starter_learning_rate'] - args = ClassArg()\ - .add('start_pref_e', float, default = 0.02)\ - .add('limit_pref_e', float, default = 1.00)\ - .add('start_pref_f', float, default = 1000)\ - .add('limit_pref_f', float, default = 1.00)\ - .add('start_pref_v', float, default = 0)\ - .add('limit_pref_v', float, default = 0)\ - .add('start_pref_ae', float, default = 0)\ - .add('limit_pref_ae', float, default = 0)\ - .add('start_pref_pf', float, default = 0)\ - .add('limit_pref_pf', float, default = 0)\ - .add('relative_f', float) - class_data = args.parse(jdata) - self.start_pref_e = class_data['start_pref_e'] - self.limit_pref_e = class_data['limit_pref_e'] - self.start_pref_f = class_data['start_pref_f'] - self.limit_pref_f = class_data['limit_pref_f'] - self.start_pref_v = class_data['start_pref_v'] - self.limit_pref_v = class_data['limit_pref_v'] - self.start_pref_ae = class_data['start_pref_ae'] - self.limit_pref_ae = class_data['limit_pref_ae'] - self.start_pref_pf = class_data['start_pref_pf'] - self.limit_pref_pf = class_data['limit_pref_pf'] - self.relative_f = class_data['relative_f'] - self.has_e = (self.start_pref_e != 0 or self.limit_pref_e != 0) - self.has_f = (self.start_pref_f != 0 or self.limit_pref_f != 0) - self.has_v = (self.start_pref_v != 0 or self.limit_pref_v != 0) - self.has_ae = (self.start_pref_ae != 0 or self.limit_pref_ae != 0) - self.has_pf = (self.start_pref_pf != 0 or self.limit_pref_pf != 0) + """ + Standard loss function for DP models + """ + def __init__ (self, + starter_learning_rate : float, + start_pref_e : float = 0.02, + limit_pref_e : float = 1.00, + start_pref_f : float = 1000, + limit_pref_f : float = 1.00, + start_pref_v : float = 0.0, + limit_pref_v : float = 0.0, + start_pref_ae : float = 0.0, + limit_pref_ae : float = 0.0, + start_pref_pf : float = 0.0, + limit_pref_pf : float = 0.0, + relative_f : float = None + ) -> None: + self.starter_learning_rate = starter_learning_rate + self.start_pref_e = start_pref_e + self.limit_pref_e = limit_pref_e + self.start_pref_f = start_pref_f + self.limit_pref_f = limit_pref_f + self.start_pref_v = start_pref_v + self.limit_pref_v = limit_pref_v + self.start_pref_ae = start_pref_ae + self.limit_pref_ae = limit_pref_ae + self.start_pref_pf = start_pref_pf + self.limit_pref_pf = limit_pref_pf + self.relative_f = relative_f + self.has_e = (self.start_pref_e != 0.0 or self.limit_pref_e != 0.0) + self.has_f = (self.start_pref_f != 0.0 or self.limit_pref_f != 0.0) + self.has_v = (self.start_pref_v != 0.0 or self.limit_pref_v != 0.0) + self.has_ae = (self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0) + self.has_pf = (self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0) # data required - add_data_requirement('energy', 1, atomic=False, must=self.has_e, high_prec=True) - add_data_requirement('force', 3, atomic=True, must=self.has_f, high_prec=False) - add_data_requirement('virial', 9, atomic=False, must=self.has_v, high_prec=False) - add_data_requirement('atom_ener', 1, atomic=True, must=self.has_ae, high_prec=False) - add_data_requirement('atom_pref', 1, atomic=True, must=self.has_pf, high_prec=False, repeat=3) + add_data_requirement('energy', 1, atomic=False, must=False, high_prec=True) + add_data_requirement('force', 3, atomic=True, must=False, high_prec=False) + add_data_requirement('virial', 9, atomic=False, must=False, high_prec=False) + add_data_requirement('atom_ener', 1, atomic=True, must=False, high_prec=False) + add_data_requirement('atom_pref', 1, atomic=True, must=False, high_prec=False, repeat=3) def build (self, learning_rate, @@ -114,51 +117,93 @@ def build (self, l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss) more_loss['l2_pref_force_loss'] = l2_pref_force_loss + # only used when tensorboard was set as true + self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) + self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', global_cvt_2_tf_float(tf.sqrt(l2_ener_loss)) / global_cvt_2_tf_float(natoms[0])) + self.l2_loss_force_summary = tf.summary.scalar('l2_force_loss', tf.sqrt(l2_force_loss)) + self.l2_loss_virial_summary = tf.summary.scalar('l2_virial_loss', tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0])) + self.l2_l = l2_loss self.l2_more = more_loss return l2_loss, more_loss + def eval(self, sess, feed_dict, natoms): + run_data = [ + self.l2_l, + self.l2_more['l2_ener_loss'], + self.l2_more['l2_force_loss'], + self.l2_more['l2_virial_loss'], + self.l2_more['l2_atom_ener_loss'], + self.l2_more['l2_pref_force_loss'] + ] + error, error_e, error_f, error_v, error_ae, error_pf = sess.run(run_data, feed_dict=feed_dict) + results = {"natoms": natoms[0], "rmse": np.sqrt(error)} + if self.has_e: + results["rmse_e"] = np.sqrt(error_e) / natoms[0] + if self.has_ae: + results["rmse_ae"] = np.sqrt(error_ae) + if self.has_f: + results["rmse_f"] = np.sqrt(error_f) + if self.has_v: + results["rmse_v"] = np.sqrt(error_v) / natoms[0] + if self.has_pf: + results["rmse_pf"] = np.sqrt(error_pf) + return results - def print_header(self) : - prop_fmt = ' %9s %9s' + def print_header(self): # depreciated + prop_fmt = ' %11s %11s' print_str = '' - print_str += prop_fmt % ('l2_tst', 'l2_trn') + print_str += prop_fmt % ('rmse_tst', 'rmse_trn') if self.has_e : - print_str += prop_fmt % ('l2_e_tst', 'l2_e_trn') + print_str += prop_fmt % ('rmse_e_tst', 'rmse_e_trn') if self.has_ae : - print_str += prop_fmt % ('l2_ae_tst', 'l2_ae_trn') + print_str += prop_fmt % ('rmse_ae_tst', 'rmse_ae_trn') if self.has_f : - print_str += prop_fmt % ('l2_f_tst', 'l2_f_trn') + print_str += prop_fmt % ('rmse_f_tst', 'rmse_f_trn') if self.has_v : - print_str += prop_fmt % ('l2_v_tst', 'l2_v_trn') + print_str += prop_fmt % ('rmse_v_tst', 'rmse_v_trn') if self.has_pf : - print_str += prop_fmt % ('l2_pf_tst', 'l2_pf_trn') + print_str += prop_fmt % ('rmse_pf_tst', 'rmse_pf_trn') return print_str - def print_on_training(self, + tb_writer, + cur_batch, sess, natoms, feed_dict_test, - feed_dict_batch) : - error_test, error_e_test, error_f_test, error_v_test, error_ae_test, error_pf_test \ - = sess.run([self.l2_l, \ - self.l2_more['l2_ener_loss'], \ - self.l2_more['l2_force_loss'], \ - self.l2_more['l2_virial_loss'], \ - self.l2_more['l2_atom_ener_loss'],\ - self.l2_more['l2_pref_force_loss']], - feed_dict=feed_dict_test) - error_train, error_e_train, error_f_train, error_v_train, error_ae_train, error_pf_train \ - = sess.run([self.l2_l, \ - self.l2_more['l2_ener_loss'], \ - self.l2_more['l2_force_loss'], \ - self.l2_more['l2_virial_loss'], \ - self.l2_more['l2_atom_ener_loss'],\ - self.l2_more['l2_pref_force_loss']], - feed_dict=feed_dict_batch) + feed_dict_batch): # depreciated + + run_data = [ + self.l2_l, + self.l2_more['l2_ener_loss'], + self.l2_more['l2_force_loss'], + self.l2_more['l2_virial_loss'], + self.l2_more['l2_atom_ener_loss'], + self.l2_more['l2_pref_force_loss'] + ] + + # first train data + train_out = sess.run(run_data, feed_dict=feed_dict_batch) + error_train, error_e_train, error_f_train, error_v_train, error_ae_train, error_pf_train = train_out + + # than test data, if tensorboard log writter is present, commpute summary + # and write tensorboard logs + if tb_writer: + summary_merged_op = tf.summary.merge([self.l2_loss_summary, self.l2_loss_ener_summary, self.l2_loss_force_summary, self.l2_loss_virial_summary]) + run_data.insert(0, summary_merged_op) + + test_out = sess.run(run_data, feed_dict=feed_dict_test) + + if tb_writer: + summary = test_out.pop(0) + tb_writer.add_summary(summary, cur_batch) + + error_test, error_e_test, error_f_test, error_v_test, error_ae_test, error_pf_test = test_out + + print_str = "" - prop_fmt = " %9.2e %9.2e" + prop_fmt = " %11.2e %11.2e" print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train)) if self.has_e : print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0]) @@ -171,11 +216,17 @@ def print_on_training(self, if self.has_pf: print_str += prop_fmt % (np.sqrt(error_pf_test), np.sqrt(error_pf_train)) - return print_str + return print_str class EnerDipoleLoss () : - def __init__ (self, jdata, **kwarg) : + def __init__ (self, + starter_learning_rate : float, + start_pref_e : float = 0.1, + limit_pref_e : float = 1.0, + start_pref_ed : float = 1.0, + limit_pref_ed : float = 1.0 + ) -> None : self.starter_learning_rate = kwarg['starter_learning_rate'] args = ClassArg()\ .add('start_pref_e', float, must = True, default = 0.1) \ @@ -232,105 +283,80 @@ def build (self, more_loss['l2_ener_loss'] = l2_ener_loss more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss + self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) + self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', tf.sqrt(l2_ener_loss) / global_cvt_2_tf_float(natoms[0])) + self.l2_ener_dipole_loss_summary = tf.summary.scalar('l2_ener_dipole_loss', tf.sqrt(l2_ener_dipole_loss)) + self.l2_l = l2_loss self.l2_more = more_loss return l2_loss, more_loss - - def print_header(self) : + def eval(self, sess, feed_dict, natoms): + run_data = [ + self.l2_l, + self.l2_more['l2_ener_loss'], + self.l2_more['l2_ener_dipole_loss'] + ] + error, error_e, error_ed = sess.run(run_data, feed_dict=feed_dict) + results = { + 'natoms': natoms[0], + 'rmse': np.sqrt(error), + 'rmse_e': np.sqrt(error_e) / natoms[0], + 'rmse_ed': np.sqrt(error_ed) + } + return results + + @staticmethod + def print_header() : # depreciated prop_fmt = ' %9s %9s' print_str = '' print_str += prop_fmt % ('l2_tst', 'l2_trn') print_str += prop_fmt % ('l2_e_tst', 'l2_e_trn') print_str += prop_fmt % ('l2_ed_tst', 'l2_ed_trn') - return print_str - + return print_str - def print_on_training(self, + def print_on_training(self, + tb_writer, + cur_batch, sess, natoms, feed_dict_test, - feed_dict_batch) : - error_test, error_e_test, error_ed_test\ - = sess.run([self.l2_l, \ - self.l2_more['l2_ener_loss'], \ - self.l2_more['l2_ener_dipole_loss']], - feed_dict=feed_dict_test) - error_train, error_e_train, error_ed_train\ - = sess.run([self.l2_l, \ - self.l2_more['l2_ener_loss'], \ - self.l2_more['l2_ener_dipole_loss']], - feed_dict=feed_dict_batch) - print_str = "" - prop_fmt = " %9.2e %9.2e" - print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train)) - print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0]) - print_str += prop_fmt % (np.sqrt(error_ed_test), np.sqrt(error_ed_train)) - return print_str - - -class TensorLoss () : - def __init__ (self, jdata, **kwarg) : - try: - model = kwarg['model'] - type_sel = model.get_sel_type() - except : - type_sel = None - self.tensor_name = kwarg['tensor_name'] - self.tensor_size = kwarg['tensor_size'] - self.label_name = kwarg['label_name'] - self.atomic = kwarg.get('atomic', True) - if jdata is not None: - self.scale = jdata.get('scale', 1.0) - else: - self.scale = 1.0 - # data required - add_data_requirement(self.label_name, - self.tensor_size, - atomic=self.atomic, - must=True, - high_prec=False, - type_sel = type_sel) + feed_dict_batch): # depreciated - def build (self, - learning_rate, - natoms, - model_dict, - label_dict, - suffix): - polar_hat = label_dict[self.label_name] - polar = model_dict[self.tensor_name] - l2_loss = tf.reduce_mean( tf.square(self.scale*(polar - polar_hat)), name='l2_'+suffix) - more_loss = {'nonorm': l2_loss} - if not self.atomic : - atom_norm = 1./ global_cvt_2_tf_float(natoms[0]) - l2_loss = l2_loss * atom_norm - self.l2_l = l2_loss - self.l2_more = more_loss['nonorm'] + run_data = [ + self.l2_l, + self.l2_more['l2_ener_loss'], + self.l2_more['l2_ener_dipole_loss'] + ] - return l2_loss, more_loss + # first train data + train_out = sess.run(run_data, feed_dict=feed_dict_batch) + error_train, error_e_train, error_ed_train = train_out - def print_header(self) : - prop_fmt = ' %9s %9s' - print_str = '' - print_str += prop_fmt % ('l2_tst', 'l2_trn') - return print_str + # than test data, if tensorboard log writter is present, commpute summary + # and write tensorboard logs + if tb_writer: + summary_merged_op = tf.summary.merge([ + self.l2_loss_summary, + self.l2_loss_ener_summary, + self.l2_ener_dipole_loss_summary + ]) + run_data.insert(0, summary_merged_op) + + test_out = sess.run(run_data, feed_dict=feed_dict_test) + + if tb_writer: + summary = test_out.pop(0) + tb_writer.add_summary(summary, cur_batch) + + error_test, error_e_test, error_ed_test = test_out - def print_on_training(self, - sess, - natoms, - feed_dict_test, - feed_dict_batch) : - error_test\ - = sess.run([self.l2_more], \ - feed_dict=feed_dict_test) - error_train\ - = sess.run([self.l2_more], \ - feed_dict=feed_dict_batch) print_str = "" prop_fmt = " %9.2e %9.2e" print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train)) + print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0]) + print_str += prop_fmt % (np.sqrt(error_ed_test), np.sqrt(error_ed_train)) + return print_str - return print_str diff --git a/deepmd/loss/tensor.py b/deepmd/loss/tensor.py new file mode 100644 index 0000000000..72a7046d10 --- /dev/null +++ b/deepmd/loss/tensor.py @@ -0,0 +1,194 @@ +import numpy as np +from deepmd.env import tf +from deepmd.common import ClassArg, add_data_requirement + +from deepmd.env import global_cvt_2_tf_float +from deepmd.env import global_cvt_2_ener_float + +class TensorLoss () : + """ + Loss function for tensorial properties. + """ + def __init__ (self, jdata, **kwarg) : + try: + model = kwarg['model'] + self.type_sel = model.get_sel_type() + except : + self.type_sel = None + self.tensor_name = kwarg['tensor_name'] + self.tensor_size = kwarg['tensor_size'] + self.label_name = kwarg['label_name'] + if jdata is not None: + self.scale = jdata.get('scale', 1.0) + else: + self.scale = 1.0 + + # YHT: added for global / local dipole combination + assert jdata is not None, "Please provide loss parameters!" + # YWolfeee: modify, use pref / pref_atomic, instead of pref_weight / pref_atomic_weight + self.local_weight = jdata.get('pref_atomic', None) + self.global_weight = jdata.get('pref', None) + + assert (self.local_weight is not None and self.global_weight is not None), "Both `pref` and `pref_atomic` should be provided." + assert self.local_weight >= 0.0 and self.global_weight >= 0.0, "Can not assign negative weight to `pref` and `pref_atomic`" + assert (self.local_weight >0.0) or (self.global_weight>0.0), AssertionError('Can not assian zero weight both to `pref` and `pref_atomic`') + + # data required + add_data_requirement("atomic_" + self.label_name, + self.tensor_size, + atomic=True, + must=False, + high_prec=False, + type_sel = self.type_sel) + add_data_requirement(self.label_name, + self.tensor_size, + atomic=False, + must=False, + high_prec=False, + type_sel = self.type_sel) + + def build (self, + learning_rate, + natoms, + model_dict, + label_dict, + suffix): + polar_hat = label_dict[self.label_name] + atomic_polar_hat = label_dict["atomic_" + self.label_name] + polar = model_dict[self.tensor_name] + + find_global = label_dict['find_' + self.label_name] + find_atomic = label_dict['find_atomic_' + self.label_name] + + + + # YHT: added for global / local dipole combination + l2_loss = global_cvt_2_tf_float(0.0) + more_loss = { + "local_loss":global_cvt_2_tf_float(0.0), + "global_loss":global_cvt_2_tf_float(0.0) + } + + + if self.local_weight > 0.0: + 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 + l2_loss += self.local_weight * local_loss + self.l2_loss_local_summary = tf.summary.scalar('l2_local_loss', + tf.sqrt(more_loss['local_loss'])) + + + if self.global_weight > 0.0: # Need global loss + atoms = 0 + if self.type_sel is not None: + for w in self.type_sel: + atoms += natoms[2+w] + else: + atoms = natoms[0] + nframes = tf.shape(polar)[0] // self.tensor_size // atoms + # get global results + global_polar = tf.reshape(tf.reduce_sum(tf.reshape( + polar, [nframes, -1, self.tensor_size]), axis=1),[-1]) + #if self.atomic: # If label is local, however + # global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape( + # polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1]) + #else: + # global_polar_hat = polar_hat + + global_loss = global_cvt_2_tf_float(find_global) * tf.reduce_mean( tf.square(self.scale*(global_polar - polar_hat)), name='l2_'+suffix) + + more_loss['global_loss'] = global_loss + self.l2_loss_global_summary = tf.summary.scalar('l2_global_loss', + tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms)) + + # YWolfeee: should only consider atoms with dipole, i.e. atoms + # atom_norm = 1./ global_cvt_2_tf_float(natoms[0]) + atom_norm = 1./ global_cvt_2_tf_float(atoms) + global_loss *= atom_norm + + l2_loss += self.global_weight * global_loss + + self.l2_more = more_loss + self.l2_l = l2_loss + + self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss)) + return l2_loss, more_loss + + def eval(self, sess, feed_dict, natoms): + atoms = 0 + if self.type_sel is not None: + for w in self.type_sel: + atoms += natoms[2+w] + else: + atoms = natoms[0] + + run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']] + error, error_lc, error_gl = sess.run(run_data, feed_dict=feed_dict) + + results = {"natoms": atoms, "rmse": np.sqrt(error)} + if self.local_weight > 0.0: + results["rmse_lc"] = np.sqrt(error_lc) + if self.global_weight > 0.0: + results["rmse_gl"] = np.sqrt(error_gl) / atoms + return results + + def print_header(self): # depreciated + prop_fmt = ' %11s %11s' + print_str = '' + print_str += prop_fmt % ('rmse_tst', 'rmse_trn') + if self.local_weight > 0.0: + print_str += prop_fmt % ('rmse_lc_tst', 'rmse_lc_trn') + if self.global_weight > 0.0: + print_str += prop_fmt % ('rmse_gl_tst', 'rmse_gl_trn') + return print_str + + def print_on_training(self, + tb_writer, + cur_batch, + sess, + natoms, + feed_dict_test, + feed_dict_batch) : # depreciated + + # YHT: added to calculate the atoms number + atoms = 0 + if self.type_sel is not None: + for w in self.type_sel: + atoms += natoms[2+w] + else: + atoms = natoms[0] + + run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']] + summary_list = [self.l2_loss_summary] + if self.local_weight > 0.0: + summary_list.append(self.l2_loss_local_summary) + if self.global_weight > 0.0: + summary_list.append(self.l2_loss_global_summary) + + # first train data + error_train = sess.run(run_data, feed_dict=feed_dict_batch) + + # than test data, if tensorboard log writter is present, commpute summary + # and write tensorboard logs + if tb_writer: + #summary_merged_op = tf.summary.merge([self.l2_loss_summary]) + summary_merged_op = tf.summary.merge(summary_list) + run_data.insert(0, summary_merged_op) + + test_out = sess.run(run_data, feed_dict=feed_dict_test) + + if tb_writer: + summary = test_out.pop(0) + tb_writer.add_summary(summary, cur_batch) + + error_test = test_out + + print_str = "" + prop_fmt = " %11.2e %11.2e" + print_str += prop_fmt % (np.sqrt(error_test[0]), np.sqrt(error_train[0])) + if self.local_weight > 0.0: + print_str += prop_fmt % (np.sqrt(error_test[1]), np.sqrt(error_train[1]) ) + if self.global_weight > 0.0: + print_str += prop_fmt % (np.sqrt(error_test[2])/atoms, np.sqrt(error_train[2])/atoms) + + return print_str diff --git a/deepmd/model/__init__.py b/deepmd/model/__init__.py new file mode 100644 index 0000000000..e50a138aec --- /dev/null +++ b/deepmd/model/__init__.py @@ -0,0 +1,5 @@ +from .ener import EnerModel +from .tensor import WFCModel +from .tensor import DipoleModel +from .tensor import PolarModel +from .tensor import GlobalPolarModel diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py new file mode 100644 index 0000000000..6d7230b6cd --- /dev/null +++ b/deepmd/model/ener.py @@ -0,0 +1,251 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.utils.pair_tab import PairTab +from deepmd.common import ClassArg +from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION +from deepmd.env import op_module +from .model_stat import make_stat_input, merge_sys_stat + +class EnerModel() : + model_type = 'ener' + + def __init__ ( + self, + descrpt, + fitting, + typeebd = None, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2, + use_srtab : str = None, + smin_alpha : float = None, + sw_rmin : float = None, + sw_rmax : float = None + ) -> None: + """ + Constructor + + Parameters + ---------- + descrpt + Descriptor + fitting + Fitting net + type_map + Mapping atom type to the name (str) of the type. + For example `type_map[1]` gives the name of the type 1. + data_stat_nbatch + Number of frames used for data statistic + data_stat_protect + Protect parameter for atomic energy regression + 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. + 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. + sw_rmin + The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + sw_rmin + The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + """ + # descriptor + self.descrpt = descrpt + self.rcut = self.descrpt.get_rcut() + self.ntypes = self.descrpt.get_ntypes() + # fitting + self.fitting = fitting + self.numb_fparam = self.fitting.get_numb_fparam() + # type embedding + self.typeebd = typeebd + # other inputs + if type_map is None: + self.type_map = [] + else: + self.type_map = type_map + self.data_stat_nbatch = data_stat_nbatch + self.data_stat_protect = data_stat_protect + self.srtab_name = use_srtab + if self.srtab_name is not None : + self.srtab = PairTab(self.srtab_name) + self.smin_alpha = smin_alpha + self.sw_rmin = sw_rmin + self.sw_rmax = sw_rmax + else : + self.srtab = None + + + def get_rcut (self) : + return self.rcut + + def get_ntypes (self) : + return self.ntypes + + def get_type_map (self) : + return self.type_map + + def data_stat(self, data): + all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False) + m_all_stat = merge_sys_stat(all_stat) + self._compute_input_stat(m_all_stat, protection = self.data_stat_protect) + self._compute_output_stat(all_stat) + # self.bias_atom_e = data.compute_energy_shift(self.rcond) + + def _compute_input_stat (self, all_stat, protection = 1e-2) : + self.descrpt.compute_input_stats(all_stat['coord'], + all_stat['box'], + all_stat['type'], + all_stat['natoms_vec'], + all_stat['default_mesh'], + all_stat) + self.fitting.compute_input_stats(all_stat, protection = protection) + + def _compute_output_stat (self, all_stat) : + self.fitting.compute_output_stats(all_stat) + + + def build (self, + coord_, + atype_, + natoms, + box, + mesh, + input_dict, + suffix = '', + reuse = None): + + with tf.variable_scope('model_attr' + suffix, reuse = reuse) : + 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) + + if self.srtab is not None : + tab_info, tab_data = self.srtab.get() + 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)) + + coord = tf.reshape (coord_, [-1, natoms[1] * 3]) + atype = tf.reshape (atype_, [-1, natoms[1]]) + + # type embedding if any + if self.typeebd is not None: + type_embedding = self.typeebd.build( + self.ntypes, + reuse = reuse, + suffix = suffix, + ) + input_dict['type_embedding'] = type_embedding + + dout \ + = self.descrpt.build(coord_, + atype_, + natoms, + box, + mesh, + input_dict, + suffix = suffix, + reuse = reuse) + dout = tf.identity(dout, name='o_descriptor') + + if self.srtab is not None : + nlist, rij, sel_a, sel_r = self.descrpt.get_nlist() + nnei_a = np.cumsum(sel_a)[-1] + nnei_r = np.cumsum(sel_r)[-1] + + atom_ener = self.fitting.build (dout, + natoms, + input_dict, + reuse = reuse, + suffix = suffix) + + if self.srtab is not None : + sw_lambda, sw_deriv \ + = op_module.soft_min_switch(atype, + rij, + nlist, + natoms, + sel_a = sel_a, + sel_r = sel_r, + alpha = self.smin_alpha, + rmin = self.sw_rmin, + rmax = self.sw_rmax) + inv_sw_lambda = 1.0 - sw_lambda + # NOTICE: + # atom energy is not scaled, + # force and virial are scaled + tab_atom_ener, tab_force, tab_atom_virial \ + = op_module.pair_tab(self.tab_info, + self.tab_data, + atype, + rij, + nlist, + natoms, + sw_lambda, + sel_a = sel_a, + sel_r = sel_r) + energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]]) + tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape(tab_atom_ener, [-1]) + atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener + energy_raw = tab_atom_ener + atom_ener + else : + energy_raw = atom_ener + + energy_raw = tf.reshape(energy_raw, [-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, virial, atom_virial \ + = self.descrpt.prod_force_virial (atom_ener, natoms) + + if self.srtab is not None : + sw_force \ + = op_module.soft_min_force(energy_diff, + sw_deriv, + nlist, + natoms, + n_a_sel = nnei_a, + n_r_sel = nnei_r) + force = force + sw_force + tab_force + + force = tf.reshape (force, [-1, 3 * natoms[1]], name = "o_force"+suffix) + + if self.srtab is not None : + sw_virial, sw_atom_virial \ + = op_module.soft_min_virial (energy_diff, + sw_deriv, + rij, + nlist, + natoms, + n_a_sel = nnei_a, + n_r_sel = nnei_r) + atom_virial = atom_virial + sw_atom_virial + tab_atom_virial + virial = virial + sw_virial \ + + tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis = 1) + + virial = tf.reshape (virial, [-1, 9], name = "o_virial"+suffix) + atom_virial = tf.reshape (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 + diff --git a/deepmd/model/model_stat.py b/deepmd/model/model_stat.py new file mode 100644 index 0000000000..61f151f27b --- /dev/null +++ b/deepmd/model/model_stat.py @@ -0,0 +1,60 @@ +import numpy as np +from collections import defaultdict + +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 + 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 = list(all_stat.keys())[0] + 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 + diff --git a/deepmd/model/tensor.py b/deepmd/model/tensor.py new file mode 100644 index 0000000000..114a319edd --- /dev/null +++ b/deepmd/model/tensor.py @@ -0,0 +1,182 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.common import ClassArg +from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION +from deepmd.env import op_module +from .model_stat import make_stat_input, merge_sys_stat + +class TensorModel() : + def __init__ ( + self, + tensor_name : str, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2, + )->None: + """ + Constructor + + Parameters + ---------- + tensor_name + Name of the tensor. + descrpt + Descriptor + fitting + Fitting net + type_map + Mapping atom type to the name (str) of the type. + For example `type_map[1]` gives the name of the type 1. + data_stat_nbatch + Number of frames used for data statistic + data_stat_protect + Protect parameter for atomic energy regression + """ + self.model_type = tensor_name + # descriptor + self.descrpt = descrpt + self.rcut = self.descrpt.get_rcut() + self.ntypes = self.descrpt.get_ntypes() + # fitting + self.fitting = fitting + # other params + if type_map is None: + self.type_map = [] + else: + self.type_map = type_map + self.data_stat_nbatch = data_stat_nbatch + self.data_stat_protect = data_stat_protect + + def get_rcut (self) : + return self.rcut + + def get_ntypes (self) : + return self.ntypes + + def get_type_map (self) : + return self.type_map + + def get_sel_type(self): + return self.fitting.get_sel_type() + + def get_out_size (self) : + return self.fitting.get_out_size() + + def data_stat(self, data): + all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False) + m_all_stat = merge_sys_stat(all_stat) + self._compute_input_stat (m_all_stat, protection = self.data_stat_protect) + self._compute_output_stat(all_stat) + + def _compute_input_stat(self, all_stat, protection = 1e-2) : + self.descrpt.compute_input_stats(all_stat['coord'], + all_stat['box'], + all_stat['type'], + all_stat['natoms_vec'], + all_stat['default_mesh'], + all_stat) + if hasattr(self.fitting, 'compute_input_stats'): + self.fitting.compute_input_stats(all_stat, protection = protection) + + def _compute_output_stat (self, all_stat) : + if hasattr(self.fitting, 'compute_output_stats'): + self.fitting.compute_output_stats(all_stat) + + def build (self, + coord_, + atype_, + natoms, + box, + mesh, + input_dict, + suffix = '', + reuse = None): + with tf.variable_scope('model_attr' + suffix, reuse = reuse) : + t_tmap = tf.constant(' '.join(self.type_map), + name = 'tmap', + dtype = tf.string) + t_st = tf.constant(self.get_sel_type(), + name = 'sel_type', + dtype = tf.int32) + 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) + t_od = tf.constant(self.get_out_size(), + name = 'output_dim', + dtype = tf.int32) + + dout \ + = self.descrpt.build(coord_, + atype_, + natoms, + box, + mesh, + input_dict, + suffix = suffix, + reuse = reuse) + dout = tf.identity(dout, name='o_descriptor') + rot_mat = self.descrpt.get_rot_mat() + rot_mat = tf.identity(rot_mat, name = 'o_rot_mat') + + output = self.fitting.build (dout, + rot_mat, + natoms, + reuse = reuse, + suffix = suffix) + output = tf.identity(output, name = 'o_' + self.model_type) + + return {self.model_type: output} + + +class WFCModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'wfc', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + +class DipoleModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'dipole', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + +class PolarModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + +class GlobalPolarModel(TensorModel): + def __init__( + self, + descrpt, + fitting, + type_map : List[str] = None, + data_stat_nbatch : int = 10, + data_stat_protect : float = 1e-2 + ) -> None: + TensorModel.__init__(self, 'global_polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect) + + diff --git a/deepmd/op/__init__.py b/deepmd/op/__init__.py new file mode 100644 index 0000000000..1771422a40 --- /dev/null +++ b/deepmd/op/__init__.py @@ -0,0 +1,28 @@ +"""This module will house cust Tf OPs after CMake installation.""" + +from pathlib import Path +import importlib +import logging + +NOT_LOADABLE = ("__init__.py") +PACKAGE_BASE = "deepmd.op" + +log = logging.getLogger(__name__) + + +def import_ops(): + """Import all custom TF ops that are present in this submodule. + + Note + ---- + Initialy this subdir is unpopulated. CMake will install all the op module python + files and shared libs. + """ + for module_file in Path(__file__).parent.glob("*.py"): + if module_file.name not in NOT_LOADABLE: + module_name = f".{module_file.stem}" + log.debug(f"importing op module: {module_name}") + importlib.import_module(module_name, PACKAGE_BASE) + + +import_ops() diff --git a/deepmd/train/run_options.py b/deepmd/train/run_options.py new file mode 100644 index 0000000000..25029c4308 --- /dev/null +++ b/deepmd/train/run_options.py @@ -0,0 +1,353 @@ +"""Module taking care of important package constants.""" + +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import numpy as np +from deepmd.cluster import get_resource +from deepmd.env import get_tf_default_nthreads, tf, GLOBAL_CONFIG, global_float_prec +from deepmd.loggers import set_log_handles + +if TYPE_CHECKING: + from mpi4py import MPI + + try: + from typing import Protocol # python >=3.8 + except ImportError: + from typing_extensions import Protocol # type: ignore + + class TFServerV1(Protocol): + """Prococol mimicking parser object.""" + + server_def: tf.train.ServerDef + target: str + + +__all__ = [ + "WELCOME", + "CITATION", + "BUILD", + "RunOptions", +] + +log = logging.getLogger(__name__) + + +# http://patorjk.com/software/taag. Font:Big" +WELCOME = ( # noqa + " _____ _____ __ __ _____ _ _ _ ", + "| __ \ | __ \ | \/ || __ \ | | (_)| | ", + "| | | | ___ ___ | |__) || \ / || | | | ______ | | __ _ | |_ ", + "| | | | / _ \ / _ \| ___/ | |\/| || | | ||______|| |/ /| || __|", + "| |__| || __/| __/| | | | | || |__| | | < | || |_ ", + "|_____/ \___| \___||_| |_| |_||_____/ |_|\_\|_| \__|", +) + +CITATION = ( + "Please read and cite:", + "Wang, Zhang, Han and E, Comput.Phys.Comm. 228, 178-184 (2018)", +) + +_sep = "\n " +BUILD = ( + f"installed to: {GLOBAL_CONFIG['install_prefix']}", + f"source : {GLOBAL_CONFIG['git_summ']}", + f"source brach: {GLOBAL_CONFIG['git_branch']}", + f"source commit: {GLOBAL_CONFIG['git_hash']}", + f"source commit at: {GLOBAL_CONFIG['git_date']}", + f"build float prec: {global_float_prec}", + f"build with tf inc: {GLOBAL_CONFIG['tf_include_dir']}", + f"build with tf lib: {GLOBAL_CONFIG['tf_libs'].replace(';', _sep)}" # noqa +) + + +def _is_distributed(MPI: "MPI") -> bool: + """Check if there are more than one MPI processes. + + Parameters + ---------- + MPI : MPI + MPI object + + Returns + ------- + bool + True if we have more than 1 MPI process + """ + return MPI.COMM_WORLD.Get_size() > 1 + + +def _distributed_task_config( + MPI: "MPI", + node_name: str, + node_list_: List[str], + gpu_list: Optional[List[int]] = None, + default_port: int = 2222, +) -> Tuple[Dict[str, List[str]], str, int, str, str]: + """Create configuration for distributed tensorflow session. + + Parameters + ---------- + MPI : mpi4py.MPI + MPI module + node_name : str + the name of current node + node_list_ : List[str] + the list of nodes of the current mpirun + gpu_list : Optional[List[int]], optional + the list of GPUs on each node, by default None + default_port : int, optional + the default port for socket communication, by default 2222 + + Returns + ------- + Tuple[Dict[str, List[str]], str, int, str, str] + cluster specification, job name of this task, index of this task, + hostname:port socket of this task, the device for this task + """ + # setup cluster + node_list = list(set(node_list_)) + node_list.sort() + node_color = node_list.index(node_name) + world_idx = MPI.COMM_WORLD.Get_rank() + node_comm = MPI.COMM_WORLD.Split(node_color, world_idx) + node_task_idx = node_comm.Get_rank() + node_numb_task = node_comm.Get_size() + + socket_list = [] + for ii in node_list: + for jj in range(node_numb_task): + socket_list.append(f"{ii}:{default_port + jj}") + ps_map = socket_list[0:1] + worker_map = socket_list[1:] + + if node_color == 0 and node_task_idx == 0: + my_job = "ps" + my_socket = ps_map[0] + my_task_idx = ps_map.index(my_socket) + else: + my_job = "worker" + my_socket = f"{node_name}:{default_port - node_task_idx}" + assert my_socket in worker_map + my_task_idx = worker_map.index(my_socket) + + # setup gpu/cpu devices + if gpu_list is not None: + numb_gpu = len(gpu_list) + gpu_idx = node_numb_task - node_task_idx - 1 + if gpu_idx >= numb_gpu: + my_device = "cpu:0" # "cpu:%d" % node_task_idx + else: + my_device = f"gpu:{gpu_idx:d}" + else: + my_device = "cpu:0" # "cpu:%d" % node_task_idx + + cluster = {"worker": worker_map, "ps": ps_map} + return cluster, my_job, my_task_idx, my_socket, my_device + + +class RunOptions: + """Class with inf oon how to run training (cluster, MPI and GPU config). + + Attributes + ---------- + cluster: Optional[Dict[str, List[str]]] + cluster informations as dict + cluster_spec: Optional[tf.train.ClusterSpec] + `tf.train.ClusterSpec` or None if training is serial + gpus: Optional[List[int]] + list of GPUs if any are present else None + is_chief: bool + in distribured training it is true for tha main MPI process in serail it is + always true + my_job_name: str + name of the training job + my_socket: Optional[str] + communication socket for distributed training + my_task_index: int + index of the MPI task + nodename: str + name of the node + num_ps: Optional[int] + number of ps + num_workers: Optional[int] + number of workers + server: Optional[tf.train.Server] + `tf.train.Server` or `None` for serial training + my_device: str + deviice type - gpu or cpu + """ + + cluster: Optional[Dict[str, List[str]]] + cluster_spec: Optional[tf.train.ClusterSpec] + gpus: Optional[List[int]] + is_chief: bool + my_job_name: str + my_socket: Optional[str] + my_task_index: int + nodename: str + num_ps: Optional[int] + num_workers: Optional[int] + server: Optional["TFServerV1"] + my_device: str + + _MPI: Optional["MPI"] + _log_handles_already_set: bool = False + + def __init__( + self, + init_model: Optional[str] = None, + restart: Optional[str] = None, + log_path: Optional[str] = None, + log_level: int = 0, + mpi_log: str = "master", + try_distrib: bool = False + ): + # distributed tasks + if try_distrib: + self._try_init_mpi() + else: + self.is_distrib = False + self._init_serial() + + if all((init_model, restart)): + raise RuntimeError( + "--init-model and --restart should not be set at the same time" + ) + + # model init options + self.restart = restart + self.init_model = init_model + self.init_mode = "init_from_scratch" + + if restart is not None: + self.restart = os.path.abspath(restart) + self.init_mode = "restart" + elif init_model is not None: + self.init_model = os.path.abspath(init_model) + self.init_mode = "init_from_model" + + self._setup_logger(Path(log_path) if log_path else None, log_level, mpi_log) + + def print_resource_summary(self): + """Print build and current running cluster configuration summary.""" + log.info("---Summary of the training---------------------------------------") + if self.is_distrib: + log.info("distributed") + log.info(f"ps list: {self.cluster['ps']}") + log.info(f"worker list: {self.cluster['worker']}") + log.info(f"chief on: {self.nodename}") + else: + log.info(f"running on: {self.nodename}") + if self.gpus is None: + log.info(f"CUDA_VISIBLE_DEVICES: unset") + else: + log.info(f"CUDA_VISIBLE_DEVICES: {self.gpus}") + intra, inter = get_tf_default_nthreads() + log.info(f"num_intra_threads: {intra:d}") + log.info(f"num_inter_threads: {inter:d}") + log.info("-----------------------------------------------------------------") + + def _setup_logger( + self, + log_path: Optional[Path], + log_level: int, + mpi_log: Optional[str], + ): + """Set up package loggers. + + Parameters + ---------- + log_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 than also `MPI` object must + be passed in. by default None + """ + if not self._log_handles_already_set: + if not self._MPI: + mpi_log = None + set_log_handles(log_level, log_path, mpi_log=mpi_log, MPI=self._MPI) + self._log_handles_already_set = True + log.debug("Log handles were successfully set") + else: + log.warning( + f"Log handles have already been set. It is not advisable to " + f"reset them{', especially when runnig with MPI!' if self._MPI else ''}" + ) + + def _try_init_mpi(self): + try: + from mpi4py import MPI + except ImportError: + raise RuntimeError( + "cannot import mpi4py module, cannot do distributed simulation" + ) + else: + self.is_distrib = _is_distributed(MPI) + if self.is_distrib: + self._init_distributed(MPI) + self._MPI = MPI + else: + self._init_serial() + self._MPI = None + + def _init_distributed(self, MPI: "MPI"): + """Initialize settings for distributed training. + + Parameters + ---------- + MPI : MPI + MPI object + """ + nodename, nodelist, gpus = get_resource() + self.nodename = nodename + self.gpus = gpus + ( + self.cluster, + self.my_job_name, + self.my_task_index, + self.my_socket, + self.my_device, + ) = _distributed_task_config(MPI, nodename, nodelist, gpus) + self.is_chief = self.my_job_name == "worker" and self.my_task_index == 0 + self.num_ps = len(self.cluster["ps"]) + self.num_workers = len(self.cluster["worker"]) + self.cluster_spec = tf.train.ClusterSpec(self.cluster) + self.server = tf.train.Server( + server_or_cluster_def=self.cluster_spec, + job_name=self.my_job_name, + task_index=self.my_task_index, + ) + + def _init_serial(self): + """Initialize setting for serial training.""" + nodename, _, gpus = get_resource() + + self.cluster = None + self.cluster_spec = None + self.gpus = gpus + self.is_chief = True + self.my_job_name = nodename + self.my_socket = None + self.my_task_index = 0 + self.nodename = nodename + self.num_ps = None + self.num_workers = None + self.server = None + + if gpus is not None: + self.my_device = "gpu:" + str(gpus[0]) + else: + self.my_device = "cpu:0" + + self._MPI = None diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py new file mode 100644 index 0000000000..36fc60f029 --- /dev/null +++ b/deepmd/train/trainer.py @@ -0,0 +1,647 @@ +#!/usr/bin/env python3 +import logging +import os +import time +import shutil +import numpy as np +from deepmd.env import tf +from deepmd.env import default_tf_session_config +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION +from deepmd.fit import EnerFitting, WFCFitting, PolarFittingLocFrame, PolarFittingSeA, GlobalPolarFittingSeA, DipoleFittingSeA +from deepmd.descriptor import DescrptLocFrame +from deepmd.descriptor import DescrptSeA +from deepmd.descriptor import DescrptSeT +from deepmd.descriptor import DescrptSeAEbd +from deepmd.descriptor import DescrptSeAEf +from deepmd.descriptor import DescrptSeR +from deepmd.descriptor import DescrptSeAR +from deepmd.descriptor import DescrptHybrid +from deepmd.model import EnerModel, WFCModel, DipoleModel, PolarModel, GlobalPolarModel +from deepmd.loss import EnerStdLoss, EnerDipoleLoss, TensorLoss +from deepmd.utils.learning_rate import LearningRateExp +from deepmd.utils.neighbor_stat import NeighborStat +from deepmd.utils.type_embed import TypeEmbedNet + +from tensorflow.python.client import timeline +from deepmd.env import op_module + +# load grad of force module +import deepmd.op + +from deepmd.common import j_must_have, ClassArg + +log = logging.getLogger(__name__) + + +def _is_subdir(path, directory): + path = os.path.realpath(path) + directory = os.path.realpath(directory) + if path == directory: + return False + relative = os.path.relpath(path, directory) + os.sep + return not relative.startswith(os.pardir + os.sep) + +def _generate_descrpt_from_param_dict(descrpt_param): + try: + descrpt_type = descrpt_param['type'] + except KeyError: + raise KeyError('the type of descriptor should be set by `type`') + descrpt_param.pop('type', None) + to_pop = [] + for kk in descrpt_param: + if kk[0] == '_': + to_pop.append(kk) + for kk in to_pop: + descrpt_param.pop(kk, None) + if descrpt_type == 'loc_frame': + descrpt = DescrptLocFrame(**descrpt_param) + elif descrpt_type == 'se_e2_a' or descrpt_type == 'se_a' : + descrpt = DescrptSeA(**descrpt_param) + elif descrpt_type == 'se_e2_r' or descrpt_type == 'se_r' : + descrpt = DescrptSeR(**descrpt_param) + elif descrpt_type == 'se_e3' or descrpt_type == 'se_at' or descrpt_type == 'se_a_3be' : + descrpt = DescrptSeT(**descrpt_param) + elif descrpt_type == 'se_a_tpe' or descrpt_type == 'se_a_ebd' : + descrpt = DescrptSeAEbd(**descrpt_param) + elif descrpt_type == 'se_a_ef' : + descrpt = DescrptSeAEf(**descrpt_param) + elif descrpt_type == 'se_ar' : + descrpt = DescrptSeAR(descrpt_param) + else : + raise RuntimeError('unknow model type ' + descrpt_type) + return descrpt + + +class DPTrainer (object): + def __init__(self, + jdata, + run_opt): + self.run_opt = run_opt + self._init_param(jdata) + + def _init_param(self, jdata): + # model config + model_param = j_must_have(jdata, 'model') + descrpt_param = j_must_have(model_param, 'descriptor') + fitting_param = j_must_have(model_param, 'fitting_net') + typeebd_param = model_param.get('type_embedding', None) + self.model_param = model_param + self.descrpt_param = descrpt_param + + # descriptor + try: + descrpt_type = descrpt_param['type'] + except KeyError: + raise KeyError('the type of descriptor should be set by `type`') + + if descrpt_type != 'hybrid': + self.descrpt = _generate_descrpt_from_param_dict(descrpt_param) + else : + descrpt_list = [] + for ii in descrpt_param.get('list', []): + descrpt_list.append(_generate_descrpt_from_param_dict(ii)) + self.descrpt = DescrptHybrid(descrpt_list) + + # fitting net + try: + fitting_type = fitting_param['type'] + except: + fitting_type = 'ener' + fitting_param.pop('type', None) + fitting_param['descrpt'] = self.descrpt + if fitting_type == 'ener': + self.fitting = EnerFitting(**fitting_param) + # elif fitting_type == 'wfc': + # self.fitting = WFCFitting(fitting_param, self.descrpt) + elif fitting_type == 'dipole': + if descrpt_type == 'se_e2_a': + self.fitting = DipoleFittingSeA(**fitting_param) + else : + raise RuntimeError('fitting dipole only supports descrptors: se_e2_a') + elif fitting_type == 'polar': + # if descrpt_type == 'loc_frame': + # self.fitting = PolarFittingLocFrame(fitting_param, self.descrpt) + if descrpt_type == 'se_e2_a': + self.fitting = PolarFittingSeA(**fitting_param) + else : + raise RuntimeError('fitting polar only supports descrptors: loc_frame and se_e2_a') + elif fitting_type == 'global_polar': + if descrpt_type == 'se_e2_a': + self.fitting = GlobalPolarFittingSeA(**fitting_param) + else : + raise RuntimeError('fitting global_polar only supports descrptors: loc_frame and se_e2_a') + else : + raise RuntimeError('unknow fitting type ' + fitting_type) + + # type embedding + if typeebd_param is not None: + self.typeebd = TypeEmbedNet( + neuron=typeebd_param['neuron'], + resnet_dt=typeebd_param['resnet_dt'], + activation_function=typeebd_param['activation_function'], + precision=typeebd_param['precision'], + trainable=typeebd_param['trainable'], + seed=typeebd_param['seed'] + ) + else: + self.typeebd = None + + # init model + # infer model type by fitting_type + if fitting_type == 'ener': + self.model = EnerModel( + self.descrpt, + self.fitting, + self.typeebd, + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2), + model_param.get('use_srtab'), + model_param.get('smin_alpha'), + model_param.get('sw_rmin'), + model_param.get('sw_rmax') + ) + # elif fitting_type == 'wfc': + # self.model = WFCModel(model_param, self.descrpt, self.fitting) + elif fitting_type == 'dipole': + self.model = DipoleModel( + self.descrpt, + self.fitting, + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2) + ) + elif fitting_type == 'polar': + self.model = PolarModel( + self.descrpt, + self.fitting, + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2) + ) + elif fitting_type == 'global_polar': + self.model = GlobalPolarModel( + self.descrpt, + self.fitting, + model_param.get('type_map'), + model_param.get('data_stat_nbatch', 10), + model_param.get('data_stat_protect', 1e-2) + ) + else : + raise RuntimeError('get unknown fitting type when building model') + + # learning rate + lr_param = j_must_have(jdata, 'learning_rate') + try: + lr_type = lr_param['type'] + except: + lr_type = 'exp' + if lr_type == 'exp': + self.lr = LearningRateExp(lr_param['start_lr'], + lr_param['stop_lr'], + lr_param['decay_steps']) + else : + raise RuntimeError('unknown learning_rate type ' + lr_type) + + # loss + # infer loss type by fitting_type + try : + loss_param = jdata['loss'] + loss_type = loss_param.get('type', 'ener') + except: + loss_param = None + loss_type = 'ener' + + if fitting_type == 'ener': + loss_param.pop('type', None) + loss_param['starter_learning_rate'] = self.lr.start_lr() + if loss_type == 'ener': + self.loss = EnerStdLoss(**loss_param) + elif loss_type == 'ener_dipole': + self.loss = EnerDipoleLoss(**loss_param) + else: + raise RuntimeError('unknow loss type') + elif fitting_type == 'wfc': + self.loss = TensorLoss(loss_param, + model = self.model, + tensor_name = 'wfc', + tensor_size = self.model.get_out_size(), + label_name = 'wfc') + elif fitting_type == 'dipole': + self.loss = TensorLoss(loss_param, + model = self.model, + tensor_name = 'dipole', + tensor_size = 3, + label_name = 'dipole') + elif fitting_type == 'polar': + self.loss = TensorLoss(loss_param, + model = self.model, + tensor_name = 'polar', + tensor_size = 9, + label_name = 'polarizability') + elif fitting_type == 'global_polar': + self.loss = TensorLoss(loss_param, + model = self.model, + tensor_name = 'global_polar', + tensor_size = 9, + atomic = False, + label_name = 'polarizability') + else : + raise RuntimeError('get unknown fitting type when building loss function') + + # training + tr_data = jdata['training'] + self.disp_file = tr_data.get('disp_file', 'lcurve.out') + self.disp_freq = tr_data.get('disp_freq', 1000) + self.save_freq = tr_data.get('save_freq', 1000) + self.save_ckpt = tr_data.get('save_ckpt', 'model.ckpt') + self.display_in_training = tr_data.get('disp_training', True) + self.timing_in_training = tr_data.get('time_training', True) + self.profiling = tr_data.get('profiling', False) + self.profiling_file = tr_data.get('profiling_file', 'timeline.json') + self.tensorboard = tr_data.get('tensorboard', False) + self.tensorboard_log_dir = tr_data.get('tensorboard_log_dir', 'log') + # self.sys_probs = tr_data['sys_probs'] + # self.auto_prob_style = tr_data['auto_prob'] + self.useBN = False + if fitting_type == 'ener' and self.fitting.get_numb_fparam() > 0 : + self.numb_fparam = self.fitting.get_numb_fparam() + else : + self.numb_fparam = 0 + + if tr_data.get("validation_data", None) is not None: + self.valid_numb_batch = tr_data["validation_data"].get("numb_btch", 1) + else: + self.valid_numb_batch = 1 + + + def build (self, + data, + stop_batch = 0) : + self.ntypes = self.model.get_ntypes() + # Usually, the type number of the model should be equal to that of the data + # However, nt_model > nt_data should be allowed, since users may only want to + # train using a dataset that only have some of elements + assert (self.ntypes >= data.get_ntypes()), "ntypes should match that found in data" + self.stop_batch = stop_batch + + self.batch_size = data.get_batch_size() + + if self.numb_fparam > 0 : + log.info("training with %d frame parameter(s)" % self.numb_fparam) + else: + log.info("training without frame parameter") + + self.type_map = data.get_type_map() + + self.model.data_stat(data) + + if 'compress' in self.model_param and self.model_param['compress']['compress']: + assert 'rcut' in self.descrpt_param, "Error: descriptor must have attr rcut!" + self.neighbor_stat \ + = NeighborStat(self.ntypes, self.descrpt_param['rcut']) + self.min_nbor_dist, self.max_nbor_size \ + = self.neighbor_stat.get_stat(data) + self.descrpt.enable_compression(self.min_nbor_dist, self.model_param['compress']['model_file'], self.model_param['compress']['table_config'][0], self.model_param['compress']['table_config'][1], self.model_param['compress']['table_config'][2], self.model_param['compress']['table_config'][3]) + + worker_device = "/job:%s/task:%d/%s" % (self.run_opt.my_job_name, + self.run_opt.my_task_index, + self.run_opt.my_device) + + with tf.device(tf.train.replica_device_setter(worker_device = worker_device, + cluster = self.run_opt.cluster_spec)): + self._build_lr() + self._build_network(data) + self._build_training() + + + def _build_lr(self): + self._extra_train_ops = [] + self.global_step = tf.train.get_or_create_global_step() + self.learning_rate = self.lr.build(self.global_step, self.stop_batch) + log.info("built lr") + + def _build_network(self, data): + self.place_holders = {} + data_dict = data.get_data_dict() + for kk in data_dict.keys(): + if kk == 'type': + continue + prec = GLOBAL_TF_FLOAT_PRECISION + if data_dict[kk]['high_prec'] : + prec = GLOBAL_ENER_FLOAT_PRECISION + self.place_holders[kk] = tf.placeholder(prec, [None], name = 't_' + kk) + self.place_holders['find_'+kk] = tf.placeholder(tf.float32, name = 't_find_' + kk) + + self.place_holders['type'] = tf.placeholder(tf.int32, [None], name='t_type') + self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name='t_natoms') + self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name='t_mesh') + self.place_holders['is_training'] = tf.placeholder(tf.bool) + self.model_pred\ + = self.model.build (self.place_holders['coord'], + self.place_holders['type'], + self.place_holders['natoms_vec'], + self.place_holders['box'], + self.place_holders['default_mesh'], + self.place_holders, + suffix = "", + reuse = False) + + self.l2_l, self.l2_more\ + = self.loss.build (self.learning_rate, + self.place_holders['natoms_vec'], + self.model_pred, + self.place_holders, + suffix = "test") + + log.info("built network") + + def _build_training(self): + trainable_variables = tf.trainable_variables() + optimizer = tf.train.AdamOptimizer(learning_rate = self.learning_rate) + if self.run_opt.is_distrib : + optimizer = tf.train.SyncReplicasOptimizer( + optimizer, + replicas_to_aggregate = self.run_opt.cluster_spec.num_tasks("worker"), + total_num_replicas = self.run_opt.cluster_spec.num_tasks("worker"), + name = "sync_replicas") + self.sync_replicas_hook = optimizer.make_session_run_hook(self.run_opt.is_chief) + grads = tf.gradients(self.l2_l, trainable_variables) + apply_op = optimizer.apply_gradients (zip (grads, trainable_variables), + global_step=self.global_step, + name='train_step') + train_ops = [apply_op] + self._extra_train_ops + self.train_op = tf.group(*train_ops) + log.info("built training") + + def _init_sess_serial(self) : + self.sess = tf.Session(config=default_tf_session_config) + self.saver = tf.train.Saver() + saver = self.saver + if self.run_opt.init_mode == 'init_from_scratch' : + log.info("initialize model from scratch") + init_op = tf.global_variables_initializer() + self.sess.run(init_op) + fp = open(self.disp_file, "w") + fp.close () + elif self.run_opt.init_mode == 'init_from_model' : + log.info("initialize from model %s" % self.run_opt.init_model) + init_op = tf.global_variables_initializer() + self.sess.run(init_op) + saver.restore (self.sess, self.run_opt.init_model) + self.sess.run(self.global_step.assign(0)) + fp = open(self.disp_file, "w") + fp.close () + elif self.run_opt.init_mode == 'restart' : + log.info("restart from model %s" % self.run_opt.restart) + init_op = tf.global_variables_initializer() + self.sess.run(init_op) + saver.restore (self.sess, self.run_opt.restart) + else : + raise RuntimeError ("unkown init mode") + + def _init_sess_distrib(self): + ckpt_dir = os.path.join(os.getcwd(), self.save_ckpt) + assert(_is_subdir(ckpt_dir, os.getcwd())), "the checkpoint dir must be a subdir of the current dir" + if self.run_opt.init_mode == 'init_from_scratch' : + log.info("initialize model from scratch") + if self.run_opt.is_chief : + if os.path.exists(ckpt_dir): + shutil.rmtree(ckpt_dir) + if not os.path.exists(ckpt_dir) : + os.makedirs(ckpt_dir) + fp = open(self.disp_file, "w") + fp.close () + elif self.run_opt.init_mode == 'init_from_model' : + raise RuntimeError("distributed training does not support %s" % self.run_opt.init_mode) + elif self.run_opt.init_mode == 'restart' : + log.info("restart from model %s" % ckpt_dir) + if self.run_opt.is_chief : + assert(os.path.isdir(ckpt_dir)), "the checkpoint dir %s should exists" % ckpt_dir + else : + raise RuntimeError ("unkown init mode") + + saver = tf.train.Saver(max_to_keep = 1) + self.saver = None + # gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.5) + # config = tf.ConfigProto(allow_soft_placement=True, + # gpu_options = gpu_options, + # intra_op_parallelism_threads=self.run_opt.num_intra_threads, + # inter_op_parallelism_threads=self.run_opt.num_inter_threads) + config = tf.ConfigProto(intra_op_parallelism_threads=self.run_opt.num_intra_threads, + inter_op_parallelism_threads=self.run_opt.num_inter_threads) + # The stop_hook handles stopping after running given steps + # stop_hook = tf.train.StopAtStepHook(last_step = stop_batch) + # hooks = [self.sync_replicas_hook, stop_hook] + hooks = [self.sync_replicas_hook] + scaffold = tf.train.Scaffold(saver=saver) + # Use monitor session for distributed computation + self.sess = tf.train.MonitoredTrainingSession(master = self.run_opt.server.target, + is_chief = self.run_opt.is_chief, + config = config, + hooks = hooks, + scaffold = scaffold, + checkpoint_dir = ckpt_dir) + # , + # save_checkpoint_steps = self.save_freq) + + def train (self, train_data, valid_data=None) : + + # if valid_data is None: # no validation set specified. + # valid_data = train_data # using training set as validation set. + + stop_batch = self.stop_batch + if self.run_opt.is_distrib : + self._init_sess_distrib() + else : + self._init_sess_serial() + + # self.print_head() + fp = None + if self.run_opt.is_chief : + fp = open(self.disp_file, "a") + + cur_batch = self.sess.run(self.global_step) + is_first_step = True + self.cur_batch = cur_batch + log.info("start training at lr %.2e (== %.2e), decay_step %d, decay_rate %f, final lr will be %.2e" % + (self.sess.run(self.learning_rate), + self.lr.value(cur_batch), + self.lr.decay_steps_, + self.lr.decay_rate_, + self.lr.value(stop_batch)) + ) + + prf_options = None + prf_run_metadata = None + if self.profiling : + prf_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) + prf_run_metadata = tf.RunMetadata() + + # set tensorboard execution environment + if self.tensorboard : + summary_merged_op = tf.summary.merge_all() + # Remove TB old logging directory from previous run + try: + shutil.rmtree(self.tensorboard_log_dir) + except FileNotFoundError: + pass # directory does not exist, this is OK + except Exception as e: + # general error when removing directory, warn user + log.exception( + f"Could not remove old tensorboard logging directory: " + f"{self.tensorboard_log_dir}. Error: {e}" + ) + else: + log.debug("Removing old tensorboard log directory.") + tb_train_writer = tf.summary.FileWriter(self.tensorboard_log_dir + '/train', self.sess.graph) + tb_valid_writer = tf.summary.FileWriter(self.tensorboard_log_dir + '/test') + else: + tb_train_writer = None + tb_valid_writer = None + + train_time = 0 + while cur_batch < stop_batch : + + # first round validation: + train_batch = train_data.get_batch() + if self.display_in_training and is_first_step: + valid_batches = [valid_data.get_batch() for ii in range(self.valid_numb_batch)] if valid_data is not None else None + self.valid_on_the_fly(fp, [train_batch], valid_batches, print_header=True) + is_first_step = False + + if self.timing_in_training: tic = time.time() + train_feed_dict = self.get_feed_dict(train_batch, is_training=True) + # use tensorboard to visualize the training of deepmd-kit + # it will takes some extra execution time to generate the tensorboard data + if self.tensorboard : + summary, _ = self.sess.run([summary_merged_op, self.train_op], feed_dict=train_feed_dict, + options=prf_options, run_metadata=prf_run_metadata) + tb_train_writer.add_summary(summary, cur_batch) + else : + self.sess.run([self.train_op], feed_dict=train_feed_dict, + options=prf_options, run_metadata=prf_run_metadata) + if self.timing_in_training: toc = time.time() + if self.timing_in_training: train_time += toc - tic + cur_batch = self.sess.run(self.global_step) + self.cur_batch = cur_batch + + # on-the-fly validation + if self.display_in_training and (cur_batch % self.disp_freq == 0): + if self.timing_in_training: + tic = time.time() + valid_batches = [valid_data.get_batch() for ii in range(self.valid_numb_batch)] if valid_data is not None else None + self.valid_on_the_fly(fp, [train_batch], valid_batches) + if self.timing_in_training: + toc = time.time() + test_time = toc - tic + log.info("batch %7d training time %.2f s, testing time %.2f s" + % (cur_batch, train_time, test_time)) + train_time = 0 + if self.save_freq > 0 and cur_batch % self.save_freq == 0 and self.run_opt.is_chief : + if self.saver is not None : + self.saver.save (self.sess, os.getcwd() + "/" + self.save_ckpt) + log.info("saved checkpoint %s" % self.save_ckpt) + if self.run_opt.is_chief: + fp.close () + if self.profiling and self.run_opt.is_chief : + fetched_timeline = timeline.Timeline(prf_run_metadata.step_stats) + chrome_trace = fetched_timeline.generate_chrome_trace_format() + with open(self.profiling_file, 'w') as f: + f.write(chrome_trace) + + def get_feed_dict(self, batch, is_training): + feed_dict = {} + for kk in batch.keys(): + if kk == 'find_type' or kk == 'type': + continue + if 'find_' in kk: + feed_dict[self.place_holders[kk]] = batch[kk] + else: + feed_dict[self.place_holders[kk]] = np.reshape(batch[kk], [-1]) + for ii in ['type']: + feed_dict[self.place_holders[ii]] = np.reshape(batch[ii], [-1]) + for ii in ['natoms_vec', 'default_mesh']: + feed_dict[self.place_holders[ii]] = batch[ii] + feed_dict[self.place_holders['is_training']] = is_training + return feed_dict + + def get_global_step(self): + return self.sess.run(self.global_step) + + # def print_head (self) : # depreciated + # if self.run_opt.is_chief: + # fp = open(self.disp_file, "a") + # print_str = "# %5s" % 'batch' + # print_str += self.loss.print_header() + # print_str += ' %8s\n' % 'lr' + # fp.write(print_str) + # fp.close () + + def valid_on_the_fly(self, + fp, + train_batches, + valid_batches, + print_header=False): + train_results = self.get_evaluation_results(train_batches) + valid_results = self.get_evaluation_results(valid_batches) + + cur_batch = self.cur_batch + current_lr = self.sess.run(self.learning_rate) + if print_header: + self.print_header(fp, train_results, valid_results) + self.print_on_training(fp, train_results, valid_results, cur_batch, current_lr) + + @staticmethod + def print_header(fp, train_results, valid_results): + print_str = '' + print_str += "# %5s" % 'step' + if valid_results is not None: + prop_fmt = ' %11s %11s' + for k in train_results.keys(): + print_str += prop_fmt % (k + '_val', k + '_trn') + else: + prop_fmt = ' %11s' + for k in train_results.keys(): + print_str += prop_fmt % (k + '_trn') + print_str += ' %8s\n' % 'lr' + fp.write(print_str) + fp.flush() + + @staticmethod + def print_on_training(fp, train_results, valid_results, cur_batch, cur_lr): + print_str = '' + print_str += "%7d" % cur_batch + if valid_results is not None: + prop_fmt = " %11.2e %11.2e" + for k in valid_results.keys(): + # assert k in train_results.keys() + print_str += prop_fmt % (valid_results[k], train_results[k]) + else: + prop_fmt = " %11.2e" + for k in train_results.keys(): + print_str += prop_fmt % (train_results[k]) + print_str += " %8.1e\n" % cur_lr + fp.write(print_str) + fp.flush() + + def get_evaluation_results(self, batch_list): + if batch_list is None: return None + numb_batch = len(batch_list) + + sum_results = {} # sum of losses on all atoms + sum_natoms = 0 + for i in range(numb_batch): + batch = batch_list[i] + natoms = batch["natoms_vec"] + feed_dict = self.get_feed_dict(batch, is_training=False) + results = self.loss.eval(self.sess, feed_dict, natoms) + + for k, v in results.items(): + if k == "natoms": + sum_natoms += v + else: + sum_results[k] = sum_results.get(k, 0.) + v * results["natoms"] + avg_results = {k: v / sum_natoms for k, v in sum_results.items() if not k == "natoms"} + return avg_results diff --git a/deepmd/utils/__init__.py b/deepmd/utils/__init__.py new file mode 100644 index 0000000000..a54f69b853 --- /dev/null +++ b/deepmd/utils/__init__.py @@ -0,0 +1,9 @@ +# +from .data import DeepmdData +from .data_system import DeepmdDataSystem + +# out-of-dated +from .data import DataSets +from .data_system import DataSystem +from .pair_tab import PairTab +from .learning_rate import LearningRateExp diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py new file mode 100644 index 0000000000..684b0eacbc --- /dev/null +++ b/deepmd/utils/argcheck.py @@ -0,0 +1,649 @@ +from dargs import dargs, Argument, Variant +from deepmd.common import ACTIVATION_FN_DICT, PRECISION_DICT + + +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())}' + doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' + doc_trainable = 'If the parameters in the embedding net are trainable' + + return [ + Argument("neuron", list, optional = True, default = [2, 4, 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 = "float64", doc = doc_precision), + Argument("trainable", bool, optional = True, default = True, doc = doc_trainable), + Argument("seed", [int,None], optional = True, doc = doc_seed), + ] + + +# --- Descriptor configurations: --- # +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 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+4]: type of the atom defining the second axis of type-i atom.\n\n\ +- axis_rule[i*6+5]: 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.' + + return [ + Argument("sel_a", list, optional = False, doc = doc_sel_a), + Argument("sel_r", list, optional = False, doc = doc_sel_r), + Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut), + Argument("axis_rule", list, optional = False, doc = doc_axis_rule) + ] + + +def descrpt_se_a_args(): + doc_sel = 'A list of integers. 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.' + 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())}' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_type_one_side = 'Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets' + doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' + doc_trainable = 'If the parameters in the embedding net is trainable' + doc_seed = 'Random seed for parameter initialization' + doc_exclude_types = 'The Excluded types' + 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, optional = False, 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, optional = True, default = [10,20,40], doc = doc_neuron), + Argument("axis_neuron", int, optional = True, default = 4, 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 = "float64", 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, optional = True, default = [], doc = doc_exclude_types), + Argument("set_davg_zero", bool, optional = True, default = False, doc = doc_set_davg_zero) + ] + + +def descrpt_se_t_args(): + doc_sel = 'A list of integers. 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.' + 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())}' + 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())}' + 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, optional = False, 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, 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 = "float64", 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) + ] + + + +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) + ] + + +def descrpt_se_r_args(): + doc_sel = 'A list of integers. 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.' + 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())}' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_type_one_side = 'Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets' + doc_precision = f'The precision of the embedding net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' + doc_trainable = 'If the parameters in the embedding net are trainable' + doc_seed = 'Random seed for parameter initialization' + doc_exclude_types = 'The Excluded types' + 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, optional = False, 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, 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 = "float64", 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, optional = True, default = [], doc = doc_exclude_types), + Argument("set_davg_zero", bool, optional = True, default = False, doc = doc_set_davg_zero) + ] + + +def descrpt_se_ar_args(): + link = make_link('se_a', 'model/descriptor[se_a]') + doc_a = f'The parameters of descriptor {link}' + link = make_link('se_r', 'model/descriptor[se_r]') + doc_r = f'The parameters of descriptor {link}' + + return [ + Argument("a", dict, optional = False, doc = doc_a), + Argument("r", dict, optional = False, doc = doc_r), + ] + + +def descrpt_hybrid_args(): + doc_list = f'A list of descriptor definitions' + + return [ + Argument("list", list, optional = False, doc = doc_list) + ] + + +def descrpt_variant_type_args(): + 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]') + doc_descrpt_type = f'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\ +- `hybrid`: Concatenate of a list of descriptors as a new descriptor.' + + return Variant("type", [ + Argument("loc_frame", dict, descrpt_local_frame_args()), + Argument("se_e2_a", dict, descrpt_se_a_args(), alias = ['se_a']), + Argument("se_e2_r", dict, descrpt_se_r_args(), alias = ['se_r']), + Argument("se_e3", dict, descrpt_se_t_args(), alias = ['se_at', 'se_a_3be', 'se_t']), + Argument("se_a_tpe", dict, descrpt_se_a_tpe_args(), alias = ['se_a_ebd']), + Argument("hybrid", dict, descrpt_hybrid_args()), + ], doc = doc_descrpt_type) + + +# --- Fitting net configurations: --- # +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())}' + doc_precision = f'The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())}' + 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.' + doc_seed = 'Random seed for parameter initialization of the fitting net' + doc_atom_ener = 'Specify the atomic energy in vacuum for each type' + + 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, 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], optional = True, default = True, doc = doc_trainable), + Argument("rcond", float, optional = True, default = 1e-3, doc = doc_rcond), + Argument("seed", [int,None], optional = True, doc = doc_seed), + Argument("atom_ener", list, optional = True, default = [], doc = doc_atom_ener) + ] + + +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())}' + 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())}' + 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, optional = True, default = [120,120,120], 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 = 'float64', doc = doc_precision), + Argument("fit_diag", bool, optional = True, default = True, doc = doc_fit_diag), + Argument("scale", [list,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,None], optional = True, doc = doc_sel_type), + Argument("seed", [int,None], optional = True, doc = doc_seed) + ] + + +#def fitting_global_polar(): +# return fitting_polar() + + +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())}' + 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())}' + 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, optional = True, default = [120,120,120], 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 = 'float64', doc = doc_precision), + Argument("sel_type", [list,int,None], optional = True, 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\ +- `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", [Argument("ener", dict, fitting_ener()), + Argument("dipole", dict, fitting_dipole()), + Argument("polar", dict, fitting_polar()), + ], + 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 = f"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, optional = False, doc = doc_model_charge_map), + Argument("sys_charge_map", list, 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_compress = "The name of the frozen model file." + doc_model_file = f"The input model file, which will be compressed by the DeePMD-kit." + doc_table_config = f"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)." + + return [ + Argument("compress", bool, optional = False, default = True, doc = doc_compress), + Argument("model_file", str, optional = False, default = 'frozen_model.pb', doc = doc_model_file), + Argument("table_config", list, optional = False, default = [5, 0.01, 0.1, -1], doc = doc_table_config), + ] + +# --- 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 (): + 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.' + 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_type_embedding = "The type embedding." + doc_descrpt = 'The descriptor of atomic environment.' + doc_fitting = 'The fitting of physical properties.' + 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_compress_config = 'Model compression configurations' + + ca = Argument("model", dict, + [Argument("type_map", list, 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("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("type_embedding", dict, type_embedding_args(), [], optional = True, doc = doc_type_embedding), + Argument("descriptor", dict, [], [descrpt_variant_type_args()], doc = doc_descrpt), + Argument("fitting_net", dict, [], [fitting_variant_type_args()], doc = doc_fitting), + 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) + ]) + # print(ca.gen_doc()) + return ca + + +# --- Learning rate configurations: --- # +def learning_rate_exp(): + doc_start_lr = 'The learning rate 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_lr = "The definitio of learning rate" + return Argument("learning_rate", dict, [], + [learning_rate_variant_type_args()], + doc = doc_lr) + + +# --- Loss configurations: --- # +def start_pref(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 {item} label should be provided by file {item}.npy in each data system. If both start_pref_{item} and limit_pref_{item} 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.' + + +def loss_ener(): + doc_start_pref_e = start_pref('energy') + doc_limit_pref_e = limit_pref('energy') + doc_start_pref_f = start_pref('force') + doc_limit_pref_f = limit_pref('force') + 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_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.' + 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("relative_f", [float,None], optional = True, doc = doc_relative_f) + ] + +# YWolfeee: Modified to support tensor type of loss args. +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`. \n\.' + + + return Variant("type", + [Argument("ener", dict, loss_ener()), + Argument("tensor", dict, loss_tensor()), + #Argument("polar", dict, loss_tensor()), + #Argument("global_polar", dict, loss_tensor("global")) + ], + 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.\n\.' + ca = Argument('loss', dict, [], + [loss_variant_type_args()], + optional = True, + doc = doc_loss) + 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.' + 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], optional=False, default=".", doc=doc_systems), + Argument("set_prefix", str, optional=True, default='set', doc=doc_set_prefix), + Argument("batch_size", [list, 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, optional=True, default=None, doc=doc_sys_probs, alias=["sys_weights"]), + ] + + doc_training_data = "Configurations of training data." + return Argument("training_data", dict, optional=False, + 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 systems to be sampled for each validation period." + + args = [ + Argument("systems", [list, str], optional=False, default=".", doc=doc_systems), + Argument("set_prefix", str, optional=True, default='set', doc=doc_set_prefix), + Argument("batch_size", [list, 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, 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 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_numb_test = 'Number of frames used for the test during training.' + doc_save_freq = 'The frequency of saving check point.' + doc_save_ckpt = 'The file name of saving check point.' + 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_tensorboard = 'Enable tensorboard' + doc_tensorboard_log_dir = 'The log directory of tensorboard outputs' + + arg_training_data = training_data_args() + arg_validation_data = validation_data_args() + + args = [ + arg_training_data, + arg_validation_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='lcueve.out', doc=doc_disp_file), + Argument("disp_freq", int, optional=True, default=1000, doc=doc_disp_freq), + Argument("numb_test", [list,int,str], optional=True, default=1, doc=doc_numb_test), + 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("tensorboard", bool, optional=True, default=False, doc=doc_tensorboard), + Argument("tensorboard_log_dir", str, optional=True, default='log', doc=doc_tensorboard_log_dir), + ] + + 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 + ma = model_args() + lra = learning_rate_args() + la = loss_args() + ta = training_args() + ptr = [] + ptr.append(ma.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs)) + ptr.append(la.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs)) + ptr.append(lra.gen_doc(make_anchor=make_anchor, make_link=make_link, **kwargs)) + ptr.append(ta.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 normalize_hybrid_list(hy_list): + new_list = [] + base = Argument("base", dict, [], [descrpt_variant_type_args()], doc = "") + for ii in range(len(hy_list)): + data = base.normalize_value(hy_list[ii], trim_pattern="_*") + base.check_value(data, strict=True) + new_list.append(data) + return new_list + + +def normalize(data): + if "hybrid" == data["model"]["descriptor"]["type"]: + data["model"]["descriptor"]["list"] \ + = normalize_hybrid_list(data["model"]["descriptor"]["list"]) + + ma = model_args() + lra = learning_rate_args() + la = loss_args() + ta = training_args() + + base = Argument("base", dict, [ma, lra, la, ta]) + data = base.normalize_value(data, trim_pattern="_*") + base.check_value(data, strict=True) + + return data + + +if __name__ == '__main__': + gen_doc() + diff --git a/deepmd/utils/compat.py b/deepmd/utils/compat.py new file mode 100644 index 0000000000..861a00439c --- /dev/null +++ b/deepmd/utils/compat.py @@ -0,0 +1,298 @@ +"""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 + +from deepmd.common import j_must_have + + +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 = {} + if "with_distrib" in jdata: + output["with_distrib"] = jdata["with_distrib"] + 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 + must_have : bool + ensure that the source dictionary contains the copyyied keys + """ + for k in keys: + dst[k] = src[k] + + +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 + + jdata["training"] = new_tr_cfg + + 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 updata_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, dump) + elif is_deepmd_v1_input(jdata): + jdata = convert_input_v1_v2(jdata, warning, dump) + else: + pass + + return jdata diff --git a/source/train/Data.py b/deepmd/utils/data.py similarity index 76% rename from source/train/Data.py rename to deepmd/utils/data.py index b38839a546..c75e312ad1 100644 --- a/source/train/Data.py +++ b/deepmd/utils/data.py @@ -4,16 +4,44 @@ import glob import numpy as np import os.path -from deepmd.RunOptions import global_np_float_precision -from deepmd.RunOptions import global_ener_float_precision +from typing import Tuple, List +import logging + +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION + +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` + """ def __init__ (self, - sys_path, - set_prefix = 'set', - shuffle_test = True, - type_map = None, - modifier = None) : + sys_path : str, + set_prefix : str = 'set', + shuffle_test : bool = True, + type_map : List[str] = None, + modifier = None, + trn_all_set : bool = False) : + """ + Constructor + + 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 + 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. + """ self.dirs = glob.glob (os.path.join(sys_path, set_prefix + ".*")) self.dirs.sort() # load atom type @@ -35,10 +63,13 @@ def __init__ (self, self.idx_map = self._make_idx_map(self.atom_type) # train dirs self.test_dir = self.dirs[-1] - if len(self.dirs) == 1 : + if trn_all_set: self.train_dirs = self.dirs - else : - self.train_dirs = self.dirs[:-1] + 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 = True) @@ -52,13 +83,37 @@ def __init__ (self, def add(self, - key, - ndof, - atomic = False, - must = False, - high_prec = False, - type_sel = None, - repeat = 1) : + key : str, + ndof : int, + atomic : bool = False, + must : bool = False, + high_prec : bool = False, + type_sel : List[int] = None, + repeat : int = 1 + ) : + """ + 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. + """ self.data_dict[key] = {'ndof': ndof, 'atomic': atomic, 'must': must, @@ -71,8 +126,19 @@ def add(self, def reduce(self, - key_out, - key_in) : + 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 (not(key_out in self.data_dict)), 'output key should not have been added' @@ -88,15 +154,21 @@ def reduce(self, } return self - def get_data_dict(self): + def get_data_dict(self) -> dict: + """ + Get the `data_dict` + """ return self.data_dict - def check_batch_size (self, batch_size) : + def check_batch_size (self, batch_size) : + """ + 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 = np.load(os.path.join(ii, "coord.npy")).astype(global_ener_float_precision) + tmpe = np.load(os.path.join(ii, "coord.npy")).astype(GLOBAL_ENER_FLOAT_PRECISION) else: - tmpe = np.load(os.path.join(ii, "coord.npy")).astype(global_np_float_precision) + tmpe = np.load(os.path.join(ii, "coord.npy")).astype(GLOBAL_NP_FLOAT_PRECISION) if tmpe.ndim == 1: tmpe = tmpe.reshape([1,-1]) if tmpe.shape[0] < batch_size : @@ -104,10 +176,13 @@ def check_batch_size (self, batch_size) : 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 = np.load(os.path.join(self.test_dir, "coord.npy")).astype(global_ener_float_precision) + tmpe = np.load(os.path.join(self.test_dir, "coord.npy")).astype(GLOBAL_ENER_FLOAT_PRECISION) else: - tmpe = np.load(os.path.join(self.test_dir, "coord.npy")).astype(global_np_float_precision) + tmpe = np.load(os.path.join(self.test_dir, "coord.npy")).astype(GLOBAL_NP_FLOAT_PRECISION) if tmpe.ndim == 1: tmpe = tmpe.reshape([1,-1]) if tmpe.shape[0] < test_size : @@ -115,7 +190,17 @@ def check_test_size (self, test_size) : else : return None - def get_batch(self, batch_size) : + 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 : @@ -134,7 +219,17 @@ def get_batch(self, batch_size) : ret = self._get_subdata(self.batch_set, idx) return ret - def get_test (self, ntests = -1) : + 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: @@ -148,44 +243,89 @@ def get_test (self, ntests = -1) : self.modifier.modify_data(ret) return ret - def get_ntypes(self) : + 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) : + def get_type_map(self) -> List[str]: + """ + Get the type map + """ return self.type_map - def get_atom_type(self) : + def get_atom_type(self) -> List[int]: + """ + Get atom types + """ return self.atom_type - def get_numb_set (self) : + def get_numb_set (self) -> int: + """ + Get number of training sets + """ return len (self.train_dirs) - def get_numb_batch (self, batch_size, set_idx) : + 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) : + 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) : + 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] @@ -272,9 +412,9 @@ def _load_set(self, set_name) : # get nframes path = os.path.join(set_name, "coord.npy") if self.data_dict['coord']['high_prec'] : - coord = np.load(path).astype(global_ener_float_precision) + coord = np.load(path).astype(GLOBAL_ENER_FLOAT_PRECISION) else: - coord = np.load(path).astype(global_np_float_precision) + coord = np.load(path).astype(GLOBAL_NP_FLOAT_PRECISION) if coord.ndim == 1: coord = coord.reshape([1,-1]) nframes = coord.shape[0] @@ -299,7 +439,7 @@ def _load_set(self, set_name) : 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) + tmp_in = data[k_in].astype(GLOBAL_ENER_FLOAT_PRECISION) data[kk] = np.sum(np.reshape(tmp_in, [nframes, self.natoms, ndof]), axis = 1) return data @@ -321,14 +461,20 @@ def _load_data(self, set_name, key, nframes, ndof_, atomic = False, must = True, path = os.path.join(set_name, key+".npy") if os.path.isfile (path) : if high_prec : - data = np.load(path).astype(global_ener_float_precision) + data = np.load(path).astype(GLOBAL_ENER_FLOAT_PRECISION) else: - data = np.load(path).astype(global_np_float_precision) - if atomic : - data = data.reshape([nframes, natoms, -1]) - data = data[:,idx_map,:] - data = data.reshape([nframes, -1]) - data = np.reshape(data, [nframes, ndof]) + data = np.load(path).astype(GLOBAL_NP_FLOAT_PRECISION) + 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 @@ -336,9 +482,9 @@ def _load_data(self, set_name, key, nframes, ndof_, atomic = False, must = True, raise RuntimeError("%s not found!" % path) else: if high_prec : - data = np.zeros([nframes,ndof]).astype(global_ener_float_precision) + data = np.zeros([nframes,ndof]).astype(GLOBAL_ENER_FLOAT_PRECISION) else : - data = np.zeros([nframes,ndof]).astype(global_np_float_precision) + data = np.zeros([nframes,ndof]).astype(GLOBAL_NP_FLOAT_PRECISION) if repeat != 1: data = np.repeat(data, repeat).reshape([nframes, -1]) return np.float32(0.0), data @@ -370,6 +516,9 @@ def _check_pbc(self, sys_path): class DataSets (object): + """ + Outdated class for one data system. Not maintained anymore. + """ def __init__ (self, sys_path, set_prefix, @@ -562,7 +711,7 @@ def get_set(self, data, idx = None) : if ii == "type": new_data[ii] = dd else: - new_data[ii] = dd.astype(global_np_float_precision) + new_data[ii] = dd.astype(GLOBAL_NP_FLOAT_PRECISION) return new_data def get_test (self) : diff --git a/source/train/DataSystem.py b/deepmd/utils/data_system.py similarity index 71% rename from source/train/DataSystem.py rename to deepmd/utils/data_system.py index 5aa866d550..8b84319eb2 100644 --- a/source/train/DataSystem.py +++ b/deepmd/utils/data_system.py @@ -1,37 +1,90 @@ #!/usr/bin/env python3 +import logging import os import collections import warnings import numpy as np -from deepmd.Data import DataSets -from deepmd.Data import DeepmdData +from typing import Tuple, List + +from deepmd.utils.data import DataSets +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, - batch_size, - test_size, - rcut, - set_prefix = 'set', - shuffle_test = True, - type_map = None, - modifier = None) : + systems : List[str], + batch_size : int, + test_size : int, + rcut : float, + set_prefix : str = 'set', + shuffle_test : bool = True, + type_map : List[str] = None, + modifier = None, + trn_all_set = False, + sys_probs = None, + auto_prob_style ="prob_sys_size") : + """ + 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 + 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.""" # 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, - modifier = modifier)) + self.data_systems.append( + DeepmdData( + ii, + set_prefix=set_prefix, + shuffle_test=shuffle_test, + type_map = type_map, + modifier = modifier, + trn_all_set = trn_all_set + )) # batch size self.batch_size = batch_size - if isinstance(self.batch_size, int) : + 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(':') @@ -92,6 +145,10 @@ def __init__ (self, 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]) @@ -142,8 +199,20 @@ def compute_energy_shift(self, rcond = 1e-3, key = 'energy') : = np.linalg.lstsq(sys_tynatom, sys_ener, rcond = rcond) return energy_shift - - def add_dict(self, adict) : + def add_dict(self, adict: dict) -> None: + """ + Add items to the data system by a `dict`. + `adict` should have items like + 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'], @@ -154,32 +223,80 @@ def add_dict(self, adict) : repeat=adict[kk]['repeat']) def add(self, - key, - ndof, - atomic = False, - must = False, - high_prec = False, - type_sel = None, - repeat = 1) : + key : str, + ndof : int, + atomic : bool = False, + must : bool = False, + high_prec : bool = False, + type_sel : List[int] = None, + repeat : int = 1 + ) : + """ + 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. + """ for ii in self.data_systems: ii.add(key, ndof, atomic=atomic, must=must, high_prec=high_prec, repeat=repeat, type_sel=type_sel) - def reduce(self, - key_out, - key_in) : + 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, k_in) + ii.reduce(key_out, key_in) - def get_data_dict(self) : - return self.data_systems[0].get_data_dict() + 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./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 = self._prob_sys_size_ext(auto_prob_style) + else: + raise RuntimeError("Unknown auto prob style: " + auto_prob_style) + else: + probs = self._process_sys_probs(sys_probs) + self.sys_probs = probs def _get_sys_probs(self, sys_probs, - auto_prob_style) : + auto_prob_style) : # depreciated if sys_probs is None : if auto_prob_style == "prob_uniform" : - prob = None + prob_v = 1./float(self.nsystems) + prob = [prob_v for ii in range(self.nsystems)] elif auto_prob_style == "prob_sys_size" : prob = self.prob_nbatches elif auto_prob_style[:14] == "prob_sys_size;" : @@ -190,13 +307,14 @@ def _get_sys_probs(self, prob = self._process_sys_probs(sys_probs) return prob - - def get_batch (self, - sys_idx = None, - sys_probs = None, - auto_prob_style = "prob_sys_size") : + def get_batch(self, sys_idx : int = None): + # 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 system + Get a batch of data from the data systems Parameters ---------- @@ -204,27 +322,14 @@ def get_batch (self, 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. - 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: float - 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. """ if not hasattr(self, 'default_mesh') : self._make_default_mesh() 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 = np.random.choice(np.arange(self.nsystems), p = prob) + # prob = self._get_sys_probs(sys_probs, auto_prob_style) + self.pick_idx = np.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] @@ -232,9 +337,19 @@ def get_batch (self, # ! altered by Marián Rynik def get_test (self, - sys_idx = None, - n_test = -1) : + sys_idx : 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, 'default_mesh') : self._make_default_mesh() if not hasattr(self, 'test_data') : @@ -252,29 +367,49 @@ def get_test (self, 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.""" + """ + 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): + def get_type_map(self) -> List[str]: + """ + Get the type map + """ return self.type_map - def get_nbatches (self) : + def get_nbatches (self) -> int: + """ + Get the total number of batches + """ return self.nbatches - def get_ntypes (self) : + def get_ntypes (self) -> int: + """ + Get the number of types + """ return self.sys_ntypes - def get_nsystems (self) : + def get_nsystems (self) -> int: + """ + Get the number of data systems + """ return self.nsystems - def get_sys (self, idx) : + def get_sys (self, idx : int) -> DeepmdData: + """ + Get a certain data system + """ return self.data_systems[idx] - def get_batch_size(self) : + def get_batch_size(self) -> int: + """ + Get the batch size + """ return self.batch_size def _format_name_length(self, name, width) : @@ -285,29 +420,24 @@ def _format_name_length(self, name, width) : name = '-- ' + name return name - def print_summary(self, - run_opt, - sys_probs = None, - auto_prob_style = "prob_sys_size") : - prob = self._get_sys_probs(sys_probs, auto_prob_style) - tmp_msg = "" + def print_summary(self, name) : # width 65 sys_width = 42 - tmp_msg += "---Summary of DataSystem------------------------------------------------\n" - tmp_msg += "found %d system(s):\n" % self.nsystems - tmp_msg += "%s " % self._format_name_length('system', sys_width) - tmp_msg += "%s %s %s %s %5s\n" % ('natoms', 'bch_sz', 'n_bch', "n_test", 'prob') + 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) : - tmp_msg += ("%s %6d %6d %6d %6d %5.3f\n" % - (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.test_size[ii], - prob[ii]) ) - tmp_msg += "------------------------------------------------------------------------\n" - run_opt.message(tmp_msg) + 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 = [] @@ -347,9 +477,12 @@ def _process_sys_probs(self, sys_probs) : assigned_sum_prob = np.sum(type_filter * sys_probs) assert assigned_sum_prob <= 1, "the sum of assigned probability should be less than 1" rest_sum_prob = 1. - assigned_sum_prob - rest_nbatch = (1 - type_filter) * self.nbatches - rest_prob = rest_sum_prob * rest_nbatch / np.sum(rest_nbatch) - ret_prob = rest_prob + type_filter * sys_probs + if rest_sum_prob != 0 : + rest_nbatch = (1 - type_filter) * self.nbatches + 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.sum(ret_prob) == 1, "sum of probs should be 1" return ret_prob @@ -378,6 +511,9 @@ def _prob_sys_size_ext(self, keywords): class DataSystem (object) : + """ + Outdated class for the data systems. Not maintained anymore. + """ def __init__ (self, systems, set_prefix, @@ -475,7 +611,7 @@ def format_name_length(self, name, width) : name = '-- ' + name return name - def print_summary(self, run_opt) : + def print_summary(self) : tmp_msg = "" # width 65 sys_width = 42 @@ -490,7 +626,7 @@ def print_summary(self, run_opt) : self.batch_size[ii], self.nbatches[ii]) ) tmp_msg += "-----------------------------------------------------------------\n" - run_opt.message(tmp_msg) + log.info(tmp_msg) def compute_energy_shift(self) : sys_ener = np.array([]) diff --git a/deepmd/utils/learning_rate.py b/deepmd/utils/learning_rate.py new file mode 100644 index 0000000000..572f317a92 --- /dev/null +++ b/deepmd/utils/learning_rate.py @@ -0,0 +1,95 @@ +import numpy as np +from deepmd.env import tf +from deepmd.common import ClassArg + +class LearningRateExp (object) : + """ + The exponentially decaying learning rate. + + The learning rate at step t is given by + + lr(t) = start_lr * decay_rate ^ ( t / decay_steps ) + """ + def __init__ (self, + start_lr : float, + stop_lr : float = 5e-8, + decay_steps : int = 5000, + decay_rate : float = 0.95 + ) -> None : + """ + Constructor + + Parameters + ---------- + start_lr + Starting learning rate + stop_lr + Stop learning rate + decay_steps + Learning rate decay every this number of steps + decay_rate + The decay rate. + If `stop_step` is provided in `build`, then it will be determined automatically and overwritten. + """ + # args = ClassArg()\ + # .add('decay_steps', int, must = False)\ + # .add('decay_rate', float, must = False)\ + # .add('start_lr', float, must = True)\ + # .add('stop_lr', float, must = False) + # self.cd = args.parse(jdata) + self.cd = {} + self.cd['start_lr'] = start_lr + self.cd['stop_lr'] = stop_lr + self.cd['decay_steps'] = decay_steps + self.cd['decay_rate'] = decay_rate + self.start_lr_ = self.cd['start_lr'] + + def build(self, + global_step : tf.Tensor, + stop_step : int = None + ) -> tf.Tensor : + """ + Build the learning rate + + Parameters + ---------- + global_step + The tf Tensor prividing the global training step + stop_step + The stop step. If provided, the decay_rate will be determined automatically and overwritten. + + Returns + ------- + learning_rate + The learning rate + """ + if stop_step is None: + self.decay_steps_ = self.cd['decay_steps'] if self.cd['decay_steps'] is not None else 5000 + self.decay_rate_ = self.cd['decay_rate'] if self.cd['decay_rate'] is not None else 0.95 + else: + self.stop_lr_ = self.cd['stop_lr'] if self.cd['stop_lr'] is not None else 5e-8 + default_ds = 100 if stop_step // 10 > 100 else stop_step // 100 + 1 + self.decay_steps_ = self.cd['decay_steps'] if self.cd['decay_steps'] is not None else default_ds + if self.decay_steps_ >= stop_step: + self.decay_steps_ = default_ds + self.decay_rate_ = np.exp(np.log(self.stop_lr_ / self.start_lr_) / (stop_step / self.decay_steps_)) + + return tf.train.exponential_decay(self.start_lr_, + global_step, + self.decay_steps_, + self.decay_rate_, + staircase=True) + def start_lr(self) -> float: + """ + Get the start lr + """ + return self.start_lr_ + + def value (self, + step : int + ) -> float: + """ + Get the lr at a certain step + """ + return self.start_lr_ * np.power (self.decay_rate_, (step // self.decay_steps_)) + diff --git a/deepmd/utils/neighbor_stat.py b/deepmd/utils/neighbor_stat.py new file mode 100644 index 0000000000..991993e364 --- /dev/null +++ b/deepmd/utils/neighbor_stat.py @@ -0,0 +1,95 @@ +import math +import logging +import numpy as np +from tqdm import tqdm +from deepmd.env import tf +from typing import Tuple, List +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.utils.data_system import DeepmdDataSystem + +log = logging.getLogger(__name__) + +class NeighborStat(): + """ + Class for getting training data information. + It loads data from DeepmdData object, and measures the data info, including neareest nbor distance between atoms, max nbor size of atoms and the output data range of the environment matrix. + """ + def __init__(self, + ntypes : int, + rcut: float) -> None: + """ + Constructor + + Parameters + ---------- + ntypes + The num of atom types + rcut + The cut-off radius + """ + self.rcut = rcut + self.ntypes = ntypes + self.place_holders = {} + sub_graph = tf.Graph() + with sub_graph.as_default(): + for ii in ['coord', 'box']: + self.place_holders[ii] = tf.placeholder(GLOBAL_NP_FLOAT_PRECISION, [None, None], name='t_'+ii) + self.place_holders['type'] = tf.placeholder(tf.int32, [None, None], name='t_type') + self.place_holders['natoms_vec'] = tf.placeholder(tf.int32, [self.ntypes+2], name='t_natoms') + self.place_holders['default_mesh'] = tf.placeholder(tf.int32, [None], name='t_mesh') + self._max_nbor_size, self._min_nbor_dist \ + = op_module.neighbor_stat(self.place_holders['coord'], + self.place_holders['type'], + self.place_holders['natoms_vec'], + self.place_holders['box'], + self.place_holders['default_mesh'], + rcut = self.rcut) + self.sub_sess = tf.Session(graph = sub_graph, config=default_tf_session_config) + + def get_stat(self, + data : DeepmdDataSystem) -> Tuple[float, List[int]]: + """ + get the data statistics of the training data, including nearest nbor distance between atoms, max nbor size of atoms + + Parameters + ---------- + data + Class for manipulating many data systems. It is implemented with the help of DeepmdData. + + Returns + ------- + min_nbor_dist + The nearest distance between neighbor atoms + max_nbor_size + A list with ntypes integers, denotes the actual achieved max sel + """ + self.min_nbor_dist = 100.0 + self.max_nbor_size = [0] * self.ntypes + + # for ii in tqdm(range(len(data.system_dirs)), desc = 'DEEPMD INFO |-> deepmd.utils.neighbor_stat\t\t\tgetting neighbor status'): + for ii in range(len(data.system_dirs)): + for jj in data.data_systems[ii].dirs: + data_set = data.data_systems[ii]._load_set(jj) + for kk in range(np.array(data_set['type']).shape[0]): + mn, dt \ + = self.sub_sess.run([self._max_nbor_size, self._min_nbor_dist], + feed_dict = { + self.place_holders['coord']: np.array(data_set['coord'])[kk].reshape([-1, data.natoms[ii] * 3]), + self.place_holders['type']: np.array(data_set['type'])[kk].reshape([-1, data.natoms[ii]]), + self.place_holders['natoms_vec']: np.array(data.natoms_vec[ii]), + self.place_holders['box']: np.array(data_set['box'])[kk].reshape([-1, 9]), + self.place_holders['default_mesh']: np.array(data.default_mesh[ii]), + }) + dt = np.min(dt) + if dt < self.min_nbor_dist: + self.min_nbor_dist = dt + for ww in range(self.ntypes): + var = np.max(mn[:, ww]) + if var > self.max_nbor_size[ww]: + self.max_nbor_size[ww] = var + + log.info('training data with min nbor dist: ' + str(self.min_nbor_dist)) + log.info('training data with max nbor size: ' + str(self.max_nbor_size)) + return self.min_nbor_dist, self.max_nbor_size diff --git a/deepmd/utils/network.py b/deepmd/utils/network.py new file mode 100644 index 0000000000..b7d2fb24c2 --- /dev/null +++ b/deepmd/utils/network.py @@ -0,0 +1,186 @@ +import numpy as np + +from deepmd.env import tf +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION + +def one_layer_rand_seed_shift(): + return 3 + +def one_layer(inputs, + outputs_size, + activation_fn=tf.nn.tanh, + precision = GLOBAL_TF_FLOAT_PRECISION, + stddev=1.0, + bavg=0.0, + name='linear', + reuse=None, + seed=None, + use_timestep = False, + trainable = True, + useBN = False, + uniform_seed = False): + with tf.variable_scope(name, reuse=reuse): + shape = inputs.get_shape().as_list() + w = tf.get_variable('matrix', + [shape[1], outputs_size], + precision, + tf.random_normal_initializer( + stddev=stddev/np.sqrt(shape[1]+outputs_size), + seed = seed if (seed is None or uniform_seed) else seed + 0 + ), + trainable = trainable) + variable_summaries(w, 'matrix') + b = tf.get_variable('bias', + [outputs_size], + precision, + tf.random_normal_initializer( + stddev=stddev, + mean = bavg, + seed = seed if (seed is None or uniform_seed) else seed + 1 + ), + trainable = trainable) + variable_summaries(b, 'bias') + hidden = tf.matmul(inputs, w) + b + if activation_fn != None and use_timestep : + idt = tf.get_variable('idt', + [outputs_size], + precision, + tf.random_normal_initializer( + stddev=0.001, + mean = 0.1, + seed = seed if (seed is None or uniform_seed) else seed + 2 + ), + trainable = trainable) + variable_summaries(idt, 'idt') + if activation_fn != None: + if useBN: + None + # hidden_bn = self._batch_norm(hidden, name=name+'_normalization', reuse=reuse) + # return activation_fn(hidden_bn) + else: + if use_timestep : + return tf.reshape(activation_fn(hidden), [-1, outputs_size]) * idt + else : + return tf.reshape(activation_fn(hidden), [-1, outputs_size]) + else: + if useBN: + None + # return self._batch_norm(hidden, name=name+'_normalization', reuse=reuse) + else: + return hidden + + +def embedding_net_rand_seed_shift( + network_size +): + shift = 3 * (len(network_size) + 1) + return shift + +def embedding_net(xx, + network_size, + precision, + activation_fn = tf.nn.tanh, + resnet_dt = False, + name_suffix = '', + stddev = 1.0, + bavg = 0.0, + seed = None, + trainable = True, + uniform_seed = False): + """ + Parameters + ---------- + xx : Tensor + Input tensor of shape [-1,1] + network_size: list of int + Size of the embedding network. For example [16,32,64] + precision: + Precision of network weights. For example, tf.float64 + activation_fn: + Activation function + resnet_dt: boolean + Using time-step in the ResNet construction + name_suffix: str + The name suffix append to each variable. + stddev: float + Standard deviation of initializing network parameters + bavg: float + Mean of network intial bias + seed: int + Random seed for initializing network parameters + trainable: boolean + If the netowk is trainable + """ + input_shape = xx.get_shape().as_list() + outputs_size = [input_shape[1]] + network_size + + for ii in range(1, len(outputs_size)): + w = tf.get_variable('matrix_'+str(ii)+name_suffix, + [outputs_size[ii - 1], outputs_size[ii]], + precision, + tf.random_normal_initializer( + stddev=stddev/np.sqrt(outputs_size[ii]+outputs_size[ii-1]), + seed = seed if (seed is None or uniform_seed) else seed + ii*3+0 + ), + trainable = trainable) + variable_summaries(w, 'matrix_'+str(ii)+name_suffix) + + b = tf.get_variable('bias_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer( + stddev=stddev, + mean = bavg, + seed = seed if (seed is None or uniform_seed) else seed + 3*ii+1 + ), + trainable = trainable) + variable_summaries(b, 'bias_'+str(ii)+name_suffix) + + hidden = tf.reshape(activation_fn(tf.matmul(xx, w) + b), [-1, outputs_size[ii]]) + if resnet_dt : + idt = tf.get_variable('idt_'+str(ii)+name_suffix, + [1, outputs_size[ii]], + precision, + tf.random_normal_initializer( + stddev=0.001, + mean = 1.0, + seed = seed if (seed is None or uniform_seed) else seed + 3*ii+2 + ), + trainable = trainable) + variable_summaries(idt, 'idt_'+str(ii)+name_suffix) + + if outputs_size[ii] == outputs_size[ii-1]: + if resnet_dt : + xx += hidden * idt + else : + xx += hidden + elif outputs_size[ii] == outputs_size[ii-1] * 2: + if resnet_dt : + xx = tf.concat([xx,xx], 1) + hidden * idt + else : + xx = tf.concat([xx,xx], 1) + hidden + else: + xx = hidden + + return xx + +def variable_summaries(var: tf.Variable, name: str): + """Attach a lot of summaries to a Tensor (for TensorBoard visualization). + + Parameters + ---------- + var : tf.Variable + [description] + name : str + variable name + """ + with tf.name_scope(name): + mean = tf.reduce_mean(var) + tf.summary.scalar('mean', mean) + + with tf.name_scope('stddev'): + stddev = tf.sqrt(tf.reduce_mean(tf.square(var - mean))) + tf.summary.scalar('stddev', stddev) + tf.summary.scalar('max', tf.reduce_max(var)) + tf.summary.scalar('min', tf.reduce_min(var)) + tf.summary.histogram('histogram', var) diff --git a/source/train/TabInter.py b/deepmd/utils/pair_tab.py similarity index 57% rename from source/train/TabInter.py rename to deepmd/utils/pair_tab.py index 99190e1191..4e22033a3b 100644 --- a/source/train/TabInter.py +++ b/deepmd/utils/pair_tab.py @@ -1,15 +1,45 @@ #!/usr/bin/env python3 import numpy as np +from typing import Tuple, List + from scipy.interpolate import CubicSpline -class TabInter (object): +class PairTab (object): def __init__(self, - filename): + filename : str + ) -> None: + """ + Constructor + + 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.reinit(filename) def reinit(self, - filename): + 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] @@ -22,7 +52,10 @@ def reinit(self, self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) self.tab_data = self._make_data() - def get(self) : + def get(self) -> Tuple[np.array, np.array]: + """ + Get the serialized table. + """ return self.tab_info, self.tab_data def _make_data(self) : diff --git a/deepmd/utils/tabulate.py b/deepmd/utils/tabulate.py new file mode 100644 index 0000000000..1590d459dd --- /dev/null +++ b/deepmd/utils/tabulate.py @@ -0,0 +1,253 @@ +import re +import math +import logging +import numpy as np +from tqdm import tqdm +from typing import Tuple, List +from deepmd.env import tf +from deepmd.env import op_module +from tensorflow.python.platform import gfile +from tensorflow.python.framework import tensor_util + +log = logging.getLogger(__name__) + +class DeepTabulate(): + """ + Class for tabulation. + Compress a model, which including tabulating the embedding-net. + The table is composed of fifth-order polynomial coefficients and is assembled from two sub-tables. The first table takes the stride(parameter) as it\'s uniform stride, while the second table takes 10 * stride as it\s uniform stride + The range of the first table is automatically detected by deepmd-kit, while the second table ranges from the first table\'s upper boundary(upper) to the extrapolate(parameter) * upper. + """ + def __init__(self, + model_file : str, + type_one_side : bool = False) -> None: + """ + Constructor + + Parameters + ---------- + model_file + The frozen model + type_one_side + Try to build N_types tables. Otherwise, building N_types^2 tables + """ + + self.model_file = model_file + self.type_one_side = type_one_side + + self.graph, self.graph_def = self._load_graph() + self.sess = tf.Session(graph = self.graph) + + self.sub_graph, self.sub_graph_def = self._load_sub_graph() + self.sub_sess = tf.Session(graph = self.sub_graph) + + try: + self.sel_a = self.graph.get_operation_by_name('ProdEnvMatA').get_attr('sel_a') + self.descrpt = self.graph.get_operation_by_name ('ProdEnvMatA') + except Exception: + self.sel_a = self.graph.get_operation_by_name('DescrptSeA').get_attr('sel_a') + self.descrpt = self.graph.get_operation_by_name ('DescrptSeA') + self.ntypes = self._get_tensor_value(self.graph.get_tensor_by_name ('descrpt_attr/ntypes:0')) + self.davg = self._get_tensor_value(self.graph.get_tensor_by_name ('descrpt_attr/t_avg:0')) + self.dstd = self._get_tensor_value(self.graph.get_tensor_by_name ('descrpt_attr/t_std:0')) + + + self.rcut = self.descrpt.get_attr('rcut_r') + self.rcut_smth = self.descrpt.get_attr('rcut_r_smth') + + self.filter_variable_nodes = self._load_matrix_node() + self.layer_size = int(len(self.filter_variable_nodes) / (self.ntypes * self.ntypes * 2)) + self.table_size = self.ntypes * self.ntypes + if type_one_side : + self.layer_size = int(len(self.filter_variable_nodes) / (self.ntypes * 2)) + self.table_size = self.ntypes + # self.value_type = self.filter_variable_nodes["filter_type_0/matrix_1_0"].dtype #"filter_type_0/matrix_1_0" must exit~ + # get trained variables + self.bias = self._get_bias() + self.matrix = self._get_matrix() + + self.data_type = type(self.matrix["layer_1"][0][0][0]) + assert self.matrix["layer_1"][0].size > 0, "no matrix exist in matrix array!" + self.last_layer_size = self.matrix["layer_" + str(self.layer_size)][0].shape[1] + # define tables + self.data = {} + + # TODO: Need a check function to determine if the current model is properly + + def build(self, + min_nbor_dist : float, + extrapolate : float, + stride0 : float, + stride1 : float) -> Tuple[int, int]: + """ + Build the tables for model compression + + Parameters + ---------- + min_nbor_dist + The nearest distance between neighbor atoms + extrapolate + The scale of model extrapolation + stride0 + The uniform stride of the first table + stride1 + The uniform stride of the second table + + Returns + ---------- + lower + The lower boundary of environment matrix + upper + The upper boundary of environment matrix + """ + # tabulate range [lower, upper] with stride0 'stride0' + lower, upper = self._get_env_mat_range(min_nbor_dist) + xx = np.arange(lower, upper, stride0, dtype = self.data_type) + xx = np.append(xx, np.arange(upper, extrapolate * upper, stride1, dtype = self.data_type)) + xx = np.append(xx, np.array([extrapolate * upper], dtype = self.data_type)) + self.nspline = int((upper - lower) / stride0 + (extrapolate * upper - upper) / stride1) + for ii in range(self.table_size): + vv, dd, d2 = self._make_data(xx, ii) + if self.type_one_side: + net = "filter_-1_net_" + str(int(ii)) + else: + net = "filter_" + str(int(ii / self.ntypes)) + "_net_" + str(int(ii % self.ntypes)) + self.data[net] = np.zeros([self.nspline, 6 * self.last_layer_size], dtype = self.data_type) + # for jj in tqdm(range(self.nspline), desc = 'DEEPMD INFO |-> deepmd.utils.tabulate\t\t\t' + net + ', tabulating'): + for jj in range(self.nspline): + for kk in range(self.last_layer_size): + if jj < int((upper - lower) / stride0): + tt = stride0 + else: + tt = stride1 + hh = vv[jj + 1][kk] - vv[jj][kk] + self.data[net][jj][kk * 6 + 0] = vv[jj][kk] + self.data[net][jj][kk * 6 + 1] = dd[jj][kk] + self.data[net][jj][kk * 6 + 2] = 0.5 * d2[jj][kk] + self.data[net][jj][kk * 6 + 3] = (1 / (2 * tt * tt * tt)) * (20 * hh - (8 * dd[jj + 1][kk] + 12 * dd[jj][kk]) * tt - (3 * d2[jj][kk] - d2[jj + 1][kk]) * tt * tt) + self.data[net][jj][kk * 6 + 4] = (1 / (2 * tt * tt * tt * tt)) * (-30 * hh + (14 * dd[jj + 1][kk] + 16 * dd[jj][kk]) * tt + (3 * d2[jj][kk] - 2 * d2[jj + 1][kk]) * tt * tt) + self.data[net][jj][kk * 6 + 5] = (1 / (2 * tt * tt * tt * tt * tt)) * (12 * hh - 6 * (dd[jj + 1][kk] + dd[jj][kk]) * tt + (d2[jj + 1][kk] - d2[jj][kk]) * tt * tt) + self.data[net] + return lower, upper + + def _load_graph(self): + graph_def = tf.GraphDef() + with open(self.model_file, "rb") as f: + graph_def.ParseFromString(f.read()) + with tf.Graph().as_default() as graph: + tf.import_graph_def(graph_def, name = "") + return graph, graph_def + + def _load_sub_graph(self): + sub_graph_def = tf.GraphDef() + with tf.Graph().as_default() as sub_graph: + tf.import_graph_def(sub_graph_def, name = "") + return sub_graph, sub_graph_def + + def _get_tensor_value(self, tensor) : + with self.sess.as_default(): + self.sess.run(tensor) + value = tensor.eval() + return value + + def _load_matrix_node(self): + matrix_node = {} + matrix_node_pattern = "filter_type_\d+/matrix_\d+_\d+|filter_type_\d+/bias_\d+_\d+|filter_type_\d+/idt_\d+_\d+|filter_type_all/matrix_\d+_\d+|filter_type_all/bias_\d+_\d+|filter_type_all/idt_\d+_\d" + for node in self.graph_def.node: + if re.fullmatch(matrix_node_pattern, node.name) != None: + matrix_node[node.name] = node.attr["value"].tensor + for key in matrix_node.keys() : + assert key.find('bias') > 0 or key.find('matrix') > 0, "currently, only support weight matrix and bias matrix at the tabulation op!" + return matrix_node + + def _get_bias(self): + bias = {} + for layer in range(1, self.layer_size + 1): + bias["layer_" + str(layer)] = [] + if self.type_one_side: + for ii in range(0, self.ntypes): + tensor_value = np.frombuffer (self.filter_variable_nodes["filter_type_all/bias_" + str(layer) + "_" + str(int(ii))].tensor_content) + tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_all/bias_" + str(layer) + "_" + str(int(ii))].tensor_shape).as_list() + bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) + else: + for ii in range(0, self.ntypes * self.ntypes): + tensor_value = np.frombuffer(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content) + tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/bias_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list() + bias["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) + return bias + + def _get_matrix(self): + matrix = {} + for layer in range(1, self.layer_size + 1): + matrix["layer_" + str(layer)] = [] + if self.type_one_side: + for ii in range(0, self.ntypes): + tensor_value = np.frombuffer (self.filter_variable_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(int(ii))].tensor_content) + tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_all/matrix_" + str(layer) + "_" + str(int(ii))].tensor_shape).as_list() + matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) + else: + for ii in range(0, self.ntypes * self.ntypes): + tensor_value = np.frombuffer(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_content) + tensor_shape = tf.TensorShape(self.filter_variable_nodes["filter_type_" + str(int(ii / self.ntypes)) + "/matrix_" + str(layer) + "_" + str(int(ii % self.ntypes))].tensor_shape).as_list() + matrix["layer_" + str(layer)].append(np.reshape(tensor_value, tensor_shape)) + return matrix + + # one-by-one executions + def _make_data(self, xx, idx): + with self.sub_graph.as_default(): + with self.sub_sess.as_default(): + xx = tf.reshape(xx, [xx.size, -1]) + for layer in range(self.layer_size): + if layer == 0: + yy = self._layer_0(xx, self.matrix["layer_" + str(layer + 1)][idx], self.bias["layer_" + str(layer + 1)][idx]) + dy = op_module.unaggregated_dy_dx_s(yy, self.matrix["layer_" + str(layer + 1)][idx]) + dy2 = op_module.unaggregated_dy2_dx_s(yy, dy, self.matrix["layer_" + str(layer + 1)][idx]) + else: + tt, yy = self._layer_1(yy, self.matrix["layer_" + str(layer + 1)][idx], self.bias["layer_" + str(layer + 1)][idx]) + dz = op_module.unaggregated_dy_dx(yy - tt, self.matrix["layer_" + str(layer + 1)][idx], dy) + dy2 = op_module.unaggregated_dy2_dx(yy - tt, self.matrix["layer_" + str(layer + 1)][idx], dz, dy, dy2) + dy = dz + + vv = yy.eval() + dd = dy.eval() + d2 = dy2.eval() + return vv, dd, d2 + + def _layer_0(self, x, w, b): + return tf.nn.tanh(tf.matmul(x, w) + b) + + def _layer_1(self, x, w, b): + t = tf.concat([x, x], axis = 1) + return t, tf.nn.tanh(tf.matmul(x, w) + b) + t + + def _save_data(self): + for ii in range(self.ntypes * self.ntypes): + net = "filter_" + str(int(ii / self.ntypes)) + "_net_" + str(int(ii % self.ntypes)) + np.savetxt('data_' + str(int(ii)), self.data[net]) + + def _get_env_mat_range(self, + min_nbor_dist): + lower = 100.0 + upper = -10.0 + sw = self._spline5_switch(min_nbor_dist, self.rcut_smth, self.rcut) + for ii in range(self.ntypes): + if lower > -self.davg[ii][0] / self.dstd[ii][0]: + lower = -self.davg[ii][0] / self.dstd[ii][0] + if upper < ((1 / min_nbor_dist) * sw - self.davg[ii][0]) / self.dstd[ii][0]: + upper = ((1 / min_nbor_dist) * sw - self.davg[ii][0]) / self.dstd[ii][0] + log.info('training data with lower boundary: ' + str(lower)) + log.info('training data with upper boundary: ' + str(upper)) + return math.floor(lower), math.ceil(upper) + + def _spline5_switch(self, + xx, + rmin, + rmax): + if xx < rmin: + vv = 1 + elif xx < rmax: + uu = (xx - rmin) / (rmax - rmin) + vv = uu*uu*uu * (-6 * uu*uu + 15 * uu - 10) + 1 + else: + vv = 0 + return vv \ No newline at end of file diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py new file mode 100644 index 0000000000..1c36d522b4 --- /dev/null +++ b/deepmd/utils/type_embed.py @@ -0,0 +1,142 @@ +import numpy as np +from typing import Tuple, List + +from deepmd.env import tf +from deepmd.utils.network import one_layer +from deepmd.env import GLOBAL_TF_FLOAT_PRECISION +from deepmd.env import GLOBAL_NP_FLOAT_PRECISION +from deepmd.env import op_module +from deepmd.env import default_tf_session_config +from deepmd.utils.network import embedding_net + +import math +from deepmd.common import get_activation_func, get_precision, ACTIVATION_FN_DICT, PRECISION_DICT, docstring_parameter, get_np_precision +from deepmd.utils.argcheck import list_to_doc +from deepmd.utils.tabulate import DeepTabulate + + +def embed_atom_type( + ntypes : int, + natoms : tf.Tensor, + type_embedding : tf.Tensor, +): + """ + Make the embedded type for the atoms in system. + The atoms are assumed to be sorted according to the type, + thus their types are described by a `tf.Tensor` natoms, see explanation below. + + Parameters + ---------- + ntypes: + Number of types. + natoms: + The number of atoms. This tensor has the length of Ntypes + 2 + 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 + type_embedding: + The type embedding. + It has the shape of [ntypes, embedding_dim] + + Returns + ------- + atom_embedding + The embedded type of each atom. + It has the shape of [numb_atoms, embedding_dim] + """ + te_out_dim = type_embedding.get_shape().as_list()[-1] + atype = [] + for ii in range(ntypes): + atype.append(tf.tile([ii], [natoms[2+ii]])) + atype = tf.concat(atype, axis = 0) + atm_embed = tf.nn.embedding_lookup(type_embedding,tf.cast(atype,dtype=tf.int32)) #(nf*natom)*nchnl + atm_embed = tf.reshape(atm_embed,[-1,te_out_dim]) + return atm_embed + + +class TypeEmbedNet(): + @docstring_parameter(list_to_doc(ACTIVATION_FN_DICT.keys()), list_to_doc(PRECISION_DICT.keys())) + def __init__( + self, + neuron: List[int]=[], + resnet_dt: bool = False, + activation_function: str = 'tanh', + precision: str = 'default', + trainable: bool = True, + seed: int = None, + uniform_seed: bool = False, + )->None: + """ + Constructor + Parameters + ---------- + neuron : list[int] + Number of neurons in each hidden layers of the embedding net + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + activation_function + The activation function in the embedding net. Supported options are {0} + precision + The precision of the embedding net parameters. Supported options are {1} + trainable + If the weights of embedding net are trainable. + seed + Random seed for initializing the network parameters. + uniform_seed + Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed + """ + self.neuron = neuron + self.seed = seed + self.filter_resnet_dt = resnet_dt + self.filter_precision = get_precision(precision) + self.filter_activation_fn = get_activation_func(activation_function) + self.trainable = trainable + self.uniform_seed = uniform_seed + + + def build( + self, + ntypes: int, + reuse = None, + suffix = '', + ): + """ + Build the computational graph for the descriptor + + Parameters + ---------- + ntypes + Number of atom types. + reuse + The weights in the networks should be reused when get the variable. + suffix + Name suffix to identify this descriptor + + Returns + ------- + embedded_types + The computational graph for embedded types + """ + types = tf.convert_to_tensor( + [ii for ii in range(ntypes)], + dtype = tf.int32 + ) + ebd_type = tf.cast(tf.one_hot(tf.cast(types,dtype=tf.int32),int(ntypes)), self.filter_precision) + ebd_type = tf.reshape(ebd_type, [-1, ntypes]) + name = 'type_embed_net' + suffix + with tf.variable_scope(name, reuse=reuse): + ebd_type = embedding_net( + ebd_type, + self.neuron, + activation_fn = self.filter_activation_fn, + precision = self.filter_precision, + resnet_dt = self.filter_resnet_dt, + seed = self.seed, + trainable = self.trainable, + uniform_seed = self.uniform_seed) + ebd_type = tf.reshape(ebd_type, [-1, self.neuron[-1]]) # nnei * neuron[-1] + self.ebd_type = tf.identity(ebd_type, name ='t_typeebd') + return self.ebd_type + + diff --git a/deepmd/utils/weight_avg.py b/deepmd/utils/weight_avg.py new file mode 100644 index 0000000000..aec5026ae4 --- /dev/null +++ b/deepmd/utils/weight_avg.py @@ -0,0 +1,34 @@ +from typing import TYPE_CHECKING, List, Dict, Optional, Tuple +import numpy as np + + +def weighted_average( + errors: List[Dict[str, Tuple[float, float]]] +) -> Dict: + """Compute wighted average of prediction errors 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 + Tuple: (error, weight) + + Returns + ------- + Dict + weighted averages + """ + sum_err = {} + sum_siz = {} + for err in errors: + for kk, (ee, ss) in err.items(): + if kk in sum_err: + sum_err[kk] += ee * ee * ss + sum_siz[kk] += ss + else : + sum_err[kk] = ee * ee * ss + sum_siz[kk] = ss + for kk in sum_err.keys(): + sum_err[kk] = np.sqrt(sum_err[kk] / sum_siz[kk]) + return sum_err diff --git a/doc/Doxyfile b/doc/Doxyfile new file mode 100644 index 0000000000..c7af10d41a --- /dev/null +++ b/doc/Doxyfile @@ -0,0 +1,2662 @@ +# Doxyfile 1.9.2 + +# This file describes the settings to be used by the documentation system +# doxygen (www.doxygen.org) for a project. +# +# All text after a double hash (##) is considered a comment and is placed in +# front of the TAG it is preceding. +# +# All text after a single hash (#) is considered a comment and will be ignored. +# The format is: +# TAG = value [value, ...] +# For lists, items can also be appended using: +# TAG += value [value, ...] +# Values that contain spaces should be placed between quotes (\" \"). + +#--------------------------------------------------------------------------- +# Project related configuration options +#--------------------------------------------------------------------------- + +# This tag specifies the encoding used for all characters in the configuration +# file that follow. The default is UTF-8 which is also the encoding used for all +# text before the first occurrence of this tag. Doxygen uses libiconv (or the +# iconv built into libc) for the transcoding. See +# https://www.gnu.org/software/libiconv/ for the list of possible encodings. +# The default value is: UTF-8. + +DOXYFILE_ENCODING = UTF-8 + +# The PROJECT_NAME tag is a single word (or a sequence of words surrounded by +# double-quotes, unless you are using Doxywizard) that should identify the +# project for which the documentation is generated. This name is used in the +# title of most generated pages and in a few other places. +# The default value is: My Project. + +PROJECT_NAME = "My Project" + +# The PROJECT_NUMBER tag can be used to enter a project or revision number. This +# could be handy for archiving the generated documentation or if some version +# control system is used. + +PROJECT_NUMBER = + +# Using the PROJECT_BRIEF tag one can provide an optional one line description +# for a project that appears at the top of each page and should give viewer a +# quick idea about the purpose of the project. Keep the description short. + +PROJECT_BRIEF = + +# With the PROJECT_LOGO tag one can specify a logo or an icon that is included +# in the documentation. The maximum height of the logo should not exceed 55 +# pixels and the maximum width should not exceed 200 pixels. Doxygen will copy +# the logo to the output directory. + +PROJECT_LOGO = + +# The OUTPUT_DIRECTORY tag is used to specify the (relative or absolute) path +# into which the generated documentation will be written. If a relative path is +# entered, it will be relative to the location where doxygen was started. If +# left blank the current directory will be used. + +OUTPUT_DIRECTORY = _build/ + +# If the CREATE_SUBDIRS tag is set to YES then doxygen will create 4096 sub- +# directories (in 2 levels) under the output directory of each output format and +# will distribute the generated files over these directories. Enabling this +# option can be useful when feeding doxygen a huge amount of source files, where +# putting all generated files in the same directory would otherwise causes +# performance problems for the file system. +# The default value is: NO. + +CREATE_SUBDIRS = NO + +# If the ALLOW_UNICODE_NAMES tag is set to YES, doxygen will allow non-ASCII +# characters to appear in the names of generated files. If set to NO, non-ASCII +# characters will be escaped, for example _xE3_x81_x84 will be used for Unicode +# U+3044. +# The default value is: NO. + +ALLOW_UNICODE_NAMES = NO + +# The OUTPUT_LANGUAGE tag is used to specify the language in which all +# documentation generated by doxygen is written. Doxygen will use this +# information to generate all constant output in the proper language. +# Possible values are: Afrikaans, Arabic, Armenian, Brazilian, Catalan, Chinese, +# Chinese-Traditional, Croatian, Czech, Danish, Dutch, English (United States), +# Esperanto, Farsi (Persian), Finnish, French, German, Greek, Hungarian, +# Indonesian, Italian, Japanese, Japanese-en (Japanese with English messages), +# Korean, Korean-en (Korean with English messages), Latvian, Lithuanian, +# Macedonian, Norwegian, Persian (Farsi), Polish, Portuguese, Romanian, Russian, +# Serbian, Serbian-Cyrillic, Slovak, Slovene, Spanish, Swedish, Turkish, +# Ukrainian and Vietnamese. +# The default value is: English. + +OUTPUT_LANGUAGE = English + +# If the BRIEF_MEMBER_DESC tag is set to YES, doxygen will include brief member +# descriptions after the members that are listed in the file and class +# documentation (similar to Javadoc). Set to NO to disable this. +# The default value is: YES. + +BRIEF_MEMBER_DESC = YES + +# If the REPEAT_BRIEF tag is set to YES, doxygen will prepend the brief +# description of a member or function before the detailed description +# +# Note: If both HIDE_UNDOC_MEMBERS and BRIEF_MEMBER_DESC are set to NO, the +# brief descriptions will be completely suppressed. +# The default value is: YES. + +REPEAT_BRIEF = YES + +# This tag implements a quasi-intelligent brief description abbreviator that is +# used to form the text in various listings. Each string in this list, if found +# as the leading text of the brief description, will be stripped from the text +# and the result, after processing the whole list, is used as the annotated +# text. Otherwise, the brief description is used as-is. If left blank, the +# following values are used ($name is automatically replaced with the name of +# the entity):The $name class, The $name widget, The $name file, is, provides, +# specifies, contains, represents, a, an and the. + +ABBREVIATE_BRIEF = "The $name class" \ + "The $name widget" \ + "The $name file" \ + is \ + provides \ + specifies \ + contains \ + represents \ + a \ + an \ + the + +# If the ALWAYS_DETAILED_SEC and REPEAT_BRIEF tags are both set to YES then +# doxygen will generate a detailed section even if there is only a brief +# description. +# The default value is: NO. + +ALWAYS_DETAILED_SEC = NO + +# If the INLINE_INHERITED_MEMB tag is set to YES, doxygen will show all +# inherited members of a class in the documentation of that class as if those +# members were ordinary class members. Constructors, destructors and assignment +# operators of the base classes will not be shown. +# The default value is: NO. + +INLINE_INHERITED_MEMB = NO + +# If the FULL_PATH_NAMES tag is set to YES, doxygen will prepend the full path +# before files name in the file list and in the header files. If set to NO the +# shortest path that makes the file name unique will be used +# The default value is: YES. + +FULL_PATH_NAMES = YES + +# The STRIP_FROM_PATH tag can be used to strip a user-defined part of the path. +# Stripping is only done if one of the specified strings matches the left-hand +# part of the path. The tag can be used to show relative paths in the file list. +# If left blank the directory from which doxygen is run is used as the path to +# strip. +# +# Note that you can specify absolute paths here, but also relative paths, which +# will be relative from the directory where doxygen is started. +# This tag requires that the tag FULL_PATH_NAMES is set to YES. + +STRIP_FROM_PATH = + +# The STRIP_FROM_INC_PATH tag can be used to strip a user-defined part of the +# path mentioned in the documentation of a class, which tells the reader which +# header file to include in order to use a class. If left blank only the name of +# the header file containing the class definition is used. Otherwise one should +# specify the list of include paths that are normally passed to the compiler +# using the -I flag. + +STRIP_FROM_INC_PATH = + +# If the SHORT_NAMES tag is set to YES, doxygen will generate much shorter (but +# less readable) file names. This can be useful is your file systems doesn't +# support long names like on DOS, Mac, or CD-ROM. +# The default value is: NO. + +SHORT_NAMES = NO + +# If the JAVADOC_AUTOBRIEF tag is set to YES then doxygen will interpret the +# first line (until the first dot) of a Javadoc-style comment as the brief +# description. If set to NO, the Javadoc-style will behave just like regular Qt- +# style comments (thus requiring an explicit @brief command for a brief +# description.) +# The default value is: NO. + +JAVADOC_AUTOBRIEF = NO + +# If the JAVADOC_BANNER tag is set to YES then doxygen will interpret a line +# such as +# /*************** +# as being the beginning of a Javadoc-style comment "banner". If set to NO, the +# Javadoc-style will behave just like regular comments and it will not be +# interpreted by doxygen. +# The default value is: NO. + +JAVADOC_BANNER = NO + +# If the QT_AUTOBRIEF tag is set to YES then doxygen will interpret the first +# line (until the first dot) of a Qt-style comment as the brief description. If +# set to NO, the Qt-style will behave just like regular Qt-style comments (thus +# requiring an explicit \brief command for a brief description.) +# The default value is: NO. + +QT_AUTOBRIEF = NO + +# The MULTILINE_CPP_IS_BRIEF tag can be set to YES to make doxygen treat a +# multi-line C++ special comment block (i.e. a block of //! or /// comments) as +# a brief description. This used to be the default behavior. The new default is +# to treat a multi-line C++ comment block as a detailed description. Set this +# tag to YES if you prefer the old behavior instead. +# +# Note that setting this tag to YES also means that rational rose comments are +# not recognized any more. +# The default value is: NO. + +MULTILINE_CPP_IS_BRIEF = NO + +# By default Python docstrings are displayed as preformatted text and doxygen's +# special commands cannot be used. By setting PYTHON_DOCSTRING to NO the +# doxygen's special commands can be used and the contents of the docstring +# documentation blocks is shown as doxygen documentation. +# The default value is: YES. + +PYTHON_DOCSTRING = YES + +# If the INHERIT_DOCS tag is set to YES then an undocumented member inherits the +# documentation from any documented member that it re-implements. +# The default value is: YES. + +INHERIT_DOCS = YES + +# If the SEPARATE_MEMBER_PAGES tag is set to YES then doxygen will produce a new +# page for each member. If set to NO, the documentation of a member will be part +# of the file/class/namespace that contains it. +# The default value is: NO. + +SEPARATE_MEMBER_PAGES = NO + +# The TAB_SIZE tag can be used to set the number of spaces in a tab. Doxygen +# uses this value to replace tabs by spaces in code fragments. +# Minimum value: 1, maximum value: 16, default value: 4. + +TAB_SIZE = 4 + +# This tag can be used to specify a number of aliases that act as commands in +# the documentation. An alias has the form: +# name=value +# For example adding +# "sideeffect=@par Side Effects:^^" +# will allow you to put the command \sideeffect (or @sideeffect) in the +# documentation, which will result in a user-defined paragraph with heading +# "Side Effects:". Note that you cannot put \n's in the value part of an alias +# to insert newlines (in the resulting output). You can put ^^ in the value part +# of an alias to insert a newline as if a physical newline was in the original +# file. When you need a literal { or } or , in the value part of an alias you +# have to escape them by means of a backslash (\), this can lead to conflicts +# with the commands \{ and \} for these it is advised to use the version @{ and +# @} or use a double escape (\\{ and \\}) + +ALIASES = + +# Set the OPTIMIZE_OUTPUT_FOR_C tag to YES if your project consists of C sources +# only. Doxygen will then generate output that is more tailored for C. For +# instance, some of the names that are used will be different. The list of all +# members will be omitted, etc. +# The default value is: NO. + +OPTIMIZE_OUTPUT_FOR_C = NO + +# Set the OPTIMIZE_OUTPUT_JAVA tag to YES if your project consists of Java or +# Python sources only. Doxygen will then generate output that is more tailored +# for that language. For instance, namespaces will be presented as packages, +# qualified scopes will look different, etc. +# The default value is: NO. + +OPTIMIZE_OUTPUT_JAVA = NO + +# Set the OPTIMIZE_FOR_FORTRAN tag to YES if your project consists of Fortran +# sources. Doxygen will then generate output that is tailored for Fortran. +# The default value is: NO. + +OPTIMIZE_FOR_FORTRAN = NO + +# Set the OPTIMIZE_OUTPUT_VHDL tag to YES if your project consists of VHDL +# sources. Doxygen will then generate output that is tailored for VHDL. +# The default value is: NO. + +OPTIMIZE_OUTPUT_VHDL = NO + +# Set the OPTIMIZE_OUTPUT_SLICE tag to YES if your project consists of Slice +# sources only. Doxygen will then generate output that is more tailored for that +# language. For instance, namespaces will be presented as modules, types will be +# separated into more groups, etc. +# The default value is: NO. + +OPTIMIZE_OUTPUT_SLICE = NO + +# Doxygen selects the parser to use depending on the extension of the files it +# parses. With this tag you can assign which parser to use for a given +# extension. Doxygen has a built-in mapping, but you can override or extend it +# using this tag. The format is ext=language, where ext is a file extension, and +# language is one of the parsers supported by doxygen: IDL, Java, JavaScript, +# Csharp (C#), C, C++, Lex, D, PHP, md (Markdown), Objective-C, Python, Slice, +# VHDL, Fortran (fixed format Fortran: FortranFixed, free formatted Fortran: +# FortranFree, unknown formatted Fortran: Fortran. In the later case the parser +# tries to guess whether the code is fixed or free formatted code, this is the +# default for Fortran type files). For instance to make doxygen treat .inc files +# as Fortran files (default is PHP), and .f files as C (default is Fortran), +# use: inc=Fortran f=C. +# +# Note: For files without extension you can use no_extension as a placeholder. +# +# Note that for custom extensions you also need to set FILE_PATTERNS otherwise +# the files are not read by doxygen. When specifying no_extension you should add +# * to the FILE_PATTERNS. +# +# Note see also the list of default file extension mappings. + +EXTENSION_MAPPING = + +# If the MARKDOWN_SUPPORT tag is enabled then doxygen pre-processes all comments +# according to the Markdown format, which allows for more readable +# documentation. See https://daringfireball.net/projects/markdown/ for details. +# The output of markdown processing is further processed by doxygen, so you can +# mix doxygen, HTML, and XML commands with Markdown formatting. Disable only in +# case of backward compatibilities issues. +# The default value is: YES. + +MARKDOWN_SUPPORT = YES + +# When the TOC_INCLUDE_HEADINGS tag is set to a non-zero value, all headings up +# to that level are automatically included in the table of contents, even if +# they do not have an id attribute. +# Note: This feature currently applies only to Markdown headings. +# Minimum value: 0, maximum value: 99, default value: 5. +# This tag requires that the tag MARKDOWN_SUPPORT is set to YES. + +TOC_INCLUDE_HEADINGS = 5 + +# When enabled doxygen tries to link words that correspond to documented +# classes, or namespaces to their corresponding documentation. Such a link can +# be prevented in individual cases by putting a % sign in front of the word or +# globally by setting AUTOLINK_SUPPORT to NO. +# The default value is: YES. + +AUTOLINK_SUPPORT = YES + +# If you use STL classes (i.e. std::string, std::vector, etc.) but do not want +# to include (a tag file for) the STL sources as input, then you should set this +# tag to YES in order to let doxygen match functions declarations and +# definitions whose arguments contain STL classes (e.g. func(std::string); +# versus func(std::string) {}). This also make the inheritance and collaboration +# diagrams that involve STL classes more complete and accurate. +# The default value is: NO. + +BUILTIN_STL_SUPPORT = NO + +# If you use Microsoft's C++/CLI language, you should set this option to YES to +# enable parsing support. +# The default value is: NO. + +CPP_CLI_SUPPORT = NO + +# Set the SIP_SUPPORT tag to YES if your project consists of sip (see: +# https://www.riverbankcomputing.com/software/sip/intro) sources only. Doxygen +# will parse them like normal C++ but will assume all classes use public instead +# of private inheritance when no explicit protection keyword is present. +# The default value is: NO. + +SIP_SUPPORT = NO + +# For Microsoft's IDL there are propget and propput attributes to indicate +# getter and setter methods for a property. Setting this option to YES will make +# doxygen to replace the get and set methods by a property in the documentation. +# This will only work if the methods are indeed getting or setting a simple +# type. If this is not the case, or you want to show the methods anyway, you +# should set this option to NO. +# The default value is: YES. + +IDL_PROPERTY_SUPPORT = YES + +# If member grouping is used in the documentation and the DISTRIBUTE_GROUP_DOC +# tag is set to YES then doxygen will reuse the documentation of the first +# member in the group (if any) for the other members of the group. By default +# all members of a group must be documented explicitly. +# The default value is: NO. + +DISTRIBUTE_GROUP_DOC = NO + +# If one adds a struct or class to a group and this option is enabled, then also +# any nested class or struct is added to the same group. By default this option +# is disabled and one has to add nested compounds explicitly via \ingroup. +# The default value is: NO. + +GROUP_NESTED_COMPOUNDS = NO + +# Set the SUBGROUPING tag to YES to allow class member groups of the same type +# (for instance a group of public functions) to be put as a subgroup of that +# type (e.g. under the Public Functions section). Set it to NO to prevent +# subgrouping. Alternatively, this can be done per class using the +# \nosubgrouping command. +# The default value is: YES. + +SUBGROUPING = YES + +# When the INLINE_GROUPED_CLASSES tag is set to YES, classes, structs and unions +# are shown inside the group in which they are included (e.g. using \ingroup) +# instead of on a separate page (for HTML and Man pages) or section (for LaTeX +# and RTF). +# +# Note that this feature does not work in combination with +# SEPARATE_MEMBER_PAGES. +# The default value is: NO. + +INLINE_GROUPED_CLASSES = NO + +# When the INLINE_SIMPLE_STRUCTS tag is set to YES, structs, classes, and unions +# with only public data fields or simple typedef fields will be shown inline in +# the documentation of the scope in which they are defined (i.e. file, +# namespace, or group documentation), provided this scope is documented. If set +# to NO, structs, classes, and unions are shown on a separate page (for HTML and +# Man pages) or section (for LaTeX and RTF). +# The default value is: NO. + +INLINE_SIMPLE_STRUCTS = NO + +# When TYPEDEF_HIDES_STRUCT tag is enabled, a typedef of a struct, union, or +# enum is documented as struct, union, or enum with the name of the typedef. So +# typedef struct TypeS {} TypeT, will appear in the documentation as a struct +# with name TypeT. When disabled the typedef will appear as a member of a file, +# namespace, or class. And the struct will be named TypeS. This can typically be +# useful for C code in case the coding convention dictates that all compound +# types are typedef'ed and only the typedef is referenced, never the tag name. +# The default value is: NO. + +TYPEDEF_HIDES_STRUCT = NO + +# The size of the symbol lookup cache can be set using LOOKUP_CACHE_SIZE. This +# cache is used to resolve symbols given their name and scope. Since this can be +# an expensive process and often the same symbol appears multiple times in the +# code, doxygen keeps a cache of pre-resolved symbols. If the cache is too small +# doxygen will become slower. If the cache is too large, memory is wasted. The +# cache size is given by this formula: 2^(16+LOOKUP_CACHE_SIZE). The valid range +# is 0..9, the default is 0, corresponding to a cache size of 2^16=65536 +# symbols. At the end of a run doxygen will report the cache usage and suggest +# the optimal cache size from a speed point of view. +# Minimum value: 0, maximum value: 9, default value: 0. + +LOOKUP_CACHE_SIZE = 0 + +# The NUM_PROC_THREADS specifies the number threads doxygen is allowed to use +# during processing. When set to 0 doxygen will based this on the number of +# cores available in the system. You can set it explicitly to a value larger +# than 0 to get more control over the balance between CPU load and processing +# speed. At this moment only the input processing can be done using multiple +# threads. Since this is still an experimental feature the default is set to 1, +# which effectively disables parallel processing. Please report any issues you +# encounter. Generating dot graphs in parallel is controlled by the +# DOT_NUM_THREADS setting. +# Minimum value: 0, maximum value: 32, default value: 1. + +NUM_PROC_THREADS = 1 + +#--------------------------------------------------------------------------- +# Build related configuration options +#--------------------------------------------------------------------------- + +# If the EXTRACT_ALL tag is set to YES, doxygen will assume all entities in +# documentation are documented, even if no documentation was available. Private +# class members and static file members will be hidden unless the +# EXTRACT_PRIVATE respectively EXTRACT_STATIC tags are set to YES. +# Note: This will also disable the warnings about undocumented members that are +# normally produced when WARNINGS is set to YES. +# The default value is: NO. + +EXTRACT_ALL = NO + +# If the EXTRACT_PRIVATE tag is set to YES, all private members of a class will +# be included in the documentation. +# The default value is: NO. + +EXTRACT_PRIVATE = NO + +# If the EXTRACT_PRIV_VIRTUAL tag is set to YES, documented private virtual +# methods of a class will be included in the documentation. +# The default value is: NO. + +EXTRACT_PRIV_VIRTUAL = NO + +# If the EXTRACT_PACKAGE tag is set to YES, all members with package or internal +# scope will be included in the documentation. +# The default value is: NO. + +EXTRACT_PACKAGE = NO + +# If the EXTRACT_STATIC tag is set to YES, all static members of a file will be +# included in the documentation. +# The default value is: NO. + +EXTRACT_STATIC = NO + +# If the EXTRACT_LOCAL_CLASSES tag is set to YES, classes (and structs) defined +# locally in source files will be included in the documentation. If set to NO, +# only classes defined in header files are included. Does not have any effect +# for Java sources. +# The default value is: YES. + +EXTRACT_LOCAL_CLASSES = YES + +# This flag is only useful for Objective-C code. If set to YES, local methods, +# which are defined in the implementation section but not in the interface are +# included in the documentation. If set to NO, only methods in the interface are +# included. +# The default value is: NO. + +EXTRACT_LOCAL_METHODS = NO + +# If this flag is set to YES, the members of anonymous namespaces will be +# extracted and appear in the documentation as a namespace called +# 'anonymous_namespace{file}', where file will be replaced with the base name of +# the file that contains the anonymous namespace. By default anonymous namespace +# are hidden. +# The default value is: NO. + +EXTRACT_ANON_NSPACES = NO + +# If this flag is set to YES, the name of an unnamed parameter in a declaration +# will be determined by the corresponding definition. By default unnamed +# parameters remain unnamed in the output. +# The default value is: YES. + +RESOLVE_UNNAMED_PARAMS = YES + +# If the HIDE_UNDOC_MEMBERS tag is set to YES, doxygen will hide all +# undocumented members inside documented classes or files. If set to NO these +# members will be included in the various overviews, but no documentation +# section is generated. This option has no effect if EXTRACT_ALL is enabled. +# The default value is: NO. + +HIDE_UNDOC_MEMBERS = NO + +# If the HIDE_UNDOC_CLASSES tag is set to YES, doxygen will hide all +# undocumented classes that are normally visible in the class hierarchy. If set +# to NO, these classes will be included in the various overviews. This option +# has no effect if EXTRACT_ALL is enabled. +# The default value is: NO. + +HIDE_UNDOC_CLASSES = NO + +# If the HIDE_FRIEND_COMPOUNDS tag is set to YES, doxygen will hide all friend +# declarations. If set to NO, these declarations will be included in the +# documentation. +# The default value is: NO. + +HIDE_FRIEND_COMPOUNDS = NO + +# If the HIDE_IN_BODY_DOCS tag is set to YES, doxygen will hide any +# documentation blocks found inside the body of a function. If set to NO, these +# blocks will be appended to the function's detailed documentation block. +# The default value is: NO. + +HIDE_IN_BODY_DOCS = NO + +# The INTERNAL_DOCS tag determines if documentation that is typed after a +# \internal command is included. If the tag is set to NO then the documentation +# will be excluded. Set it to YES to include the internal documentation. +# The default value is: NO. + +INTERNAL_DOCS = NO + +# With the correct setting of option CASE_SENSE_NAMES doxygen will better be +# able to match the capabilities of the underlying filesystem. In case the +# filesystem is case sensitive (i.e. it supports files in the same directory +# whose names only differ in casing), the option must be set to YES to properly +# deal with such files in case they appear in the input. For filesystems that +# are not case sensitive the option should be be set to NO to properly deal with +# output files written for symbols that only differ in casing, such as for two +# classes, one named CLASS and the other named Class, and to also support +# references to files without having to specify the exact matching casing. On +# Windows (including Cygwin) and MacOS, users should typically set this option +# to NO, whereas on Linux or other Unix flavors it should typically be set to +# YES. +# The default value is: system dependent. + +CASE_SENSE_NAMES = YES + +# If the HIDE_SCOPE_NAMES tag is set to NO then doxygen will show members with +# their full class and namespace scopes in the documentation. If set to YES, the +# scope will be hidden. +# The default value is: NO. + +HIDE_SCOPE_NAMES = NO + +# If the HIDE_COMPOUND_REFERENCE tag is set to NO (default) then doxygen will +# append additional text to a page's title, such as Class Reference. If set to +# YES the compound reference will be hidden. +# The default value is: NO. + +HIDE_COMPOUND_REFERENCE= NO + +# If the SHOW_INCLUDE_FILES tag is set to YES then doxygen will put a list of +# the files that are included by a file in the documentation of that file. +# The default value is: YES. + +SHOW_INCLUDE_FILES = YES + +# If the SHOW_GROUPED_MEMB_INC tag is set to YES then Doxygen will add for each +# grouped member an include statement to the documentation, telling the reader +# which file to include in order to use the member. +# The default value is: NO. + +SHOW_GROUPED_MEMB_INC = NO + +# If the FORCE_LOCAL_INCLUDES tag is set to YES then doxygen will list include +# files with double quotes in the documentation rather than with sharp brackets. +# The default value is: NO. + +FORCE_LOCAL_INCLUDES = NO + +# If the INLINE_INFO tag is set to YES then a tag [inline] is inserted in the +# documentation for inline members. +# The default value is: YES. + +INLINE_INFO = YES + +# If the SORT_MEMBER_DOCS tag is set to YES then doxygen will sort the +# (detailed) documentation of file and class members alphabetically by member +# name. If set to NO, the members will appear in declaration order. +# The default value is: YES. + +SORT_MEMBER_DOCS = YES + +# If the SORT_BRIEF_DOCS tag is set to YES then doxygen will sort the brief +# descriptions of file, namespace and class members alphabetically by member +# name. If set to NO, the members will appear in declaration order. Note that +# this will also influence the order of the classes in the class list. +# The default value is: NO. + +SORT_BRIEF_DOCS = NO + +# If the SORT_MEMBERS_CTORS_1ST tag is set to YES then doxygen will sort the +# (brief and detailed) documentation of class members so that constructors and +# destructors are listed first. If set to NO the constructors will appear in the +# respective orders defined by SORT_BRIEF_DOCS and SORT_MEMBER_DOCS. +# Note: If SORT_BRIEF_DOCS is set to NO this option is ignored for sorting brief +# member documentation. +# Note: If SORT_MEMBER_DOCS is set to NO this option is ignored for sorting +# detailed member documentation. +# The default value is: NO. + +SORT_MEMBERS_CTORS_1ST = NO + +# If the SORT_GROUP_NAMES tag is set to YES then doxygen will sort the hierarchy +# of group names into alphabetical order. If set to NO the group names will +# appear in their defined order. +# The default value is: NO. + +SORT_GROUP_NAMES = NO + +# If the SORT_BY_SCOPE_NAME tag is set to YES, the class list will be sorted by +# fully-qualified names, including namespaces. If set to NO, the class list will +# be sorted only by class name, not including the namespace part. +# Note: This option is not very useful if HIDE_SCOPE_NAMES is set to YES. +# Note: This option applies only to the class list, not to the alphabetical +# list. +# The default value is: NO. + +SORT_BY_SCOPE_NAME = NO + +# If the STRICT_PROTO_MATCHING option is enabled and doxygen fails to do proper +# type resolution of all parameters of a function it will reject a match between +# the prototype and the implementation of a member function even if there is +# only one candidate or it is obvious which candidate to choose by doing a +# simple string match. By disabling STRICT_PROTO_MATCHING doxygen will still +# accept a match between prototype and implementation in such cases. +# The default value is: NO. + +STRICT_PROTO_MATCHING = NO + +# The GENERATE_TODOLIST tag can be used to enable (YES) or disable (NO) the todo +# list. This list is created by putting \todo commands in the documentation. +# The default value is: YES. + +GENERATE_TODOLIST = YES + +# The GENERATE_TESTLIST tag can be used to enable (YES) or disable (NO) the test +# list. This list is created by putting \test commands in the documentation. +# The default value is: YES. + +GENERATE_TESTLIST = YES + +# The GENERATE_BUGLIST tag can be used to enable (YES) or disable (NO) the bug +# list. This list is created by putting \bug commands in the documentation. +# The default value is: YES. + +GENERATE_BUGLIST = YES + +# The GENERATE_DEPRECATEDLIST tag can be used to enable (YES) or disable (NO) +# the deprecated list. This list is created by putting \deprecated commands in +# the documentation. +# The default value is: YES. + +GENERATE_DEPRECATEDLIST= YES + +# The ENABLED_SECTIONS tag can be used to enable conditional documentation +# sections, marked by \if ... \endif and \cond +# ... \endcond blocks. + +ENABLED_SECTIONS = + +# The MAX_INITIALIZER_LINES tag determines the maximum number of lines that the +# initial value of a variable or macro / define can have for it to appear in the +# documentation. If the initializer consists of more lines than specified here +# it will be hidden. Use a value of 0 to hide initializers completely. The +# appearance of the value of individual variables and macros / defines can be +# controlled using \showinitializer or \hideinitializer command in the +# documentation regardless of this setting. +# Minimum value: 0, maximum value: 10000, default value: 30. + +MAX_INITIALIZER_LINES = 30 + +# Set the SHOW_USED_FILES tag to NO to disable the list of files generated at +# the bottom of the documentation of classes and structs. If set to YES, the +# list will mention the files that were used to generate the documentation. +# The default value is: YES. + +SHOW_USED_FILES = YES + +# Set the SHOW_FILES tag to NO to disable the generation of the Files page. This +# will remove the Files entry from the Quick Index and from the Folder Tree View +# (if specified). +# The default value is: YES. + +SHOW_FILES = YES + +# Set the SHOW_NAMESPACES tag to NO to disable the generation of the Namespaces +# page. This will remove the Namespaces entry from the Quick Index and from the +# Folder Tree View (if specified). +# The default value is: YES. + +SHOW_NAMESPACES = YES + +# The FILE_VERSION_FILTER tag can be used to specify a program or script that +# doxygen should invoke to get the current version for each file (typically from +# the version control system). Doxygen will invoke the program by executing (via +# popen()) the command command input-file, where command is the value of the +# FILE_VERSION_FILTER tag, and input-file is the name of an input file provided +# by doxygen. Whatever the program writes to standard output is used as the file +# version. For an example see the documentation. + +FILE_VERSION_FILTER = + +# The LAYOUT_FILE tag can be used to specify a layout file which will be parsed +# by doxygen. The layout file controls the global structure of the generated +# output files in an output format independent way. To create the layout file +# that represents doxygen's defaults, run doxygen with the -l option. You can +# optionally specify a file name after the option, if omitted DoxygenLayout.xml +# will be used as the name of the layout file. See also section "Changing the +# layout of pages" for information. +# +# Note that if you run doxygen from a directory containing a file called +# DoxygenLayout.xml, doxygen will parse it automatically even if the LAYOUT_FILE +# tag is left empty. + +LAYOUT_FILE = + +# The CITE_BIB_FILES tag can be used to specify one or more bib files containing +# the reference definitions. This must be a list of .bib files. The .bib +# extension is automatically appended if omitted. This requires the bibtex tool +# to be installed. See also https://en.wikipedia.org/wiki/BibTeX for more info. +# For LaTeX the style of the bibliography can be controlled using +# LATEX_BIB_STYLE. To use this feature you need bibtex and perl available in the +# search path. See also \cite for info how to create references. + +CITE_BIB_FILES = + +#--------------------------------------------------------------------------- +# Configuration options related to warning and progress messages +#--------------------------------------------------------------------------- + +# The QUIET tag can be used to turn on/off the messages that are generated to +# standard output by doxygen. If QUIET is set to YES this implies that the +# messages are off. +# The default value is: NO. + +QUIET = NO + +# The WARNINGS tag can be used to turn on/off the warning messages that are +# generated to standard error (stderr) by doxygen. If WARNINGS is set to YES +# this implies that the warnings are on. +# +# Tip: Turn warnings on while writing the documentation. +# The default value is: YES. + +WARNINGS = YES + +# If the WARN_IF_UNDOCUMENTED tag is set to YES then doxygen will generate +# warnings for undocumented members. If EXTRACT_ALL is set to YES then this flag +# will automatically be disabled. +# The default value is: YES. + +WARN_IF_UNDOCUMENTED = YES + +# If the WARN_IF_DOC_ERROR tag is set to YES, doxygen will generate warnings for +# potential errors in the documentation, such as documenting some parameters in +# a documented function twice, or documenting parameters that don't exist or +# using markup commands wrongly. +# The default value is: YES. + +WARN_IF_DOC_ERROR = YES + +# If WARN_IF_INCOMPLETE_DOC is set to YES, doxygen will warn about incomplete +# function parameter documentation. If set to NO, doxygen will accept that some +# parameters have no documentation without warning. +# The default value is: YES. + +WARN_IF_INCOMPLETE_DOC = YES + +# This WARN_NO_PARAMDOC option can be enabled to get warnings for functions that +# are documented, but have no documentation for their parameters or return +# value. If set to NO, doxygen will only warn about wrong parameter +# documentation, but not about the absence of documentation. If EXTRACT_ALL is +# set to YES then this flag will automatically be disabled. See also +# WARN_IF_INCOMPLETE_DOC +# The default value is: NO. + +WARN_NO_PARAMDOC = NO + +# If the WARN_AS_ERROR tag is set to YES then doxygen will immediately stop when +# a warning is encountered. If the WARN_AS_ERROR tag is set to FAIL_ON_WARNINGS +# then doxygen will continue running as if WARN_AS_ERROR tag is set to NO, but +# at the end of the doxygen process doxygen will return with a non-zero status. +# Possible values are: NO, YES and FAIL_ON_WARNINGS. +# The default value is: NO. + +WARN_AS_ERROR = NO + +# The WARN_FORMAT tag determines the format of the warning messages that doxygen +# can produce. The string should contain the $file, $line, and $text tags, which +# will be replaced by the file and line number from which the warning originated +# and the warning text. Optionally the format may contain $version, which will +# be replaced by the version of the file (if it could be obtained via +# FILE_VERSION_FILTER) +# The default value is: $file:$line: $text. + +WARN_FORMAT = "$file:$line: $text" + +# The WARN_LOGFILE tag can be used to specify a file to which warning and error +# messages should be written. If left blank the output is written to standard +# error (stderr). + +WARN_LOGFILE = + +#--------------------------------------------------------------------------- +# Configuration options related to the input files +#--------------------------------------------------------------------------- + +# The INPUT tag is used to specify the files and/or directories that contain +# documented source files. You may enter file names like myfile.cpp or +# directories like /usr/src/myproject. Separate the files or directories with +# spaces. See also FILE_PATTERNS and EXTENSION_MAPPING +# Note: If this tag is empty the current directory is searched. + +INPUT = ../source/api_cc/include/ + +# This tag can be used to specify the character encoding of the source files +# that doxygen parses. Internally doxygen uses the UTF-8 encoding. Doxygen uses +# libiconv (or the iconv built into libc) for the transcoding. See the libiconv +# documentation (see: +# https://www.gnu.org/software/libiconv/) for the list of possible encodings. +# The default value is: UTF-8. + +INPUT_ENCODING = UTF-8 + +# If the value of the INPUT tag contains directories, you can use the +# FILE_PATTERNS tag to specify one or more wildcard patterns (like *.cpp and +# *.h) to filter out the source-files in the directories. +# +# Note that for custom extensions or not directly supported extensions you also +# need to set EXTENSION_MAPPING for the extension otherwise the files are not +# read by doxygen. +# +# Note the list of default checked file patterns might differ from the list of +# default file extension mappings. +# +# If left blank the following patterns are tested:*.c, *.cc, *.cxx, *.cpp, +# *.c++, *.java, *.ii, *.ixx, *.ipp, *.i++, *.inl, *.idl, *.ddl, *.odl, *.h, +# *.hh, *.hxx, *.hpp, *.h++, *.l, *.cs, *.d, *.php, *.php4, *.php5, *.phtml, +# *.inc, *.m, *.markdown, *.md, *.mm, *.dox (to be provided as doxygen C +# comment), *.py, *.pyw, *.f90, *.f95, *.f03, *.f08, *.f18, *.f, *.for, *.vhd, +# *.vhdl, *.ucf, *.qsf and *.ice. + +FILE_PATTERNS = *.c \ + *.cc \ + *.cxx \ + *.cpp \ + *.c++ \ + *.java \ + *.ii \ + *.ixx \ + *.ipp \ + *.i++ \ + *.inl \ + *.idl \ + *.ddl \ + *.odl \ + *.h \ + *.hh \ + *.hxx \ + *.hpp \ + *.h++ \ + *.l \ + *.cs \ + *.d \ + *.php \ + *.php4 \ + *.php5 \ + *.phtml \ + *.inc \ + *.m \ + *.markdown \ + *.md \ + *.mm \ + *.dox \ + *.py \ + *.pyw \ + *.f90 \ + *.f95 \ + *.f03 \ + *.f08 \ + *.f18 \ + *.f \ + *.for \ + *.vhd \ + *.vhdl \ + *.ucf \ + *.qsf \ + *.ice + +# The RECURSIVE tag can be used to specify whether or not subdirectories should +# be searched for input files as well. +# The default value is: NO. + +RECURSIVE = NO + +# The EXCLUDE tag can be used to specify files and/or directories that should be +# excluded from the INPUT source files. This way you can easily exclude a +# subdirectory from a directory tree whose root is specified with the INPUT tag. +# +# Note that relative paths are relative to the directory from which doxygen is +# run. + +EXCLUDE = + +# The EXCLUDE_SYMLINKS tag can be used to select whether or not files or +# directories that are symbolic links (a Unix file system feature) are excluded +# from the input. +# The default value is: NO. + +EXCLUDE_SYMLINKS = NO + +# If the value of the INPUT tag contains directories, you can use the +# EXCLUDE_PATTERNS tag to specify one or more wildcard patterns to exclude +# certain files from those directories. +# +# Note that the wildcards are matched against the file with absolute path, so to +# exclude all test directories for example use the pattern */test/* + +EXCLUDE_PATTERNS = + +# The EXCLUDE_SYMBOLS tag can be used to specify one or more symbol names +# (namespaces, classes, functions, etc.) that should be excluded from the +# output. The symbol name can be a fully qualified name, a word, or if the +# wildcard * is used, a substring. Examples: ANamespace, AClass, +# AClass::ANamespace, ANamespace::*Test +# +# Note that the wildcards are matched against the file with absolute path, so to +# exclude all test directories use the pattern */test/* + +EXCLUDE_SYMBOLS = + +# The EXAMPLE_PATH tag can be used to specify one or more files or directories +# that contain example code fragments that are included (see the \include +# command). + +EXAMPLE_PATH = + +# If the value of the EXAMPLE_PATH tag contains directories, you can use the +# EXAMPLE_PATTERNS tag to specify one or more wildcard pattern (like *.cpp and +# *.h) to filter out the source-files in the directories. If left blank all +# files are included. + +EXAMPLE_PATTERNS = * + +# If the EXAMPLE_RECURSIVE tag is set to YES then subdirectories will be +# searched for input files to be used with the \include or \dontinclude commands +# irrespective of the value of the RECURSIVE tag. +# The default value is: NO. + +EXAMPLE_RECURSIVE = NO + +# The IMAGE_PATH tag can be used to specify one or more files or directories +# that contain images that are to be included in the documentation (see the +# \image command). + +IMAGE_PATH = + +# The INPUT_FILTER tag can be used to specify a program that doxygen should +# invoke to filter for each input file. Doxygen will invoke the filter program +# by executing (via popen()) the command: +# +# +# +# where is the value of the INPUT_FILTER tag, and is the +# name of an input file. Doxygen will then use the output that the filter +# program writes to standard output. If FILTER_PATTERNS is specified, this tag +# will be ignored. +# +# Note that the filter must not add or remove lines; it is applied before the +# code is scanned, but not when the output code is generated. If lines are added +# or removed, the anchors will not be placed correctly. +# +# Note that for custom extensions or not directly supported extensions you also +# need to set EXTENSION_MAPPING for the extension otherwise the files are not +# properly processed by doxygen. + +INPUT_FILTER = + +# The FILTER_PATTERNS tag can be used to specify filters on a per file pattern +# basis. Doxygen will compare the file name with each pattern and apply the +# filter if there is a match. The filters are a list of the form: pattern=filter +# (like *.cpp=my_cpp_filter). See INPUT_FILTER for further information on how +# filters are used. If the FILTER_PATTERNS tag is empty or if none of the +# patterns match the file name, INPUT_FILTER is applied. +# +# Note that for custom extensions or not directly supported extensions you also +# need to set EXTENSION_MAPPING for the extension otherwise the files are not +# properly processed by doxygen. + +FILTER_PATTERNS = + +# If the FILTER_SOURCE_FILES tag is set to YES, the input filter (if set using +# INPUT_FILTER) will also be used to filter the input files that are used for +# producing the source files to browse (i.e. when SOURCE_BROWSER is set to YES). +# The default value is: NO. + +FILTER_SOURCE_FILES = NO + +# The FILTER_SOURCE_PATTERNS tag can be used to specify source filters per file +# pattern. A pattern will override the setting for FILTER_PATTERN (if any) and +# it is also possible to disable source filtering for a specific pattern using +# *.ext= (so without naming a filter). +# This tag requires that the tag FILTER_SOURCE_FILES is set to YES. + +FILTER_SOURCE_PATTERNS = + +# If the USE_MDFILE_AS_MAINPAGE tag refers to the name of a markdown file that +# is part of the input, its contents will be placed on the main page +# (index.html). This can be useful if you have a project on for instance GitHub +# and want to reuse the introduction page also for the doxygen output. + +USE_MDFILE_AS_MAINPAGE = + +#--------------------------------------------------------------------------- +# Configuration options related to source browsing +#--------------------------------------------------------------------------- + +# If the SOURCE_BROWSER tag is set to YES then a list of source files will be +# generated. Documented entities will be cross-referenced with these sources. +# +# Note: To get rid of all source code in the generated output, make sure that +# also VERBATIM_HEADERS is set to NO. +# The default value is: NO. + +SOURCE_BROWSER = NO + +# Setting the INLINE_SOURCES tag to YES will include the body of functions, +# classes and enums directly into the documentation. +# The default value is: NO. + +INLINE_SOURCES = NO + +# Setting the STRIP_CODE_COMMENTS tag to YES will instruct doxygen to hide any +# special comment blocks from generated source code fragments. Normal C, C++ and +# Fortran comments will always remain visible. +# The default value is: YES. + +STRIP_CODE_COMMENTS = YES + +# If the REFERENCED_BY_RELATION tag is set to YES then for each documented +# entity all documented functions referencing it will be listed. +# The default value is: NO. + +REFERENCED_BY_RELATION = NO + +# If the REFERENCES_RELATION tag is set to YES then for each documented function +# all documented entities called/used by that function will be listed. +# The default value is: NO. + +REFERENCES_RELATION = NO + +# If the REFERENCES_LINK_SOURCE tag is set to YES and SOURCE_BROWSER tag is set +# to YES then the hyperlinks from functions in REFERENCES_RELATION and +# REFERENCED_BY_RELATION lists will link to the source code. Otherwise they will +# link to the documentation. +# The default value is: YES. + +REFERENCES_LINK_SOURCE = YES + +# If SOURCE_TOOLTIPS is enabled (the default) then hovering a hyperlink in the +# source code will show a tooltip with additional information such as prototype, +# brief description and links to the definition and documentation. Since this +# will make the HTML file larger and loading of large files a bit slower, you +# can opt to disable this feature. +# The default value is: YES. +# This tag requires that the tag SOURCE_BROWSER is set to YES. + +SOURCE_TOOLTIPS = YES + +# If the USE_HTAGS tag is set to YES then the references to source code will +# point to the HTML generated by the htags(1) tool instead of doxygen built-in +# source browser. The htags tool is part of GNU's global source tagging system +# (see https://www.gnu.org/software/global/global.html). You will need version +# 4.8.6 or higher. +# +# To use it do the following: +# - Install the latest version of global +# - Enable SOURCE_BROWSER and USE_HTAGS in the configuration file +# - Make sure the INPUT points to the root of the source tree +# - Run doxygen as normal +# +# Doxygen will invoke htags (and that will in turn invoke gtags), so these +# tools must be available from the command line (i.e. in the search path). +# +# The result: instead of the source browser generated by doxygen, the links to +# source code will now point to the output of htags. +# The default value is: NO. +# This tag requires that the tag SOURCE_BROWSER is set to YES. + +USE_HTAGS = NO + +# If the VERBATIM_HEADERS tag is set the YES then doxygen will generate a +# verbatim copy of the header file for each class for which an include is +# specified. Set to NO to disable this. +# See also: Section \class. +# The default value is: YES. + +VERBATIM_HEADERS = YES + +#--------------------------------------------------------------------------- +# Configuration options related to the alphabetical class index +#--------------------------------------------------------------------------- + +# If the ALPHABETICAL_INDEX tag is set to YES, an alphabetical index of all +# compounds will be generated. Enable this if the project contains a lot of +# classes, structs, unions or interfaces. +# The default value is: YES. + +ALPHABETICAL_INDEX = YES + +# In case all classes in a project start with a common prefix, all classes will +# be put under the same header in the alphabetical index. The IGNORE_PREFIX tag +# can be used to specify a prefix (or a list of prefixes) that should be ignored +# while generating the index headers. +# This tag requires that the tag ALPHABETICAL_INDEX is set to YES. + +IGNORE_PREFIX = + +#--------------------------------------------------------------------------- +# Configuration options related to the HTML output +#--------------------------------------------------------------------------- + +# If the GENERATE_HTML tag is set to YES, doxygen will generate HTML output +# The default value is: YES. + +GENERATE_HTML = NO + +# The HTML_OUTPUT tag is used to specify where the HTML docs will be put. If a +# relative path is entered the value of OUTPUT_DIRECTORY will be put in front of +# it. +# The default directory is: html. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_OUTPUT = html + +# The HTML_FILE_EXTENSION tag can be used to specify the file extension for each +# generated HTML page (for example: .htm, .php, .asp). +# The default value is: .html. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_FILE_EXTENSION = .html + +# The HTML_HEADER tag can be used to specify a user-defined HTML header file for +# each generated HTML page. If the tag is left blank doxygen will generate a +# standard header. +# +# To get valid HTML the header file that includes any scripts and style sheets +# that doxygen needs, which is dependent on the configuration options used (e.g. +# the setting GENERATE_TREEVIEW). It is highly recommended to start with a +# default header using +# doxygen -w html new_header.html new_footer.html new_stylesheet.css +# YourConfigFile +# and then modify the file new_header.html. See also section "Doxygen usage" +# for information on how to generate the default header that doxygen normally +# uses. +# Note: The header is subject to change so you typically have to regenerate the +# default header when upgrading to a newer version of doxygen. For a description +# of the possible markers and block names see the documentation. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_HEADER = + +# The HTML_FOOTER tag can be used to specify a user-defined HTML footer for each +# generated HTML page. If the tag is left blank doxygen will generate a standard +# footer. See HTML_HEADER for more information on how to generate a default +# footer and what special commands can be used inside the footer. See also +# section "Doxygen usage" for information on how to generate the default footer +# that doxygen normally uses. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_FOOTER = + +# The HTML_STYLESHEET tag can be used to specify a user-defined cascading style +# sheet that is used by each HTML page. It can be used to fine-tune the look of +# the HTML output. If left blank doxygen will generate a default style sheet. +# See also section "Doxygen usage" for information on how to generate the style +# sheet that doxygen normally uses. +# Note: It is recommended to use HTML_EXTRA_STYLESHEET instead of this tag, as +# it is more robust and this tag (HTML_STYLESHEET) will in the future become +# obsolete. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_STYLESHEET = + +# The HTML_EXTRA_STYLESHEET tag can be used to specify additional user-defined +# cascading style sheets that are included after the standard style sheets +# created by doxygen. Using this option one can overrule certain style aspects. +# This is preferred over using HTML_STYLESHEET since it does not replace the +# standard style sheet and is therefore more robust against future updates. +# Doxygen will copy the style sheet files to the output directory. +# Note: The order of the extra style sheet files is of importance (e.g. the last +# style sheet in the list overrules the setting of the previous ones in the +# list). For an example see the documentation. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_EXTRA_STYLESHEET = + +# The HTML_EXTRA_FILES tag can be used to specify one or more extra images or +# other source files which should be copied to the HTML output directory. Note +# that these files will be copied to the base HTML output directory. Use the +# $relpath^ marker in the HTML_HEADER and/or HTML_FOOTER files to load these +# files. In the HTML_STYLESHEET file, use the file name only. Also note that the +# files will be copied as-is; there are no commands or markers available. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_EXTRA_FILES = + +# The HTML_COLORSTYLE_HUE tag controls the color of the HTML output. Doxygen +# will adjust the colors in the style sheet and background images according to +# this color. Hue is specified as an angle on a color-wheel, see +# https://en.wikipedia.org/wiki/Hue for more information. For instance the value +# 0 represents red, 60 is yellow, 120 is green, 180 is cyan, 240 is blue, 300 +# purple, and 360 is red again. +# Minimum value: 0, maximum value: 359, default value: 220. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_COLORSTYLE_HUE = 220 + +# The HTML_COLORSTYLE_SAT tag controls the purity (or saturation) of the colors +# in the HTML output. For a value of 0 the output will use gray-scales only. A +# value of 255 will produce the most vivid colors. +# Minimum value: 0, maximum value: 255, default value: 100. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_COLORSTYLE_SAT = 100 + +# The HTML_COLORSTYLE_GAMMA tag controls the gamma correction applied to the +# luminance component of the colors in the HTML output. Values below 100 +# gradually make the output lighter, whereas values above 100 make the output +# darker. The value divided by 100 is the actual gamma applied, so 80 represents +# a gamma of 0.8, The value 220 represents a gamma of 2.2, and 100 does not +# change the gamma. +# Minimum value: 40, maximum value: 240, default value: 80. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_COLORSTYLE_GAMMA = 80 + +# If the HTML_TIMESTAMP tag is set to YES then the footer of each generated HTML +# page will contain the date and time when the page was generated. Setting this +# to YES can help to show when doxygen was last run and thus if the +# documentation is up to date. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_TIMESTAMP = NO + +# If the HTML_DYNAMIC_MENUS tag is set to YES then the generated HTML +# documentation will contain a main index with vertical navigation menus that +# are dynamically created via JavaScript. If disabled, the navigation index will +# consists of multiple levels of tabs that are statically embedded in every HTML +# page. Disable this option to support browsers that do not have JavaScript, +# like the Qt help browser. +# The default value is: YES. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_DYNAMIC_MENUS = YES + +# If the HTML_DYNAMIC_SECTIONS tag is set to YES then the generated HTML +# documentation will contain sections that can be hidden and shown after the +# page has loaded. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_DYNAMIC_SECTIONS = NO + +# With HTML_INDEX_NUM_ENTRIES one can control the preferred number of entries +# shown in the various tree structured indices initially; the user can expand +# and collapse entries dynamically later on. Doxygen will expand the tree to +# such a level that at most the specified number of entries are visible (unless +# a fully collapsed tree already exceeds this amount). So setting the number of +# entries 1 will produce a full collapsed tree by default. 0 is a special value +# representing an infinite number of entries and will result in a full expanded +# tree by default. +# Minimum value: 0, maximum value: 9999, default value: 100. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_INDEX_NUM_ENTRIES = 100 + +# If the GENERATE_DOCSET tag is set to YES, additional index files will be +# generated that can be used as input for Apple's Xcode 3 integrated development +# environment (see: +# https://developer.apple.com/xcode/), introduced with OSX 10.5 (Leopard). To +# create a documentation set, doxygen will generate a Makefile in the HTML +# output directory. Running make will produce the docset in that directory and +# running make install will install the docset in +# ~/Library/Developer/Shared/Documentation/DocSets so that Xcode will find it at +# startup. See https://developer.apple.com/library/archive/featuredarticles/Doxy +# genXcode/_index.html for more information. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +GENERATE_DOCSET = NO + +# This tag determines the name of the docset feed. A documentation feed provides +# an umbrella under which multiple documentation sets from a single provider +# (such as a company or product suite) can be grouped. +# The default value is: Doxygen generated docs. +# This tag requires that the tag GENERATE_DOCSET is set to YES. + +DOCSET_FEEDNAME = "Doxygen generated docs" + +# This tag specifies a string that should uniquely identify the documentation +# set bundle. This should be a reverse domain-name style string, e.g. +# com.mycompany.MyDocSet. Doxygen will append .docset to the name. +# The default value is: org.doxygen.Project. +# This tag requires that the tag GENERATE_DOCSET is set to YES. + +DOCSET_BUNDLE_ID = org.doxygen.Project + +# The DOCSET_PUBLISHER_ID tag specifies a string that should uniquely identify +# the documentation publisher. This should be a reverse domain-name style +# string, e.g. com.mycompany.MyDocSet.documentation. +# The default value is: org.doxygen.Publisher. +# This tag requires that the tag GENERATE_DOCSET is set to YES. + +DOCSET_PUBLISHER_ID = org.doxygen.Publisher + +# The DOCSET_PUBLISHER_NAME tag identifies the documentation publisher. +# The default value is: Publisher. +# This tag requires that the tag GENERATE_DOCSET is set to YES. + +DOCSET_PUBLISHER_NAME = Publisher + +# If the GENERATE_HTMLHELP tag is set to YES then doxygen generates three +# additional HTML index files: index.hhp, index.hhc, and index.hhk. The +# index.hhp is a project file that can be read by Microsoft's HTML Help Workshop +# on Windows. In the beginning of 2021 Microsoft took the original page, with +# a.o. the download links, offline the HTML help workshop was already many years +# in maintenance mode). You can download the HTML help workshop from the web +# archives at Installation executable (see: +# http://web.archive.org/web/20160201063255/http://download.microsoft.com/downlo +# ad/0/A/9/0A939EF6-E31C-430F-A3DF-DFAE7960D564/htmlhelp.exe). +# +# The HTML Help Workshop contains a compiler that can convert all HTML output +# generated by doxygen into a single compiled HTML file (.chm). Compiled HTML +# files are now used as the Windows 98 help format, and will replace the old +# Windows help format (.hlp) on all Windows platforms in the future. Compressed +# HTML files also contain an index, a table of contents, and you can search for +# words in the documentation. The HTML workshop also contains a viewer for +# compressed HTML files. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +GENERATE_HTMLHELP = NO + +# The CHM_FILE tag can be used to specify the file name of the resulting .chm +# file. You can add a path in front of the file if the result should not be +# written to the html output directory. +# This tag requires that the tag GENERATE_HTMLHELP is set to YES. + +CHM_FILE = + +# The HHC_LOCATION tag can be used to specify the location (absolute path +# including file name) of the HTML help compiler (hhc.exe). If non-empty, +# doxygen will try to run the HTML help compiler on the generated index.hhp. +# The file has to be specified with full path. +# This tag requires that the tag GENERATE_HTMLHELP is set to YES. + +HHC_LOCATION = + +# The GENERATE_CHI flag controls if a separate .chi index file is generated +# (YES) or that it should be included in the main .chm file (NO). +# The default value is: NO. +# This tag requires that the tag GENERATE_HTMLHELP is set to YES. + +GENERATE_CHI = NO + +# The CHM_INDEX_ENCODING is used to encode HtmlHelp index (hhk), content (hhc) +# and project file content. +# This tag requires that the tag GENERATE_HTMLHELP is set to YES. + +CHM_INDEX_ENCODING = + +# The BINARY_TOC flag controls whether a binary table of contents is generated +# (YES) or a normal table of contents (NO) in the .chm file. Furthermore it +# enables the Previous and Next buttons. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTMLHELP is set to YES. + +BINARY_TOC = NO + +# The TOC_EXPAND flag can be set to YES to add extra items for group members to +# the table of contents of the HTML help documentation and to the tree view. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTMLHELP is set to YES. + +TOC_EXPAND = NO + +# If the GENERATE_QHP tag is set to YES and both QHP_NAMESPACE and +# QHP_VIRTUAL_FOLDER are set, an additional index file will be generated that +# can be used as input for Qt's qhelpgenerator to generate a Qt Compressed Help +# (.qch) of the generated HTML documentation. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +GENERATE_QHP = NO + +# If the QHG_LOCATION tag is specified, the QCH_FILE tag can be used to specify +# the file name of the resulting .qch file. The path specified is relative to +# the HTML output folder. +# This tag requires that the tag GENERATE_QHP is set to YES. + +QCH_FILE = + +# The QHP_NAMESPACE tag specifies the namespace to use when generating Qt Help +# Project output. For more information please see Qt Help Project / Namespace +# (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#namespace). +# The default value is: org.doxygen.Project. +# This tag requires that the tag GENERATE_QHP is set to YES. + +QHP_NAMESPACE = org.doxygen.Project + +# The QHP_VIRTUAL_FOLDER tag specifies the namespace to use when generating Qt +# Help Project output. For more information please see Qt Help Project / Virtual +# Folders (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#virtual-folders). +# The default value is: doc. +# This tag requires that the tag GENERATE_QHP is set to YES. + +QHP_VIRTUAL_FOLDER = doc + +# If the QHP_CUST_FILTER_NAME tag is set, it specifies the name of a custom +# filter to add. For more information please see Qt Help Project / Custom +# Filters (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#custom-filters). +# This tag requires that the tag GENERATE_QHP is set to YES. + +QHP_CUST_FILTER_NAME = + +# The QHP_CUST_FILTER_ATTRS tag specifies the list of the attributes of the +# custom filter to add. For more information please see Qt Help Project / Custom +# Filters (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#custom-filters). +# This tag requires that the tag GENERATE_QHP is set to YES. + +QHP_CUST_FILTER_ATTRS = + +# The QHP_SECT_FILTER_ATTRS tag specifies the list of the attributes this +# project's filter section matches. Qt Help Project / Filter Attributes (see: +# https://doc.qt.io/archives/qt-4.8/qthelpproject.html#filter-attributes). +# This tag requires that the tag GENERATE_QHP is set to YES. + +QHP_SECT_FILTER_ATTRS = + +# The QHG_LOCATION tag can be used to specify the location (absolute path +# including file name) of Qt's qhelpgenerator. If non-empty doxygen will try to +# run qhelpgenerator on the generated .qhp file. +# This tag requires that the tag GENERATE_QHP is set to YES. + +QHG_LOCATION = + +# If the GENERATE_ECLIPSEHELP tag is set to YES, additional index files will be +# generated, together with the HTML files, they form an Eclipse help plugin. To +# install this plugin and make it available under the help contents menu in +# Eclipse, the contents of the directory containing the HTML and XML files needs +# to be copied into the plugins directory of eclipse. The name of the directory +# within the plugins directory should be the same as the ECLIPSE_DOC_ID value. +# After copying Eclipse needs to be restarted before the help appears. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +GENERATE_ECLIPSEHELP = NO + +# A unique identifier for the Eclipse help plugin. When installing the plugin +# the directory name containing the HTML and XML files should also have this +# name. Each documentation set should have its own identifier. +# The default value is: org.doxygen.Project. +# This tag requires that the tag GENERATE_ECLIPSEHELP is set to YES. + +ECLIPSE_DOC_ID = org.doxygen.Project + +# If you want full control over the layout of the generated HTML pages it might +# be necessary to disable the index and replace it with your own. The +# DISABLE_INDEX tag can be used to turn on/off the condensed index (tabs) at top +# of each HTML page. A value of NO enables the index and the value YES disables +# it. Since the tabs in the index contain the same information as the navigation +# tree, you can set this option to YES if you also set GENERATE_TREEVIEW to YES. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +DISABLE_INDEX = NO + +# The GENERATE_TREEVIEW tag is used to specify whether a tree-like index +# structure should be generated to display hierarchical information. If the tag +# value is set to YES, a side panel will be generated containing a tree-like +# index structure (just like the one that is generated for HTML Help). For this +# to work a browser that supports JavaScript, DHTML, CSS and frames is required +# (i.e. any modern browser). Windows users are probably better off using the +# HTML help feature. Via custom style sheets (see HTML_EXTRA_STYLESHEET) one can +# further fine tune the look of the index (see "Fine-tuning the output"). As an +# example, the default style sheet generated by doxygen has an example that +# shows how to put an image at the root of the tree instead of the PROJECT_NAME. +# Since the tree basically has the same information as the tab index, you could +# consider setting DISABLE_INDEX to YES when enabling this option. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +GENERATE_TREEVIEW = NO + +# When both GENERATE_TREEVIEW and DISABLE_INDEX are set to YES, then the +# FULL_SIDEBAR option determines if the side bar is limited to only the treeview +# area (value NO) or if it should extend to the full height of the window (value +# YES). Setting this to YES gives a layout similar to +# https://docs.readthedocs.io with more room for contents, but less room for the +# project logo, title, and description. If either GENERATOR_TREEVIEW or +# DISABLE_INDEX is set to NO, this option has no effect. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +FULL_SIDEBAR = NO + +# The ENUM_VALUES_PER_LINE tag can be used to set the number of enum values that +# doxygen will group on one line in the generated HTML documentation. +# +# Note that a value of 0 will completely suppress the enum values from appearing +# in the overview section. +# Minimum value: 0, maximum value: 20, default value: 4. +# This tag requires that the tag GENERATE_HTML is set to YES. + +ENUM_VALUES_PER_LINE = 4 + +# If the treeview is enabled (see GENERATE_TREEVIEW) then this tag can be used +# to set the initial width (in pixels) of the frame in which the tree is shown. +# Minimum value: 0, maximum value: 1500, default value: 250. +# This tag requires that the tag GENERATE_HTML is set to YES. + +TREEVIEW_WIDTH = 250 + +# If the EXT_LINKS_IN_WINDOW option is set to YES, doxygen will open links to +# external symbols imported via tag files in a separate window. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +EXT_LINKS_IN_WINDOW = NO + +# If the HTML_FORMULA_FORMAT option is set to svg, doxygen will use the pdf2svg +# tool (see https://github.com/dawbarton/pdf2svg) or inkscape (see +# https://inkscape.org) to generate formulas as SVG images instead of PNGs for +# the HTML output. These images will generally look nicer at scaled resolutions. +# Possible values are: png (the default) and svg (looks nicer but requires the +# pdf2svg or inkscape tool). +# The default value is: png. +# This tag requires that the tag GENERATE_HTML is set to YES. + +HTML_FORMULA_FORMAT = png + +# Use this tag to change the font size of LaTeX formulas included as images in +# the HTML documentation. When you change the font size after a successful +# doxygen run you need to manually remove any form_*.png images from the HTML +# output directory to force them to be regenerated. +# Minimum value: 8, maximum value: 50, default value: 10. +# This tag requires that the tag GENERATE_HTML is set to YES. + +FORMULA_FONTSIZE = 10 + +# Use the FORMULA_TRANSPARENT tag to determine whether or not the images +# generated for formulas are transparent PNGs. Transparent PNGs are not +# supported properly for IE 6.0, but are supported on all modern browsers. +# +# Note that when changing this option you need to delete any form_*.png files in +# the HTML output directory before the changes have effect. +# The default value is: YES. +# This tag requires that the tag GENERATE_HTML is set to YES. + +FORMULA_TRANSPARENT = YES + +# The FORMULA_MACROFILE can contain LaTeX \newcommand and \renewcommand commands +# to create new LaTeX commands to be used in formulas as building blocks. See +# the section "Including formulas" for details. + +FORMULA_MACROFILE = + +# Enable the USE_MATHJAX option to render LaTeX formulas using MathJax (see +# https://www.mathjax.org) which uses client side JavaScript for the rendering +# instead of using pre-rendered bitmaps. Use this if you do not have LaTeX +# installed or if you want to formulas look prettier in the HTML output. When +# enabled you may also need to install MathJax separately and configure the path +# to it using the MATHJAX_RELPATH option. +# The default value is: NO. +# This tag requires that the tag GENERATE_HTML is set to YES. + +USE_MATHJAX = NO + +# With MATHJAX_VERSION it is possible to specify the MathJax version to be used. +# Note that the different versions of MathJax have different requirements with +# regards to the different settings, so it is possible that also other MathJax +# settings have to be changed when switching between the different MathJax +# versions. +# Possible values are: MathJax_2 and MathJax_3. +# The default value is: MathJax_2. +# This tag requires that the tag USE_MATHJAX is set to YES. + +MATHJAX_VERSION = MathJax_2 + +# When MathJax is enabled you can set the default output format to be used for +# the MathJax output. For more details about the output format see MathJax +# version 2 (see: +# http://docs.mathjax.org/en/v2.7-latest/output.html) and MathJax version 3 +# (see: +# http://docs.mathjax.org/en/latest/web/components/output.html). +# Possible values are: HTML-CSS (which is slower, but has the best +# compatibility. This is the name for Mathjax version 2, for MathJax version 3 +# this will be translated into chtml), NativeMML (i.e. MathML. Only supported +# for NathJax 2. For MathJax version 3 chtml will be used instead.), chtml (This +# is the name for Mathjax version 3, for MathJax version 2 this will be +# translated into HTML-CSS) and SVG. +# The default value is: HTML-CSS. +# This tag requires that the tag USE_MATHJAX is set to YES. + +MATHJAX_FORMAT = HTML-CSS + +# When MathJax is enabled you need to specify the location relative to the HTML +# output directory using the MATHJAX_RELPATH option. The destination directory +# should contain the MathJax.js script. For instance, if the mathjax directory +# is located at the same level as the HTML output directory, then +# MATHJAX_RELPATH should be ../mathjax. The default value points to the MathJax +# Content Delivery Network so you can quickly see the result without installing +# MathJax. However, it is strongly recommended to install a local copy of +# MathJax from https://www.mathjax.org before deployment. The default value is: +# - in case of MathJax version 2: https://cdn.jsdelivr.net/npm/mathjax@2 +# - in case of MathJax version 3: https://cdn.jsdelivr.net/npm/mathjax@3 +# This tag requires that the tag USE_MATHJAX is set to YES. + +MATHJAX_RELPATH = + +# The MATHJAX_EXTENSIONS tag can be used to specify one or more MathJax +# extension names that should be enabled during MathJax rendering. For example +# for MathJax version 2 (see +# https://docs.mathjax.org/en/v2.7-latest/tex.html#tex-and-latex-extensions): +# MATHJAX_EXTENSIONS = TeX/AMSmath TeX/AMSsymbols +# For example for MathJax version 3 (see +# http://docs.mathjax.org/en/latest/input/tex/extensions/index.html): +# MATHJAX_EXTENSIONS = ams +# This tag requires that the tag USE_MATHJAX is set to YES. + +MATHJAX_EXTENSIONS = + +# The MATHJAX_CODEFILE tag can be used to specify a file with javascript pieces +# of code that will be used on startup of the MathJax code. See the MathJax site +# (see: +# http://docs.mathjax.org/en/v2.7-latest/output.html) for more details. For an +# example see the documentation. +# This tag requires that the tag USE_MATHJAX is set to YES. + +MATHJAX_CODEFILE = + +# When the SEARCHENGINE tag is enabled doxygen will generate a search box for +# the HTML output. The underlying search engine uses javascript and DHTML and +# should work on any modern browser. Note that when using HTML help +# (GENERATE_HTMLHELP), Qt help (GENERATE_QHP), or docsets (GENERATE_DOCSET) +# there is already a search function so this one should typically be disabled. +# For large projects the javascript based search engine can be slow, then +# enabling SERVER_BASED_SEARCH may provide a better solution. It is possible to +# search using the keyboard; to jump to the search box use + S +# (what the is depends on the OS and browser, but it is typically +# , /