diff --git a/.gitignore b/.gitignore index 916dd4c35..4bb008d93 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,9 @@ pwscf.save/ /tests/test_ase/__pycache__/ /tests/test_outputs/* test_files/qe_output_*.out +/tests/*log* +/tests/*.pickle +/tests/*.npy /tests/*.xyz /tests/*.gp /tests/*-bak diff --git a/.travis.yml b/.travis.yml index 53e87185a..feb9dcc2a 100644 --- a/.travis.yml +++ b/.travis.yml @@ -14,17 +14,19 @@ before_install: - pwd - wget http://folk.uio.no/anjohan/lmp - "wget https://github.com/cp2k/cp2k/releases/download/\ - v6.1.0/cp2k-6.1-Linux-x86_64.sopt" + v7.1.0/cp2k-7.1-Linux-x86_64.sopt" - chmod u+x lmp - - chmod u+x cp2k-6.1-Linux-x86_64.sopt + - chmod u+x cp2k-7.1-Linux-x86_64.sopt - pip install -r requirements.txt script: - pwd - cd tests - - PWSCF_COMMAND=pw.x CP2K_COMMAND=../cp2k-6.1-Linux-x86_64.sopt + - ls test_files + - PWSCF_COMMAND=pw.x lmp=$(pwd)/../lmp - pytest --show-capture=all -vv --durations=0 --cov=../flare/ + CP2K_COMMAND=../cp2k-7.1-Linux-x86_64.sopt + pytest -vv --durations=0 --cov=../flare/ - coverage xml after_success: diff --git a/docs/images/GPFA_tutorial.png b/docs/images/GPFA_tutorial.png new file mode 100644 index 000000000..a88b041b4 Binary files /dev/null and b/docs/images/GPFA_tutorial.png differ diff --git a/docs/source/conf.py b/docs/source/conf.py index dc97b95a7..80d5002e9 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -47,7 +47,8 @@ 'sphinx.ext.autodoc', 'sphinx.ext.imgmath', 'sphinx_rtd_theme', - 'sphinx.ext.napoleon' + 'sphinx.ext.napoleon', + 'nbsphinx' ] napoleon_use_param = False # Add any paths that contain templates here, relative to this directory. diff --git a/docs/source/flare/ase/ase.rst b/docs/source/flare/ase/ase.rst index f57966db9..9fd1b4b24 100644 --- a/docs/source/flare/ase/ase.rst +++ b/docs/source/flare/ase/ase.rst @@ -3,11 +3,9 @@ ASE Interface We provide an interface to do the OTF training coupling with ASE. Including wrapping the FLARE's GaussianProcess and MappedGaussianProcess into an ASE calculator: :class:`FLARE_Calculator`, -the on-the-fly training module coupled with ASE molecular dynamics engines :module:`otf_md.py` .. toctree:: :maxdepth: 2 calculator otf - otf_md diff --git a/docs/source/flare/ase/otf_md.rst b/docs/source/flare/ase/otf_md.rst deleted file mode 100644 index e8ee4bf8a..000000000 --- a/docs/source/flare/ase/otf_md.rst +++ /dev/null @@ -1,5 +0,0 @@ -ASE MD engines -============== - -.. automodule:: flare.ase.otf_md - :members: diff --git a/docs/source/flare/utils/env_getarray.rst b/docs/source/flare/utils/env_getarray.rst new file mode 100644 index 000000000..554c509e7 --- /dev/null +++ b/docs/source/flare/utils/env_getarray.rst @@ -0,0 +1,5 @@ +Construct Atomic Environment +=================== + +.. automodule:: flare.utils.env_getarray + :members: diff --git a/docs/source/flare/utils/flare_io.rst b/docs/source/flare/utils/flare_io.rst new file mode 100644 index 000000000..e8c236404 --- /dev/null +++ b/docs/source/flare/utils/flare_io.rst @@ -0,0 +1,6 @@ +I/O for trajectories +=================== + +.. automodule:: flare.utils.flare_io + :members: + diff --git a/docs/source/flare/utils/mask_helper.rst b/docs/source/flare/utils/mask_helper.rst index 7acc96ea5..32a0c58b2 100644 --- a/docs/source/flare/utils/mask_helper.rst +++ b/docs/source/flare/utils/mask_helper.rst @@ -1,5 +1,5 @@ Advanced Hyperparameters Set Up =================== -.. automodule:: flare.utils.mask_helper +.. automodule:: flare.utils.parameter_helper :members: diff --git a/docs/source/flare/utils/utils.rst b/docs/source/flare/utils/utils.rst index 6b167a37c..a1562385c 100644 --- a/docs/source/flare/utils/utils.rst +++ b/docs/source/flare/utils/utils.rst @@ -7,3 +7,6 @@ Utility element_coder learner mask_helper + env_getarray + md_helper + flare_io diff --git a/docs/source/tutorials/after_training.ipynb b/docs/source/tutorials/after_training.ipynb new file mode 100644 index 000000000..f10468105 --- /dev/null +++ b/docs/source/tutorials/after_training.ipynb @@ -0,0 +1,265 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# After Training\n", + "\n", + "After the on-the-fly training is complete, we can play with the force field we obtained. \n", + "We are going to do the following things:\n", + "\n", + "1. Parse the on-the-fly training trajectory to collect training data\n", + "2. Reconstruct the GP model from the training trajectory\n", + "3. Build up Mapped GP (MGP) for accelerated force field, and save coefficient file for LAMMPS\n", + "4. Use LAMMPS to run fast simulation using MGP pair style\n", + "\n", + "## Parse OTF log file\n", + "\n", + "After the on-the-fly training is complete, we have a log file and can use the `otf_parser` module to parse the trajectory. " + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from flare import otf_parser\n", + "\n", + "logdir = '../../../tests/test_files'\n", + "file_name = f'{logdir}/AgI_snippet.out'\n", + "hyp_no = 2 # use the hyperparameters from the 2nd training step\n", + "otf_object = otf_parser.OtfAnalysis(file_name)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Construct GP model from log file\n", + "\n", + "We can reconstruct GP model from the parsed log file (the on-the-fly training trajectory). Here we build up the GP model with 2+3 body kernel from the on-the-fly log file. " + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Final name of the gp instance is default_gp_2\n" + ] + } + ], + "source": [ + "gp_model = otf_object.make_gp(hyp_no=hyp_no)\n", + "gp_model.parallel = True\n", + "gp_model.hyp_labels = ['sig2', 'ls2', 'sig3', 'ls3', 'noise']\n", + "\n", + "# write model to a binary file\n", + "gp_model.write_model('AgI.gp', format='pickle')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The last step `write_model` is to write this GP model into a binary file, \n", + "so next time we can directly load the model from the pickle file as" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Final name of the gp instance is default_gp_2_2\n" + ] + } + ], + "source": [ + "from flare.gp import GaussianProcess\n", + "\n", + "gp_model = GaussianProcess.from_file('AgI.gp.pickle')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Map the GP force field & Dump LAMMPS coefficient file\n", + "\n", + "To use the trained force field with accelerated version MGP, or in LAMMPS, we need to build MGP from GP model. \n", + "Since 2-body and 3-body are both included, we need to set up the number of grid points for 2-body and 3-body in `grid_params`.\n", + "We build up energy mapping, thus set `map_force=False`.\n", + "See [MGP tutorial](https://flare.readthedocs.io/en/latest/tutorials/mgp.html) for more explanation of the MGP settings." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "from flare.mgp import MappedGaussianProcess\n", + "\n", + "grid_params = {'twobody': {'grid_num': [64]}, \n", + " 'threebody': {'grid_num': [20, 20, 20]}}\n", + "\n", + "data = gp_model.training_statistics\n", + "lammps_location = 'AgI_Molten_15.txt'\n", + "\n", + "mgp_model = MappedGaussianProcess(grid_params, data['species'], \n", + " map_force=False, lmp_file_name='AgI_Molten_15.txt', n_cpus=1)\n", + "mgp_model.build_map(gp_model)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The coefficient file for LAMMPS mgp pair_style is automatically saved once the mapping is done. \n", + "Saved as `lmp_file_name`. \n", + "\n", + "## Run LAMMPS with MGP pair style\n", + "\n", + "With the above coefficient file, we can run LAMMPS simulation with the mgp pair style. \n", + "First download our mgp pair style files, compile your lammps executable with mgp pair style following our [instruction](https://flare.readthedocs.io/en/latest/tutorials/lammps.html).\n", + "\n", + "1. One way to use it is running `lmp_executable < in.lammps > log.lammps` \n", + "with the executable provided in our repository. \n", + "When creating the input file, please note to set\n", + "\n", + "```\n", + "newton off\n", + "pair_style mgp\n", + "pair_coeff * * yes/no yes/no\n", + "```\n", + "\n", + "An example is using coefficient file `AgI_Molten_15.txt` for AgI system, \n", + "with two-body (the 1st `yes`) together with three-body (the 2nd `yes`).\n", + "\n", + "```\n", + "pair_coeff * * AgI_Molten_15.txt Ag I yes yes\n", + "```\n", + "\n", + "**Note**: if you build force mapping (`map_force=True`) instead of energy mapping, please use\n", + "```\n", + "pair_style mgpf\n", + "```\n", + "\n", + "2. The third way is to use the ASE LAMMPS interface" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from flare.ase.calculator import FLARE_Calculator\n", + "\n", + "# get chemical symbols, masses etc.\n", + "species = gp_model.training_statistics['species']\n", + "specie_symbol_list = \" \".join(species)\n", + "masses=[f\"{i} {_Z_to_mass[_element_to_Z[species[i]]]}\" for i in range(len(species))]\n", + "\n", + "# set up input params\n", + "parameters = {'command': os.environ.get('lmp'), # set up executable for ASE\n", + " 'newton': 'off',\n", + " 'pair_style': 'mgp',\n", + " 'pair_coeff': [f'* * {lammps_location} {specie_symbol_list} yes yes'],\n", + " 'mass': masses}\n", + "files = [lammps_location]\n", + "\n", + "# create ASE calc\n", + "lmp_calc = LAMMPS(label=f'tmp_AgI', keep_tmp_files=True, tmp_dir='./tmp/',\n", + " parameters=parameters, files=files, specorder=species)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "3. The second way to run LAMMPS is using our LAMMPS interface, please set the\n", + "environment variable `$lmp` to the executable." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from flare import struc\n", + "from flare.lammps import lammps_calculator\n", + "\n", + "# lmp coef file is automatically written now every time MGP is constructed\n", + "\n", + "# create test structure\n", + "species = otf_object.gp_species_list[-1]\n", + "positions = otf_object.position_list[-1]\n", + "forces = otf_object.force_list[-1]\n", + "otf_cell = otf_object.header['cell']\n", + "structure = struc.Structure(otf_cell, species, positions)\n", + "\n", + "atom_types = [1, 2]\n", + "atom_masses = [108, 127]\n", + "atom_species = [1, 2] * 27\n", + "\n", + "# create data file\n", + "data_file_name = 'tmp.data'\n", + "data_text = lammps_calculator.lammps_dat(structure, atom_types,\n", + " atom_masses, atom_species)\n", + "lammps_calculator.write_text(data_file_name, data_text)\n", + "\n", + "# create lammps input\n", + "style_string = 'mgp'\n", + "coeff_string = '* * {} Ag I yes yes'.format(lammps_location)\n", + "lammps_executable = '$lmp'\n", + "dump_file_name = 'tmp.dump'\n", + "input_file_name = 'tmp.in'\n", + "output_file_name = 'tmp.out'\n", + "input_text = \\\n", + " lammps_calculator.generic_lammps_input(data_file_name, style_string,\n", + " coeff_string, dump_file_name)\n", + "lammps_calculator.write_text(input_file_name, input_text)\n", + "\n", + "lammps_calculator.run_lammps(lammps_executable, input_file_name,\n", + " output_file_name)\n", + "\n", + "lammps_forces = lammps_calculator.lammps_parser(dump_file_name)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/source/tutorials/after_training.py b/docs/source/tutorials/after_training.py deleted file mode 100644 index eb45325db..000000000 --- a/docs/source/tutorials/after_training.py +++ /dev/null @@ -1,107 +0,0 @@ -import numpy as np -from flare import otf_parser - -file_name = 'test_files/AgI_snippet.out' -hyp_no = 2 # use the hyperparameters from the 2nd training step -otf_object = otf_parser.OtfAnalysis(file_name) - -# ------------------------------------------------------------------------- -# reconstruct gp model from otf snippet -# ------------------------------------------------------------------------- - - - - -gp_model = otf_object.make_gp(kernel_name="2+3_mc", - hyp_no=hyp_no) -gp_model.parallel = True -gp_model.hyp_labels = ['sig2', 'ls2', 'sig3', 'ls3', 'noise'] - -# write model to a binary file -gp_model.write_model('AgI.gp', format='pickle') -# ------------------------------------------------------------------------- -# map the potential -# ------------------------------------------------------------------------- -from flare.mgp.mgp import MappedGaussianProcess - -grid_num_2 = 64 -grid_num_3 = 20 -lower_cut = 2.5 -two_cut = 7. -three_cut = 5. -lammps_location = 'AgI_Molten_15.txt' - -# set struc params. cell and masses arbitrary? -mapped_cell = np.eye(3) * 100 # just use a sufficiently large one -struc_params = {'species': [47, 53], - 'cube_lat': mapped_cell} - -# grid parameters -grid_params = {'bounds_2': [[lower_cut], [two_cut]], - 'bounds_3': [[lower_cut, lower_cut, -1], - [three_cut, three_cut, 1]], - 'grid_num_2': grid_num_2, - 'grid_num_3': [grid_num_3, grid_num_3, grid_num_3], - 'svd_rank_2': 64, - 'svd_rank_3': 90, - 'bodies': [2, 3], - 'load_grid': None, - 'update': True} - -mgp_model = MappedGaussianProcess(gp_model.hyps, gp_model.cutoffs, - grid_params, struc_params, mean_only=True, container_only=False, - GP=gp_model, lmp_file_name=lammps_location) -# ------------------------------------------------------------------------- -# test the mapped potential -# ------------------------------------------------------------------------- -gp_pred_x = gp_model.predict(environ, 1) -mgp_pred = mgp_model.predict(environ, mean_only=True) - -# check mgp is within 1 meV/A of the gp -assert(np.abs(mgp_pred[0][0] - gp_pred_x[0]) < 1e-3) - -# ------------------------------------------------------------------------- -# check lammps potential -# ------------------------------------------------------------------------- -from flare import struc, env -from flare.lammps import lammps_calculator - -# lmp coef file is automatically written now every time MGP is constructed - -# create test structure -species = otf_object.gp_species_list[-1] -positions = otf_object.position_list[-1] -forces = otf_object.force_list[-1] -otf_cell = otf_object.header['cell'] -structure = struc.Structure(otf_cell, species, positions) - -atom_types = [1, 2] -atom_masses = [108, 127] -atom_species = [1, 2] * 27 - -# create data file -data_file_name = 'tmp.data' -data_text = lammps_calculator.lammps_dat(structure, atom_types, - atom_masses, atom_species) -lammps_calculator.write_text(data_file_name, data_text) - -# create lammps input -style_string = 'mgp' -coeff_string = '* * {} Ag I yes yes'.format(lammps_location) -lammps_executable = '$lmp' -dump_file_name = 'tmp.dump' -input_file_name = 'tmp.in' -output_file_name = 'tmp.out' -input_text = \ - lammps_calculator.generic_lammps_input(data_file_name, style_string, - coeff_string, dump_file_name) -lammps_calculator.write_text(input_file_name, input_text) - -lammps_calculator.run_lammps(lammps_executable, input_file_name, - output_file_name) - -lammps_forces = lammps_calculator.lammps_parser(dump_file_name) - -# check that lammps agrees with gp to within 1 meV/A -assert(np.abs(lammps_forces[0, 1] - forces[0, 1]) < 1e-3) - diff --git a/docs/source/tutorials/after_training.rst b/docs/source/tutorials/after_training.rst deleted file mode 100644 index 6747874e2..000000000 --- a/docs/source/tutorials/after_training.rst +++ /dev/null @@ -1,72 +0,0 @@ -After Training -============== - -.. toctree:: - :maxdepth: 2 - -After the on-the-fly training is complete, we can play with the force field we obtained. -We are going to do the following things: - -1. Parse the on-the-fly training trajectory to collect training data -2. Reconstruct the GP model from the training trajectory -3. Build up Mapped GP (MGP) for accelerated force field, and save coefficient file for LAMMPS -4. Use LAMMPS to run fast simulation using MGP pair style - -Parse OTF log file ------------------- -After the on-the-fly training is complete, we have a log file and can use the `otf_parser` module to parse the trajectory. - -.. literalinclude:: after_training.py - :lines: 1-6 - -Construct GP model from log file --------------------------------- -We can reconstruct GP model from the parsed log file (the on-the-fly training trajectory) - -.. literalinclude:: after_training.py - :lines: 11-21 - -The last step `write_model` is to write this GP model into a binary file, -so next time we can directly load the model from the pickle file as - -.. code-block:: python - - gp_model = pickle.load(open('AgI.gp.pickle', 'rb')) - - -Map the GP force field & Dump LAMMPS coefficient file ------------------------------------------------------ -To use the trained force field with accelerated version MGP, or in LAMMPS, we need to build MGP from GP model - -.. literalinclude:: after_training.py - :lines: 25-53 - -The coefficient file for LAMMPS mgp pair_style is automatically saved once the mapping is done. -Saved as `lmp_file_name`. - -Run LAMMPS with MGP pair style ------------------------------- -With the above coefficient file, we can run LAMMPS simulation with the mgp pair style. - -1. One way to use it is running `lmp_executable < in.lammps > log.lammps` -with the executable provided in our repository. -When creating the input file, please note to set - -.. code-block:: C - - newton off - pair_style mgp - pair_coeff * * yes/no yes/no - -An example is using coefficient file `AgI_Molten_15.txt` for AgI system, -with two-body (the 1st `yes`) together with three-body (the 2nd `yes`). - -.. code-block:: C - - pair_coeff * * AgI_Molten_15.txt Ag I yes yes - -2. Another way to run LAMMPS is using our LAMMPS interface, please set the -environment variable `$lmp` to the executable. - -.. literalinclude:: after_training.py - :lines: 66-106 diff --git a/docs/source/tutorials/ase.ipynb b/docs/source/tutorials/ase.ipynb new file mode 100644 index 000000000..0f5bc0073 --- /dev/null +++ b/docs/source/tutorials/ase.ipynb @@ -0,0 +1,490 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# On-the-fly training using ASE\n", + "\n", + "This is a quick introduction of how to set up our ASE-OTF interface to train a force field. We will train a force field model for diamond. To run the on-the-fly training, we will need to\n", + "\n", + "1. Create a supercell with ASE Atoms object\n", + "2. Set up FLARE ASE calculator, including the kernel functions, hyperparameters, cutoffs for Gaussian process, and mapping parameters (if Mapped Gaussian Process is used)\n", + "3. Set up DFT ASE calculator. Here we will give an example of Quantum Espresso\n", + "4. Set up on-the-fly training with ASE MD engine\n", + "\n", + "Please make sure you are using the LATEST FLARE code in our master branch.\n", + "\n", + "## Step 1: Set up supercell with ASE\n", + "\n", + "Here we create a 2x1x1 supercell with lattice constant 3.855, and randomly perturb the positions of the atoms, so that they will start MD with non-zero forces." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from ase import units\n", + "from ase.spacegroup import crystal\n", + "from ase.build import bulk\n", + "\n", + "np.random.seed(12345)\n", + "\n", + "a = 3.52678\n", + "super_cell = bulk('C', 'diamond', a=a, cubic=True) " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Set up FLARE calculator\n", + "\n", + "Now let’s set up our Gaussian process model in the same way as introduced before" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "twobody0 cutoff is not define. it's going to use the universal cutoff.\n", + "threebody0 cutoff is not define. it's going to use the universal cutoff.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "hyps [0.92961609 0.31637555 0.18391881 0.20456028 0.05 ]\n" + ] + } + ], + "source": [ + "from flare.gp import GaussianProcess\n", + "from flare.utils.parameter_helper import ParameterHelper\n", + "\n", + "# set up GP hyperparameters\n", + "kernels = ['twobody', 'threebody'] # use 2+3 body kernel\n", + "parameters = {'cutoff_twobody': 5.0, \n", + " 'cutoff_threebody': 3.5}\n", + "pm = ParameterHelper(\n", + " kernels = kernels, \n", + " random = True,\n", + " parameters=parameters\n", + ")\n", + "\n", + "hm = pm.as_dict()\n", + "hyps = hm['hyps']\n", + "cut = hm['cutoffs']\n", + "print('hyps', hyps)\n", + "\n", + "gp_model = GaussianProcess(\n", + " kernels = kernels,\n", + " component = 'sc', # single-component. For multi-comp, use 'mc'\n", + " hyps = hyps,\n", + " cutoffs = cut,\n", + " hyp_labels = ['sig2','ls2','sig3','ls3','noise'],\n", + " opt_algorithm = 'L-BFGS-B',\n", + " n_cpus = 1\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Optional\n", + "\n", + "If you want to use Mapped Gaussian Process (MGP), then set up MGP as follows" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "from flare.mgp import MappedGaussianProcess\n", + "\n", + "grid_params = {'twobody': {'grid_num': [64]},\n", + " 'threebody': {'grid_num': [16, 16, 16]}}\n", + "\n", + "mgp_model = MappedGaussianProcess(grid_params, \n", + " unique_species = [6], \n", + " n_cpus = 1,\n", + " map_force = False, \n", + " mean_only = False)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now let's set up FLARE's ASE calculator. If you want to use MGP model, then set `use_mapping = True` and `mgp_model = mgp_model` below." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "from flare.ase.calculator import FLARE_Calculator\n", + "\n", + "flare_calculator = FLARE_Calculator(gp_model, \n", + " par = True, \n", + " mgp_model = None,\n", + " use_mapping = False)\n", + "\n", + "super_cell.set_calculator(flare_calculator)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 3: Set up DFT calculator\n", + "For DFT calculator, you can use any calculator provided by ASE, e.g. [Quantum Espresso (QE)](https://wiki.fysik.dtu.dk/ase/ase/calculators/espresso.html), [VASP](https://wiki.fysik.dtu.dk/ase/ase/calculators/vasp.html), etc. \n", + "\n", + "For a quick illustration of our interface, we use the [Lennard-Jones (LJ)](https://wiki.fysik.dtu.dk/ase/ase/calculators/others.html?highlight=lj#module-ase.calculators.lj) potential as an example. " + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "from ase.calculators.lj import LennardJones\n", + "lj_calc = LennardJones() " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Optional: alternatively, set up Quantum Espresso calculator\n", + "\n", + "We also give the code below for setting up the ASE quantum espresso calculator, following the [instruction](https://wiki.fysik.dtu.dk/ase/ase/calculators/espresso.html) on ASE website.\n", + "\n", + "First, we need to set up our environment variable `ASE_ESPRESSO_COMMAND` to our QE executable, so that ASE can find this calculator. Then set up our input parameters of QE and create an ASE calculator" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from ase.calculators.espresso import Espresso\n", + "\n", + "# ---------------- set up executable ---------------- \n", + "label = 'C'\n", + "input_file = label+'.pwi'\n", + "output_file = label+'.pwo'\n", + "no_cpus = 32\n", + "npool = 32\n", + "pw_loc = 'path/to/pw.x'\n", + "\n", + "# serial\n", + "os.environ['ASE_ESPRESSO_COMMAND'] = f'{pw_loc} < {input_file} > {output_file}'\n", + "\n", + "## parallel qe using mpirun\n", + "# os.environ['ASE_ESPRESSO_COMMAND'] = f'mpirun -np {no_cpus} {pw_loc} -npool {npool} < {input_file} > {output_file}'\n", + "\n", + "## parallel qe using srun (for slurm system)\n", + "# os.environ['ASE_ESPRESSO_COMMAND'] = 'srun -n {no_cpus} --mpi=pmi2 {pw_loc} -npool {npool} < {input_file} > {output_file}'\n", + "\n", + "\n", + "# -------------- set up input parameters -------------- \n", + "input_data = {'control': {'prefix': label, \n", + " 'pseudo_dir': './',\n", + " 'outdir': './out',\n", + " 'calculation': 'scf'},\n", + " 'system': {'ibrav': 0, \n", + " 'ecutwfc': 60,\n", + " 'ecutrho': 360},\n", + " 'electrons': {'conv_thr': 1.0e-9,\n", + " 'electron_maxstep': 100,\n", + " 'mixing_beta': 0.7}}\n", + "\n", + "# ---------------- pseudo-potentials ----------------- \n", + "ion_pseudo = {'C': 'C.pz-rrkjus.UPF'}\n", + "\n", + "# -------------- create ASE calculator ---------------- \n", + "dft_calc = Espresso(pseudopotentials=ion_pseudo, label=label, \n", + " tstress=True, tprnfor=True, nosym=True, \n", + " input_data=input_data, kpts=(8, 8, 8)) " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 4: Set up On-The-Fly MD engine\n", + "\n", + "Finally, our OTF is compatible with 5 MD engines that ASE supports: VelocityVerlet, NVTBerendsen, NPTBerendsen, NPT and Langevin. We can choose any of them, and set up the parameters based on [ASE requirements](https://wiki.fysik.dtu.dk/ase/ase/md.html). After everything is set up, we can run the on-the-fly training.\n", + "\n", + "**Note**: Currently, only VelocityVerlet is tested on real system, NPT may have issue with pressure and stress.\n", + "\n", + "Set up ASE_OTF training engine:\n", + "1. Initialize the velocities of atoms as 500K\n", + "2. Set up MD arguments as a dictionary based on [ASE MD parameters](https://wiki.fysik.dtu.dk/ase/ase/md.html). For VelocityVerlet, we don't need to set up extra parameters.\n", + " \n", + " E.g. for NVTBerendsen, we can set `md_kwargs = {'temperature': 500, 'taut': 0.5e3 * units.fs}`" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "from ase import units\n", + "from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution,\n", + " Stationary, ZeroRotation)\n", + "\n", + "temperature = 500\n", + "MaxwellBoltzmannDistribution(super_cell, temperature * units.kB)\n", + "Stationary(super_cell) # zero linear momentum\n", + "ZeroRotation(super_cell) # zero angular momentum\n", + "\n", + "md_engine = 'VelocityVerlet'\n", + "md_kwargs = {}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "3. Set up parameters for On-The-Fly (OTF) training. The descriptions of the parameters are in [ASE OTF module](https://flare-yuuuu.readthedocs.io/en/development/flare/ase/otf.html).\n", + "4. Set up the ASE_OTF training engine, and run\n", + "5. Check `otf.out` after the training is done." + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO:otflog:2020-06-10 14:53:00.574573\n", + "INFO:otflog:\n", + "GaussianProcess Object\n", + "Number of cpu cores: 1\n", + "Kernel: ['twobody', 'threebody']\n", + "Training points: 0\n", + "Cutoffs: {'twobody': 5.0, 'threebody': 3.5}\n", + "Model Likelihood: None\n", + "Number of hyperparameters: 5\n", + "Hyperparameters_array: [0.92961609 0.31637555 0.18391881 0.20456028 0.05 ]\n", + "Hyperparameters: \n", + "sig2: 0.9296160928171479 \n", + "ls2: 0.3163755545817859 \n", + "sig3: 0.18391881167709445 \n", + "ls3: 0.2045602785530397 \n", + "noise: 0.05 \n", + "Hyps_mask train_noise: True \n", + "Hyps_mask nspecie: 1 \n", + "Hyps_mask twobody_start: 0 \n", + "Hyps_mask ntwobody: 1 \n", + "Hyps_mask threebody_start: 2 \n", + "Hyps_mask nthreebody: 1 \n", + "Hyps_mask kernels: ['twobody', 'threebody'] \n", + "Hyps_mask cutoffs: {'twobody': 5.0, 'threebody': 3.5} \n", + "\n", + "uncertainty tolerance: 2 times noise hyperparameter \n", + "timestep (ps): 0.09822694788464063\n", + "number of frames: 3\n", + "number of atoms: 8\n", + "system species: {'C'}\n", + "periodic cell: \n", + "[[3.52678 0. 0. ]\n", + " [0. 3.52678 0. ]\n", + " [0. 0. 3.52678]]\n", + "\n", + "previous positions (A):\n", + "C 0.0000 0.0000 0.0000\n", + "C 0.8817 0.8817 0.8817\n", + "C 0.0000 1.7634 1.7634\n", + "C 0.8817 2.6451 2.6451\n", + "C 1.7634 0.0000 1.7634\n", + "C 2.6451 0.8817 2.6451\n", + "C 1.7634 1.7634 0.0000\n", + "C 2.6451 2.6451 0.8817\n", + "--------------------------------------------------------------------------------\n", + "\n", + "INFO:otflog:\n", + "Calling DFT...\n", + "\n", + "INFO:otflog:DFT run complete.\n", + "INFO:otflog:number of DFT calls: 1\n", + "INFO:otflog:wall time from start: 0.03 s\n", + "INFO:otflog:Adding atom [0, 1, 2, 3] to the training set.\n", + "INFO:otflog:Uncertainty: [0. 0. 0.]\n", + "INFO:otflog:\n", + "GP hyperparameters: \n", + "INFO:otflog:Hyp0 : sig2 = 0.9296\n", + "INFO:otflog:Hyp1 : ls2 = 0.3164\n", + "INFO:otflog:Hyp2 : sig3 = 0.1839\n", + "INFO:otflog:Hyp3 : ls3 = 0.2046\n", + "INFO:otflog:Hyp4 : noise = 0.0010\n", + "INFO:otflog:likelihood: 71.8658\n", + "INFO:otflog:likelihood gradient: [-1.79487637e-06 -6.53005436e-06 -8.57610391e-06 -1.76345959e-05\n", + " -1.19999904e+04]\n", + "INFO:otflog:wall time from start: 7.23 s\n", + "INFO:otflog:\n", + "*-Frame: 0 \n", + "Simulation Time: 0.0 ps \n", + "El Position (A) DFT Force (ev/A) Std. Dev (ev/A) Velocities (A/ps) \n", + "C 0.0000 0.0000 0.0000 0.0000 0.0000 0.0000 0.0000 0.0000 0.0000 0.0417 0.0382 0.0035\n", + "C 0.8817 0.8817 0.8817 0.0000 0.0000 0.0000 0.0000 0.0000 0.0000 -0.0135 0.0222 0.0320\n", + "C 0.0000 1.7634 1.7634 0.0000 0.0000 -0.0000 0.0000 0.0000 0.0000 0.0204 -0.0705 0.0194\n", + "C 0.8817 2.6451 2.6451 -0.0000 -0.0000 -0.0000 0.0000 0.0000 0.0000 -0.0013 0.0346 0.0278\n", + "C 1.7634 0.0000 1.7634 0.0000 0.0000 -0.0000 0.0000 0.0000 0.0000 -0.0676 -0.0129 0.0242\n", + "C 2.6451 0.8817 2.6451 -0.0000 -0.0000 -0.0000 0.0000 0.0000 0.0000 -0.0027 -0.0121 -0.0340\n", + "C 1.7634 1.7634 0.0000 0.0000 -0.0000 0.0000 0.0000 0.0000 0.0000 0.0658 -0.0278 -0.0497\n", + "C 2.6451 2.6451 0.8817 -0.0000 -0.0000 0.0000 0.0000 0.0000 0.0000 -0.0427 0.0283 -0.0231\n", + "\n", + "temperature: 199.00 K \n", + "kinetic energy: 0.180060 eV \n", + "\n", + "INFO:otflog:wall time from start: 7.23 s\n", + "INFO:otflog:--------------------------------------------------------------------------------\n", + "-Frame: 1 \n", + "Simulation Time: 0.0982 ps \n", + "El Position (A) GP Force (ev/A) Std. Dev (ev/A) Velocities (A/ps) \n", + "C 0.0082 0.0075 0.0007 -0.0000 -0.0000 -0.0000 16.1332 16.9877 6.7169 0.0417 0.0382 0.0035\n", + "C 0.8790 0.8861 0.8880 0.0000 -0.0000 -0.0000 15.3645 9.5079 17.2434 -0.0135 0.0222 0.0320\n", + "C 0.0040 1.7495 1.7672 -0.0000 0.0000 -0.0000 10.8208 37.7660 9.7448 0.0204 -0.0705 0.0194\n", + "C 0.8814 2.6519 2.6505 -0.0000 -0.0000 0.0000 9.5129 20.9930 4.7440 -0.0013 0.0346 0.0278\n", + "C 1.7501 -0.0025 1.7682 0.0000 0.0000 -0.0000 23.4305 10.7241 12.5271 -0.0676 -0.0129 0.0242\n", + "C 2.6446 0.8793 2.6384 0.0000 -0.0000 0.0000 3.1513 13.4187 7.3315 -0.0027 -0.0121 -0.0340\n", + "C 1.7763 1.7579 -0.0098 -0.0000 0.0000 0.0000 20.5584 9.0554 17.2228 0.0658 -0.0278 -0.0497\n", + "C 2.6367 2.6506 0.8772 0.0000 -0.0000 0.0000 17.8540 7.7304 12.6641 -0.0427 0.0283 -0.0231\n", + "\n", + "temperature: 199.00 K \n", + "kinetic energy: 0.180060 eV \n", + "\n", + "INFO:otflog:wall time from start: 9.65 s\n", + "INFO:otflog:\n", + "Calling DFT...\n", + "\n", + "INFO:otflog:DFT run complete.\n", + "INFO:otflog:number of DFT calls: 2\n", + "INFO:otflog:wall time from start: 9.67 s\n", + "INFO:otflog:\n", + "*-Frame: 1 \n", + "Simulation Time: 0.0982 ps \n", + "El Position (A) DFT Force (ev/A) Std. Dev (ev/A) Velocities (A/ps) \n", + "C 0.0164 0.0150 0.0013 0.0372 -0.0001 -0.0205 16.1332 16.9877 6.7169 0.0835 0.0764 0.0069\n", + "C 0.8763 0.8904 0.8943 -0.0669 -0.0327 0.0578 15.3645 9.5079 17.2434 -0.0274 0.0442 0.0641\n", + "C 0.0080 1.7356 1.7710 0.0322 -0.1194 0.0093 10.8208 37.7660 9.7448 0.0409 -0.1414 0.0388\n", + "C 0.8812 2.6587 2.6560 0.0353 0.0737 -0.0165 9.5129 20.9930 4.7440 -0.0025 0.0695 0.0555\n", + "C 1.7368 -0.0051 1.7729 -0.0330 0.0001 0.0250 23.4305 10.7241 12.5271 -0.1353 -0.0259 0.0486\n", + "C 2.6440 0.8770 2.6317 -0.0014 0.0643 0.0114 3.1513 13.4187 7.3315 -0.0053 -0.0238 -0.0680\n", + "C 1.7893 1.7525 -0.0195 0.0479 0.0129 -0.0207 20.5584 9.0554 17.2228 0.1317 -0.0555 -0.0994\n", + "C 2.6283 2.6562 0.8726 -0.0513 0.0013 -0.0459 17.8540 7.7304 12.6641 -0.0856 0.0565 -0.0465\n", + "\n", + "temperature: 798.57 K \n", + "kinetic energy: 0.722563 eV \n", + "\n", + "INFO:otflog:wall time from start: 9.68 s\n", + "INFO:otflog:mean absolute error: 0.0340 eV/A\n", + "INFO:otflog:mean absolute dft component: 0.0340 eV/A\n", + "INFO:otflog:Adding atom [6, 3, 4, 2] to the training set.\n", + "INFO:otflog:Uncertainty: [20.55839508 9.05540846 17.22284583]\n", + "INFO:otflog:\n", + "GP hyperparameters: \n", + "INFO:otflog:Hyp0 : sig2 = 0.7038\n", + "INFO:otflog:Hyp1 : ls2 = 2.0405\n", + "INFO:otflog:Hyp2 : sig3 = 0.0000\n", + "INFO:otflog:Hyp3 : ls3 = 9.6547\n", + "INFO:otflog:Hyp4 : noise = 0.0010\n", + "INFO:otflog:likelihood: 122.4930\n", + "INFO:otflog:likelihood gradient: [-1.34065483e+00 -1.79554908e-01 -4.94110742e-02 1.54534584e-10\n", + " -1.82026091e+04]\n", + "INFO:otflog:wall time from start: 30.46 s\n", + "INFO:otflog:--------------------------------------------------------------------------------\n", + "-Frame: 2 \n", + "Simulation Time: 0.196 ps \n", + "El Position (A) GP Force (ev/A) Std. Dev (ev/A) Velocities (A/ps) \n", + "C 0.0247 0.0225 0.0020 0.0748 -0.0003 -0.0400 0.0008 0.0015 0.0008 0.0000 0.0000 0.0000\n", + "C 0.8735 0.8947 0.9007 -0.1357 -0.0645 0.1173 0.0014 0.0015 0.0010 0.0000 0.0000 0.0000\n", + "C 0.0121 1.7215 1.7748 0.0632 -0.2385 0.0151 0.0010 0.0015 0.0016 0.0000 0.0000 0.0000\n", + "C 0.8810 2.6657 2.6614 0.0692 0.1497 -0.0328 0.0011 0.0013 0.0010 0.0000 0.0000 0.0000\n", + "C 1.7235 -0.0076 1.7778 -0.0661 -0.0006 0.0515 0.0015 0.0013 0.0008 0.0000 0.0000 0.0000\n", + "C 2.6435 0.8748 2.6251 -0.0019 0.1297 0.0235 0.0009 0.0017 0.0010 0.0000 0.0000 0.0000\n", + "C 1.8023 1.7471 -0.0293 0.0980 0.0221 -0.0451 0.0015 0.0018 0.0012 0.0000 0.0000 0.0000\n", + "C 2.6197 2.6617 0.8679 -0.1015 0.0024 -0.0895 0.0013 0.0012 0.0012 0.0000 0.0000 0.0000\n", + "\n", + "temperature: 0.00 K \n", + "kinetic energy: 0.000000 eV \n", + "\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "INFO:otflog:wall time from start: 33.38 s\n", + "INFO:otflog:--------------------\n", + "INFO:otflog:Run complete.\n" + ] + } + ], + "source": [ + "from flare.ase.otf import ASE_OTF\n", + "\n", + "otf_params = {'init_atoms': [0, 1, 2, 3],\n", + " 'output_name': 'otf',\n", + " 'std_tolerance_factor': 2,\n", + " 'max_atoms_added' : 4,\n", + " 'freeze_hyps': 10}\n", + "\n", + "test_otf = ASE_OTF(super_cell, \n", + " timestep = 1 * units.fs,\n", + " number_of_steps = 3,\n", + " dft_calc = lj_calc,\n", + " md_engine = md_engine,\n", + " md_kwargs = md_kwargs,\n", + " **otf_params)\n", + "\n", + "test_otf.run()" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.7.5" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/source/tutorials/ase.rst b/docs/source/tutorials/ase.rst deleted file mode 100644 index ead929590..000000000 --- a/docs/source/tutorials/ase.rst +++ /dev/null @@ -1,95 +0,0 @@ -On-the-fly training using ASE -============================= - -.. toctree:: - :maxdepth: 2 - -This is a quick introduction of how to set up our ASE-OTF interface to train a force field. We will train a force field model for bulk `AgI `_. To run the on-the-fly training, we will need to - - 1. Create a supercell with ASE Atoms object - - 2. Set up FLARE ASE calculator, including the kernel functions, hyperparameters, cutoffs for Gaussian process, and mapping parameters (if Mapped Gaussian Process is used) - - 3. Set up DFT ASE calculator. Here we will give an example of Quantum Espresso - - 4. Set up on-the-fly training with ASE MD engine - -To increase the flexibility and make it easier for testing, we suggest creating four ".py" files for the 4 steps above. And please make sure you are using the LATEST FLARE code in our master branch. - -Setup supercell with ASE ------------------------- -Here we create a 2x1x1 supercell with lattice constant 3.855, and randomly perturb the positions of the atoms, so that they will start MD with non-zero forces. - -.. literalinclude:: ../../../tests/ase_otf/atom_setup.py - -Setup FLARE calculator ----------------------- -Now let's set up our Gaussian process model in the same way as introduced before - -.. literalinclude:: ../../../tests/ase_otf/flare_setup.py - :lines: 1-27 - -**Optional:** if you want to use the LAMMPS interface with the trained force field, you need to construct Mapped Gaussian Process (MGP). Accelerated on-the-fly training with MGP is also enabled, but not thoroughly tested. You can set up MGP in FLARE calculator as below: - -.. literalinclude:: ../../../tests/ase_otf/flare_setup.py - :lines: 28-51 - -Create a ``Calculator`` object - -.. literalinclude:: ../../../tests/ase_otf/flare_setup.py - :lines: 52-53 - - -Setup DFT calculator --------------------- -For DFT calculator, here we use `Quantum Espresso (QE) `_ as an example. -First, we need to set up our environment variable ``ASE_ESPRESSO_COMMAND`` to our QE executable, -so that ASE can find this calculator. Then set up our input parameters of QE and create an ASE calculator - -.. literalinclude:: ../../../tests/ase_otf/qe_setup.py - -Setup On-The-Fly MD engine --------------------------- -Finally, our OTF is compatible with 4 MD engines that ASE supports: -VelocityVerlet, NVTBerendsen, NPTBerendsen and NPT. -We can choose any of them, and set up the parameters based on -`ASE requirements `_. -After everything is set up, we can run the on-the-fly training by method ``otf_run(number_of_steps)`` - -**Note:** Currently, only ``VelocityVerlet`` is tested on real system, ``NPT`` may have issue with pressure and stress. - -.. literalinclude:: ../../../tests/ase_otf/otf_setup.py - -When the OTF training is finished, there will be data files saved including: - -1. A log file ``otf_run.log`` of the information in training. -If ``data_in_logfile=True``, then the data in ``otf_data`` folder -(described below) will also be written in this log file. - -2. A folder ``otf_data`` containing: - - * positions.xyz: the trajectory of the on-the-fly MD run - * velocities.dat: the velocities of the frames in the trajectory - * forces.dat: the forces of the frames in trajectory predicted by FLARE - * uncertainties.dat: the uncertainties of the frames in trajectory predicted by FLARE - * dft_positions.xyz: the DFT calculated frames - * dft_forces.dat: the DFT forces correspond to frames in dft_positions.xyz - * added_atoms.dat: the list of atoms added to the training set of FLARE in each DFT calculated frame - -3. Kernel matrix and alpha vector used in GP: ``ky_mat_inv.npy`` and ``alpha.npy`` - -4. If MGP is used, i.e. ``use_mapping=True``, and 3-body kernel is used and mapped, -then there will be two files saving grid values: ``grid3_mean.npy`` and ``grid3_var.npy``. - -Restart from previous training ------------------------------- -We have an option for continuing from a finished training. - -1. Move all the saved files mentioned above to one folder, e.g. ``restart_data``. -(All the ``.xyz``, ``.dat``, ``.npy`` files should be in one folder) - -2. Set ``otf_params['restart_from'] = 'restart_data'`` - -3. Run as mentioned in above sections - - diff --git a/docs/source/tutorials/gpfa.rst b/docs/source/tutorials/gpfa.rst index 19c3834e2..c376eb1d4 100644 --- a/docs/source/tutorials/gpfa.rst +++ b/docs/source/tutorials/gpfa.rst @@ -8,8 +8,15 @@ Molecular Dynamics (AIMD) trajectory can be used to train a Gaussian Process mo We can use a very short trajectory for a very simple molecule which is already included in the test files in order to demonstrate how to set up and run the code. The trajectory this tutorial focuses on involves a few frames of the -molecule Methanol vibrating about it's equilibrium configuration ran in VASP. +molecule Methanol vibrating about its equilibrium configuration ran in VASP. +Roadmap Figure +-------------- +In this tutorial, we will walk through the first two steps contained in the below figure. the GP from AIMD module is designed to give you the tools necessary to extract FLARE structures from a previously existing molecular dynamics run. + +.. figure:: ../../images/GPFA_tutorial.png + :figwidth: 800 % + :align: center Step 1: Setting up a Gaussian Process Object @@ -54,9 +61,9 @@ We will then set up the ``GaussianProcess`` object. .. code-block:: python - gp = GaussianProcess(kernel_name="2+3mc", + gp = GaussianProcess(kernels=['twobody', 'threebody'], hyps=[0.01, 0.01, 0.01, 0.01, 0.01], - cutoffs = (7,7), + cutoffs = {'twobody':7, 'threebody':3}, hyp_labels=['Two-Body Signal Variance','Two-Body Length Scale','Three-Body Signal Variance', 'Three-Body Length Scale', 'Noise Variance'] ) @@ -93,7 +100,7 @@ You can open it via the command from json import loads from flare.struc import Structure with open('path-to-methanol-frames','r') as f: - loaded_dicts = [loads(line) for line in f.readlines()] + loaded_dicts = [loads(line) for line in f.readlines()] trajectory = [Structure.from_dict(d) for d in loaded_dicts] Our trajectory is a list of FLARE structures, each of which is decorated with diff --git a/docs/source/tutorials/lammps.rst b/docs/source/tutorials/lammps.rst new file mode 100644 index 000000000..4f008d6fb --- /dev/null +++ b/docs/source/tutorials/lammps.rst @@ -0,0 +1,137 @@ +Compile LAMMPS with MGP Pair Style +================================== +Anders Johansson + +MPI +--- + +For when you can't get Kokkos+OpenMP to work. + +Compiling +********* + +.. code-block:: bash + + cp lammps_plugin/pair_mgp.* /path/to/lammps/src + cd /path/to/lammps/src + make -j$(nproc) mpi + +You can replace ``mpi`` with your favourite Makefile, e.g. ``intel_cpu_intelmpi``, or use the CMake build system. + + +Running +******* + +.. code-block:: bash + + mpirun /path/to/lammps/src/lmp_mpi -in in.lammps + +as usual, but your LAMMPS script ``in.lammps`` needs to specify ``newton off``. + +MPI+OpenMP through Kokkos +------------------------- + +For OpenMP parallelisation on your laptop or on one node, or for hybrid parallelisation on multiple nodes. + +Compiling +********* + +.. code-block:: bash + + cp lammps_plugin/pair_mgp*.* /path/to/lammps/src + cd /path/to/lammps/src + make yes-kokkos + make -j$(nproc) kokkos_omp + +You can change the compiler flags etc. in ``/path/to/lammps/src/MAKE/OPTIONS/Makefile.kokkos_omp``. +As of writing, the pair style is not detected by CMake. + +Running +******* + +With ``newton on`` in your LAMMPS script: + +.. code-block:: bash + + mpirun /path/to/lammps/src/lmp_kokkos_omp -k on t 4 -sf kk -package kokkos newton on neigh half -in in.lammps + +With ``newton off`` in your LAMMPS script: + +.. code-block:: bash + + mpirun /path/to/lammps/src/lmp_kokkos_omp -k on t 4 -sf kk -package kokkos newton off neigh full -in in.lammps + +Replace 4 with the desired number of threads *per MPI task*. Skip ``mpirun`` when running on one machine. + +If you are running on multiple nodes on a cluster, you would typically launch one MPI task per node, +and then set the number of threads equal to the number of cores (or hyperthreads) per node. +A sample SLURM job script for 4 nodes, each with 48 cores, may look something like this: + +.. code-block:: bash + + #SBATCH --nodes=4 + #SBATCH --ntasks-per-node=1 + #SBATCH --cpus-per-task=48 + mpirun -np $SLURM_NTASKS /path/to/lammps/src/lmp_kokkos_omp -k on t $SLURM_CPUS_PER_TASK -sf kk -package kokkos newton off neigh full -in in.lammps + +When running on Knight's Landing or other heavily hyperthreaded systems, you may want to try using more than one thread per CPU. + +MPI+CUDA through Kokkos +----------------------- + +For running on the GPU on your laptop, or for multiple GPUs on one or more nodes. + +Compiling +********* + +.. code-block:: bash + + cp lammps_plugin/pair_mgp*.* /path/to/lammps/src + cd /path/to/lammps/src + make yes-kokkos + make -j$(nproc) KOKKOS_ARCH=Volta70 kokkos_cuda_mpi + +The ``KOKKOS_ARCH`` must be changed according to your GPU model. ``Volta70`` is for V100, ``Pascal60`` is for P100, etc. + +You can change the compiler flags etc. in ``/path/to/lammps/src/MAKE/OPTIONS/Makefile.kokkos_cuda_mpi``. +As of writing, the pair style is not detected by CMake. + +Running +******* + +With ``newton on`` in your LAMMPS script: + +.. code-block:: bash + + mpirun /path/to/lammps/src/lmp_kokkos_cuda_mpi -k on g 4 -sf kk -package kokkos newton on neigh half -in in.lammps + +With ``newton off`` in your LAMMPS script: + +.. code-block:: bash + + mpirun /path/to/lammps/src/lmp_kokkos_cuda_mpi -k on g 4 -sf kk -package kokkos newton off neigh full -in in.lammps + +Replace 4 with the desired number of GPUs *per node*, skip ``mpirun`` if you are using 1 GPU. +The number of MPI tasks should be set equal to the total number of GPUs. + +If you are running on multiple nodes on a cluster, you would typically launch one MPI task per GPU. +A sample SLURM job script for 4 nodes, each with 2 GPUs, may look something like this: + +.. code-block:: bash + + #SBATCH --nodes=4 + #SBATCH --ntasks-per-node=2 + #SBATCH --cpus-per-task=1 + #SBATCH --gpus-per-node=2 + mpirun -np $SLURM_NTASKS /path/to/lammps/src/lmp_kokkos_cuda_mpi -k on g $SLURM_GPUS_PER_NODE -sf kk -package kokkos newton off neigh full -in in.lammps + +Notes on Newton (only relevant with Kokkos) +------------------------------------------- + +There are defaults which will kick in if you don't specify anything in the input +script and/or skip the ``-package kokkos newton ... neigh ...`` flag. +You can try these at your own risk, but it is safest to specify everything. +See also the `documentation `_. + +``newton on`` will probably be faster if you have a 2-body potential, +otherwise the alternatives should give roughly equal performance. diff --git a/docs/source/tutorials/otf_al.rst b/docs/source/tutorials/otf_al.rst index 296a9bc99..4a1cf7ddb 100644 --- a/docs/source/tutorials/otf_al.rst +++ b/docs/source/tutorials/otf_al.rst @@ -17,10 +17,10 @@ or :doc:`mc_simple.py <../flare/kernels/mc_simple>` (multi-component) for more o # make gp model hyps = np.array([0.1, 1, 0.01]) hyp_labels = ['Signal Std', 'Length Scale', 'Noise Std'] - cutoffs = np.array([3.9, 3.9]) + cutoffs = {'threebody':3.9} gp = \ - GaussianProcess(kernel_name='3b', + GaussianProcess(kernels=['threebody'], hyps=hyps, cutoffs=cutoffs, hyp_labels=hyp_labels, @@ -29,22 +29,20 @@ or :doc:`mc_simple.py <../flare/kernels/mc_simple>` (multi-component) for more o **Some Explanation about the parameters:** -* ``kernel_name``: set to be the name of kernel functions +* ``kernels``: set to be the name list of kernel functions to use - * import from :doc:`sc.py <../flare/kernels/sc>` (single-component system) - or :doc:`mc_simple.py <../flare/kernels/mc_simple>` (multi-component system). - * Currently we have the choices of two-body, three-body and two-plus-three-body kernel functions. - * Two-plus-three-body kernel function is simply the summation of two-body and three-body kernels, - and is tested to have best performance. + * Currently we have the choices of ``twobody``, ``threebody`` and ``manybody`` kernel functions. + * If multiple kernels are listed, the resulted kernel is simply the summation of all listed kernels, * ``hyps``: the array of hyperparameters, whose names are shown in ``hyp_labels``. * For two-body kernel function, an array of length 3 is needed, ``hyps=[sigma_2, ls_2, sigma_n]``; * For three-body, ``hyps=[sigma_3, ls_3, sigma_n]``; - * For two-plus-three-body, ``hyps=[sigma_2, ls_2, sigma_3, ls_3, sigma_n]``. + * For twobody plus threebody, ``hyps=[sigma_2, ls_2, sigma_3, ls_3, sigma_n]``. + * For twobody, threebody plus manybody, ``hyps=[sigma_2, ls_2, sigma_3, ls_3, sigma_m, ls_m, sigma_n]``. -* ``cutoffs``: consists of two values. The 1st is the cutoff of two-body and the 2nd is for three-body kernel. - Usually we will set a larger one for two-body. +* ``cutoffs``: a dictionary consists of corresponding cutoff values for each kernel. + Usually we will set a larger one for two-body, and smaller one for threebody and manybody * ``maxiter``: set to constrain the number of steps in training hyperparameters. @@ -52,6 +50,7 @@ or :doc:`mc_simple.py <../flare/kernels/mc_simple>` (multi-component) for more o **Note:** 1. See :doc:`GaussianProcess <../flare/gp>` for complete description of arguments of ``GaussianProcess`` class. +2. See :doc:`AdvancedHyperparametersSetUp <../flare/utils/mask_helper>` for more complicated hyper-parameters set up. Step 2: Set up DFT Calculator diff --git a/docs/source/tutorials/prepare_data.rst b/docs/source/tutorials/prepare_data.rst new file mode 100644 index 000000000..6dc76b923 --- /dev/null +++ b/docs/source/tutorials/prepare_data.rst @@ -0,0 +1,105 @@ +Prepare your data +================= + +If you have collected data for training, including atomic positions, chemical +species, cell etc., you need to convert it into a list of ``Structure`` objects. +Below we provide a few examples. + + +VASP data +--------- + +If you have AIMD data from VASP, you can follow +`the step 2 of this instruction `_ +to generate ``Structure`` with the ``vasprun.xml`` file. + + +Data from Quantum Espresso, LAMMPS, etc. +---------------------------------------- + +If you have collected data from any +`calculator that ASE supports `_, +or have dumped data file of `format that ASE supports `_, +you can convert your data into ASE ``Atoms``, then from ``Atoms`` to +``Structure`` via ``Structure.from_ase_atoms``. + +For example, if you have collected data from QE, and obtained the QE output file ``.pwo``, +you can parse it with ASE, and convert ASE ``Atoms`` into ``Structure``. + + +.. code-block:: python + + from ase.io import read + from flare.struc import Structure + + frames = read('data.pwo', index=':', format='espresso-out') # read the whole traj + trajectory = [] + for atoms in frames: + trajectory.append(Structure.from_ase_atoms(atoms)) + + +If the data is from the LAMMPS dump file, use + +.. code-block:: python + + # if it's text file + frames = read('data.dump', index=':', format='lammps-dump-text') + + # if it's binary file + frames = read('data.dump', index=':', format='lammps-dump-binary') + + +Then the ``trajectory`` can be used to +`train GP from AIMD data `_. + + +Try building GP from data +------------------------- + +To have a more complete and better monitored training process, please use our +`GPFA module `_. + +Here we are not going to use this module, but only provide a simple example on +how the GP is constructed from the data. + +.. code-block:: python + + from flare.gp import GaussianProcess + from flare.utils.parameter_helper import ParameterHelper + + # set up hyperparameters, cutoffs + kernels = ['twobody', 'threebody'] + parameters = {'cutoff_twobody': 4.0, 'cutoff_threebody': 3.0} + pm = ParameterHelper(kernels=kernels, + random=True, + parameters=parameters) + hm = pm.as_dict() + hyps = hm['hyps'] + cutoffs = hm['cutoffs'] + hl = hm['hyp_labels'] + + kernel_type = 'mc' # multi-component. use 'sc' for single component system + + # build up GP model + gp_model = \ + GaussianProcess(kernels=kernels, + component=kernel_type, + hyps=hyps, + hyp_labels=hl, + cutoffs=cutoffs, + hyps_mask=hm, + parallel=False, + n_cpus=1) + + # feed training data into GP + # use the "trajectory" as from above, a list of Structure objects + for train_struc in trajectory: + gp_model.update_db(train_struc, forces) + gp_model.check_L_alpha() # build kernel matrix from training data + + # make a prediction with gp, test on a training data + test_env = gp_model.training_data[0] + gp_pred = gp_model.predict(test_env, 1) # obtain the x-component + # (force_x, var_x) + # x: 1, y: 2, z: 3 + print(gp_pred) diff --git a/docs/source/tutorials/tutorials.rst b/docs/source/tutorials/tutorials.rst index c7a692795..7e50588b1 100644 --- a/docs/source/tutorials/tutorials.rst +++ b/docs/source/tutorials/tutorials.rst @@ -4,7 +4,9 @@ Tutorials .. toctree:: :maxdepth: 3 + prepare_data + gpfa otf_al ase - gpfa after_training + lammps diff --git a/flare/ase/calculator.py b/flare/ase/calculator.py index 269324a3a..c69659ede 100644 --- a/flare/ase/calculator.py +++ b/flare/ase/calculator.py @@ -1,23 +1,29 @@ -""":class:`FLARE_Calculator` is a calculator compatible with `ASE`. You can build up `ASE Atoms` for your atomic structure, and use `get_forces`, `get_potential_energy` as general `ASE Calculators`, and use it in `ASE Molecular Dynamics` and our `ASE OTF` training module.""" +''':class:`FLARE_Calculator` is a calculator compatible with `ASE`. +You can build up `ASE Atoms` for your atomic structure, and use `get_forces`, +`get_potential_energy` as general `ASE Calculators`, and use it in +`ASE Molecular Dynamics` and our `ASE OTF` training module. For the usage +users can refer to `ASE Calculator module `_ +and `ASE Calculator tutorial `_.''' + +import warnings import numpy as np import multiprocessing as mp from flare.env import AtomicEnvironment from flare.struc import Structure -from flare.mgp.mgp import MappedGaussianProcess +from flare.mgp import MappedGaussianProcess from flare.predict import predict_on_structure_par_en, predict_on_structure_en from ase.calculators.calculator import Calculator class FLARE_Calculator(Calculator): - """Build FLARE as an ASE Calculator, which is compatible with ASE Atoms and Molecular Dynamics. - - :param gp_model: FLARE's Gaussian process object - :type gp_model: GaussianProcess - :param mgp_model: FLARE's Mapped Gaussian Process object. `None` by default. MGP will only be used if `use_mapping` is set to True - :type mgp_model: MappedGaussianProcess - :param par: set to `True` if parallelize the prediction. `False` by default. - :type par: Bool - :param use_mapping: set to `True` if use MGP for prediction. `False` by default. - :type use_mapping: Bool + """ + Build FLARE as an ASE Calculator, which is compatible with ASE Atoms and + Molecular Dynamics. + Args: + gp_model (GaussianProcess): FLARE's Gaussian process object + mgp_model (MappedGaussianProcess): FLARE's Mapped Gaussian Process object. + `None` by default. MGP will only be used if `use_mapping` is set to True + par (Bool): set to `True` if parallelize the prediction. `False` by default. + use_mapping (Bool): set to `True` if use MGP for prediction. `False` by default. """ def __init__(self, gp_model, mgp_model=None, par=False, use_mapping=False): @@ -46,7 +52,7 @@ def get_forces(self, atoms): def get_stress(self, atoms): if not self.use_mapping: - raise NotImplementedError("Stress is only supported in MGP") + warnings.warn('Stress is only implemented in MGP, not in GP. Will return zeros.') return self.get_property('stress', atoms) @@ -88,7 +94,13 @@ def calculate_gp(self, atoms): self.results['local_energies'] = local_energies self.results['energy'] = np.sum(local_energies) atoms.get_uncertainties = self.get_uncertainties - return forces + + # GP stress not implemented yet + self.results['stresses'] = np.zeros((nat, 6)) + volume = atoms.get_volume() + total_stress = np.sum(self.results['stresses'], axis=0) + self.results['stress'] = total_stress / volume + def calculate_mgp_serial(self, atoms): @@ -101,11 +113,19 @@ def calculate_mgp_serial(self, atoms): self.results['stresses'] = np.zeros((nat, 6)) self.results['stds'] = np.zeros((nat, 3)) self.results['local_energies'] = np.zeros(nat) + for n in range(nat): chemenv = AtomicEnvironment(struc_curr, n, - self.mgp_model.cutoffs, + self.gp_model.cutoffs, cutoffs_mask = self.mgp_model.hyps_mask) - f, v, vir, e = self.mgp_model.predict(chemenv, mean_only=False) + + try: + f, v, vir, e = self.mgp_model.predict(chemenv, mean_only=False) + except ValueError: # if lower_bound error is raised + warnings.warn('Re-build map with a new lower bound') + self.mgp_model.build_map(self.gp_model) + f, v, vir, e = self.mgp_model.predict(chemenv, mean_only=False) + self.results['forces'][n] = f self.results['stresses'][n] = vir self.results['stds'][n] = np.sqrt(np.absolute(v)) @@ -127,47 +147,3 @@ def calculate_mgp_par(self, atoms): def calculation_required(self, atoms, quantities): return True - - def train_gp(self, **kwargs): - """ - The same function of training GP hyperparameters as `train()` in :class:`GaussianProcess` - """ - self.gp_model.train(**kwargs) - - - def build_mgp(self, skip=True): - """ - Construct :class:`MappedGaussianProcess` based on the current GP - TODO: change to the re-build method - - :param skip: if `True`, then it will not construct MGP - :type skip: Bool - """ - # l_bound not implemented - - if skip: - return 1 - - # set svd rank based on the training set, grid number and threshold 1000 - grid_params = self.mgp_model.grid_params - struc_params = self.mgp_model.struc_params - map_force = self.mgp_model.map_force - lmp_file_name = self.mgp_model.lmp_file_name - mean_only = self.mgp_model.mean_only - n_cpus = self.mgp_model.n_cpus - container_only = False - - train_size = len(self.gp_model.training_data) - rank_2 = np.min([1000, grid_params['grid_num_2'], train_size*3]) - rank_3 = np.min([1000, grid_params['grid_num_3'][0]**3, train_size*3]) - grid_params['svd_rank_2'] = rank_2 - grid_params['svd_rank_3'] = rank_3 - - self.mgp_model = MappedGaussianProcess(grid_params, - struc_params, - map_force=map_force, - GP=self.gp_model, - mean_only=mean_only, - container_only=container_only, - lmp_file_name=lmp_file_name, - n_cpus=n_cpus) diff --git a/flare/ase/dft.py b/flare/ase/dft.py new file mode 100644 index 000000000..a72491fed --- /dev/null +++ b/flare/ase/dft.py @@ -0,0 +1,33 @@ +''' +This module is to provide the same interface as the module `dft_interface`, so we can use ASE atoms and calculators to run OTF +''' + +import numpy as np +from copy import deepcopy + +def parse_dft_input(atoms): + pos = np.copy(atoms.positions) + spc = atoms.get_chemical_symbols() + cell = np.array(atoms.get_cell()) + + # build mass dict + mass = atoms.get_masses() + mass_dict = {} + for i in range(len(spc)): + spec_ind = str(spc[i]) + if spec_ind not in mass_dict.keys(): + mass_dict[spec_ind] = mass[i] + return pos, spc, cell, mass_dict + +def run_dft_par(atoms, structure, dft_calc, **dft_kwargs): + ''' + Assume that the atoms have been updated + ''' + # change from FLARE to DFT calculator + calc = deepcopy(dft_calc) + atoms.set_calculator(calc) + + # calculate DFT forces + forces = atoms.get_forces() + + return forces diff --git a/flare/ase/logger.py b/flare/ase/logger.py deleted file mode 100644 index 84b2cdc75..000000000 --- a/flare/ase/logger.py +++ /dev/null @@ -1,203 +0,0 @@ -import numpy as np -import datetime -import time -import os - -from ase.md import MDLogger -from ase import units - -from flare.ase.calculator import FLARE_Calculator - - -class OTFLogger(MDLogger): - - def __init__(self, dyn, atoms, logfile, header=False, stress=False, - peratom=False, mode="w", data_folder='otf_data', - data_in_logfile=False): - - super().__init__(dyn, atoms, logfile, header, stress, - peratom, mode) - - self.natoms = self.atoms.get_number_of_atoms() - self.write_header_info() - self.start_time = time.time() - if data_folder not in os.listdir(): - os.mkdir(data_folder) - - self.positions_xyz = open(data_folder+'/positions.xyz', mode=mode) - self.velocities_dat = open(data_folder+'/velocities.dat', mode=mode) - self.forces_dat = open(data_folder+'/forces.dat', mode=mode) - self.uncertainties_dat = open(data_folder+'/uncertainties.dat', - mode=mode) - self.dft_positions_xyz = open(data_folder+'/dft_positions.xyz', - mode=mode) - self.dft_forces_dat = open(data_folder+'/dft_forces.dat', mode=mode) - self.added_atoms_dat = open(data_folder+'/added_atoms.dat', mode=mode) - - self.traj_files = [self.positions_xyz, self.velocities_dat, - self.forces_dat, self.uncertainties_dat] - self.dft_data_files = [self.dft_positions_xyz, self.dft_forces_dat] - self.data_in_logfile = data_in_logfile - - def write_header_info(self): - gp_model = self.atoms.calc.gp_model - self.logfile.write(str(datetime.datetime.now())) - self.logfile.write('\nnumber of cpu cores: ') # TODO - self.logfile.write('\ncutoffs: '+str(gp_model.cutoffs)) - self.logfile.write('\nkernel_name: '+gp_model.kernel.__name__) - self.logfile.write('\nnumber of hyperparameters: '+str(len(gp_model.hyps))) - self.logfile.write('\nhyperparameters: '+str(gp_model.hyps)) - self.logfile.write('\nhyperparameter optimization algorithm: ' + - gp_model.opt_algorithm) - self.logfile.write('\nuncertainty tolerance: {} times noise'.format( - str(self.dyn.std_tolerance_factor))) - self.logfile.write('\ntimestep (ps): {}'.format(self.dyn.dt/1000)) - self.logfile.write('\nnumber of frames: {}'.format(0)) - self.logfile.write('\nnumber of atoms: {}'.format( - len(self.atoms.positions))) - self.logfile.write('\nsystem species: {}'.format( - self.atoms.get_chemical_symbols())) - self.logfile.write('\nperiodic cell:\n'+str(np.array(self.atoms.cell))) - self.logfile.write('\n') - self.write_prev_positions() - - - def write_hyps(self, hyp_labels, hyps, like, like_grad): - self.logfile.write('\n\nGP hyperparameters: \n') - for i, label in enumerate(hyp_labels): - self.logfile.write('Hyp{} : {} = {}\n'.format(i, label, hyps[i])) - - self.logfile.write('likelihood: '+str(like)+'\n') - self.logfile.write('likelihood gradient: '+str(like_grad)) - - def write_wall_time(self): - self.logfile.write('\nwall time from start: %.2f s \n' - % (time.time()-self.start_time)) - - def write_prev_positions(self): - v = self.atoms.get_velocities() - pos = self.atoms.get_positions() - dt = self.dyn.dt - prev_pos = pos - v * dt - species = self.atoms.get_chemical_symbols() - nat = len(pos) - self.logfile.write('previous positions (A):\n') - template = '{} {:9f} {:9f} {:9f}\n' - for n in range(nat): - self.logfile.write(template.format(species[n], - prev_pos[n,0], prev_pos[n,1], prev_pos[n,2])) - - def __call__(self): - self.write_logfile() - self.write_datafiles() - - def write_datafiles(self): - template = '{} {:9f} {:9f} {:9f}' - steps = self.dyn.nsteps - t = steps / 1000 - - species = self.atoms.get_chemical_symbols() - positions = self.atoms.get_positions() - forces = self.atoms.get_forces() - if type(self.atoms.calc) == FLARE_Calculator: - velocities = self.atoms.get_velocities() - stds = self.atoms.get_uncertainties(self.atoms) - data_files = self.traj_files - data = [positions, velocities, forces, stds] - else: - data_files = self.dft_data_files - data = [positions, forces] - self.added_atoms_dat.write('Frame '+str(steps)+'\n') - - for ind, f in enumerate(data_files): - f.write(str(self.natoms)) - f.write('\nFrame '+str(steps)+'\n') - for atom in range(self.natoms): - dat = template.format(species[atom], - data[ind][atom][0], - data[ind][atom][1], - data[ind][atom][2]) - f.write(dat+'\n') - f.flush() - - - def write_logfile(self): - self.logfile.write(50*'-') - if self.dyn is not None: - steps = self.dyn.nsteps - t = steps / 1000 - if type(self.atoms.calc) != FLARE_Calculator: - self.logfile.write('\n*-Frame: '+str(steps)) - else: - self.logfile.write('\n-Frame: '+str(steps)) - self.logfile.write('\nSimulation time: '+str(t)+' ps') - self.logfile.write('\n') - - if self.data_in_logfile: - self.write_data_to_logfile() - - # write energy, temperature info - epot = self.atoms.get_potential_energy() - ekin = self.atoms.get_kinetic_energy() - temp = ekin / (1.5 * units.kB * self.natoms) - if self.peratom: - epot /= self.natoms - ekin /= self.natoms -# self.logfile.write('\ntotal energy: '+str(epot+ekin)) - self.logfile.write('\ntemperature: '+str(temp)+' K') - self.logfile.write('\nkinetic energy: '+str(ekin)+' eV') - self.write_wall_time() - - self.logfile.flush() - - def write_data_to_logfile(self): - # add positions, forces and stds to be written - species = self.atoms.get_chemical_symbols() - positions = self.atoms.get_positions() - forces = self.atoms.get_forces() - velocities = self.atoms.get_velocities() - if type(self.atoms.calc) == FLARE_Calculator: - stds = self.atoms.get_uncertainties(self.atoms) - force_str = 'GP Forces' - else: - stds = np.zeros(positions.shape) - force_str = 'DFT Forces' - self.logfile.write('Type'+(7*' ') + 'Positions'+(23*' ') + - force_str+(22*' ') + 'Uncertainties'+(22*' ') + - 'Velocities\n') - - template = '{} {:9f} {:9f} {:9f} {:9f} {:9f} {:9f} '\ - '{:9f} {:9f} {:9f} {:9f} {:9f} {:9f}' - for atom in range(len(positions)): - dat = \ - template.format(species[atom], positions[atom][0], - positions[atom][1], positions[atom][2], - forces[atom][0], forces[atom][1], - forces[atom][2], stds[atom][0], - stds[atom][1], stds[atom][2], - velocities[atom][0], velocities[atom][1], - velocities[atom][2]) - self.logfile.write(dat+'\n') - - def add_atom_info(self, target_atom, uncertainty): - if not isinstance(target_atom, list): - target_atom = [target_atom] - self.logfile.write('\nAdding atom {} to the training set.\ - \nUncertainty: {}.'.format(target_atom, uncertainty)) - self.added_atoms_dat.write(str(target_atom[0])+' ') # temporarily support 1 atom - - def write_mgp_train(self, mgp_model, train_time): - train_size = len(mgp_model.GP.training_data) - self.logfile.write('\ntraining set size: {}\n'.format(train_size)) - self.logfile.write('lower bound: {}\n'.format(mgp_model.l_bound)) - self.logfile.write('mgp l_bound: {}\n'.format(mgp_model.grid_params - ['bounds_2'][0, 0])) - self.logfile.write('building mapping time: {}'.format(train_time)) - - def run_complete(self): - self.logfile.write('Run complete.') - for f in self.traj_files: - f.close() - for f in self.dft_data_files: - f.close() - self.added_atoms_dat.close() diff --git a/flare/ase/otf.py b/flare/ase/otf.py index fa059b2d7..33feb3e63 100644 --- a/flare/ase/otf.py +++ b/flare/ase/otf.py @@ -1,358 +1,172 @@ ''' -:class:`OTF` is the on-the-fly training module for ASE, WITHOUT molecular dynamics engine. -It needs to be used adjointly with ASE MD engine. Please refer to our -`OTF MD module `_ for the -complete training module with OTF and MD. +:class:`OTF` is the on-the-fly training module for ASE, WITHOUT molecular dynamics engine. +It needs to be used adjointly with ASE MD engine. ''' import os import sys import inspect +from time import time from copy import deepcopy +import numpy as np +from ase.md.npt import NPT +from ase.md.nvtberendsen import NVTBerendsen +from ase.md.nptberendsen import NPTBerendsen +from ase.md.verlet import VelocityVerlet +from ase.md.langevin import Langevin + from flare.struc import Structure from flare.gp import GaussianProcess +from flare.otf import OTF from flare.utils.learner import is_std_in_bound -from flare.mgp.utils import get_l_bound -import numpy as np -from ase import units -from ase.calculators.espresso import Espresso +from flare.ase.calculator import FLARE_Calculator +import flare.ase.dft as dft_source -class OTF: - """ - OTF (on-the-fly) training with the ASE interface. +class ASE_OTF(OTF): - Note: Dft calculator is set outside of the otf module, and input as - dft_calc, so that different calculators can be used + ''' + On-the-fly training module using ASE MD engine, a subclass of OTF. Args: - dft_calc (ASE Calculator): the ASE DFT calculator (see ASE documentaion) - dft_count (int): initial number of DFT calls - std_tolerance_factor (float): the threshold of calling DFT = noise * - std_tolerance_factor - init_atoms (list): the list of atoms in the first DFT call to add to - the training set, since there's no uncertainty prediction initially - calculate_energy (bool): if True, the energy will be calculated; - otherwise, only forces will be predicted - max_atoms_added (int): the maximal number of atoms to add to the - training set after each DFT calculation - freeze_hyps (int or None): the hyperparameters will only be trained for - the first `freeze_hyps` DFT calls, and will be fixed after that - restart_from (str or None): the path of the directory that stores the - training data from last OTF run, and this OTF will restart from it - - Other Parameters: - use_mapping (bool): if True, the MGP will be used - non_mapping_steps (list): a list of steps that MGP will not be - constructed and used - l_bound (float): the lower bound of the interatomic distance, used for - MGP construction - two_d (bool): used in the calculation of l_bound. If 2-D material is - considered, set to True, then the atomic environment construction - will only search the x & y periodic boundaries to save time - """ + atoms (ASE Atoms): the ASE Atoms object for the on-the-fly MD run, + with calculator set as FLARE_Calculator. + timestep: the timestep in MD. Please use ASE units, e.g. if the timestep + is 1 fs, then set `timestep = 1 * units.fs` + number_of_steps (int): the total number of steps for MD. + dft_calc (ASE Calculator): any ASE calculator is supported, + e.g. Espresso, VASP etc. + md_engine (str): the name of MD thermostat, only `VelocityVerlet`, + `NVTBerendsen`, `NPTBerendsen`, `NPT` and `Langevin` are supported. + md_kwargs (dict): Specify the args for MD as a dictionary, the args are + as required by the ASE MD modules consistent with the `md_engine`. + trajectory (ASE Trajectory): default `None`, not recommended, currently + in experiment. + + The following arguments are for on-the-fly training, the user can also + refer to :class:`OTF` - def __init__(self, - # on-the-fly parameters - dft_calc=None, dft_count=None, std_tolerance_factor: float=1, - skip: int=0, init_atoms: list=[], calculate_energy=False, - max_atoms_added=1, freeze_hyps=1, restart_from=None, - # mgp parameters - use_mapping: bool=False, non_mapping_steps: list=[], - l_bound: float=None, two_d: bool=False): - - # get all arguments as attributes - arg_dict = inspect.getargvalues(inspect.currentframe())[3] - del arg_dict['self'] - self.__dict__.update(arg_dict) - - if dft_count is None: - self.dft_count = 0 - self.noa = len(self.atoms.positions) - - # initialize local energies - if calculate_energy: - self.local_energies = np.zeros(self.noa) + Args: + prev_pos_init ([type], optional): Previous positions. Defaults + to None. + rescale_steps (List[int], optional): List of frames for which the + velocities of the atoms are rescaled. Defaults to []. + rescale_temps (List[int], optional): List of rescaled temperatures. + Defaults to []. + + calculate_energy (bool, optional): If True, the energy of each + frame is calculated with the GP. Defaults to False. + write_model (int, optional): If 0, write never. If 1, write at + end of run. If 2, write after each training and end of run. + If 3, write after each time atoms are added and end of run. + + std_tolerance_factor (float, optional): Threshold that determines + when DFT is called. Specifies a multiple of the current noise + hyperparameter. If the epistemic uncertainty on a force + component exceeds this value, DFT is called. Defaults to 1. + skip (int, optional): Number of frames that are skipped when + dumping to the output file. Defaults to 0. + init_atoms (List[int], optional): List of atoms from the input + structure whose local environments and force components are + used to train the initial GP model. If None is specified, all + atoms are used to train the initial GP. Defaults to None. + output_name (str, optional): Name of the output file. Defaults to + 'otf_run'. + max_atoms_added (int, optional): Number of atoms added each time + DFT is called. Defaults to 1. + freeze_hyps (int, optional): Specifies the number of times the + hyperparameters of the GP are optimized. After this many + updates to the GP, the hyperparameters are frozen. + Defaults to 10. + + n_cpus (int, optional): Number of cpus used during training. + Defaults to 1. + ''' + + def __init__(self, atoms, timestep, number_of_steps, dft_calc, + md_engine, md_kwargs, trajectory=None, **otf_kwargs): + + self.atoms = atoms + self.md_engine = md_engine + + if md_engine == 'VelocityVerlet': + MD = VelocityVerlet + elif md_engine == 'NVTBerendsen': + MD = NVTBerendsen + elif md_engine == 'NPTBerendsen': + MD = NPTBerendsen + elif md_engine == 'NPT': + MD = NPT + elif md_engine == 'Langevin': + MD = Langevin else: - self.local_energies = None - - # initialize otf - if init_atoms is None: - self.init_atoms = [int(n) for n in range(self.noa)] - - def otf_run(self, steps, rescale_temp=[], rescale_steps=[]): - """ - Use `otf_run` intead of `run` to perform a number of time steps. - - Args: - steps (int): the number of time steps - - Other Parameters: - rescale_temp (list): a list of temepratures that rescale the system - rescale_steps (list): a list of step numbers that the temperature - rescaling in `rescale_temp` is done - - Example: - # rescale temperature to 500K and 1000K at the 100th and 200th step - rescale_temp = [500, 1000] - rescale_steps = [100, 200] - """ - - # observers - for i, obs in enumerate(self.observers): - if obs[0].__class__.__name__ == "OTFLogger": - self.logger_ind = i - break - - # initialize gp by a dft calculation - calc = self.atoms.calc - calc.mgp_updated = False - - # restart from previous OTF training - if self.restart_from is not None: - self.restart() - f = self.atoms.calc.results['forces'] + raise NotImplementedError(md_engine+' is not implemented in ASE') - if not calc.gp_model.training_data: - self.dft_count = 0 - self.stds = np.zeros((self.noa, 3)) - dft_forces = self.call_DFT() - f = dft_forces + self.md = MD(atoms=atoms, timestep=timestep, trajectory=trajectory, + **md_kwargs) - # update gp model - curr_struc = Structure.from_ase_atoms(self.atoms) - self.l_bound = get_l_bound(100, curr_struc, self.two_d) - print('l_bound:', self.l_bound) + self.atoms = atoms + force_source = dft_source + self.flare_calc = self.atoms.calc - calc.gp_model.update_db(curr_struc, dft_forces, - custom_range=self.init_atoms) + super().__init__( + dt=timestep, number_of_steps=number_of_steps, + gp=self.flare_calc.gp_model, force_source=force_source, + dft_loc=dft_calc, dft_input=self.atoms, **otf_kwargs) - # train calculator - for atom in self.init_atoms: - # the observers[0][0] is the logger - self.observers[self.logger_ind][0].add_atom_info(atom, - self.stds[atom]) - self.train() - - if self.use_mapping: - self.build_mgp() - - self.observers[self.logger_ind][0].write_wall_time() + def initialize_train(self): + super().initialize_train() if self.md_engine == 'NPT': - if not self.initialized: - self.initialize() + if not self.md.initialized: + self.md.initialize() else: - if self.have_the_atoms_been_changed(): + if self.md.have_the_atoms_been_changed(): raise NotImplementedError( "You have modified the atoms since the last timestep.") - step_0 = self.nsteps - for i in range(step_0, steps): - print('step:', i) - - calc.results = {} # clear the calculation from last step - self.stds = np.zeros((self.noa, 3)) - - # temperature rescaling - if self.nsteps in rescale_steps: - temp = rescale_temp[rescale_steps.index(self.nsteps)] - curr_velocities = self.atoms.get_velocities() - curr_temp = self.atoms.get_temperature() - self.atoms.set_velocities(curr_velocities *\ - np.sqrt(temp/curr_temp)) - - if self.md_engine == 'NPT': - self.step() - else: - f = self.step(f) - self.nsteps += 1 - self.stds = self.atoms.get_uncertainties(self.atoms) - - # figure out if std above the threshold - self.call_observers() - curr_struc = Structure.from_ase_atoms(self.atoms) - self.l_bound = get_l_bound(self.l_bound, curr_struc, self.two_d) - print('l_bound:', self.l_bound) - curr_struc.stds = np.copy(self.stds) - noise = calc.gp_model.hyps[-1] - self.std_in_bound, self.target_atoms = is_std_in_bound(\ - self.std_tolerance_factor, noise, curr_struc, self.max_atoms_added) - - print('std in bound:', self.std_in_bound, self.target_atoms) - - if not self.std_in_bound: - # call dft/eam - print('calling dft') - dft_forces = self.call_DFT() - - # update gp - print('updating gp') - self.update_GP(dft_forces) - calc.mgp_updated = False - - if self.use_mapping: - self.build_mgp() - - self.observers[self.logger_ind][0].run_complete() - - - def build_mgp(self): - # build mgp - calc = self.atoms.calc - if self.nsteps in self.non_mapping_steps: - calc.use_mapping = False - skip = True - else: - calc.use_mapping = True - - if calc.mgp_updated: - skip = True - else: - skip = False - calc.mgp_updated = True - - calc.build_mgp(skip) - - - def call_DFT(self): - self.dft_calc.nsteps = self.nsteps - prev_calc = self.atoms.calc - calc = deepcopy(self.dft_calc) - self.atoms.set_calculator(calc) - forces = self.atoms.get_forces() - self.call_observers() - self.atoms.set_calculator(prev_calc) - self.dft_count += 1 - return forces - - def update_GP(self, dft_forces): - atom_count = 0 - atom_list = [] - gp_model = self.atoms.calc.gp_model - - # build gp structure from atoms - atom_struc = Structure.from_ase_atoms(self.atoms) - - while (not self.std_in_bound and atom_count < - np.min([self.max_atoms_added, len(self.target_atoms)])): - - target_atom = self.target_atoms[atom_count] - - # update gp model - gp_model.update_db(atom_struc, dft_forces, - custom_range=[target_atom]) - - if gp_model.alpha is None: - gp_model.set_L_alpha() - else: - gp_model.update_L_alpha() - - # atom_list.append(target_atom) - ## force calculation needed before get_uncertainties - # forces = self.atoms.calc.get_forces_gp(self.atoms) - # self.stds = self.atoms.get_uncertainties() - - # write added atom to the log file, - # refer to ase.optimize.optimize.Dynamics - self.observers[self.logger_ind][0].add_atom_info(target_atom, - self.stds[target_atom]) - - #self.is_std_in_bound(atom_list) - atom_count += 1 - - self.train() - self.observers[self.logger_ind][0].added_atoms_dat.write('\n') - self.observers[self.logger_ind][0].write_wall_time() - - def train(self, output=None, skip=False): - calc = self.atoms.calc - if (self.dft_count-1) < self.freeze_hyps: - #TODO: add other args to train() - calc.gp_model.train(output=output) - self.observers[self.logger_ind][0].write_hyps(calc.gp_model.hyp_labels, - calc.gp_model.hyps, calc.gp_model.likelihood, - calc.gp_model.likelihood_gradient) - else: - #TODO: change to update_L_alpha() - calc.gp_model.set_L_alpha() - - # save gp_model everytime after training - calc.gp_model.write_model('otf_data/gp_model', format='pickle') - - def restart(self): - # Recover atomic configuration: positions, velocities, forces - positions, self.nsteps = self.read_frame('positions.xyz', -1) - self.atoms.set_positions(positions) - self.atoms.set_velocities(self.read_frame('velocities.dat', -1)[0]) - self.atoms.calc.results['forces'] = self.read_frame('forces.dat', -1)[0] - print('Last frame recovered') - -# # Recover training data set -# gp_model = self.atoms.calc.gp_model -# atoms = deepcopy(self.atoms) -# nat = len(self.atoms.positions) -# dft_positions = self.read_all_frames('dft_positions.xyz', nat) -# dft_forces = self.read_all_frames('dft_forces.dat', nat) -# added_atoms = self.read_all_frames('added_atoms.dat', 1, 1, 'int') -# for i, frame in enumerate(dft_positions): -# atoms.set_positions(frame) -# curr_struc = Structure.from_ase_atoms(atoms) -# gp_model.update_db(curr_struc, dft_forces[i], added_atoms[i].tolist()) -# gp_model.set_L_alpha() -# print('GP training set ready') - - # Recover FLARE calculator - self.atoms.calc.gp_model = GaussianProcess.from_file(self.restart_from+'/gp_model.pickle') -# gp_model.ky_mat_inv = np.load(self.restart_from+'/ky_mat_inv.npy') -# gp_model.alpha = np.load(self.restart_from+'/alpha.npy') - if self.atoms.calc.use_mapping: - for map_3 in self.atoms.calc.mgp_model.maps_3: - map_3.load_grid = self.restart_from + '/' - self.atoms.calc.build_mgp(skip=False) - self.atoms.calc.mgp_updated = True - print('GP and MGP ready') - - self.l_bound = 10 - - def read_all_frames(self, filename, nat, header=2, elem_type='xyz'): - frames = [] - with open(self.restart_from+'/'+filename) as f: - lines = f.readlines() - frame_num = len(lines) // (nat+header) - for i in range(frame_num): - start = (nat+header) * i + header - curr_frame = lines[start:start+nat] - properties = [] - for line in curr_frame: - line = line.split() - if elem_type == 'xyz': - xyz = [float(l) for l in line[1:]] - properties.append(xyz) - elif elem_type == 'int': - properties = [int(l) for l in line] - frames.append(properties) - return np.array(frames) - - - def read_frame(self, filename, frame_num): - nat = len(self.atoms.positions) - with open(self.restart_from+'/'+filename) as f: - lines = f.readlines() - if frame_num == -1: # read the last frame - start_line = - (nat+2) - frame = lines[start_line:] - else: - start_line = frame_num * (nat+2) - end_line = (frame_num+1) * (nat+2) - frame = f.lines[start_line:end_line] - - properties = [] - for line in frame[2:]: - line = line.split() - properties.append([float(d) for d in line[1:]]) - return np.array(properties), len(lines)//(nat+2) - - - + def compute_properties(self): + ''' + compute forces and stds with FLARE_Calculator + ''' + if not isinstance(self.atoms.calc, FLARE_Calculator): + self.atoms.set_calculator(self.flare_calc) + + self.atoms.calc.results = {} + f = self.atoms.get_forces(self.atoms) + stds = self.atoms.get_uncertainties(self.atoms) + self.structure.forces = deepcopy(f) + self.structure.stds = deepcopy(stds) + + def md_step(self): + ''' + Get new position in molecular dynamics based on the forces predicted by + FLARE_Calculator or DFT calculator + ''' + self.md.step() + + # Return a copy so that future updates to atoms.positions doesn't also + # update structure.positions. + return np.copy(self.atoms.positions) + + # TODO: fix the temperature output in the log file + + def update_positions(self, new_pos): + # call OTF method + super().update_positions(new_pos) + + # update ASE atoms + if self.curr_step in self.rescale_steps: + rescale_ind = self.rescale_steps.index(self.curr_step) + temp_fac = self.rescale_temps[rescale_ind] / self.temperature + vel_fac = np.sqrt(temp_fac) + curr_velocities = self.atoms.get_velocities() + self.atoms.set_velocities(curr_velocities * vel_fac) + + def update_gp(self, train_atoms, dft_frcs): + + super().update_gp(train_atoms, dft_frcs) + + if self.flare_calc.use_mapping: + self.flare_calc.mgp_model.build_map(self.flare_calc.gp_model) diff --git a/flare/ase/otf_md.py b/flare/ase/otf_md.py deleted file mode 100644 index 201d9dae3..000000000 --- a/flare/ase/otf_md.py +++ /dev/null @@ -1,198 +0,0 @@ -''' -This module provides OTF training with ASE MD engines: VerlocityVerlet, NVTBerendsen, NPTBerendsen, NPT and Langevin. -Please see the function `otf_md` below for usage -''' -import os -import sys -from flare.struc import Structure -from flare.ase.otf import OTF - -import numpy as np -from ase.calculators.espresso import Espresso -from ase.calculators.eam import EAM -from ase.md.npt import NPT -from ase.md.nvtberendsen import NVTBerendsen -from ase.md.nptberendsen import NPTBerendsen -from ase.md.verlet import VelocityVerlet -from ase.md.md import MolecularDynamics -from ase.md.langevin import Langevin -from ase import units - -class OTF_VelocityVerlet(VelocityVerlet, OTF): - """ - On-the-fly training with ASE's VelocityVerlet molecular dynamics engine. - Inherit from ASE `VelocityVerlet `_ class and our ASE-coupled on-the-fly training engine `flare.ase.OTF` - - Args: - atoms, timestep, trajectory, dt: - see `VelocityVerlet `_ - kwargs: same parameters as :class:`flare.ase.OTF` - """ - - def __init__(self, atoms, timestep, trajectory=None, dt=None, - **kwargs): - - VelocityVerlet.__init__(self, atoms=atoms, timestep=timestep, - trajectory=trajectory, dt=dt) - - OTF.__init__(self, **kwargs) - self.md_engine = 'VelocityVerlet' - -class OTF_NVTBerendsen(NVTBerendsen, OTF): - """ - On-the-fly training with ASE's NVTBerendsen molecular dynamics engine. \ - Inherit from ASE `NVTBerendsen `_ class and our ASE-coupled on-the-fly training engine `flare.ase.OTF` - - Args: - atoms, timestep, temperature, taut, fixcm: see\ - `NVTBerendsen `_. - kwargs: same parameters as :class:`flare.ase.OTF` - """ - - - def __init__(self, atoms, timestep, temperature, taut, fixcm=True, - trajectory=None, **kwargs): - - NVTBerendsen.__init__(self, atoms, timestep, temperature, taut, - fixcm, trajectory) - - OTF.__init__(self, **kwargs) - - self.md_engine = 'NVTBerendsen' - -class OTF_NPTBerendsen(NPTBerendsen, OTF): - """ - On-the-fly training with ASE's Langevin molecular dynamics engine. \ - Inherit from ASE `Langevin `_ class and our ASE-coupled on-the-fly training engine `flare.ase.OTF` - - Args: - atoms, timestep, temperature, taut, pressure, taup, compressibility, fixcm:\ - see `NPTBerendsen `_. - kwargs: same parameters as :class:`flare.ase.OTF` - """ - - - def __init__(self, atoms, timestep, temperature, taut=0.5e3, - pressure=1.01325, taup=1e3, - compressibility=4.57e-5, fixcm=True, trajectory=None, - **kwargs): - - NPTBerendsen.__init__(self, atoms, timestep, temperature, taut, - pressure, taup, - compressibility, fixcm, trajectory) - - OTF.__init__(self, **kwargs) - - self.md_engine = 'NPTBerendsen' - -class OTF_NPT(NPT, OTF): - """ - On-the-fly training with ASE's Langevin molecular dynamics engine. \ - Inherit from ASE `NPT `_ class and our ASE-coupled on-the-fly training engine `flare.ase.OTF` - - Args: - atoms, timestep, temperature, friction:\ - see `NPT `_ - kwargs: same parameters as :class:`flare.ase.OTF` - """ - - - def __init__(self, atoms, timestep, temperature, externalstress, - ttime, pfactor, mask=None, trajectory=None, **kwargs): - - NPT.__init__(self, atoms, timestep, temperature, - externalstress, ttime, pfactor, mask, - trajectory) - - OTF.__init__(self, **kwargs) - - self.md_engine = 'NPT' - -class OTF_Langevin(Langevin, OTF): - """ - On-the-fly training with ASE's Langevin molecular dynamics engine. \ - Inherit from ASE `Langevin `_ class and our ASE-coupled on-the-fly training engine `flare.ase.OTF` - - Args: - atoms, timestep, temperature, friction:\ - see `Langevin `_. - kwargs: same parameters as :class:`flare.ase.OTF` - """ - - def __init__(self, atoms, timestep=None, temperature=None, friction=None, - fixcm=True, trajectory=None, **kwargs): - - Langevin.__init__(self, atoms, timestep, temperature, friction, - fixcm, trajectory) - - OTF.__init__(self, **kwargs) - - self.md_engine = 'Langevin' - - - -def otf_md(md_engine: str, atoms, md_params: dict, otf_params: dict): - ''' - Create an OTF MD engine - - Args: - md_engine (str): the name of md engine, including `VelocityVerlet`, - `NVTBerendsen`, `NPTBerendsen`, `NPT`, `Langevin` - atoms (Atoms): ASE Atoms to apply this md engine - md_params (dict): parameters used in MD engines, - must include: `timestep`, `trajectory` (usually set to None). - Also include those parameters required for ASE MD engine, - please look at ASE website to find out parameters for different engines - otf_params (dict): parameters used in OTF module - - Return: - An OTF MD class object - - Example: - >>> from ase import units - >>> from ase.spacegroup import crystal - >>> super_cell = crystal(['Ag', 'I'], - basis=[(0, 0, 0), (0.5, 0.5, 0.5)], - size=(2, 1, 1), - cellpar=[3.85, 3.85, 3.85, 90, 90, 90]) - >>> md_engine = 'VelocityVerlet' - >>> md_params = {'timestep': 1 * units.fs, 'trajectory': None, - 'dt': None} - >>> otf_params = {'dft_calc': dft_calc, - 'init_atoms': [0], - 'std_tolerance_factor': 1, - 'max_atoms_added' : len(super_cell.positions), - 'freeze_hyps': 10, - 'use_mapping': False} - >>> test_otf = otf_md(md_engine, super_cell, md_params, otf_params) - ''' - - md = md_params - timestep = md['timestep'] - trajectory = md['trajectory'] - - if md_engine == 'VelocityVerlet': - return OTF_VelocityVerlet(atoms, timestep, trajectory=trajectory, - dt=md['dt'], **otf_params) - - elif md_engine == 'NVTBerendsen': - return OTF_NVTBerendsen(atoms, timestep, md['temperature'], - md['taut'], md['fixcm'], trajectory, **otf_params) - - elif md_engine == 'NPTBerendsen': - return OTF_NPTBerendsen(atoms, timestep, md['temperature'], - md['taut'], md['pressure'], md['taup'], - md['compressibility'], md['fixcm'], trajectory, **otf_params) - - elif md_engine == 'NPT': - return OTF_NPT(atoms, timestep, md['temperature'], - md['externalstress'], md['ttime'], md['pfactor'], - md['mask'], trajectory, **otf_params) - - elif md_engine == 'Langevin': - return OTF_Langevin(atoms, timestep, md['temperature'], - md['friction'], md['fixcm'], trajectory, **otf_params) - - else: - raise NotImplementedError(md_engine+' is not implemented') - diff --git a/flare/dft_interface/vasp_util.py b/flare/dft_interface/vasp_util.py index b6ffa368c..a788abcde 100644 --- a/flare/dft_interface/vasp_util.py +++ b/flare/dft_interface/vasp_util.py @@ -1,15 +1,17 @@ import numpy as np -import os, shutil +import os +import shutil + +from json import dump, load +from subprocess import call +from typing import List, Union from pymatgen.io.vasp.inputs import Poscar from pymatgen.io.vasp.outputs import Vasprun from pymatgen.io.vasp.sets import VaspInputSet from pymatgen.core.periodic_table import Element -from subprocess import call from flare.struc import Structure -from typing import List, Union -from json import dump, load from flare.utils.element_coder import NumpyEncoder name="VASP" @@ -86,9 +88,10 @@ def run_dft(calc_dir: str, dft_loc: str, try: forces = parse_func("vasprun.xml") except FileNotFoundError: - raise FileNotFoundError("""Could not load vasprun.xml. - The calculation may not have finished. - Current directory is %s""" % os.getcwd()) + os.chdir(currdir) + raise FileNotFoundError("Could not load vasprun.xml."\ + "The calculation may not have finished." + f"Current directory is {os.getcwd()}") os.chdir(currdir) return forces @@ -103,11 +106,12 @@ def run_dft_par(dft_input: str, structure: Structure, dft_out="vasprun.xml", parallel_prefix="mpi", mpi = None, npool = None, + screen_out='vasp.out', **dft_kwargs): # TODO Incorporate Custodian. edit_dft_input_positions(dft_input, structure) - if dft_command is None and not os.environ.get('VASP_COMMAND'): + if dft_command is None or not os.environ.get('VASP_COMMAND'): raise FileNotFoundError\ ("Warning: No VASP Command passed, or stored in " "environment as VASP_COMMAND. ") @@ -123,7 +127,8 @@ def run_dft_par(dft_input: str, structure: Structure, serial_prefix = dft_kwargs.get('serial_prefix', '') dft_command = f'{serial_prefix} {dft_command}' - call(dft_command, shell=True) + with open(screen_out, "w+") as fout: + call(dft_command.split(), stdout=fout) return parse_dft_forces(dft_out) diff --git a/flare/env.py b/flare/env.py index 2558cd124..2f0de0ce0 100644 --- a/flare/env.py +++ b/flare/env.py @@ -2,12 +2,13 @@ environment of an atom. :class:`AtomicEnvironment` objects are inputs to the 2-, 3-, and 2+3-body kernels.""" import numpy as np -from math import sqrt, ceil -from numba import njit +from copy import deepcopy +from math import ceil from flare.struc import Structure -from flare.utils.mask_helper import HyperParameterMasking -from flare.kernels.kernels import coordination_number, q_value_mc +from flare.parameters import Parameters import flare.kernels.cutoffs as cf +from flare.utils.env_getarray import get_2_body_arrays, get_3_body_arrays,\ + get_m2_body_arrays, get_m3_body_arrays class AtomicEnvironment: """Contains information about the local environment of an atom, @@ -35,16 +36,16 @@ class AtomicEnvironment: Li to group 1, and Be to group 0 (the 0th register is ignored). * nspecie: Integer, number of different species groups (equal to number of unique values in specie_mask). - * nbond: Integer, number of different hyperparameter/cutoff sets to associate with + * ntwobody: Integer, number of different hyperparameter/cutoff sets to associate with different 2-body pairings of atoms in groups defined in specie_mask. - * bond_mask: Array of length nspecie^2, which describes the cutoff to + * twobody_mask: Array of length nspecie^2, which describes the cutoff to associate with different pairings of species types. For example, if there - are atoms of type 0 and 1, then bond_mask defines which cutoff + are atoms of type 0 and 1, then twobody_mask defines which cutoff to use for parings [0-0, 0-1, 1-0, 1-1]: if we wanted cutoff0 for 0-0 parings and set 1 for 0-1 and 1-1 pairings, then we would make - bond_mask [0, 1, 1, 1]. - * cutoff_2b: Array of length nbond, which stores the cutoff used for different - types of bonds defined in bond_mask + twobody_mask [0, 1, 1, 1]. + * twobody_cutoff_list: Array of length ntwobody, which stores the cutoff used for different + types of bonds defined in twobody_mask * ncut3b: Integer, number of different cutoffs sets to associate with different 3-body pariings of atoms in groups defined in specie_mask. * cut3b_mask: Array of length nspecie^2, which describes the cutoff to @@ -53,102 +54,163 @@ class AtomicEnvironment: If C and O are associate with atom group 1 in specie_mask and H are associate with group 0 in specie_mask, the cut3b_mask[1*nspecie+0] determines the C/O-H bond cutoff, and cut3b_mask[1*nspecie+1] determines the C-O bond cutoff. If we want the - former one to use the 1st cutoff in cutoff_3b and the later to use the 2nd cutoff - in cutoff_3b, the cut3b_mask should be [0, 0, 0, 1] - * cutoff_3b: Array of length ncut3b, which stores the cutoff used for different + former one to use the 1st cutoff in threebody_cutoff_list and the later to use the 2nd cutoff + in threebody_cutoff_list, the cut3b_mask should be [0, 0, 0, 1] + * threebody_cutoff_list: Array of length ncut3b, which stores the cutoff used for different types of bonds in triplets. - * nmb : Integer, number of different cutoffs set to associate with different coordination + * nmanybody : Integer, number of different cutoffs set to associate with different coordination numbers - * mb_mask: similar to bond_mask and cut3b_mask. - * cutoff_mb: Array of length nmb, stores the cutoff used for different many body terms + * manybody_mask: similar to twobody_mask and cut3b_mask. + * manybody_cutoff_list: Array of length nmanybody, stores the cutoff used for different many body terms Examples can be found at the end of in tests/test_env.py """ + all_kernel_types = ['twobody', 'threebody', 'manybody'] + ndim = {'twobody': 2, 'threebody': 3, 'manybody': 2, 'cut3b': 2} + def __init__(self, structure: Structure, atom: int, cutoffs, cutoffs_mask=None): + self.structure = structure self.positions = structure.wrapped_positions self.cell = structure.cell self.species = structure.coded_species + # backward compatability + if not isinstance(cutoffs, dict): + newcutoffs = {'twobody':cutoffs[0]} + if len(cutoffs)>1: + newcutoffs['threebody'] = cutoffs[1] + if len(cutoffs)>2: + newcutoffs['manybody'] = cutoffs[2] + cutoffs = newcutoffs + + if cutoffs_mask is None: + cutoffs_mask = {'cutoffs': cutoffs} + elif cutoffs is not None: + cutoffs_mask['cutoffs'] = deepcopy(cutoffs) + # Set the sweep array based on the max cutoff. - sweep_val = ceil(np.max(cutoffs) / structure.max_cutoff) + sweep_val = ceil(np.max(list(cutoffs.values())) / structure.max_cutoff) self.sweep_val = sweep_val self.sweep_array = np.arange(-sweep_val, sweep_val + 1, 1) self.atom = atom self.ctype = structure.coded_species[atom] - self.cutoffs = np.copy(cutoffs) - self.cutoffs_mask = cutoffs_mask - self.setup_mask() + self.twobody_cutoff = 0 + self.threebody_cutoff = 0 + self.manybody_cutoff = 0 + + self.ntwobody = 1 + self.ncut3b = 0 + self.nmanybody = 0 + + self.nspecie = 1 + self.specie_mask = None + self.twobody_mask = None + self.threebody_mask = None + self.manybody_mask = None + self.twobody_cutoff_list = None + self.threebody_cutoff_list = None + self.manybody_cutoff_list = None + + self.setup_mask(cutoffs_mask) - assert self.scalar_cutoff_3 <= self.scalar_cutoff_2, \ + assert self.threebody_cutoff <= self.twobody_cutoff, \ "2b cutoff has to be larger than 3b cutoff" # # TO DO, once the mb function is updated to use the bond_array_2 # # this block should be activated. - # assert self.scalar_cutoff_mb <= self.scalar_cutoff_2, \ + # assert self.manybody_cutoff <= self.twobody_cutoff, \ # "mb cutoff has to be larger than mb cutoff" - self.compute_env() + self.bond_array_2 = None + self.etypes = None + self.bond_inds = None + self.bond_array_3 = None + self.cross_bond_inds = None + self.cross_bond_dists = None + self.triplet_counts = None + self.q_array = None + self.q_neigh_array = None + self.q_grads = None + self.q_neigh_grads = None + self.unique_species = None + self.etypes_mb = None - def setup_mask(self): + self.compute_env() - self.scalar_cutoff_2, self.scalar_cutoff_3, self.scalar_cutoff_mb, \ - self.cutoff_2b, self.cutoff_3b, self.cutoff_mb, \ - self.nspecie, self.n2b, self.n3b, self.nmb, self.specie_mask, \ - self.bond_mask, self.cut3b_mask, self.mb_mask = \ - HyperParameterMasking.mask2cutoff(self.cutoffs, self.cutoffs_mask) + def setup_mask(self, cutoffs_mask): + + self.cutoffs_mask = deepcopy(cutoffs_mask) + self.cutoffs = cutoffs_mask['cutoffs'] + + for kernel in AtomicEnvironment.all_kernel_types: + ndim = AtomicEnvironment.ndim[kernel] + if kernel in self.cutoffs: + setattr(self, kernel+'_cutoff', self.cutoffs[kernel]) + + if (self.twobody_cutoff == 0): + self.twobody_cutoff = np.max([self.threebody_cutoff, self.manybody_cutoff]) + self.cutoffs['twobody'] = self.twobody_cutoff + + self.nspecie = cutoffs_mask.get('nspecie', 1) + if 'specie_mask' in cutoffs_mask: + self.specie_mask = np.array(cutoffs_mask['specie_mask'], dtype=np.int) + + for kernel in AtomicEnvironment.all_kernel_types: + ndim = AtomicEnvironment.ndim[kernel] + if kernel in self.cutoffs: + setattr(self, kernel+'_cutoff', self.cutoffs[kernel]) + setattr(self, 'n'+kernel, 1) + if kernel != 'threebody': + name_list = [kernel+'_cutoff_list', + 'n'+kernel, kernel+'_mask'] + for name in name_list: + if name in cutoffs_mask: + setattr(self, name, cutoffs_mask[name]) + else: + self.ncut3b = cutoffs_mask.get('ncut3b', 1) + self.cut3b_mask = cutoffs_mask.get('cut3b_mask', None) + if 'threebody_cutoff_list' in cutoffs_mask: + self.threebody_cutoff_list = np.array(cutoffs_mask['threebody_cutoff_list'], dtype=np.float) def compute_env(self): # get 2-body arrays - if (self.n2b > 1): - bond_array_2, bond_positions_2, etypes, bond_inds = \ - get_2_body_arrays_sepcut(self.positions, self.atom, self.cell, - self.cutoff_2b, self.species, self.sweep_array, - self.nspecie, self.specie_mask, self.bond_mask) - else: + if (self.ntwobody >= 1): bond_array_2, bond_positions_2, etypes, bond_inds = \ - get_2_body_arrays(self.positions, self.atom, self.cell, - self.scalar_cutoff_2, self.species, self.sweep_array) + get_2_body_arrays(self.positions, self.atom, self.cell, self.twobody_cutoff, + self.twobody_cutoff_list, self.species, self.sweep_array, + self.nspecie, self.specie_mask, self.twobody_mask) - self.bond_array_2 = bond_array_2 - self.etypes = etypes - self.bond_inds = bond_inds + self.bond_array_2 = bond_array_2 + self.etypes = etypes + self.bond_inds = bond_inds # if 2 cutoffs are given, create 3-body arrays - if self.scalar_cutoff_3 > 0: - if (self.n3b > 1): - bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts = \ - get_3_body_arrays_sepcut(bond_array_2, bond_positions_2, - self.species[self.atom], etypes, self.cutoff_3b, - self.nspecie, self.specie_mask, self.cut3b_mask) - else: - bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts = \ - get_3_body_arrays( - bond_array_2, bond_positions_2, self.scalar_cutoff_3) - + if self.ncut3b > 0: + bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts = \ + get_3_body_arrays(bond_array_2, bond_positions_2, + self.species[self.atom], etypes, self.threebody_cutoff, + self.threebody_cutoff_list, + self.nspecie, self.specie_mask, self.cut3b_mask) self.bond_array_3 = bond_array_3 self.cross_bond_inds = cross_bond_inds self.cross_bond_dists = cross_bond_dists self.triplet_counts = triplet_counts # if 3 cutoffs are given, create many-body arrays - if self.scalar_cutoff_mb > 0: - if (self.nmb > 1): - self.q_array, self.q_neigh_array, self.q_neigh_grads, \ - self.unique_species, self.etypes_mb = \ - get_m_body_arrays_sepcut(self.positions, self.atom, self.cell, \ - self.cutoff_mb, self.species, self.sweep_array, self.nspecie, self.specie_mask, self.mb_mask,\ - cf.quadratic_cutoff) - else: - self.q_array, self.q_neigh_array, self.q_neigh_grads, \ - self.unique_species, self.etypes_mb = get_m_body_arrays(\ - self.positions, self.atom, self.cell, \ - self.scalar_cutoff_mb, self.species, self.sweep_array, cf.quadratic_cutoff) + if self.nmanybody > 0: + self.q_array, self.q_neigh_array, self.q_grads, self.q_neigh_grads, \ + self.unique_species, self.etypes_mb = \ + get_m2_body_arrays(self.positions, self.atom, self.cell, + self.manybody_cutoff, self.manybody_cutoff_list, + self.species, self.sweep_array, + self.nspecie, self.specie_mask, + self.manybody_mask, cf.quadratic_cutoff) def as_dict(self): """ @@ -189,10 +251,6 @@ def from_dict(dictionary): cutoffs = dictionary['cutoffs'] else: cutoffs = {} - for cutoff_type in ['2','3','mb']: - key = 'scalar_cutoff_'+cutoff_type - if (key in dictionary): - cutoffs[key] = dictionary[key] cutoffs_mask = dictionary.get('cutoffs_mask', None) @@ -207,511 +265,3 @@ def __str__(self): sorted(list(set(neighbor_types)))) return string - - -@njit -def get_2_body_arrays(positions: np.ndarray, atom: int, cell: np.ndarray, - cutoff_2: float, species: np.ndarray, sweep: np.ndarray): - """Returns distances, coordinates, and species of atoms in the 2-body - local environment. This method is implemented outside the AtomicEnvironment - class to allow for njit acceleration with Numba. - - :param positions: Positions of atoms in the structure. - :type positions: np.ndarray - :param atom: Index of the central atom of the local environment. - :type atom: int - :param cell: 3x3 array whose rows are the Bravais lattice vectors of the - cell. - :type cell: np.ndarray - :param cutoff_2: 2-body cutoff radius. - :type cutoff_2: float - :param species: Numpy array of species represented by their atomic numbers. - :type species: np.ndarray - :return: Tuple of arrays describing pairs of atoms in the 2-body local - environment. - - bond_array_2: Array containing the distances and relative - coordinates of atoms in the 2-body local environment. First column - contains distances, remaining columns contain Cartesian coordinates - divided by the distance (with the origin defined as the position of the - central atom). The rows are sorted by distance from the central atom. - - bond_positions_2: Coordinates of atoms in the 2-body local environment. - - etypes: Species of atoms in the 2-body local environment represented by - their atomic number. - - bond_indices: Structure indices of atoms in the local environment. - - :rtype: np.ndarray, np.ndarray, np.ndarray, np.ndarray - """ - - noa = len(positions) - pos_atom = positions[atom] - super_count = sweep.shape[0]**3 - coords = np.zeros((noa, 3, super_count)) - dists = np.zeros((noa, super_count)) - cutoff_count = 0 - - vec1 = cell[0] - vec2 = cell[1] - vec3 = cell[2] - - # record distances and positions of images - for n in range(noa): - diff_curr = positions[n] - pos_atom - im_count = 0 - for s1 in sweep: - for s2 in sweep: - for s3 in sweep: - im = diff_curr + s1 * vec1 + s2 * vec2 + s3 * vec3 - dist = sqrt(im[0] * im[0] + im[1] * im[1] - + im[2] * im[2]) - if (dist < cutoff_2) and (dist != 0): - dists[n, im_count] = dist - coords[n, :, im_count] = im - cutoff_count += 1 - im_count += 1 - - # create 2-body bond array - bond_indices = np.zeros(cutoff_count, dtype=np.int8) - bond_array_2 = np.zeros((cutoff_count, 4), dtype=np.float64) - bond_positions_2 = np.zeros((cutoff_count, 3), dtype=np.float64) - etypes = np.zeros(cutoff_count, dtype=np.int8) - bond_count = 0 - - for m in range(noa): - spec_curr = species[m] - for n in range(super_count): - dist_curr = dists[m, n] - if (dist_curr < cutoff_2) and (dist_curr != 0): - coord = coords[m, :, n] - bond_array_2[bond_count, 0] = dist_curr - bond_array_2[bond_count, 1:4] = coord / dist_curr - bond_positions_2[bond_count, :] = coord - etypes[bond_count] = spec_curr - bond_indices[bond_count] = m - bond_count += 1 - - # sort by distance - sort_inds = bond_array_2[:, 0].argsort() - bond_array_2 = bond_array_2[sort_inds] - bond_positions_2 = bond_positions_2[sort_inds] - bond_indices = bond_indices[sort_inds] - etypes = etypes[sort_inds] - - return bond_array_2, bond_positions_2, etypes, bond_indices - - -@njit -def get_3_body_arrays(bond_array_2, bond_positions_2, cutoff_3: float): - """Returns distances and coordinates of triplets of atoms in the - 3-body local environment. - - :param bond_array_2: 2-body bond array. - :type bond_array_2: np.ndarray - :param bond_positions_2: Coordinates of atoms in the 2-body local - environment. - :type bond_positions_2: np.ndarray - :param cutoff_3: 3-body cutoff radius. - :type cutoff_3: float - :return: Tuple of 4 arrays describing triplets of atoms in the 3-body local - environment. - - bond_array_3: Array containing the distances and relative - coordinates of atoms in the 3-body local environment. First column - contains distances, remaining columns contain Cartesian coordinates - divided by the distance (with the origin defined as the position of the - central atom). The rows are sorted by distance from the central atom. - - cross_bond_inds: Two dimensional array whose row m contains the indices - of atoms n > m that are within a distance cutoff_3 of both atom n and the - central atom. - - cross_bond_dists: Two dimensional array whose row m contains the - distances from atom m of atoms n > m that are within a distance cutoff_3 - of both atom n and the central atom. - - triplet_counts: One dimensional array of integers whose entry m is the - number of atoms that are within a distance cutoff_3 of atom m. - - :rtype: (np.ndarray, np.ndarray, np.ndarray, np.ndarray) - """ - - # get 3-body bond array - ind_3 = -1 - noa = bond_array_2.shape[0] - for count, dist in enumerate(bond_array_2[:, 0]): - if dist > cutoff_3: - ind_3 = count - break - if ind_3 == -1: - ind_3 = noa - - bond_array_3 = bond_array_2[0:ind_3, :] - bond_positions_3 = bond_positions_2[0:ind_3, :] - - # get cross bond array - cross_bond_inds = np.zeros((ind_3, ind_3), dtype=np.int8) - 1 - cross_bond_dists = np.zeros((ind_3, ind_3)) - triplet_counts = np.zeros(ind_3, dtype=np.int8) - for m in range(ind_3): - pos1 = bond_positions_3[m] - count = m + 1 - trips = 0 - for n in range(m + 1, ind_3): - pos2 = bond_positions_3[n] - diff = pos2 - pos1 - dist_curr = sqrt( - diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2]) - - if dist_curr < cutoff_3: - cross_bond_inds[m, count] = n - cross_bond_dists[m, count] = dist_curr - count += 1 - trips += 1 - triplet_counts[m] = trips - - return bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts - - -@njit -def get_2_body_arrays_sepcut(positions, atom: int, cell, cutoff_2, species, sweep, - nspecie, specie_mask, bond_mask): - """Returns distances, coordinates, species of atoms, and indices of neighbors - in the 2-body local environment. This method is implemented outside - the AtomicEnvironment class to allow for njit acceleration with Numba. - - :param positions: Positions of atoms in the structure. - :type positions: np.ndarray - :param atom: Index of the central atom of the local environment. - :type atom: int - :param cell: 3x3 array whose rows are the Bravais lattice vectors of the - cell. - :type cell: np.ndarray - :param cutoff_2: 2-body cutoff radius. - :type cutoff_2: np.ndarray - :param species: Numpy array of species represented by their atomic numbers. - :type species: np.ndarray - :param nspecie: number of atom types to define bonds - :type: int - :param specie_mask: mapping from atomic number to atom types - :type: np.ndarray - :param bond_mask: mapping from the types of end atoms to bond types - :type: np.ndarray - :return: Tuple of arrays describing pairs of atoms in the 2-body local - environment. - - bond_array_2: Array containing the distances and relative - coordinates of atoms in the 2-body local environment. First column - contains distances, remaining columns contain Cartesian coordinates - divided by the distance (with the origin defined as the position of the - central atom). The rows are sorted by distance from the central atom. - - bond_positions_2: Coordinates of atoms in the 2-body local environment. - - etypes: Species of atoms in the 2-body local environment represented by - their atomic number. - - bond_indices: Structure indices of atoms in the local environment. - - :rtype: np.ndarray, np.ndarray, np.ndarray, np.ndarray - """ - noa = len(positions) - pos_atom = positions[atom] - coords = np.zeros((noa, 3, 27), dtype=np.float64) - dists = np.zeros((noa, 27), dtype=np.float64) - cutoff_count = 0 - - vec1 = cell[0] - vec2 = cell[1] - vec3 = cell[2] - - bc = specie_mask[species[atom]] - bcn = nspecie * bc - - # record distances and positions of images - for n in range(noa): - diff_curr = positions[n] - pos_atom - im_count = 0 - bn = specie_mask[species[n]] - rcut = cutoff_2[bond_mask[bn+bcn]] - - for s1 in sweep: - for s2 in sweep: - for s3 in sweep: - im = diff_curr + s1 * vec1 + s2 * vec2 + s3 * vec3 - dist = sqrt(im[0] * im[0] + im[1] * im[1] + im[2] * im[2]) - if (dist < rcut) and (dist != 0): - dists[n, im_count] = dist - coords[n, :, im_count] = im - cutoff_count += 1 - im_count += 1 - - # create 2-body bond array - bond_indices = np.zeros(cutoff_count, dtype=np.int8) - bond_array_2 = np.zeros((cutoff_count, 4), dtype=np.float64) - bond_positions_2 = np.zeros((cutoff_count, 3), dtype=np.float64) - etypes = np.zeros(cutoff_count, dtype=np.int8) - bond_count = 0 - - for m in range(noa): - spec_curr = species[m] - bm = specie_mask[species[m]] - rcut = cutoff_2[bond_mask[bm+bcn]] - for n in range(27): - dist_curr = dists[m, n] - if (dist_curr < rcut) and (dist_curr != 0): - coord = coords[m, :, n] - bond_array_2[bond_count, 0] = dist_curr - bond_array_2[bond_count, 1:4] = coord / dist_curr - bond_positions_2[bond_count, :] = coord - etypes[bond_count] = spec_curr - bond_indices[bond_count] = m - bond_count += 1 - - # sort by distance - sort_inds = bond_array_2[:, 0].argsort() - bond_array_2 = bond_array_2[sort_inds] - bond_positions_2 = bond_positions_2[sort_inds] - bond_indices = bond_indices[sort_inds] - etypes = etypes[sort_inds] - - return bond_array_2, bond_positions_2, etypes, bond_indices - - -@njit -def get_3_body_arrays_sepcut(bond_array_2, bond_positions_2, ctype, - etypes, cutoff_3, - nspecie, specie_mask, cut3b_mask): - """Returns distances and coordinates of triplets of atoms in the - 3-body local environment. - - :param bond_array_2: 2-body bond array. - :type bond_array_2: np.ndarray - :param bond_positions_2: Coordinates of atoms in the 2-body local - environment. - :type bond_positions_2: np.ndarray - :param ctype: atomic number of the center atom - :type: int - :param cutoff_3: 3-body cutoff radius. - :type cutoff_3: np.ndarray - :param nspecie: number of atom types to define bonds - :type: int - :param specie_mask: mapping from atomic number to atom types - :type: np.ndarray - :param cut3b_mask: mapping from the types of end atoms to bond types - :type: np.ndarray - :return: Tuple of 4 arrays describing triplets of atoms in the 3-body local - environment. - - bond_array_3: Array containing the distances and relative - coordinates of atoms in the 3-body local environment. First column - contains distances, remaining columns contain Cartesian coordinates - divided by the distance (with the origin defined as the position of the - central atom). The rows are sorted by distance from the central atom. - - cross_bond_inds: Two dimensional array whose row m contains the indices - of atoms n > m that are within a distance cutoff_3 of both atom n and the - central atom. - - cross_bond_dists: Two dimensional array whose row m contains the - distances from atom m of atoms n > m that are within a distance cutoff_3 - of both atom n and the central atom. - - triplet_counts: One dimensional array of integers whose entry m is the - number of atoms that are within a distance cutoff_3 of atom m. - - :rtype: (np.ndarray, np.ndarray, np.ndarray, np.ndarray) - """ - - bc = specie_mask[ctype] - bcn = nspecie * bc - - cut3 = np.max(cutoff_3) - - # get 3-body bond array - ind_3_l = np.where(bond_array_2[:, 0] > cut3)[0] - if (ind_3_l.shape[0] > 0): - ind_3 = ind_3_l[0] - else: - ind_3 = bond_array_2.shape[0] - - bond_array_3 = bond_array_2[0:ind_3, :] - bond_positions_3 = bond_positions_2[0:ind_3, :] - - # get cross bond array - cross_bond_inds = np.zeros((ind_3, ind_3), dtype=np.int8) - 1 - cross_bond_dists = np.zeros((ind_3, ind_3), dtype=np.float64) - triplet_counts = np.zeros(ind_3, dtype=np.int8) - for m in range(ind_3): - pos1 = bond_positions_3[m] - count = m + 1 - trips = 0 - - # choose bond dependent bond - bm = specie_mask[etypes[m]] - btype_m = cut3b_mask[bm + bcn] # (m, c) - cut_m = cutoff_3[btype_m] - bmn = nspecie * bm # for cross_dist usage - - for n in range(m + 1, ind_3): - - bn = specie_mask[etypes[n]] - btype_n = cut3b_mask[bn + bcn] # (n, c) - cut_n = cutoff_3[btype_n] - - # for cross_dist (m,n) pair - btype_mn = cut3b_mask[bn + bmn] - cut_mn = cutoff_3[btype_mn] - - pos2 = bond_positions_3[n] - diff = pos2 - pos1 - dist_curr = sqrt( - diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2]) - - if dist_curr < cut_mn \ - and bond_array_2[m, 0] < cut_m \ - and bond_array_2[n, 0] < cut_n: - cross_bond_inds[m, count] = n - cross_bond_dists[m, count] = dist_curr - count += 1 - trips += 1 - - triplet_counts[m] = trips - - return bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts - - -@njit -def get_m_body_arrays(positions, atom: int, cell, cutoff_mb: float, species, - sweep: np.ndarray, cutoff_func=cf.quadratic_cutoff): - # TODO: - # 1. need to deal with the conflict of cutoff functions if other funcs are used - # 2. complete the docs of "Return" - # TODO: this can be probably improved using stored arrays, redundant calls to get_2_body_arrays - # Get distances, positions, species and indices of neighbouring atoms - """ - Args: - positions (np.ndarray): Positions of atoms in the structure. - atom (int): Index of the central atom of the local environment. - cell (np.ndarray): 3x3 array whose rows are the Bravais lattice vectors of the - cell. - cutoff_mb (float): 2-body cutoff radius. - species (np.ndarray): Numpy array of species represented by their atomic numbers. - - Return: - Tuple of arrays describing pairs of atoms in the 2-body local - environment. - """ - # Get distances, positions, species and indexes of neighbouring atoms - bond_array_mb, __, etypes, bond_inds = get_2_body_arrays( - positions, atom, cell, cutoff_mb, species, sweep) - - species_list = np.array(list(set(species)), dtype=np.int8) - n_bonds = len(bond_inds) - n_specs = len(species_list) - qs = np.zeros(n_specs, dtype=np.float64) - qs_neigh = np.zeros((n_bonds, n_specs), dtype=np.float64) - q_grads = np.zeros((n_bonds, 3), dtype=np.float64) - - # get coordination number of center atom for each species - for s in range(n_specs): - qs[s] = q_value_mc(bond_array_mb[:, 0], cutoff_mb, species_list[s], - etypes, cutoff_func) - - # get coordination number of all neighbor atoms for each species - for i in range(n_bonds): - neigh_bond_array, _, neigh_etypes, _ = get_2_body_arrays(positions, - bond_inds[i], cell, cutoff_mb, species, sweep) - for s in range(n_specs): - qs_neigh[i, s] = q_value_mc(neigh_bond_array[:, 0], cutoff_mb, - species_list[s], neigh_etypes, cutoff_func) - - # get grad from each neighbor atom - for i in range(n_bonds): - ri = bond_array_mb[i, 0] - for d in range(3): - ci = bond_array_mb[i, d+1] - _, q_grads[i, d] = coordination_number(ri, ci, cutoff_mb, - cutoff_func) - - return qs, qs_neigh, q_grads, species_list, etypes - - -@njit -def get_m_body_arrays_sepcut(positions, atom: int, cell, cutoff_mb, - species, sweep: np.ndarray, nspec, spec_mask, mb_mask, - cutoff_func=cf.quadratic_cutoff): - # TODO: - # 1. need to deal with the conflict of cutoff functions if other funcs are used - # 2. complete the docs of "Return" - # TODO: this can be probably improved using stored arrays, redundant calls to get_2_body_arrays - # Get distances, positions, species and indices of neighbouring atoms - """ - Args: - positions (np.ndarray): Positions of atoms in the structure. - atom (int): Index of the central atom of the local environment. - cell (np.ndarray): 3x3 array whose rows are the Bravais lattice vectors of the - cell. - cutoff_mb (float): 2-body cutoff radius. - species (np.ndarray): Numpy array of species represented by their atomic numbers. - - Return: - Tuple of arrays describing pairs of atoms in the 2-body local - environment. - """ - # Get distances, positions, species and indexes of neighbouring atoms - bond_array_mb, __, etypes, bond_inds = get_2_body_arrays_sepcut( - positions, atom, cell, cutoff_mb, species, sweep, - nspec, spec_mask, mb_mask) - - bc = spec_mask[species[atom]] - bcn = bc * nspec - - species_list = np.array(list(set(species)), dtype=np.int8) - n_bonds = len(bond_inds) - n_specs = len(species_list) - qs = np.zeros(n_specs, dtype=np.float64) - qs_neigh = np.zeros((n_bonds, n_specs), dtype=np.float64) - q_grads = np.zeros((n_bonds, 3), dtype=np.float64) - - # get coordination number of center atom for each species - for s in range(n_specs): - bs = spec_mask[species_list[s]] - mbtype = mb_mask[bcn + bs] - r_cut = cutoff_mb[mbtype] - - qs[s] = q_value_mc(bond_array_mb[:, 0], r_cut, species_list[s], - etypes, cutoff_func) - - # get coordination number of all neighbor atoms for each species - for i in range(n_bonds): - be = spec_mask[etypes[i]] - ben = be * nspec - - neigh_bond_array, _, neigh_etypes, _ = \ - get_2_body_arrays_sepcut(positions, bond_inds[i], cell, - cutoff_mb, species, sweep, nspec, spec_mask, mb_mask) - for s in range(n_specs): - bs = spec_mask[species_list[s]] - mbtype = mb_mask[bs + ben] - r_cut = cutoff_mb[mbtype] - - qs_neigh[i, s] = q_value_mc(neigh_bond_array[:, 0], r_cut, - species_list[s], neigh_etypes, cutoff_func) - - # get grad from each neighbor atom - for i in range(n_bonds): - be = spec_mask[etypes[i]] - mbtype = mb_mask[bcn + be] - r_cut = cutoff_mb[mbtype] - - ri = bond_array_mb[i, 0] - for d in range(3): - ci = bond_array_mb[i, d+1] - - _, q_grads[i, d] = coordination_number(ri, ci, r_cut, - cutoff_func) - - return qs, qs_neigh, q_grads, species_list, etypes diff --git a/flare/gp.py b/flare/gp.py index 5dadd7518..2f04f5b13 100644 --- a/flare/gp.py +++ b/flare/gp.py @@ -1,31 +1,31 @@ -import time -import math -import pickle import inspect import json +import logging +import math +import pickle +import time -import numpy as np import multiprocessing as mp +import numpy as np from collections import Counter from copy import deepcopy from numpy.random import random -from typing import List, Callable, Union from scipy.linalg import solve_triangular from scipy.optimize import minimize +from typing import List, Callable, Union, Tuple from flare.env import AtomicEnvironment -from flare.struc import Structure from flare.gp_algebra import get_like_from_mats, get_neg_like_grad, \ force_force_vector, energy_force_vector, get_force_block, \ get_ky_mat_update, _global_training_data, _global_training_labels, \ _global_training_structures, _global_energy_labels, get_Ky_mat, \ get_kernel_vector, en_kern_vec -from flare.utils.mask_helper import HyperParameterMasking - -from flare.kernels.utils import str_to_kernel_set, from_mask_to_args +from flare.kernels.utils import str_to_kernel_set, from_mask_to_args, kernel_str_to_array +from flare.output import Output, set_logger +from flare.parameters import Parameters +from flare.struc import Structure from flare.utils.element_coder import NumpyEncoder, Z_to_element -from flare.output import Output class GaussianProcess: @@ -34,18 +34,14 @@ class GaussianProcess: Williams. Args: - kernel (Callable, optional): Name of the kernel to use, or the kernel - itself. - kernel_grad (Callable, optional): Function that returns the gradient - of the GP kernel with respect to the hyperparameters. + kernels (list, optional): Determine the type of kernels. Example: + ['2', '3'], ['2', '3', 'mb'], ['2']. Defaults to ['2', '3'] + component (str, optional): Determine single- ("sc") or multi- + component ("mc") kernel to use. Defaults to "mc" hyps (np.ndarray, optional): Hyperparameters of the GP. - cutoffs (np.ndarray, optional): Cutoffs of the GP kernel. + cutoffs (Dict, optional): Cutoffs of the GP kernel. hyp_labels (List, optional): List of hyperparameter labels. Defaults to None. - energy_force_kernel (Callable, optional): Energy/force kernel of the - GP used to make energy predictions. Defaults to None. - energy_kernel (Callable, optional): Energy/energy kernel of the GP. - Defaults to None. opt_algorithm (str, optional): Hyperparameter optimization algorithm. Defaults to 'L-BFGS-B'. maxiter (int, optional): Maximum number of iterations of the @@ -58,81 +54,74 @@ class GaussianProcess: predictions. output (Output, optional): Output object used to dump hyperparameters during optimization. Defaults to None. - multihyps (bool, optional): If True, turn on multiple-group of hyper- - parameters. - hyps_mask (dict, optional): If multihyps is True, hyps_mask can set up - which hyper parameter is used for what interaction. Details see - kernels/mc_sephyps.py - kernel_name (str, optional): Determine the type of kernels. Example: - 2+3_mc, 2+3+mb_mc, 2_mc, 2_sc, 3_sc, ... + hyps_mask (dict, optional): hyps_mask can set up which hyper parameter + is used for what interaction. Details see kernels/mc_sephyps.py name (str, optional): Name for the GP instance. """ - def __init__(self, kernel: Callable = None, kernel_grad: Callable = None, - hyps: 'ndarray' = None, cutoffs: 'ndarray' = None, + def __init__(self, kernels: list = ['two', 'three'], + component: str = 'mc', + hyps: 'ndarray' = None, cutoffs={}, + hyps_mask: dict = {}, hyp_labels: List = None, opt_algorithm: str = 'L-BFGS-B', maxiter: int = 10, parallel: bool = False, per_atom_par: bool = True, n_cpus: int = 1, n_sample: int = 100, output: Output = None, - multihyps: bool = False, hyps_mask: dict = None, - kernel_name="2+3_mc", name="default_gp", + name="default_gp", energy_noise: float = 0.01, **kwargs,): """Initialize GP parameters and training data.""" # load arguments into attributes self.name = name - self.cutoffs = np.array(cutoffs, dtype=np.float64) + self.output = output self.opt_algorithm = opt_algorithm - self.output = output self.per_atom_par = per_atom_par self.maxiter = maxiter - if hyps is None: - # If no hyperparameters are passed in, assume 2 hyps for each - # cutoff, plus one noise hyperparameter, and use a guess value - hyps = np.array([0.1]*(1+2*len(cutoffs))) - - self.update_hyps(hyps, hyp_labels, multihyps, hyps_mask) - # set up parallelization self.n_cpus = n_cpus self.n_sample = n_sample self.parallel = parallel - if 'nsample' in kwargs: - DeprecationWarning("nsample is being replaced with n_sample") - self.n_sample = kwargs.get('nsample') - if 'par' in kwargs: - DeprecationWarning("par is being replaced with parallel") - self.parallel = kwargs.get('par') - if 'no_cpus' in kwargs: - DeprecationWarning("no_cpus is being replaced with n_cpu") - self.n_cpus = kwargs.get('no_cpus') - - # TO DO, clean up all the other kernel arguments - if kernel is None: - kernel, grad, ek, efk = str_to_kernel_set(kernel_name, multihyps) - self.kernel = kernel - self.kernel_grad = grad - self.energy_force_kernel = efk - self.energy_kernel = ek - self.kernel_name = kernel.__name__ + + self.component = component + self.kernels = kernels + self.cutoffs = cutoffs + self.hyp_labels = hyp_labels + self.hyps_mask = hyps_mask + self.hyps = hyps + + GaussianProcess.backward_arguments(kwargs, self.__dict__) + GaussianProcess.backward_attributes(self.__dict__) + + # ------------ "computed" attributes ------------ + + if self.output is None: + self.logger_name = self.name+"GaussianProcess" + set_logger(self.logger_name, stream=True, + fileout_name=None, verbose="info") else: - DeprecationWarning("kernel, kernel_grad, energy_force_kernel " - "and energy_kernel will be replaced by " - "kernel_name") - self.kernel_name = kernel.__name__ - self.kernel = kernel - self.kernel_grad = kernel_grad - self.energy_force_kernel = kwargs.get('energy_force_kernel') - self.energy_kernel = kwargs.get('energy_kernel') + self.logger_name = self.output.basename+'log' - self.name = name + if self.hyps is None: + # If no hyperparameters are passed in, assume 2 hyps for each + # cutoff, plus one noise hyperparameter, and use a guess value + self.hyps = np.array([0.1]*(1+2*len(self.kernels))) + else: + self.hyps = np.array(self.hyps, dtype=np.float64) + + kernel, grad, ek, efk = str_to_kernel_set( + self.kernels, self.component, self.hyps_mask) + self.kernel = kernel + self.kernel_grad = grad + self.energy_force_kernel = efk + self.energy_kernel = ek + self.kernels = kernel_str_to_array(kernel.__name__) # parallelization if self.parallel: - if n_cpus is None: + if self.n_cpus is None: self.n_cpus = mp.cpu_count() else: self.n_cpus = n_cpus @@ -141,7 +130,6 @@ def __init__(self, kernel: Callable = None, kernel_grad: Callable = None, self.training_data = [] # Atomic environments self.training_labels = [] # Forces acting on central atoms - self.training_labels_np = np.empty(0, ) self.n_envs_prev = len(self.training_data) @@ -174,82 +162,58 @@ def check_instantiation(self): :return: """ - if (self.name in _global_training_labels): + if self.logger_name is None: + if self.output is None: + self.logger_name = self.name+"GaussianProcess" + set_logger(self.logger_name, stream=True, + fileout_name=None, verbose="info") + else: + self.logger_name = self.output.basename+'log' + logger = logging.getLogger(self.logger_name) + + + # check whether it's be loaded before + loaded = False + if self.name in _global_training_labels: + if _global_training_labels.get(self.name, None) is not self.training_labels_np: + loaded = True + if self.name in _global_energy_labels: + if _global_energy_labels.get(self.name, None) is not self.energy_labels_np: + loaded = True + + if loaded: + base = f'{self.name}' count = 2 while (self.name in _global_training_labels and count < 100): time.sleep(random()) self.name = f'{base}_{count}' - print("Specified GP name is present in global memory; " - "Attempting to rename the " - f"GP instance to {self.name}") + logger.debug("Specified GP name is present in global memory; " + "Attempting to rename the " + f"GP instance to {self.name}") count += 1 if (self.name in _global_training_labels): milliseconds = int(round(time.time() * 1000) % 10000000) self.name = f"{base}_{milliseconds}" - print("Specified GP name still present in global memory: " - f"renaming the gp instance to {self.name}") - print(f"Final name of the gp instance is {self.name}") - - assert (self.name not in _global_training_labels), \ - f"the gp instance name, {self.name} is used" - assert (self.name not in _global_training_data), \ - f"the gp instance name, {self.name} is used" + logger.debug("Specified GP name still present in global memory: " + f"renaming the gp instance to {self.name}") + logger.debug(f"Final name of the gp instance is {self.name}") - _global_training_data[self.name] = self.training_data - _global_training_structures[self.name] = self.training_structures - _global_training_labels[self.name] = self.training_labels_np - _global_energy_labels[self.name] = self.energy_labels_np + self.sync_data() - assert (len(self.cutoffs) <= 3) - - if (len(self.cutoffs) > 1): - assert self.cutoffs[0] >= self.cutoffs[1], \ - "2b cutoff has to be larger than 3b cutoffs" - - if ('three' in self.kernel_name): - assert len(self.cutoffs) >= 2, \ - "3b kernel needs two cutoffs, one for building"\ - " neighbor list and one for the 3b" - if ('many' in self.kernel_name): - assert len(self.cutoffs) >= 3, \ - "many-body kernel needs three cutoffs, one for building"\ - " neighbor list and one for the 3b" - - if self.multihyps is True and self.hyps_mask is None: - raise ValueError("Warning! Multihyperparameter mode enabled," - "but no configuration hyperparameter mask was " - "passed. Did you mean to set multihyps to False?") - elif self.multihyps is False and self.hyps_mask is not None: - raise ValueError("Warning! Multihyperparameter mode disabled," - "but a configuration hyperparameter mask was " - "passed. Did you mean to set multihyps to True?") - - if self.multihyps is True: - - self.hyps_mask = HyperParameterMasking.check_instantiation( - self.hyps_mask) - HyperParameterMasking.check_matching( - self.hyps_mask, self.hyps, self.cutoffs) - self.bounds = deepcopy(self.hyps_mask.get('bounds', None)) + self.hyps_mask = Parameters.check_instantiation(self.hyps, self.cutoffs, + self.kernels, self.hyps_mask) - else: - self.multihyps = False - self.hyps_mask = None + self.bounds = deepcopy(self.hyps_mask.get('bounds', None)) - def update_kernel(self, kernel_name, multihyps=False): - kernel, grad, ek, efk = str_to_kernel_set(kernel_name, multihyps) + def update_kernel(self, kernels, component="mc", hyps_mask=None): + kernel, grad, ek, efk = str_to_kernel_set( + kernels, component, hyps_mask) self.kernel = kernel self.kernel_grad = grad self.energy_force_kernel = efk self.energy_kernel = ek - self.kernel_name = kernel.__name__ - - def update_hyps(self, hyps=None, hyp_labels=None, multihyps=False, hyps_mask=None): - self.hyps = np.array(hyps, dtype=np.float64) - self.hyp_labels = hyp_labels - self.hyps_mask = hyps_mask - self.multihyps = multihyps + self.kernels = kernel_str_to_array(kernel.__name__) def update_db(self, struc: Structure, forces: List, custom_range: List[int] = (), energy: float = None): @@ -278,7 +242,7 @@ def update_db(self, struc: Structure, forces: List, for atom in update_indices: env_curr = \ AtomicEnvironment(struc, atom, self.cutoffs, - cutoffs_mask=self.hyps_mask) + cutoffs_mask=self.hyps_mask) forces_curr = np.array(forces[atom]) self.training_data.append(env_curr) @@ -286,8 +250,6 @@ def update_db(self, struc: Structure, forces: List, # create numpy array of training labels self.training_labels_np = np.hstack(self.training_labels) - _global_training_data[self.name] = self.training_data - _global_training_labels[self.name] = self.training_labels_np # If an energy is given, update the structure list. if energy is not None: @@ -301,12 +263,11 @@ def update_db(self, struc: Structure, forces: List, self.energy_labels.append(energy) self.training_structures.append(structure_list) self.energy_labels_np = np.array(self.energy_labels) - _global_training_structures[self.name] = self.training_structures - _global_energy_labels[self.name] = self.energy_labels_np # update list of all labels self.all_labels = np.concatenate((self.training_labels_np, self.energy_labels_np)) + self.sync_data() def add_one_env(self, env: AtomicEnvironment, force, train: bool = False, **kwargs): @@ -324,8 +285,7 @@ def add_one_env(self, env: AtomicEnvironment, self.training_data.append(env) self.training_labels.append(force) self.training_labels_np = np.hstack(self.training_labels) - _global_training_data[self.name] = self.training_data - _global_training_labels[self.name] = self.training_labels_np + self.sync_data() # update list of all labels self.all_labels = np.concatenate((self.training_labels_np, @@ -334,7 +294,7 @@ def add_one_env(self, env: AtomicEnvironment, if train: self.train(**kwargs) - def train(self, output=None, custom_bounds=None, + def train(self, logger=None, custom_bounds=None, grad_tol: float = 1e-4, x_tol: float = 1e-5, line_steps: int = 20, @@ -344,7 +304,7 @@ def train(self, output=None, custom_bounds=None, (related to the covariance matrix of the training set). Args: - output (Output): Output object specifying where to write the + logger (logging.logger): logger object specifying where to write the progress of the optimization. custom_bounds (np.ndarray): Custom bounds on the hyperparameters. grad_tol (float): Tolerance of the hyperparameter gradient that @@ -355,6 +315,16 @@ def train(self, output=None, custom_bounds=None, hyperparameter optimization. """ + verbose = "warning" + if print_progress: + verbose = "info" + if logger is None: + logger = set_logger("gp_algebra", stream=True, + fileout_name="log.gp_algebra", + verbose=verbose) + + disp = print_progress + if len(self.training_data) == 0 or len(self.training_labels) == 0: raise Warning("You are attempting to train a GP with no " "training data. Add environments and forces " @@ -362,9 +332,9 @@ def train(self, output=None, custom_bounds=None, x_0 = self.hyps - args = (self.name, self.kernel_grad, output, + args = (self.name, self.kernel_grad, logger, self.cutoffs, self.hyps_mask, - self.n_cpus, self.n_sample, print_progress) + self.n_cpus, self.n_sample) res = None @@ -381,25 +351,26 @@ def train(self, output=None, custom_bounds=None, try: res = minimize(get_neg_like_grad, x_0, args, method='L-BFGS-B', jac=True, bounds=bounds, - options={'disp': False, 'gtol': grad_tol, + options={'disp': disp, 'gtol': grad_tol, 'maxls': line_steps, 'maxiter': self.maxiter}) except np.linalg.LinAlgError: - print("Warning! Algorithm for L-BFGS-B failed. Changing to " - "BFGS for remainder of run.") + logger = logging.getLogger(self.logger_name) + logger.warning("Algorithm for L-BFGS-B failed. Changing to " + "BFGS for remainder of run.") self.opt_algorithm = 'BFGS' if custom_bounds is not None: res = minimize(get_neg_like_grad, x_0, args, method='L-BFGS-B', jac=True, bounds=custom_bounds, - options={'disp': False, 'gtol': grad_tol, + options={'disp': disp, 'gtol': grad_tol, 'maxls': line_steps, 'maxiter': self.maxiter}) elif self.opt_algorithm == 'BFGS': res = minimize(get_neg_like_grad, x_0, args, method='BFGS', jac=True, - options={'disp': False, 'gtol': grad_tol, + options={'disp': disp, 'gtol': grad_tol, 'maxiter': self.maxiter}) if res is None: @@ -418,7 +389,7 @@ def check_L_alpha(self): """ # Check that alpha is up to date with training set - size3 = len(self.training_data)*3 + size3 = len(self.training_data) * 3 + len(self.training_structures) # If model is empty, then just return if size3 == 0: @@ -452,8 +423,7 @@ def predict(self, x_t: AtomicEnvironment, d: int) -> [float, float]: else: n_cpus = 1 - _global_training_data[self.name] = self.training_data - _global_training_labels[self.name] = self.training_labels_np + self.sync_data() k_v = \ get_kernel_vector(self.name, self.kernel, self.energy_force_kernel, @@ -469,7 +439,8 @@ def predict(self, x_t: AtomicEnvironment, d: int) -> [float, float]: # get predictive variance without cholesky (possibly faster) # pass args to kernel based on if mult. hyperparameters in use - args = from_mask_to_args(self.hyps, self.hyps_mask, self.cutoffs) + args = from_mask_to_args(self.hyps, self.cutoffs, self.hyps_mask) + self_kern = self.kernel(x_t, x_t, d, d, *args) pred_var = self_kern - np.matmul(np.matmul(k_v, self.ky_mat_inv), k_v) @@ -490,8 +461,7 @@ def predict_local_energy(self, x_t: AtomicEnvironment) -> float: else: n_cpus = 1 - _global_training_data[self.name] = self.training_data - _global_training_labels[self.name] = self.training_labels_np + self.sync_data() k_v = en_kern_vec(self.name, self.energy_force_kernel, self.energy_kernel, @@ -519,8 +489,7 @@ def predict_local_energy_and_var(self, x_t: AtomicEnvironment): else: n_cpus = 1 - _global_training_data[self.name] = self.training_data - _global_training_labels[self.name] = self.training_labels_np + self.sync_data() # get kernel vector k_v = en_kern_vec(self.name, self.energy_force_kernel, @@ -534,7 +503,7 @@ def predict_local_energy_and_var(self, x_t: AtomicEnvironment): # get predictive variance v_vec = solve_triangular(self.l_mat, k_v, lower=True) - args = from_mask_to_args(self.hyps, self.hyps_mask, self.cutoffs) + args = from_mask_to_args(self.hyps, self.cutoffs, self.hyps_mask) self_kern = self.energy_kernel(x_t, x_t, *args) @@ -550,10 +519,7 @@ def set_L_alpha(self): The forces and variances are later obtained using alpha. """ - _global_training_data[self.name] = self.training_data - _global_training_structures[self.name] = self.training_structures - _global_training_labels[self.name] = self.training_labels_np - _global_energy_labels[self.name] = self.energy_labels_np + self.sync_data() ky_mat = \ get_Ky_mat(self.hyps, self.name, self.kernel, @@ -587,10 +553,7 @@ def update_L_alpha(self): return # Reset global variables. - _global_training_data[self.name] = self.training_data - _global_training_structures[self.name] = self.training_structures - _global_training_labels[self.name] = self.training_labels_np - _global_energy_labels[self.name] = self.energy_labels_np + self.sync_data() ky_mat = get_ky_mat_update(self.ky_mat, self.n_envs_prev, self.hyps, self.name, self.kernel, self.energy_kernel, @@ -615,12 +578,14 @@ def __str__(self): """String representation of the GP model.""" thestr = "GaussianProcess Object\n" - thestr += f'Kernel: {self.kernel_name}\n' + thestr += f'Number of cpu cores: {self.n_cpus}\n' + thestr += f'Kernel: {self.kernels}\n' thestr += f"Training points: {len(self.training_data)}\n" thestr += f'Cutoffs: {self.cutoffs}\n' thestr += f'Model Likelihood: {self.likelihood}\n' - thestr += f'MultiHyps: {self.multihyps}\n' + thestr += f'Number of hyperparameters: {len(self.hyps)}\n' + thestr += f'Hyperparameters_array: {str(self.hyps)}\n' thestr += 'Hyperparameters: \n' if self.hyp_labels is None: # Put unlabeled hyperparameters on one line @@ -628,32 +593,10 @@ def __str__(self): thestr += str(self.hyps) + '\n' else: for hyp, label in zip(self.hyps, self.hyp_labels): - thestr += f"{label}: {hyp}\n" - - if self.multihyps: - nspecie = self.hyps_mask['nspecie'] - thestr += f'nspecie: {nspecie}\n' - thestr += f'specie_mask: \n' - thestr += str(self.hyps_mask['specie_mask']) + '\n' - - nbond = self.hyps_mask.get('nbond', 0) - thestr += f'nbond: {nbond}\n' + thestr += f"{label}: {hyp} \n" - if nbond > 1: - thestr += f'bond_mask: \n' - thestr += str(self.hyps_mask['bond_mask']) + '\n' - - ntriplet = self.hyps_mask.get('ntriplet', 0) - thestr += f'ntriplet: {ntriplet}\n' - if ntriplet > 1: - thestr += f'triplet_mask: \n' - thestr += str(self.hyps_mask['triplet_mask']) + '\n' - - nmb = self.hyps_mask.get('nmb', 0) - thestr += f'nmb: {nmb}\n' - if nmb > 1: - thestr += f'mb_mask: \n' - thestr += str(self.hyps_mask['nmb_mask']) + '\n' + for k in self.hyps_mask: + thestr += f'Hyps_mask {k}: {self.hyps_mask[k]} \n' return thestr @@ -682,33 +625,29 @@ def as_dict(self): return out_dict + def sync_data(self): + _global_training_data[self.name] = self.training_data + _global_training_labels[self.name] = self.training_labels_np + _global_training_structures[self.name] = self.training_structures + _global_energy_labels[self.name] = self.energy_labels_np + @staticmethod def from_dict(dictionary): """Create GP object from dictionary representation.""" - multihyps = dictionary.get('multihyps', False) - - new_gp = GaussianProcess(kernel_name=dictionary['kernel_name'], - cutoffs=np.array(dictionary['cutoffs']), - hyps=np.array(dictionary['hyps']), - hyp_labels=dictionary['hyp_labels'], - parallel=dictionary.get('parallel', False) or - dictionary.get('par', False), - per_atom_par=dictionary.get('per_atom_par', - True), - n_cpus=dictionary.get('n_cpus') or - dictionary.get('no_cpus'), - maxiter=dictionary['maxiter'], - opt_algorithm=dictionary.get( - 'opt_algorithm', 'L-BFGS-B'), - multihyps=multihyps, - hyps_mask=dictionary.get('hyps_mask', None), - name=dictionary.get('name', 'default_gp') - ) + GaussianProcess.backward_arguments(dictionary, dictionary) + GaussianProcess.backward_attributes(dictionary) + + new_gp = GaussianProcess(**dictionary) # Save time by attempting to load in computed attributes - new_gp.training_data = [AtomicEnvironment.from_dict(env) for env in - dictionary['training_data']] + if ('training_data' in dictionary): + new_gp.training_data = [AtomicEnvironment.from_dict(env) for env in + dictionary['training_data']] + new_gp.training_labels = deepcopy(dictionary['training_labels']) + new_gp.training_labels_np = deepcopy( + dictionary['training_labels_np']) + new_gp.sync_data() # Reconstruct training structures. if ('training_structures' in dictionary): @@ -718,31 +657,17 @@ def from_dict(dictionary): for env_curr in env_list: new_gp.training_structures[n].append( AtomicEnvironment.from_dict(env_curr)) - else: - new_gp.training_structures = [] # Environments of each structure - new_gp.energy_labels = [] # Energies of training structures - new_gp.energy_labels_np = np.empty(0, ) - new_gp.energy_noise = 0.01 - - new_gp.training_labels = deepcopy(dictionary.get('training_labels', - [])) - new_gp.training_labels_np = deepcopy(dictionary.get('training_labels_np', - np.empty(0,))) - new_gp.energy_labels = deepcopy(dictionary.get('energy_labels', - [])) - new_gp.energy_labels_np = deepcopy(dictionary.get('energy_labels_np', - np.empty(0,))) + new_gp.energy_labels = deepcopy(dictionary['energy_labels']) + new_gp.energy_labels_np = deepcopy(dictionary['energy_labels_np']) new_gp.all_labels = np.concatenate((new_gp.training_labels_np, - new_gp.energy_labels_np)) + new_gp.energy_labels_np)) + + new_gp.likelihood = dictionary.get('likelihood', None) + new_gp.likelihood_gradient = dictionary.get( + 'likelihood_gradient', None) - new_gp.likelihood = dictionary['likelihood'] - new_gp.likelihood_gradient = dictionary['likelihood_gradient'] new_gp.n_envs_prev = len(new_gp.training_data) - _global_training_data[new_gp.name] = new_gp.training_data - _global_training_structures[new_gp.name] = new_gp.training_structures - _global_training_labels[new_gp.name] = new_gp.training_labels_np - _global_energy_labels[new_gp.name] = new_gp.energy_labels_np # Save time by attempting to load in computed attributes if len(new_gp.training_data) > 5000: @@ -755,8 +680,9 @@ def from_dict(dictionary): new_gp.alpha = None new_gp.ky_mat_inv = None filename = dictionary['ky_mat_file'] - Warning("the covariance matrices are not loaded" - f"because {filename} cannot be found") + logger = logging.getLogger(self.logger_name) + logger.warning("the covariance matrices are not loaded" + f"because {filename} cannot be found") else: new_gp.ky_mat_inv = np.array(dictionary['ky_mat_inv']) \ if dictionary.get('ky_mat_inv') is not None else None @@ -801,23 +727,22 @@ def adjust_cutoffs(self, new_cutoffs: Union[list, tuple, 'np.ndarray'], if (new_hyps_mask is not None): hm = new_hyps_mask + self.hyps_mask = new_hyps_mask else: hm = self.hyps_mask # update environment nenv = len(self.training_data) for i in range(nenv): - self.training_data[i].cutoffs = np.array( - new_cutoffs, dtype=np.float) + self.training_data[i].cutoffs = new_cutoffs self.training_data[i].cutoffs_mask = hm - self.training_data[i].setup_mask() + self.training_data[i].setup_mask(hm) self.training_data[i].compute_env() # Ensure that training data and labels are still consistent - _global_training_data[self.name] = self.training_data - _global_training_labels[self.name] = self.training_labels_np + self.sync_data() - self.cutoffs = np.array(new_cutoffs) + self.cutoffs = new_cutoffs if reset_L_alpha: del self.l_mat @@ -827,6 +752,61 @@ def adjust_cutoffs(self, new_cutoffs: Union[list, tuple, 'np.ndarray'], if train: self.train() + def remove_force_data(self, indexes: Union[int, List[int]], + update_matrices: bool = True)->Tuple[List[Structure], + List['ndarray']]: + """ + Remove force components from the model. Convenience function which + deletes individual data points. + + Matrices should *always* be updated if you intend to use the GP to make + predictions afterwards. This might be time consuming for large GPs, + so, it is provided as an option, but, only do so with extreme caution. + (Undefined behavior may result if you try to make predictions and/or + add to the training set afterwards). + + Returns training data which was removed akin to a pop method, in order + of lowest to highest index passed in. + + :param indexes: Indexes of envs in training data to remove. + :param update_matrices: If false, will not update the GP's matrices + afterwards (which can be time consuming for large models). + This should essentially always be true except for niche development + applications. + :return: + """ + + # Listify input even if one integer + if isinstance(indexes, int): + indexes = [indexes] + + if max(indexes) > len(self.training_data): + raise ValueError("Index out of range of data") + + # Get in reverse order so that modifying higher indexes doesn't affect + # lower indexes + indexes.sort(reverse=True) + removed_data = [] + removed_labels = [] + for i in indexes: + removed_data.append(self.training_data.pop(i)) + removed_labels.append(self.training_labels.pop(i)) + + self.training_labels_np = np.hstack(self.training_labels) + self.all_labels = np.concatenate((self.training_labels_np, + self.energy_labels_np)) + self.sync_data() + + if update_matrices: + self.set_L_alpha() + self.compute_matrices() + + # Put removed data in order of lowest to highest index + removed_data.reverse() + removed_labels.reverse() + + return removed_data, removed_labels + def write_model(self, name: str, format: str = 'json'): """ Write model in a variety of formats to a file for later re-use. @@ -849,7 +829,6 @@ def write_model(self, name: str, format: str = 'json'): self.alpha = None self.ky_mat_inv = None - supported_formats = ['json', 'pickle', 'binary'] if format.lower() == 'json': @@ -870,8 +849,6 @@ def write_model(self, name: str, format: str = 'json'): self.alpha = temp_alpha self.ky_mat_inv = temp_ky_mat_inv - - @staticmethod def from_file(filename: str, format: str = ''): """ @@ -890,10 +867,13 @@ def from_file(filename: str, format: str = ''): elif '.pickle' in filename or 'pickle' in format: with open(filename, 'rb') as f: + gp_model = pickle.load(f) - if ('name' not in gp_model.__dict__): - gp_model.name = 'default_gp' + GaussianProcess.backward_arguments( + gp_model.__dict__, gp_model.__dict__) + + GaussianProcess.backward_attributes(gp_model.__dict__) if len(gp_model.training_data) > 5000: try: @@ -905,23 +885,14 @@ def from_file(filename: str, format: str = ''): gp_model.l_mat = None gp_model.alpha = None gp_model.ky_mat_inv = None - Warning("the covariance matrices are not loaded" + Warning("the covariance matrices are not loaded" \ f"it can take extra long time to recompute") else: raise ValueError("Warning: Format unspecieified or file is not " ".json or .pickle format.") - if ('training_structures' not in gp_model.__dict__): - gp_model.training_structures = [] # Environments of each structure - gp_model.energy_labels = [] # Energies of training structures - gp_model.energy_labels_np = np.empty(0, ) - gp_model.energy_noise = 0.01 - gp_model.all_labels = np.concatenate((gp_model.training_labels_np, - gp_model.energy_labels_np)) - gp_model.check_instantiation() - return gp_model @property @@ -962,3 +933,76 @@ def __del__(self): if (self.name in _global_training_labels): _global_training_labels.pop(self.name, None) _global_training_data.pop(self.name, None) + + @staticmethod + def backward_arguments(kwargs, new_args={}): + """ + update the initialize arguments that were renamed + """ + + if 'kernel_name' in kwargs: + DeprecationWarning( + "kernel_name is being replaced with kernels") + new_args['kernels'] = kernel_str_to_array( + kwargs['kernel_name']) + kwargs.pop('kernel_name') + if 'nsample' in kwargs: + DeprecationWarning("nsample is being replaced with n_sample") + new_args['n_sample'] = kwargs['nsample'] + kwargs.pop('nsample') + if 'par' in kwargs: + DeprecationWarning("par is being replaced with parallel") + new_args['parallel'] = kwargs['par'] + kwargs.pop('par') + if 'no_cpus' in kwargs: + DeprecationWarning("no_cpus is being replaced with n_cpu") + new_args['n_cpus'] = kwargs['no_cpus'] + kwargs.pop('no_cpus') + if 'multihyps' in kwargs: + DeprecationWarning("multihyps is removed") + kwargs.pop('multihyps') + + return new_args + + @staticmethod + def backward_attributes(dictionary): + """ + add new attributes to old instance + or update attribute types + """ + + if ('name' not in dictionary): + dictionary['name'] = 'default_gp' + if ('per_atom_par' not in dictionary): + dictionary['per_atom_par'] = True + if ('optimization_algorithm' not in dictionary): + dictionary['opt_algorithm'] = 'L-BFGS-B' + if ('hyps_mask' not in dictionary): + dictionary['hyps_mask'] = None + if ('parallel' not in dictionary): + dictionary['parallel'] = False + if ('component' not in dictionary): + dictionary['component'] = 'mc' + + if ('training_structures' not in dictionary): + # Environments of each structure + dictionary['training_structures'] = [] + dictionary['energy_labels'] = [] # Energies of training structures + dictionary['energy_labels_np'] = np.empty(0, ) + + if ('training_labels' not in dictionary): + dictionary['training_labels'] = [] + dictionary['training_labels_np'] = np.empty(0,) + + if ('energy_noise' not in dictionary): + dictionary['energy_noise'] = 0.01 + + if not isinstance(dictionary['cutoffs'], dict): + dictionary['cutoffs'] = Parameters.cutoff_array_to_dict( + dictionary['cutoffs']) + + dictionary['hyps_mask'] = Parameters.backward( + dictionary['kernels'], deepcopy(dictionary['hyps_mask'])) + + if 'logger_name' not in dictionary: + dictionary['logger_name'] = None diff --git a/flare/gp_algebra.py b/flare/gp_algebra.py index 491d9b5d7..f41d90d7a 100644 --- a/flare/gp_algebra.py +++ b/flare/gp_algebra.py @@ -1,7 +1,7 @@ import math -import time -import numpy as np import multiprocessing as mp +import numpy as np +import time from typing import List, Callable from flare.kernels.utils import from_mask_to_args, from_grad_to_mask @@ -113,7 +113,6 @@ def partition_vector(n_sample, size, n_cpus): def partition_force_energy_block(n_sample: int, size1: int, size2: int, n_cpus: int): - """Special partition method for the force/energy block. Because the number of environments in a structure can vary, we only split up the environment list, which has length size1. @@ -193,7 +192,7 @@ def obtain_noise_len(hyps, hyps_mask): if (hyps_mask is not None): train_noise = hyps_mask.get('train_noise', True) if (train_noise is False): - sigma_n = hyps_mask['original'][-1] + sigma_n = hyps_mask['original_hyps'][-1] non_noise_hyps = len(hyps) return sigma_n, non_noise_hyps, train_noise @@ -211,9 +210,9 @@ def parallel_matrix_construction(pack_function, hyps, name, kernel, cutoffs, for wid in range(nbatch): s1, e1, s2, e2 = block_id[wid] children.append(mp.Process(target=queue_wrapper, - args=(result_queue, wid, pack_function, - (hyps, name, s1, e1, s2, e2, s1 == s2, - kernel, cutoffs, hyps_mask)))) + args=(result_queue, wid, pack_function, + (hyps, name, s1, e1, s2, e2, s1 == s2, + kernel, cutoffs, hyps_mask)))) # Run child processes. for c in children: @@ -301,7 +300,7 @@ def get_force_block_pack(hyps: np.ndarray, name: str, s1: int, e1: int, ds = [1, 2, 3] # calculate elements - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) for m_index in range(size1): x_1 = training_data[int(math.floor(m_index / 3))+s1] @@ -333,7 +332,7 @@ def get_energy_block_pack(hyps: np.ndarray, name: str, s1: int, e1: int, energy_block = np.zeros([size1, size2]) # calculate elements - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) for m_index in range(size1): struc_1 = training_structures[m_index + s1] @@ -373,7 +372,7 @@ def get_force_energy_block_pack(hyps: np.ndarray, name: str, s1: int, ds = [1, 2, 3] # calculate elements - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) for m_index in range(size1): environment_1 = training_data[int(math.floor(m_index / 3)) + s1] @@ -728,7 +727,7 @@ def energy_energy_vector_unit(name, s, e, x, kernel, hyps, cutoffs=None, size = e - s energy_energy_unit = np.zeros(size, ) - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) for m_index in range(size): structure = training_structures[m_index + s] @@ -753,7 +752,7 @@ def energy_force_vector_unit(name, s, e, x, kernel, hyps, cutoffs=None, size = (e - s) * 3 k_v = np.zeros(size, ) - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) for m_index in range(size): x_2 = training_data[int(math.floor(m_index / 3))+s] @@ -770,7 +769,7 @@ def force_energy_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask, """ size = e - s - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) force_energy_unit = np.zeros(size,) for m_index in range(size): @@ -793,7 +792,7 @@ def force_force_vector_unit(name, s, e, x, kernel, hyps, cutoffs, hyps_mask, size = (e - s) ds = [1, 2, 3] - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) k_v = np.zeros(size * 3) @@ -982,7 +981,7 @@ def get_ky_and_hyp_pack(name, s1, e1, s2, e2, same: bool, hyps: np.ndarray, k_mat = np.zeros([size1, size2]) hyp_mat = np.zeros([non_noise_hyps, size1, size2]) - args = from_mask_to_args(hyps, hyps_mask, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hyps_mask) ds = [1, 2, 3] @@ -1006,6 +1005,7 @@ def get_ky_and_hyp_pack(name, s1, e1, s2, e2, same: bool, hyps: np.ndarray, # store kernel value k_mat[m_index, n_index] = cov[0] grad = from_grad_to_mask(cov[1], hyps_mask) + hyp_mat[:, m_index, n_index] = grad if (same): k_mat[n_index, m_index] = cov[0] @@ -1041,8 +1041,8 @@ def get_ky_and_hyp(hyps: np.ndarray, name, kernel_grad, cutoffs=None, n_cpus = mp.cpu_count() if (n_cpus == 1): hyp_mat0, k_mat = get_ky_and_hyp_pack( - name, 0, size, 0, size, True, - hyps, kernel_grad, cutoffs, hyps_mask) + name, 0, size, 0, size, True, + hyps, kernel_grad, cutoffs, hyps_mask) else: block_id, nbatch = partition_matrix(n_sample, size, n_cpus) @@ -1119,9 +1119,9 @@ def get_like_from_mats(ky_mat, l_mat, alpha, name): def get_neg_like_grad(hyps: np.ndarray, name: str, - kernel_grad, output=None, + kernel_grad, logger=None, cutoffs=None, hyps_mask=None, - n_cpus=1, n_sample=100, print_progress=False): + n_cpus=1, n_sample=100): """compute the log likelihood and its gradients :param hyps: list of hyper-parameters @@ -1130,7 +1130,7 @@ def get_neg_like_grad(hyps: np.ndarray, name: str, :param kernel_grad: function object of the kernel gradient :param output: Output object for dumping every hyper-parameter sets computed - :type output: class Output + :type output: logger :param cutoffs: The cutoff values used for the atomic environments :type cutoffs: list of 2 float numbers :param hyps_mask: dictionary used for multi-group hyperparmeters @@ -1141,42 +1141,24 @@ def get_neg_like_grad(hyps: np.ndarray, name: str, """ time0 = time.time() - if output is not None: - ostring = "hyps:" - for hyp in hyps: - ostring += f" {hyp}" - ostring += "\n" - output.write_to_log(ostring, name="hyps") hyp_mat, ky_mat = \ get_ky_and_hyp(hyps, name, kernel_grad, cutoffs=cutoffs, hyps_mask=hyps_mask, n_cpus=n_cpus, n_sample=n_sample) - if output is not None: - output.write_to_log(f"get_ky_and_hyp {time.time()-time0}\n", - name="hyps") + logger.debug(f"get_ky_and_hyp {time.time()-time0}") time0 = time.time() like, like_grad = \ get_like_grad_from_mats(ky_mat, hyp_mat, name) - if output is not None: - output.write_to_log(f"get_like_grad_from_mats {time.time()-time0}\n", - name="hyps") - - if output is not None: - ostring = "like grad:" - for lg in like_grad: - ostring += f" {lg}" - ostring += "\n" - output.write_to_log(ostring, name="hyps") - output.write_to_log('like: ' + str(like)+'\n', name="hyps") - - if print_progress: - print('\nHyperparameters: ', list(hyps)) - print('Likelihood: ' + str(like)) - print('Likelihood Gradient: ', list(like_grad)) + logger.debug(f"get_like_grad_from_mats {time.time()-time0}") + + logger.debug('') + logger.info(f'Hyperparameters: {list(hyps)}') + logger.info(f'Likelihood: {like}') + logger.info(f'Likelihood Gradient: {list(like_grad)}') return -like, -like_grad diff --git a/flare/gp_from_aimd.py b/flare/gp_from_aimd.py index 2e7db391c..043a5f2cd 100644 --- a/flare/gp_from_aimd.py +++ b/flare/gp_from_aimd.py @@ -33,6 +33,7 @@ """ import json as json +import logging import numpy as np import time import warnings @@ -52,7 +53,7 @@ from flare.utils.element_coder import element_to_Z, Z_to_element, NumpyEncoder from flare.utils.learner import subset_of_frame_by_element, \ is_std_in_bound_per_species, is_force_in_bound_per_species -from flare.mgp.mgp import MappedGaussianProcess +from flare.mgp import MappedGaussianProcess class TrajectoryTrainer: @@ -74,7 +75,7 @@ def __init__(self, frames: List[Structure], max_trains: int = np.inf, min_atoms_per_train: int = 1, shuffle_frames: bool = False, - verbose: int = 1, + verbose: str = "INFO", pre_train_on_skips: int = -1, pre_train_seed_frames: List[Structure] = None, pre_train_seed_envs: List[Tuple[AtomicEnvironment, @@ -112,9 +113,7 @@ def __init__(self, frames: List[Structure], :param n_cpus: Number of CPUs to parallelize over for parallelization over atoms :param shuffle_frames: Randomize order of frames for better training - :param verbose: 0: Silent, NO output written or printed at all. - 1: Minimal, - 2: Lots of information + :param verbose: same as logging level, "WARNING", "INFO", "DEBUG" :param pre_train_on_skips: Train model on every n frames before running :param pre_train_seed_frames: Frames to train on before running :param pre_train_seed_envs: Environments to train on before running @@ -155,7 +154,6 @@ def __init__(self, frames: List[Structure], self.max_atoms_from_frame = max_atoms_from_frame self.min_atoms_per_train = min_atoms_per_train self.predict_atoms_per_element = predict_atoms_per_element - self.verbose = verbose self.train_count = 0 self.calculate_energy = calculate_energy self.n_cpus = n_cpus @@ -197,9 +195,9 @@ def __init__(self, frames: List[Structure], else pre_train_seed_frames self.pre_train_env_per_species = {} if pre_train_atoms_per_element \ - is None else pre_train_atoms_per_element + is None else pre_train_atoms_per_element self.train_env_per_species = {} if train_atoms_per_element \ - is None else train_atoms_per_element + is None else train_atoms_per_element # Convert to Coded Species if self.pre_train_env_per_species: @@ -209,13 +207,10 @@ def __init__(self, frames: List[Structure], self.pre_train_env_per_species[key] # Output parameters - self.verbose = verbose - if self.verbose: - self.output = Output(output_name, always_flush=True) - else: - self.output = None + self.output = Output(output_name, verbose, always_flush=True) + self.logger_name = self.output.basename+'log' self.train_checkpoint_interval = train_checkpoint_interval or \ - checkpoint_interval + checkpoint_interval self.atom_checkpoint_interval = atom_checkpoint_interval self.model_format = model_format @@ -239,27 +234,23 @@ def pre_run(self): if self.mgp: raise NotImplementedError("Pre-running not" "yet configured for MGP") - if self.verbose: - self.output.write_header(self.gp.cutoffs, - self.gp.kernel_name, - self.gp.hyps, - self.gp.opt_algorithm, - dt=0, - Nsteps=len(self.frames), - structure=None, - std_tolerance=(self.rel_std_tolerance, - self.abs_std_tolerance), - optional={ - 'GP Statistics': - json.dumps( - self.gp.training_statistics), - 'GP Name': self.gp.name, - 'GP Write Name': - self.output_name + "_model." + self.model_format}) + self.output.write_header(str(self.gp), + dt=0, + Nsteps=len(self.frames), + structure=None, + std_tolerance=(self.rel_std_tolerance, + self.abs_std_tolerance), + optional={ + 'GP Statistics': + json.dumps( + self.gp.training_statistics), + 'GP Name': self.gp.name, + 'GP Write Name': + self.output_name + "_model." + self.model_format}) self.start_time = time.time() - if self.verbose >= 3: - print("Now beginning pre-run activity.") + logger = logging.getLogger(self.logger_name) + logger.debug("Now beginning pre-run activity.") # If seed environments were passed in, add them to the GP. for point in self.seed_envs: @@ -307,24 +298,22 @@ def pre_run(self): train_atoms=train_atoms, uncertainties=[], train=False) - if self.verbose and atom_count > 0: - self.output.write_to_log(f"Added {atom_count} atoms to " - f"pretrain.\n" - f"Pre-run GP Statistics: " - f"{json.dumps(self.gp.training_statistics)} \n", - flush=True) + logger = logging.getLogger(self.logger_name) + if atom_count > 0: + logger.info(f"Added {atom_count} atoms to " + f"pretrain.\n" + f"Pre-run GP Statistics: " + f"{json.dumps(self.gp.training_statistics)} ") if (self.seed_envs or atom_count or self.seed_frames) and \ (self.pre_train_max_iter or self.max_trains): - if self.verbose >= 3: - print("Now commencing pre-run training of GP (which has " - "non-empty training set)") + logger.debug("Now commencing pre-run training of GP (which has " + "non-empty training set)") self.train_gp(max_iter=self.pre_train_max_iter) else: - if self.verbose >= 3: - print("Now commencing pre-run set up of GP (which has " - "non-empty training set)") - self.gp.set_L_alpha() + logger.debug("Now commencing pre-run set up of GP (which has " + "non-empty training set)") + self.gp.check_L_alpha() if self.model_format and not self.mgp: self.gp.write_model(f'{self.output_name}_prerun', @@ -341,8 +330,8 @@ def run(self): """ # Perform pre-run, in which seed trames are used. - if self.verbose >= 3: - print("Commencing run with pre-run...") + logger = logging.getLogger(self.logger_name) + logger.debug("Commencing run with pre-run...") if not self.mgp: self.pre_run() @@ -358,13 +347,12 @@ def run(self): for i, cur_frame in enumerate(self.frames[::self.skip]): - if self.verbose >= 2: - print(f"=====NOW ON FRAME {i}=====") + logger.info(f"=====NOW ON FRAME {i}=====") # If no predict_atoms_per_element was specified, predict_atoms # will be equal to every atom in the frame. predict_atoms = subset_of_frame_by_element(cur_frame, - self.predict_atoms_per_element) + self.predict_atoms_per_element) # Atoms which are skipped will have NaN as their force / std values local_energies = None @@ -373,8 +361,7 @@ def run(self): if self.mgp: pred_forces, pred_stds = self.pred_func( structure=cur_frame, mgp=self.gp, write_to_structure=False, - selective_atoms=predict_atoms, skipped_atom_value= - np.nan) + selective_atoms=predict_atoms, skipped_atom_value=np.nan) elif self.calculate_energy: pred_forces, pred_stds, local_energies = self.pred_func( structure=cur_frame, gp=self.gp, n_cpus=self.n_cpus, @@ -382,11 +369,11 @@ def run(self): skipped_atom_value=np.nan) else: pred_forces, pred_stds = self.pred_func(structure=cur_frame, - gp=self.gp, - n_cpus=self.n_cpus, - write_to_structure=False, - selective_atoms=predict_atoms, - skipped_atom_value=np.nan) + gp=self.gp, + n_cpus=self.n_cpus, + write_to_structure=False, + selective_atoms=predict_atoms, + skipped_atom_value=np.nan) # Get Error dft_forces = cur_frame.forces @@ -397,13 +384,12 @@ def run(self): dummy_frame.forces = pred_forces dummy_frame.stds = pred_stds - if self.verbose: - self.output.write_gp_dft_comparison( - curr_step=i, frame=dummy_frame, - start_time=time.time(), - dft_forces=dft_forces, - error=error, - local_energies=local_energies) + self.output.write_gp_dft_comparison( + curr_step=i, frame=dummy_frame, + start_time=time.time(), + dft_forces=dft_forces, + error=error, + local_energies=local_energies) if i < train_frame: # Noise hyperparameter & relative std tolerance is not for mgp. @@ -464,14 +450,14 @@ def run(self): if self.train_checkpoint_interval and \ cur_trains_done_write and \ self.train_checkpoint_interval \ - <= cur_trains_done_write: + <= cur_trains_done_write: will_write = True cur_trains_done_write = 0 if self.atom_checkpoint_interval \ and cur_atoms_added_write \ and self.atom_checkpoint_interval \ - <= cur_atoms_added_write: + <= cur_atoms_added_write: will_write = True cur_atoms_added_write = 0 @@ -482,8 +468,7 @@ def run(self): if (i + 1) == train_frame and not self.mgp: self.gp.check_L_alpha() - if self.verbose: - self.output.conclude_run() + self.output.conclude_run() if self.model_format and not self.mgp: self.gp.write_model(f'{self.output_name}_model', @@ -511,18 +496,17 @@ def update_gp_and_print(self, frame: Structure, train_atoms: List[int], for atom, spec in zip(train_atoms, added_species): added_atoms[spec].append(atom) - if self.verbose: - self.output.write_to_log('\nAdding atom(s) ' - f'{json.dumps(added_atoms,cls=NumpyEncoder)}' - ' to the training set.\n') + logger = logging.getLogger(self.logger_name) + logger.info('Adding atom(s) ' + f'{json.dumps(added_atoms,cls=NumpyEncoder)}' + ' to the training set.') if uncertainties is None or len(uncertainties) != 0: uncertainties = frame.stds[train_atoms] - if self.verbose and len(uncertainties) != 0: - self.output.write_to_log(f'Uncertainties: ' - f'{uncertainties}.\n', - flush=True) + if len(uncertainties) != 0: + logger.info(f'Uncertainties: ' + f'{uncertainties}.') # update gp model; handling differently if it's an MGP if not self.mgp: @@ -543,8 +527,10 @@ def train_gp(self, max_iter: int = None): :type max_iter: int """ - if self.verbose >= 1: - self.output.write_to_log('Train GP\n') + logger = logging.getLogger(self.logger_name) + logger.debug('Train GP') + + logger_train = logging.getLogger(self.output.basename+'hyps') # TODO: Improve flexibility in GP training to make this next step # unnecessary, so maxiter can be passed as an argument @@ -555,17 +541,16 @@ def train_gp(self, max_iter: int = None): elif max_iter is not None: temp_maxiter = self.gp.maxiter self.gp.maxiter = max_iter - self.gp.train(output=self.output if self.verbose >= 2 else None) + self.gp.train(logger=logger_train) self.gp.maxiter = temp_maxiter else: - self.gp.train(output=self.output if self.verbose >= 2 else None) - - if self.verbose: - self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps, - self.start_time, - self.gp.likelihood, - self.gp.likelihood_gradient, - hyps_mask=self.gp.hyps_mask) + self.gp.train(logger=logger_train) + + self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps, + self.start_time, + self.gp.likelihood, + self.gp.likelihood_gradient, + hyps_mask=self.gp.hyps_mask) self.train_count += 1 diff --git a/flare/kernels/kernels.py b/flare/kernels/kernels.py index 6b5d1da8c..e80e80db5 100644 --- a/flare/kernels/kernels.py +++ b/flare/kernels/kernels.py @@ -36,15 +36,11 @@ def force_helper(A, B, C, D, fi, fj, fdi, fdj, ls1, ls2, ls3, sig2): the same type. """ E = exp(-D * ls1) - F = E * B * ls2 - G = -E * C * ls2 - H = A * E * ls2 - B * C * E * ls3 - I = E * fdi * fdj - J = F * fi * fdj - K = G * fdi * fj - L = H * fi * fj - M = sig2 * (I + J + K + L) - + I = fdi * fdj + J = B * ls2 * fi * fdj + K = - C * ls2 * fdi * fj + L = (A * ls2 - B * C * ls3) * fi * fj + M = sig2 * (I + J + K + L) * E return M diff --git a/flare/kernels/map_3b_kernel.py b/flare/kernels/map_3b_kernel.py deleted file mode 100644 index e53377fa4..000000000 --- a/flare/kernels/map_3b_kernel.py +++ /dev/null @@ -1,654 +0,0 @@ -"""Multi-element 2-, 3-, and 2+3-body kernels that restrict all signal -variance hyperparameters to a single value.""" -import numpy as np -from numba import njit, prange -from math import exp, floor -from typing import Callable - -from flare.env import AtomicEnvironment -import flare.kernels.cutoffs as cf - -def three_body_mc_en(env1: AtomicEnvironment, r1, r2, r12, c2, etypes2, - hyps: 'ndarray', cutoffs: 'ndarray', - cutoff_func: Callable = cf.quadratic_cutoff) \ - -> float: - """3-body multi-element kernel between a force component and many local - energies on the grid. - - Args: - env1 (AtomicEnvironment): First local environment. - rj1 (np.ndarray): matrix of the first edge length - rj2 (np.ndarray): matrix of the second edge length - rj12 (np.ndarray): matrix of the third edge length - c2 (int): Species of the central atom of the second local environment. - etypes2 (np.ndarray): Species of atoms in the second local - environment. - d1 (int): Force component of the first environment (1=x, 2=y, 3=z). - hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1, - sig2, ls2). - cutoffs (np.ndarray): Two-element array containing the 2- and 3-body - cutoffs. - cutoff_func (Callable): Cutoff function. - - Returns: - float: - Value of the 3-body force/energy kernel. - """ - sig = hyps[0] - ls = hyps[1] - r_cut = cutoffs[1] - - return three_body_mc_en_jit(env1.bond_array_3, env1.ctype, - env1.etypes, - env1.cross_bond_inds, - env1.cross_bond_dists, - env1.triplet_counts, - c2, etypes2, - r1, r2, r12, - sig, ls, r_cut, cutoff_func) / 9. - -def three_body_mc_en_sephyps(env1, r1, r2, r12, c2, etypes2, - cutoff_2b, cutoff_3b, nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, - cutoff_func=cf.quadratic_cutoff) -> float: - """3-body multi-element kernel between a force component and many local - energies on the grid. - - Args: - env1 (AtomicEnvironment): First local environment. - rj1 (np.ndarray): matrix of the first edge length - rj2 (np.ndarray): matrix of the second edge length - rj12 (np.ndarray): matrix of the third edge length - c2 (int): Species of the central atom of the second local environment. - etypes2 (np.ndarray): Species of atoms in the second local - environment. - d1 (int): Force component of the first environment (1=x, 2=y, 3=z). - cutoff_2b: dummy - cutoff_3b (float, np.ndarray): cutoff(s) for three-body interaction - nspec (int): number of different species groups - spec_mask (np.ndarray): 118-long integer array that determines specie group - nbond: dummy - bond_mask: dummy - ntriplet (int): number of different hyperparameter sets to associate with 3-body pairings - triplet_mask (np.ndarray): nspec^3 long integer array - ncut3b (int): number of different 3-body cutoff sets to associate with 3-body pairings - cut3b_mask (np.ndarray): nspec^2 long integer array - sig2: dummy - ls2: dummy - sig3 (np.ndarray): signal variances associates with three-body term - ls3 (np.ndarray): length scales associates with three-body term - cutoff_func (Callable): Cutoff function of the kernel. - - Returns: - float: - Value of the 3-body force/energy kernel. - """ - - ej1 = etypes2[0] - ej2 = etypes2[1] - bc1 = spec_mask[c2] - bc2 = spec_mask[ej1] - bc3 = spec_mask[ej2] - ttype = triplet_mask[nspec * nspec * bc1 + nspec*bc2 + bc3] - ls = ls3[ttype] - sig = sig3[ttype] - r_cut = cutoff_3b - - return three_body_mc_en_jit(env1.bond_array_3, env1.ctype, - env1.etypes, - env1.cross_bond_inds, - env1.cross_bond_dists, - env1.triplet_counts, - c2, etypes2, - r1, r2, r12, - sig, ls, r_cut, cutoff_func) / 9. - - -def three_body_mc_en_force(env1: AtomicEnvironment, r1, r2, r12, c2, etypes2, - d1: int, hyps: 'ndarray', cutoffs: 'ndarray', - cutoff_func: Callable = cf.quadratic_cutoff) \ - -> float: - """3-body multi-element kernel between a force component and many local - energies on the grid. - - Args: - env1 (AtomicEnvironment): First local environment. - rj1 (np.ndarray): matrix of the first edge length - rj2 (np.ndarray): matrix of the second edge length - rj12 (np.ndarray): matrix of the third edge length - c2 (int): Species of the central atom of the second local environment. - etypes2 (np.ndarray): Species of atoms in the second local - environment. - d1 (int): Force component of the first environment (1=x, 2=y, 3=z). - hyps (np.ndarray): Hyperparameters of the kernel function (sig1, ls1, - sig2, ls2). - cutoffs (np.ndarray): Two-element array containing the 2- and 3-body - cutoffs. - cutoff_func (Callable): Cutoff function. - - Returns: - float: - Value of the 3-body force/energy kernel. - """ - sig = hyps[0] - ls = hyps[1] - r_cut = cutoffs[1] - - return three_body_mc_en_force_jit(env1.bond_array_3, env1.ctype, - env1.etypes, - env1.cross_bond_inds, - env1.cross_bond_dists, - env1.triplet_counts, - c2, etypes2, - r1, r2, r12, - d1, sig, ls, r_cut, cutoff_func) / 3 - -def three_body_mc_en_force_sephyps(env1, r1, r2, r12, c2, etypes2, - d1, cutoff_2b, cutoff_3b, nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, - cutoff_func=cf.quadratic_cutoff) -> float: - """3-body multi-element kernel between a force component and many local - energies on the grid. - - Args: - env1 (AtomicEnvironment): First local environment. - rj1 (np.ndarray): matrix of the first edge length - rj2 (np.ndarray): matrix of the second edge length - rj12 (np.ndarray): matrix of the third edge length - c2 (int): Species of the central atom of the second local environment. - etypes2 (np.ndarray): Species of atoms in the second local - environment. - d1 (int): Force component of the first environment (1=x, 2=y, 3=z). - cutoff_2b: dummy - cutoff_3b (float, np.ndarray): cutoff(s) for three-body interaction - nspec (int): number of different species groups - spec_mask (np.ndarray): 118-long integer array that determines specie group - nbond: dummy - bond_mask: dummy - ntriplet (int): number of different hyperparameter sets to associate with 3-body pairings - triplet_mask (np.ndarray): nspec^3 long integer array - ncut3b (int): number of different 3-body cutoff sets to associate with 3-body pairings - cut3b_mask (np.ndarray): nspec^2 long integer array - sig2: dummy - ls2: dummy - sig3 (np.ndarray): signal variances associates with three-body term - ls3 (np.ndarray): length scales associates with three-body term - cutoff_func (Callable): Cutoff function of the kernel. - - Returns: - float: - Value of the 3-body force/energy kernel. - """ - - ej1 = etypes2[0] - ej2 = etypes2[1] - bc1 = spec_mask[c2] - bc2 = spec_mask[ej1] - bc3 = spec_mask[ej2] - ttype = triplet_mask[nspec * nspec * bc1 + nspec*bc2 + bc3] - ls = ls3[ttype] - sig = sig3[ttype] - r_cut = cutoff_3b - - return three_body_mc_en_force_jit(env1.bond_array_3, env1.ctype, - env1.etypes, - env1.cross_bond_inds, - env1.cross_bond_dists, - env1.triplet_counts, - c2, etypes2, - r1, r2, r12, - d1, sig, ls, r_cut, cutoff_func) / 3 - - -# @njit -# def three_body_mc_force_en_jit(bond_array_1, c1, etypes1, -# cross_bond_inds_1, cross_bond_dists_1, -# triplets_1, -# c2, etypes2, -# rj1, rj2, rj3, -# d1, sig, ls, r_cut, cutoff_func): -# """3-body multi-element kernel between a force component and many local -# energies on the grid. -# -# Args: -# bond_array_1 (np.ndarray): 3-body bond array of the first local -# environment. -# c1 (int): Species of the central atom of the first local environment. -# etypes1 (np.ndarray): Species of atoms in the first local -# environment. -# cross_bond_inds_1 (np.ndarray): Two dimensional array whose row m -# contains the indices of atoms n > m in the first local -# environment that are within a distance r_cut of both atom n and -# the central atom. -# cross_bond_dists_1 (np.ndarray): Two dimensional array whose row m -# contains the distances from atom m of atoms n > m in the first -# local environment that are within a distance r_cut of both atom -# n and the central atom. -# triplets_1 (np.ndarray): One dimensional array of integers whose entry -# m is the number of atoms in the first local environment that are -# within a distance r_cut of atom m. -# c2 (int): Species of the central atom of the second local environment. -# etypes2 (np.ndarray): Species of atoms in the second local -# environment. -# rj1 (np.ndarray): matrix of the first edge length -# rj2 (np.ndarray): matrix of the second edge length -# rj12 (np.ndarray): matrix of the third edge length -# d1 (int): Force component of the first environment (1=x, 2=y, 3=z). -# sig (float): 3-body signal variance hyperparameter. -# ls (float): 3-body length scale hyperparameter. -# r_cut (float): 3-body cutoff radius. -# cutoff_func (Callable): Cutoff function. -# -# Returns: -# float: -# Value of the 3-body force/energy kernel. -# """ -# -# kern = np.zeros_like(rj1, dtype=np.float64) -# -# ei1 = etypes2[0] -# ei2 = etypes2[1] -# -# all_spec = [c2, ei1, ei2] -# if (c1 not in all_spec): -# return kern -# all_spec.remove(c1) -# -# # pre-compute constants that appear in the inner loop -# sig2 = sig * sig -# ls1 = 1 / (2 * ls * ls) -# ls2 = 1 / (ls * ls) -# -# f1, fdi1 = cutoff_func(r_cut, ri1, ci1) -# -# f2, fdi2 = cutoff_func(r_cut, ri2, ci2) -# f3, fdi3 = cutoff_func(r_cut, ri3, ci3) -# fi = f1 * f2 * f3 -# fdi = fdi1 * fi2 * fi3 + fi1 * fdi2 * fi3 -# # del f1 -# # del f2 -# # del f3 -# -# for m in prange(bond_array_1.shape[0]): -# ei1 = etypes1[m] -# -# two_spec = [all_spec[0], all_spec[1]] -# if (ei1 in two_spec): -# two_spec.remove(ei1) -# one_spec = two_spec[0] -# -# rj1 = bond_array_1[m, 0] -# fj1, _ = cutoff_func(r_cut, rj1, 0) -# -# for n in prange(triplets_1[m]): -# -# ind1 = cross_bond_inds_1[m, m + n + 1] -# ej2 = etypes1[ind1] -# -# if (ej2 == one_spec): -# -# if (ei2 == ej2): -# r11 = ri1 - rj1 -# if (ei2 == ej1): -# r12 = ri1 - rj2 -# if (ei2 == c2): -# r13 = ri1 - rj3 -# -# rj2 = bond_array_1[ind1, 0] -# if (ei1 == ej2): -# r21 = ri2 - rj1 -# if (ei1 == ej1): -# r22 = ri2 - rj2 -# if (ei1 == c2): -# r23 = ri2 - rj3 -# cj2 = bond_array_1[ind1, d1] -# fj2, _ = cutoff_func(r_cut, rj2, 0) -# # del ri2 -# -# rj3 = cross_bond_dists_1[m, m + n + 1] -# if (c1 == ej2): -# r31 = ri3 - rj1 -# if (c1 == ej1): -# r32 = ri3 - rj2 -# if (c1 == c2): -# r33 = ri3 - rj3 -# fj3, _ = cutoff_func(r_cut, rj3, 0) -# # del ri3 -# -# fj = fj1 * fj2 * fj3 -# # del fj1 -# # del fj2 -# # del fj3 -# -# if (c1 == c2): -# if (ei1 == ej1) and (ei2 == ej2): -# kern += three_body_en_helper(ci1, ci2, r11, r22, -# r33, fi, fj, fdi, ls1, -# ls2, sig2) -# if (ei1 == ej2) and (ei2 == ej1): -# kern += three_body_en_helper(ci1, ci2, r12, r21, -# r33, fi, fj, fdi, ls1, -# ls2, sig2) -# if (c1 == ej1): -# if (ei1 == ej2) and (ei2 == c2): -# kern += three_body_en_helper(ci1, ci2, r13, r21, -# r32, fi, fj, fdi, ls1, -# ls2, sig2) -# if (ei1 == c2) and (ei2 == ej2): -# kern += three_body_en_helper(ci1, ci2, r11, r23, -# r32, fi, fj, fdi, ls1, -# ls2, sig2) -# if (c1 == ej2): -# if (ei1 == ej1) and (ei2 == c2): -# kern += three_body_en_helper(ci1, ci2, r13, r22, -# r31, fi, fj, fdi, ls1, -# ls2, sig2) -# if (ei1 == c2) and (ei2 == ej1): -# kern += three_body_en_helper(ci1, ci2, r12, r23, -# r31, fi, fj, fdi, ls1, -# ls2, sig2) -# return kern - -@njit -def three_body_mc_en_force_jit(bond_array_1, c1, etypes1, - cross_bond_inds_1, cross_bond_dists_1, - triplets_1, - c2, etypes2, - rj1, rj2, rj3, - d1, sig, ls, r_cut, cutoff_func): - """3-body multi-element kernel between a force component and many local - energies on the grid. - - Args: - bond_array_1 (np.ndarray): 3-body bond array of the first local - environment. - c1 (int): Species of the central atom of the first local environment. - etypes1 (np.ndarray): Species of atoms in the first local - environment. - cross_bond_inds_1 (np.ndarray): Two dimensional array whose row m - contains the indices of atoms n > m in the first local - environment that are within a distance r_cut of both atom n and - the central atom. - cross_bond_dists_1 (np.ndarray): Two dimensional array whose row m - contains the distances from atom m of atoms n > m in the first - local environment that are within a distance r_cut of both atom - n and the central atom. - triplets_1 (np.ndarray): One dimensional array of integers whose entry - m is the number of atoms in the first local environment that are - within a distance r_cut of atom m. - c2 (int): Species of the central atom of the second local environment. - etypes2 (np.ndarray): Species of atoms in the second local - environment. - rj1 (np.ndarray): matrix of the first edge length - rj2 (np.ndarray): matrix of the second edge length - rj12 (np.ndarray): matrix of the third edge length - d1 (int): Force component of the first environment (1=x, 2=y, 3=z). - sig (float): 3-body signal variance hyperparameter. - ls (float): 3-body length scale hyperparameter. - r_cut (float): 3-body cutoff radius. - cutoff_func (Callable): Cutoff function. - - Returns: - float: - Value of the 3-body force/energy kernel. - """ - - kern = np.zeros_like(rj1, dtype=np.float64) - - ej1 = etypes2[0] - ej2 = etypes2[1] - - all_spec = [c2, ej1, ej2] - if (c1 not in all_spec): - return kern - all_spec.remove(c1) - - # pre-compute constants that appear in the inner loop - sig2 = sig * sig - ls1 = 1 / (2 * ls * ls) - ls2 = 1 / (ls * ls) - - f1, _ = cutoff_func(r_cut, rj1, 0) - f2, _ = cutoff_func(r_cut, rj2, 0) - f3, _ = cutoff_func(r_cut, rj3, 0) - fj = f1 * f2 * f3 - # del f1 - # del f2 - # del f3 - - for m in prange(bond_array_1.shape[0]): - ei1 = etypes1[m] - - two_spec = [all_spec[0], all_spec[1]] - if (ei1 in two_spec): - two_spec.remove(ei1) - one_spec = two_spec[0] - - ri1 = bond_array_1[m, 0] - ci1 = bond_array_1[m, d1] - fi1, fdi1 = cutoff_func(r_cut, ri1, ci1) - - for n in prange(triplets_1[m]): - - ind1 = cross_bond_inds_1[m, m + n + 1] - ei2 = etypes1[ind1] - - if (ei2 == one_spec): - - if (ei2 == ej2): - r11 = ri1 - rj1 - if (ei2 == ej1): - r12 = ri1 - rj2 - if (ei2 == c2): - r13 = ri1 - rj3 - - ri2 = bond_array_1[ind1, 0] - if (ei1 == ej2): - r21 = ri2 - rj1 - if (ei1 == ej1): - r22 = ri2 - rj2 - if (ei1 == c2): - r23 = ri2 - rj3 - ci2 = bond_array_1[ind1, d1] - fi2, fdi2 = cutoff_func(r_cut, ri2, ci2) - # del ri2 - - ri3 = cross_bond_dists_1[m, m + n + 1] - if (c1 == ej2): - r31 = ri3 - rj1 - if (c1 == ej1): - r32 = ri3 - rj2 - if (c1 == c2): - r33 = ri3 - rj3 - fi3, _ = cutoff_func(r_cut, ri3, 0) - # del ri3 - - fi = fi1 * fi2 * fi3 - fdi = fdi1 * fi2 * fi3 + fi1 * fdi2 * fi3 - # del fi1 - # del fi2 - # del fi3 - # del fdi1 - # del fdi2 - - if (c1 == c2): - if (ei1 == ej1) and (ei2 == ej2): - kern += three_body_en_helper(ci1, ci2, r11, r22, - r33, fi, fj, fdi, ls1, - ls2, sig2) - if (ei1 == ej2) and (ei2 == ej1): - kern += three_body_en_helper(ci1, ci2, r12, r21, - r33, fi, fj, fdi, ls1, - ls2, sig2) - if (c1 == ej1): - if (ei1 == ej2) and (ei2 == c2): - kern += three_body_en_helper(ci1, ci2, r13, r21, - r32, fi, fj, fdi, ls1, - ls2, sig2) - if (ei1 == c2) and (ei2 == ej2): - kern += three_body_en_helper(ci1, ci2, r11, r23, - r32, fi, fj, fdi, ls1, - ls2, sig2) - if (c1 == ej2): - if (ei1 == ej1) and (ei2 == c2): - kern += three_body_en_helper(ci1, ci2, r13, r22, - r31, fi, fj, fdi, ls1, - ls2, sig2) - if (ei1 == c2) and (ei2 == ej1): - kern += three_body_en_helper(ci1, ci2, r12, r23, - r31, fi, fj, fdi, ls1, - ls2, sig2) - return kern - -@njit -def three_body_mc_en_jit(bond_array_1, c1, etypes1, - cross_bond_inds_1, - cross_bond_dists_1, - triplets_1, - c2, etypes2, - rj1, rj2, rj3, - sig, ls, r_cut, cutoff_func): - """3-body multi-element kernel between two local energies accelerated - with Numba. - - Args: - bond_array_1 (np.ndarray): 3-body bond array of the first local - environment. - c1 (int): Species of the central atom of the first local environment. - etypes1 (np.ndarray): Species of atoms in the first local - environment. - bond_array_2 (np.ndarray): 3-body bond array of the second local - environment. - c2 (int): Species of the central atom of the second local environment. - etypes2 (np.ndarray): Species of atoms in the second local - environment. - cross_bond_inds_1 (np.ndarray): Two dimensional array whose row m - contains the indices of atoms n > m in the first local - environment that are within a distance r_cut of both atom n and - the central atom. - cross_bond_inds_2 (np.ndarray): Two dimensional array whose row m - contains the indices of atoms n > m in the second local - environment that are within a distance r_cut of both atom n and - the central atom. - cross_bond_dists_1 (np.ndarray): Two dimensional array whose row m - contains the distances from atom m of atoms n > m in the first - local environment that are within a distance r_cut of both atom - n and the central atom. - cross_bond_dists_2 (np.ndarray): Two dimensional array whose row m - contains the distances from atom m of atoms n > m in the second - local environment that are within a distance r_cut of both atom - n and the central atom. - triplets_1 (np.ndarray): One dimensional array of integers whose entry - m is the number of atoms in the first local environment that are - within a distance r_cut of atom m. - triplets_2 (np.ndarray): One dimensional array of integers whose entry - m is the number of atoms in the second local environment that are - within a distance r_cut of atom m. - sig (float): 3-body signal variance hyperparameter. - ls (float): 3-body length scale hyperparameter. - r_cut (float): 3-body cutoff radius. - cutoff_func (Callable): Cutoff function. - - Returns: - float: - Value of the 3-body local energy kernel. - """ - - kern = np.zeros_like(rj1, dtype=np.float64) - - ej1 = etypes2[0] - ej2 = etypes2[1] - - all_spec = [c2, ej1, ej2] - if (c1 not in all_spec): - return kern - all_spec.remove(c1) - - # pre-compute constants that appear in the inner loop - sig2 = sig * sig - ls2 = 1 / (2 * ls * ls) - - - f1, _ = cutoff_func(r_cut, rj1, 0) - f2, _ = cutoff_func(r_cut, rj2, 0) - f3, _ = cutoff_func(r_cut, rj3, 0) - fj = f1 * f2 * f3 - - for m in prange(bond_array_1.shape[0]): - ri1 = bond_array_1[m, 0] - fi1, _ = cutoff_func(r_cut, ri1, 0) - ei1 = etypes1[m] - - two_spec = [all_spec[0], all_spec[1]] - if (ei1 in two_spec): - two_spec.remove(ei1) - one_spec = two_spec[0] - - for n in prange(triplets_1[m]): - ei2 = etypes1[ind1] - if (ei2 == one_spec): - - if (ei2 == ej2): - r11 = ri1 - rj1 - if (ei2 == ej1): - r12 = ri1 - rj2 - if (ei2 == c2): - r13 = ri1 - rj3 - - ri2 = bond_array_1[ind1, 0] - if (ei1 == ej2): - r21 = ri2 - rj1 - if (ei1 == ej1): - r22 = ri2 - rj2 - if (ei1 == c2): - r23 = ri2 - rj3 - ci2 = bond_array_1[ind1, d1] - fi2, _ = cutoff_func(r_cut, ri2, ci2) - # del ri2 - - ri3 = cross_bond_dists_1[m, m + n + 1] - if (c1 == ej2): - r31 = ri3 - rj1 - if (c1 == ej1): - r32 = ri3 - rj2 - if (c1 == c2): - r33 = ri3 - rj3 - fi3, _ = cutoff_func(r_cut, ri3, 0) - - fi = fi1 * fi2 * fi3 - - if (c1 == c2): - if (ei1 == ej1) and (ei2 == ej2): - C1 = r11 * r11 + r22 * r22 + r33 * r33 - kern += sig2 * np.exp(-C1 * ls2) * fi * fj - if (ei1 == ej2) and (ei2 == ej1): - C3 = r12 * r12 + r21 * r21 + r33 * r33 - kern += sig2 * np.exp(-C3 * ls2) * fi * fj - if (c1 == ej1): - if (ei1 == ej2) and (ei2 == c2): - C5 = r13 * r13 + r21 * r21 + r32 * r32 - kern += sig2 * np.exp(-C5 * ls2) * fi * fj - if (ei1 == c2) and (ei2 == ej2): - C2 = r11 * r11 + r23 * r23 + r32 * r32 - kern += sig2 * np.exp(-C2 * ls2) * fi * fj - if (c1 == ej2): - if (ei1 == ej1) and (ei2 == c2): - C6 = r13 * r13 + r22 * r22 + r31 * r31 - kern += sig2 * np.exp(-C6 * ls2) * fi * fj - if (ei1 == c2) and (ei2 == ej1): - C4 = r12 * r12 + r23 * r23 + r31 * r31 - kern += sig2 * np.exp(-C4 * ls2) * fi * fj - - return kern - - -@njit -def three_body_en_helper(ci1, ci2, r11, r22, r33, fi, fj, fdi, ls1, ls2, sig2): - - B = r11 * ci1 + r22 * ci2 - D = r11 * r11 + r22 * r22 + r33 * r33 - return -sig2 * np.exp(- D * ls1) * ( B * ls2 * fi * fj + fdi * fj) diff --git a/flare/kernels/mc_mb_sepcut.py b/flare/kernels/mc_mb_sepcut.py index 0bb2a536d..cd0bb802d 100644 --- a/flare/kernels/mc_mb_sepcut.py +++ b/flare/kernels/mc_mb_sepcut.py @@ -49,7 +49,7 @@ def many_body_mc_sepcut_jit(q_array_1, q_array_2, kern = 0 useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) bc1 = spec_mask[c1] bc1n = bc1 * nspec @@ -154,7 +154,7 @@ def many_body_mc_grad_sepcut_jit(q_array_1, q_array_2, ls_derv = np.zeros(nmb, dtype=np.float64) useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) bc1 = spec_mask[c1] bc1n = bc1 * nspec @@ -230,33 +230,37 @@ def many_body_mc_grad_sepcut_jit(q_array_1, q_array_2, dkij = 0 # c1 s and c2 s and if c1==c2 --> c1 s - kern_term_c1s = q1i_grads * q2j_grads * k12 - if (sig[mbtype1] !=0): - sig_derv[mbtype1] += kern_term_c1s * 2. / sig[mbtype1] - kern += kern_term_c1s + if k12 != 0: + kern_term_c1s = q1i_grads * q2j_grads * k12 + if sig[mbtype1] !=0: + sig_derv[mbtype1] += kern_term_c1s * 2. / sig[mbtype1] + kern += kern_term_c1s + ls_derv[mbtype1] += q1i_grads * q2j_grads * dk12 # s e1 and c2 s and c2==e1 --> c2 s - kern_term_c2s = qi1_grads * q2j_grads * ki2s - if (sig[mbtype2] !=0): - sig_derv[mbtype2] += kern_term_c2s * 2. / sig[mbtype2] - kern += kern_term_c2s + if ki2s != 0: + kern_term_c2s = qi1_grads * q2j_grads * ki2s + if sig[mbtype2] !=0: + sig_derv[mbtype2] += kern_term_c2s * 2. / sig[mbtype2] + kern += kern_term_c2s + ls_derv[mbtype2] += qi1_grads * q2j_grads * dki2s # c1 s and s e2 and c1==e2 --> c1 s - kern_term_c1s = q1i_grads * qj2_grads * k1js - if (sig[mbtype1] !=0): - sig_derv[mbtype1] += kern_term_c1s * 2. / sig[mbtype1] - kern += kern_term_c1s + if k1js != 0: + kern_term_c1s = q1i_grads * qj2_grads * k1js + if sig[mbtype1] !=0: + sig_derv[mbtype1] += kern_term_c1s * 2. / sig[mbtype1] + kern += kern_term_c1s + ls_derv[mbtype1] += q1i_grads * qj2_grads * dk1js # s e1 and s e2 and e1 == e2 -> s e - kern_term_se = qi1_grads * qj2_grads * kij - if (sig[mbtype] !=0): - sig_derv[mbtype] += kern_term_se * 2. / sig[mbtype] - kern += kern_term_se + if kij != 0: + kern_term_se = qi1_grads * qj2_grads * kij + if sig[mbtype] !=0: + sig_derv[mbtype] += kern_term_se * 2. / sig[mbtype] + kern += kern_term_se + ls_derv[mbtype] += qi1_grads * qj2_grads * dkij - ls_derv[mbtype1] += q1i_grads * q2j_grads * dk12 - ls_derv[mbtype2] += qi1_grads * q2j_grads * dki2s - ls_derv[mbtype1] += q1i_grads * qj2_grads * dk1js - ls_derv[mbtype] += qi1_grads * qj2_grads * dkij grad = np.zeros(nmb*2, dtype=np.float64) grad[:nmb] = sig_derv @@ -292,7 +296,7 @@ def many_body_mc_force_en_sepcut_jit(q_array_1, q_array_2, kern = 0 useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) bc1 = spec_mask[c1] bc1n = bc1 * nspec @@ -361,7 +365,7 @@ def many_body_mc_en_sepcut_jit(q_array_1, q_array_2, c1, c2, float: Value of the many-body kernel. """ useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) kern = 0 diff --git a/flare/kernels/mc_sephyps.py b/flare/kernels/mc_sephyps.py index 1efebc292..7c2bdbf57 100644 --- a/flare/kernels/mc_sephyps.py +++ b/flare/kernels/mc_sephyps.py @@ -7,29 +7,32 @@ To use this set of kernels, we need a hyps_mask dictionary for GaussianProcess, MappedGaussianProcess, and AtomicEnvironment (if you also set up different cutoffs). A simple example is shown below. - from flare.utils.mask_helper import HyperParameterMasking - from flare.gp import GaussianProcess - - pm = HyperParameterMasking(species=['O', 'C', 'H'], - bonds=[['*', '*'], ['O','O']], - triplets=[['*', '*', '*'], ['O','O', 'O']], - parameters={'bond0':[1, 0.5, 1], 'bond1':[2, 0.2, 2], - 'triplet0':[1, 0.5], 'triplet1':[2, 0.2], - 'cutoff2b':2, 'cutoff3b':1, 'noise': 0.05}, - constraints={'bond0':[False, True]}) - hyps_mask = pm1.generate_dict() - hyps = hyps_mask['hyps'] - cutoffs = hyps_mask['cutoffs'] - hyp_labels = hyps_mask['hyp_labels'] - gp_model = GaussianProcess(kernel_name="2+3_mc", - hyps=hyps, cutoffs=cutoffs, - hyp_labels=hyp_labels, - parallel=True, per_atom_par=False, - n_cpus=n_cpus, - multihyps=True, hyps_mask=hm) - - -In the example above, HyperParameterMasking class generates the arrays needed +Examples: + + >>> from flare.util.parameter_helper import ParameterHelper + >>> from flare.gp import GaussianProcess + + >>> pm = ParameterHelper(species=['O', 'C', 'H'], + ... kernels={'twobody':[['*', '*'], ['O','O']], + ... 'threebody':[['*', '*', '*'], ['O','O', 'O']]}, + ... parameters={'twobody0':[1, 0.5, 1], 'twobody1':[2, 0.2, 2], + ... 'triplet0':[1, 0.5], 'triplet1':[2, 0.2], + ... 'cutoff_twobody':2, 'cutoff_threebody':1, 'noise': 0.05}, + ... constraints={'twobody0':[False, True]}) + >>> hyps_mask = pm1.as_dict() + >>> hyps = hyps_mask.pop('hyps') + >>> cutoffs = hyps_mask.pop('cutoffs') + >>> hyp_labels = hyps_mask.pop('hyp_labels') + >>> kernels = hyps_mask['kernels'] + >>> gp_model = GaussianProcess(kernels=kernels, + ... hyps=hyps, cutoffs=cutoffs, + ... hyp_labels=hyp_labels, + ... parallel=True, per_atom_par=False, + ... n_cpus=n_cpus, + ... multihyps=True, hyps_mask=hm) + + +In the example above, Parameters class generates the arrays needed for these kernels and store all the grouping and mapping information in the hyps_mask dictionary. It stores following keys and values: @@ -184,12 +187,12 @@ def two_three_many_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, triplet_mask, cut3b_mask) mbmcj = many_body_mc_sepcut_jit - many_term = mbmcj(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, + many_term = mbmcj(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, + env1.ctype, env2.ctype, + env1.etypes_mb, env2.etypes_mb, + env1.unique_species, env2.unique_species, d1, d2, sigm, lsm, nspec, spec_mask, mb_mask) @@ -262,12 +265,12 @@ def two_three_many_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff ntriplet, triplet_mask, cut3b_mask) mbmcj = many_body_mc_grad_sepcut_jit - kern_many, gradm = mbmcj(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, + kern_many, gradm = mbmcj(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, + env1.ctype, env2.ctype, env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, + env1.unique_species, env2.unique_species, d1, d2, sigm, lsm, nspec, spec_mask, nmb, mb_mask) @@ -338,10 +341,10 @@ def two_three_many_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, cutoff_mb, cut3b_mask) / 3 mbmcj = many_body_mc_force_en_sepcut_jit - many_term = mbmcj(env1.q_array, env2.q_array, + many_term = mbmcj(env1.q_array, env2.q_array, env1.q_neigh_array, env1.q_neigh_grads, - env1.ctype, env2.ctype, env1.etypes_mb, - env1.unique_species, env2.unique_species, + env1.ctype, env2.ctype, env1.etypes_mb, + env1.unique_species, env2.unique_species, d1, sigm, lsm, nspec, spec_mask, mb_mask) @@ -410,8 +413,8 @@ def two_three_many_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, triplet_mask, cut3b_mask)/9. mbmcj = many_body_mc_en_sepcut_jit - many_term = mbmcj(env1.q_array, env2.q_array, - env1.ctype, env2.ctype, + many_term = mbmcj(env1.q_array, env2.q_array, + env1.ctype, env2.ctype, env1.unique_species, env2.unique_species, sigm, lsm, nspec, spec_mask, mb_mask) @@ -424,11 +427,12 @@ def two_three_many_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, # ----------------------------------------------------------------------------- -def two_plus_three_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, +def two_plus_three_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2+3-body multi-element kernel between two force components. @@ -480,11 +484,12 @@ def two_plus_three_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, return two_term + three_term -def two_plus_three_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, +def two_plus_three_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2+3-body multi-element kernel between two force components and its gradient with respect to the hyperparameters. @@ -544,11 +549,11 @@ def two_plus_three_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, return kern2 + kern3, g -def two_plus_three_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, - nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, +def two_plus_three_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, nbond, bond_mask, + ntriplet, triplet_mask, ncut3b, cut3b_mask, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2+3-body multi-element kernel between force and local energy @@ -604,11 +609,11 @@ def two_plus_three_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, return two_term + three_term -def two_plus_three_mc_en(env1, env2, cutoff_2b, cutoff_3b, - nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, +def two_plus_three_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, nbond, bond_mask, + ntriplet, triplet_mask, ncut3b, cut3b_mask, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2+3-body multi-element kernel between two local energies @@ -668,11 +673,11 @@ def two_plus_three_mc_en(env1, env2, cutoff_2b, cutoff_3b, # ----------------------------------------------------------------------------- -def three_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, - nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, +def three_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, nbond, bond_mask, + ntriplet, triplet_mask, ncut3b, cut3b_mask, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """3-body multi-element kernel between two force components. @@ -716,11 +721,11 @@ def three_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, triplet_mask, cut3b_mask) -def three_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, - nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, +def three_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, nbond, bond_mask, + ntriplet, triplet_mask, ncut3b, cut3b_mask, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """3-body multi-element kernel between two force components and its gradient with respect to the hyperparameters. @@ -767,12 +772,11 @@ def three_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, nspec, spec_mask, ntriplet, triplet_mask, cut3b_mask) -def three_body_mc_force_en( - env1, env2, d1, cutoff_2b, cutoff_3b, nspec, spec_mask, - nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, - cutoff_func=cf.quadratic_cutoff): +def three_body_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, + ncut3b, cut3b_mask, nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, + cutoff_func=cf.quadratic_cutoff): """3-body multi-element kernel between a force component and local energies Args: @@ -822,9 +826,10 @@ def three_body_mc_force_en( triplet_mask, cut3b_mask) / 3 -def three_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, nspec, spec_mask, +def three_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, sig2, ls2, sig3, ls3, + ncut3b, cut3b_mask, nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """3-body multi-element kernel between two local energies @@ -874,9 +879,9 @@ def three_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, nspec, spec_mask, def two_body_mc( - env1, env2, d1, d2, cutoff_2b, cutoff_3b, nspec, spec_mask, + env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, + nmb, mb_mask, sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2-body multi-element kernel between two force components. @@ -912,9 +917,10 @@ def two_body_mc( def two_body_mc_grad( - env1, env2, d1, d2, cutoff_2b, cutoff_3b, nspec, spec_mask, + env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, sig2, ls2, sig3, ls3, + ncut3b, cut3b_mask, nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2-body multi-element kernel between two force components and its gradient with respect to the hyperparameters. @@ -953,10 +959,11 @@ def two_body_mc_grad( nspec, spec_mask, nbond, bond_mask) -def two_body_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, +def two_body_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3, cutoff_func=cf.quadratic_cutoff): + ncut3b, cut3b_mask, nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, + cutoff_func=cf.quadratic_cutoff): """2-body multi-element kernel between a force components and local energy Args: @@ -991,9 +998,11 @@ def two_body_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, nspec, spec_mask, bond_mask) / 2 -def two_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, nspec, spec_mask, +def two_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, - ncut3b, cut3b_mask, sig2, ls2, sig3, ls3, + ncut3b, cut3b_mask, nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, cutoff_func=cf.quadratic_cutoff): """2-body multi-element kernel between two local energies @@ -1469,6 +1478,7 @@ def three_body_mc_force_en_jit(bond_array_1, c1, etypes1, return kern + @njit def three_body_mc_en_jit(bond_array_1, c1, etypes1, bond_array_2, c2, etypes2, @@ -1546,24 +1556,30 @@ def three_body_mc_en_jit(bond_array_1, c1, etypes1, if (c1 == c2): if (ei1 == ej1) and (ei2 == ej2): C1 = r11 * r11 + r22 * r22 + r33 * r33 - kern += tsig2 * exp(-C1 * tls2) * fi * fj + kern += tsig2 * \ + exp(-C1 * tls2) * fi * fj if (ei1 == ej2) and (ei2 == ej1): C3 = r12 * r12 + r21 * r21 + r33 * r33 - kern += tsig2 * exp(-C3 * tls2) * fi * fj + kern += tsig2 * \ + exp(-C3 * tls2) * fi * fj if (c1 == ej1): if (ei1 == ej2) and (ei2 == c2): C5 = r13 * r13 + r21 * r21 + r32 * r32 - kern += tsig2 * exp(-C5 * tls2) * fi * fj + kern += tsig2 * \ + exp(-C5 * tls2) * fi * fj if (ei1 == c2) and (ei2 == ej2): C2 = r11 * r11 + r23 * r23 + r32 * r32 - kern += tsig2 * exp(-C2 * tls2) * fi * fj + kern += tsig2 * \ + exp(-C2 * tls2) * fi * fj if (c1 == ej2): if (ei1 == ej1) and (ei2 == c2): C6 = r13 * r13 + r22 * r22 + r31 * r31 - kern += tsig2 * exp(-C6 * tls2) * fi * fj + kern += tsig2 * \ + exp(-C6 * tls2) * fi * fj if (ei1 == c2) and (ei2 == ej1): C4 = r12 * r12 + r23 * r23 + r31 * r31 - kern += tsig2 * exp(-C4 * tls2) * fi * fj + kern += tsig2 * \ + exp(-C4 * tls2) * fi * fj return kern @@ -1800,6 +1816,7 @@ def two_body_mc_en_jit(bond_array_1, c1, etypes1, return kern + def many_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, nspec, spec_mask, nbond, bond_mask, ntriplet, triplet_mask, @@ -1838,15 +1855,14 @@ def many_body_mc(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, Return: float: Value of the 2+3+many-body kernel. """ - return many_body_mc_sepcut_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, - env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, - d1, d2, sigm, lsm, - nspec, spec_mask, mb_mask) - + return many_body_mc_sepcut_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, + env1.q_neigh_grads, env2.q_neigh_grads, + env1.ctype, env2.ctype, + env1.etypes_mb, env2.etypes_mb, + env1.unique_species, env2.unique_species, + d1, d2, sigm, lsm, + nspec, spec_mask, mb_mask) def many_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, @@ -1891,14 +1907,14 @@ def many_body_mc_grad(env1, env2, d1, d2, cutoff_2b, cutoff_3b, cutoff_mb, with respect to the hyperparameters. """ - return many_body_mc_grad_sepcut_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, - env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, - d1, d2, sigm, lsm, - nspec, spec_mask, nmb, mb_mask) + return many_body_mc_grad_sepcut_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, + env1.q_neigh_grads, env2.q_neigh_grads, + env1.ctype, env2.ctype, + env1.etypes_mb, env2.etypes_mb, + env1.unique_species, env2.unique_species, + d1, d2, sigm, lsm, + nspec, spec_mask, nmb, mb_mask) def many_body_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, cutoff_mb, @@ -1922,14 +1938,14 @@ def many_body_mc_force_en(env1, env2, d1, cutoff_2b, cutoff_3b, cutoff_mb, float: Value of the many-body force/energy kernel. """ - return many_body_mc_force_en_sepcut_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, - env1.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, - env1.unique_species, env2.unique_species, - d1, sigm, lsm, - nspec, spec_mask, mb_mask) + return many_body_mc_force_en_sepcut_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, + env1.q_neigh_grads, + env1.ctype, env2.ctype, + env1.etypes_mb, + env1.unique_species, env2.unique_species, + d1, sigm, lsm, + nspec, spec_mask, mb_mask) def many_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, @@ -1953,11 +1969,11 @@ def many_body_mc_en(env1, env2, cutoff_2b, cutoff_3b, cutoff_mb, float: Value of the 2-body force/energy kernel. """ - return many_body_mc_en_sepcut_jit(env1.q_array, env2.q_array, - env1.ctype, env2.ctype, - env1.unique_species, env2.unique_species, - sigm, lsm, - nspec, spec_mask, mb_mask) + return many_body_mc_en_sepcut_jit(env1.q_array, env2.q_array, + env1.ctype, env2.ctype, + env1.unique_species, env2.unique_species, + sigm, lsm, + nspec, spec_mask, mb_mask) _str_to_kernel = {'2': two_body_mc, diff --git a/flare/kernels/mc_simple.py b/flare/kernels/mc_simple.py index e12696de1..4919f1071 100644 --- a/flare/kernels/mc_simple.py +++ b/flare/kernels/mc_simple.py @@ -10,8 +10,8 @@ from flare.kernels.kernels import force_helper, grad_constants, grad_helper, \ force_energy_helper, three_body_en_helper, three_body_helper_1, \ three_body_helper_2, three_body_grad_helper_1, three_body_grad_helper_2, \ - k_sq_exp_double_dev, k_sq_exp_dev, coordination_number, q_value, q_value_mc, \ - mb_grad_helper_ls_, mb_grad_helper_ls + k_sq_exp_double_dev, k_sq_exp_dev, coordination_number, q_value, \ + q_value_mc, mb_grad_helper_ls_, mb_grad_helper_ls from typing import Callable @@ -207,7 +207,8 @@ def two_plus_three_mc_en(env1: AtomicEnvironment, env2: AtomicEnvironment, # two plus three plus many body kernels # ----------------------------------------------------------------------------- -def two_plus_three_plus_many_body_mc(env1: AtomicEnvironment, env2: AtomicEnvironment, +def two_plus_three_plus_many_body_mc(env1: AtomicEnvironment, + env2: AtomicEnvironment, d1: int, d2: int, hyps, cutoffs, cutoff_func=cf.quadratic_cutoff): """2+3-body single-element kernel between two force components. @@ -250,18 +251,20 @@ def two_plus_three_plus_many_body_mc(env1: AtomicEnvironment, env2: AtomicEnviro env1.triplet_counts, env2.triplet_counts, d1, d2, sig3, ls3, r_cut_3, cutoff_func) - many_term = many_body_mc_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, - env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, - d1, d2, sigm, lsm) + many_term = \ + many_body_mc_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, + env1.q_neigh_grads, env2.q_neigh_grads, + env1.ctype, env2.ctype, + env1.etypes_mb, env2.etypes_mb, + env1.unique_species, env2.unique_species, + d1, d2, sigm, lsm) return two_term + three_term + many_term -def two_plus_three_plus_many_body_mc_grad(env1: AtomicEnvironment, env2: AtomicEnvironment, +def two_plus_three_plus_many_body_mc_grad(env1: AtomicEnvironment, + env2: AtomicEnvironment, d1: int, d2: int, hyps, cutoffs, cutoff_func=cf.quadratic_cutoff): """2+3+many-body single-element kernel between two force components. @@ -292,9 +295,10 @@ def two_plus_three_plus_many_body_mc_grad(env1: AtomicEnvironment, env2: AtomicE r_cut_3 = cutoffs[1] r_cut_m = cutoffs[2] - kern2, grad2 = two_body_mc_grad_jit(env1.bond_array_2, env1.ctype, env1.etypes, - env2.bond_array_2, env2.ctype, env2.etypes, - d1, d2, sig2, ls2, r_cut_2, cutoff_func) + kern2, grad2 = \ + two_body_mc_grad_jit(env1.bond_array_2, env1.ctype, env1.etypes, + env2.bond_array_2, env2.ctype, env2.etypes, + d1, d2, sig2, ls2, r_cut_2, cutoff_func) kern3, grad3 = \ three_body_mc_grad_jit(env1.bond_array_3, env1.ctype, env1.etypes, @@ -304,21 +308,24 @@ def two_plus_three_plus_many_body_mc_grad(env1: AtomicEnvironment, env2: AtomicE env1.triplet_counts, env2.triplet_counts, d1, d2, sig3, ls3, r_cut_3, cutoff_func) - kern_many, gradm = many_body_mc_grad_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, - env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, - d1, d2, sigm, lsm) + kern_many, gradm = \ + many_body_mc_grad_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, + env1.q_neigh_grads, env2.q_neigh_grads, + env1.ctype, env2.ctype, + env1.etypes_mb, env2.etypes_mb, + env1.unique_species, env2.unique_species, + d1, d2, sigm, lsm) return kern2 + kern3 + kern_many, np.hstack([grad2, grad3, gradm]) -def two_plus_three_plus_many_body_mc_force_en(env1: AtomicEnvironment, env2: AtomicEnvironment, +def two_plus_three_plus_many_body_mc_force_en(env1: AtomicEnvironment, + env2: AtomicEnvironment, d1: int, hyps, cutoffs, cutoff_func=cf.quadratic_cutoff): - """2+3+many-body single-element kernel between two force and energy components. + """2+3+many-body single-element kernel between two force and energy + components. Args: env1 (AtomicEnvironment): First local environment. @@ -359,11 +366,12 @@ def two_plus_three_plus_many_body_mc_force_en(env1: AtomicEnvironment, env2: Ato env1.triplet_counts, env2.triplet_counts, d1, sig3, ls3, r_cut_3, cutoff_func) / 3 - many_term = many_body_mc_force_en_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env1.q_neigh_grads, - env1.ctype, env2.ctype, env1.etypes_mb, - env1.unique_species, env2.unique_species, - d1, sigm, lsm) + many_term = \ + many_body_mc_force_en_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env1.q_neigh_grads, + env1.ctype, env2.ctype, env1.etypes_mb, + env1.unique_species, env2.unique_species, + d1, sigm, lsm) return two_term + three_term + many_term @@ -410,8 +418,8 @@ def two_plus_three_plus_many_body_mc_en(env1: AtomicEnvironment, env1.triplet_counts, env2.triplet_counts, sig3, ls3, r_cut_3, cutoff_func)/9 - many_term = many_body_mc_en_jit(env1.q_array, env2.q_array, - env1.ctype, env2.ctype, + many_term = many_body_mc_en_jit(env1.q_array, env2.q_array, + env1.ctype, env2.ctype, env1.unique_species, env2.unique_species, sigm, lsm) @@ -695,32 +703,30 @@ def many_body_mc(env1: AtomicEnvironment, env2: AtomicEnvironment, Return: float: Value of the 3-body kernel. """ - return many_body_mc_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, + return many_body_mc_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, - env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, + env1.ctype, env2.ctype, + env1.etypes_mb, env2.etypes_mb, + env1.unique_species, env2.unique_species, d1, d2, hyps[0], hyps[1]) - def many_body_mc_grad(env1: AtomicEnvironment, env2: AtomicEnvironment, d1: int, d2: int, hyps: 'ndarray', cutoffs: 'ndarray', cutoff_func: Callable = cf.quadratic_cutoff) -> float: """gradient manybody-body multi-element kernel between two force components. """ - return many_body_mc_grad_jit(env1.q_array, env2.q_array, - env1.q_neigh_array, env2.q_neigh_array, + return many_body_mc_grad_jit(env1.q_array, env2.q_array, + env1.q_neigh_array, env2.q_neigh_array, env1.q_neigh_grads, env2.q_neigh_grads, - env1.ctype, env2.ctype, + env1.ctype, env2.ctype, env1.etypes_mb, env2.etypes_mb, - env1.unique_species, env2.unique_species, + env1.unique_species, env2.unique_species, d1, d2, hyps[0], hyps[1]) - def many_body_mc_force_en(env1, env2, d1, hyps, cutoffs, cutoff_func=cf.quadratic_cutoff): """many-body single-element kernel between two local energies. @@ -737,10 +743,10 @@ def many_body_mc_force_en(env1, env2, d1, hyps, cutoffs, float: Value of the many-body force/energy kernel. """ # divide by three to account for triple counting - return many_body_mc_force_en_jit(env1.q_array, env2.q_array, + return many_body_mc_force_en_jit(env1.q_array, env2.q_array, env1.q_neigh_array, env1.q_neigh_grads, - env1.ctype, env2.ctype, env1.etypes_mb, - env1.unique_species, env2.unique_species, + env1.ctype, env2.ctype, env1.etypes_mb, + env1.unique_species, env2.unique_species, d1, hyps[0], hyps[1]) @@ -760,8 +766,8 @@ def many_body_mc_en(env1: AtomicEnvironment, env2: AtomicEnvironment, Return: float: Value of the 2-body force/energy kernel. """ - return many_body_mc_en_jit(env1.q_array, env2.q_array, - env1.ctype, env2.ctype, + return many_body_mc_en_jit(env1.q_array, env2.q_array, + env1.ctype, env2.ctype, env1.unique_species, env2.unique_species, hyps[0], hyps[1]) @@ -904,36 +910,48 @@ def three_body_mc_jit(bond_array_1, c1, etypes1, if (c1 == c2): if (ei1 == ej1) and (ei2 == ej2): kern += \ - three_body_helper_1(ci1, ci2, cj1, cj2, r11, - r22, r33, fi, fj, fdi, fdj, - ls1, ls2, ls3, sig2) + three_body_helper_1(ci1, ci2, cj1, + cj2, r11, r22, + r33, fi, fj, + fdi, fdj, ls1, + ls2, ls3, sig2) if (ei1 == ej2) and (ei2 == ej1): kern += \ - three_body_helper_1(ci1, ci2, cj2, cj1, r12, - r21, r33, fi, fj, fdi, fdj, - ls1, ls2, ls3, sig2) + three_body_helper_1(ci1, ci2, cj2, + cj1, r12, r21, + r33, fi, fj, + fdi, fdj, ls1, + ls2, ls3, sig2) if (c1 == ej1): if (ei1 == ej2) and (ei2 == c2): kern += \ - three_body_helper_2(ci2, ci1, cj2, cj1, r21, - r13, r32, fi, fj, fdi, - fdj, ls1, ls2, ls3, sig2) + three_body_helper_2(ci2, ci1, cj2, + cj1, r21, r13, + r32, fi, fj, + fdi, fdj, ls1, + ls2, ls3, sig2) if (ei1 == c2) and (ei2 == ej2): kern += \ - three_body_helper_2(ci1, ci2, cj2, cj1, r11, - r23, r32, fi, fj, fdi, - fdj, ls1, ls2, ls3, sig2) + three_body_helper_2(ci1, ci2, cj2, + cj1, r11, r23, + r32, fi, fj, + fdi, fdj, ls1, + ls2, ls3, sig2) if (c1 == ej2): if (ei1 == ej1) and (ei2 == c2): kern += \ - three_body_helper_2(ci2, ci1, cj1, cj2, r22, - r13, r31, fi, fj, fdi, - fdj, ls1, ls2, ls3, sig2) + three_body_helper_2(ci2, ci1, cj1, + cj2, r22, r13, + r31, fi, fj, + fdi, fdj, ls1, + ls2, ls3, sig2) if (ei1 == c2) and (ei2 == ej1): kern += \ - three_body_helper_2(ci1, ci2, cj1, cj2, r12, - r23, r31, fi, fj, fdi, - fdj, ls1, ls2, ls3, sig2) + three_body_helper_2(ci1, ci2, cj1, + cj2, r12, r23, + r31, fi, fj, + fdi, fdj, ls1, + ls2, ls3, sig2) return kern @@ -1679,11 +1697,11 @@ def two_body_mc_en_jit(bond_array_1, c1, etypes1, # ----------------------------------------------------------------------------- -def many_body_mc_jit(q_array_1, q_array_2, - q_neigh_array_1, q_neigh_array_2, +def many_body_mc_jit(q_array_1, q_array_2, + q_neigh_array_1, q_neigh_array_2, q_neigh_grads_1, q_neigh_grads_2, - c1, c2, etypes1, etypes2, - species1, species2, + c1, c2, etypes1, etypes2, + species1, species2, d1, d2, sig, ls): """many-body multi-element kernel between two force components accelerated with Numba. @@ -1725,25 +1743,25 @@ def many_body_mc_jit(q_array_1, q_array_2, kern = 0 useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) # loop over all possible species for s in useful_species: # Calculate many-body descriptor values for central atoms 1 and 2 - s1 = np.where(species1==s)[0][0] - s2 = np.where(species2==s)[0][0] + s1 = np.where(species1==s)[0][0] + s2 = np.where(species2==s)[0][0] q1 = q_array_1[s1] q2 = q_array_2[s2] - # compute kernel between central atoms only if central atoms are of + # compute kernel between central atoms only if central atoms are of # the same species if c1 == c2: k12 = k_sq_exp_double_dev(q1, q2, sig, ls) else: k12 = 0 - # initialize arrays of many body descriptors and gradients for the + # initialize arrays of many body descriptors and gradients for the # neighbour atoms in the two configurations # Loop over neighbours i of 1st configuration for i in range(q_neigh_array_1.shape[0]): @@ -1766,16 +1784,16 @@ def many_body_mc_jit(q_array_1, q_array_2, # Loop over neighbours j of 2 for j in range(q_neigh_array_2.shape[0]): qjs = qj2_grads = q2j_grads = k1js = 0 - + if etypes2[j] == s: q2j_grads = q_neigh_grads_2[j, d2-1] - + if c2 == s: qj2_grads = q_neigh_grads_2[j, d2-1] - + # Calculate many-body descriptor value for j qjs = q_neigh_array_2[j, s2] - + if c1 == etypes2[j]: k1js = k_sq_exp_double_dev(q1, qjs, sig, ls) @@ -1792,8 +1810,8 @@ def many_body_mc_jit(q_array_1, q_array_2, @njit -def many_body_mc_grad_jit(q_array_1, q_array_2, - q_neigh_array_1, q_neigh_array_2, +def many_body_mc_grad_jit(q_array_1, q_array_2, + q_neigh_array_1, q_neigh_array_2, q_neigh_grads_1, q_neigh_grads_2, c1, c2, etypes1, etypes2, species1, species2, d1, d2, sig, ls): @@ -1839,11 +1857,13 @@ def many_body_mc_grad_jit(q_array_1, q_array_2, ls_derv = 0.0 useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) + + print(species1, species2) for s in useful_species: - s1 = np.where(species1==s)[0][0] - s2 = np.where(species2==s)[0][0] + s1 = np.where(species1==s)[0][0] + s2 = np.where(species2==s)[0][0] q1 = q_array_1[s1] q2 = q_array_2[s2] @@ -1876,16 +1896,16 @@ def many_body_mc_grad_jit(q_array_1, q_array_2, # Loop over neighbours j of 2 for j in range(q_neigh_array_2.shape[0]): qjs = qj2_grads = q2j_grads = k1js = dk1js = 0 - + if etypes2[j] == s: q2j_grads = q_neigh_grads_2[j, d2-1] - + if c2 == s: qj2_grads = q_neigh_grads_2[j, d2-1] - + # Calculate many-body descriptor value for j qjs = q_neigh_array_2[j, s2] - + if c1 == etypes2[j]: k1js = k_sq_exp_double_dev(q1, qjs, sig, ls) q1jdiffsq = (q1 - qjs) * (q1 - qjs) @@ -1921,9 +1941,9 @@ def many_body_mc_grad_jit(q_array_1, q_array_2, @njit -def many_body_mc_force_en_jit(q_array_1, q_array_2, +def many_body_mc_force_en_jit(q_array_1, q_array_2, q_neigh_array_1, q_neigh_grads_1, - c1, c2, etypes1, + c1, c2, etypes1, species1, species2, d1, sig, ls): """many-body many-element kernel between force and energy components accelerated with Numba. @@ -1945,11 +1965,11 @@ def many_body_mc_force_en_jit(q_array_1, q_array_2, kern = 0 useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) for s in useful_species: - s1 = np.where(species1==s)[0][0] - s2 = np.where(species2==s)[0][0] + s1 = np.where(species1==s)[0][0] + s2 = np.where(species2==s)[0][0] q1 = q_array_1[s1] q2 = q_array_2[s2] @@ -1980,7 +2000,7 @@ def many_body_mc_force_en_jit(q_array_1, q_array_2, #@njit -def many_body_mc_en_jit(q_array_1, q_array_2, c1, c2, +def many_body_mc_en_jit(q_array_1, q_array_2, c1, c2, species1, species2, sig, ls): """many-body many-element kernel between energy components accelerated with Numba. @@ -2005,7 +2025,7 @@ def many_body_mc_en_jit(q_array_1, q_array_2, c1, c2, float: Value of the many-body kernel. """ useful_species = np.array( - list(set(species1).union(set(species2))), dtype=np.int8) + list(set(species1).intersection(set(species2))), dtype=np.int8) kern = 0 if c1 == c2: diff --git a/flare/kernels/utils.py b/flare/kernels/utils.py index 8e48d33e3..c26389a4c 100644 --- a/flare/kernels/utils.py +++ b/flare/kernels/utils.py @@ -1,7 +1,7 @@ import numpy as np from flare.kernels import sc, mc_simple, mc_sephyps -import flare.kernels.map_3b_kernel as map_3b +from flare.parameters import Parameters """ This module includes interface functions between kernels and gp/gp_algebra @@ -17,7 +17,9 @@ """ -def str_to_kernel_set(name: str, multihyps: bool = False): +def str_to_kernel_set(kernels: list = ['twobody', 'threebody'], + component: str = "sc", + hyps_mask: dict = None): """ return kernels and kernel gradient function base on a string. If it contains 'sc', it will use the kernel in sc module; @@ -36,39 +38,43 @@ def str_to_kernel_set(name: str, multihyps: bool = False): """ - if 'sc' in name: + # kernel name should be replace with kernel array + if component == 'sc': stk = sc._str_to_kernel else: - if (multihyps is False): - stk = mc_simple._str_to_kernel - else: + multihyps = True + if hyps_mask is None: + multihyps = False + elif hyps_mask['nspecie'] == 1: + multihyps = False + if multihyps: stk = mc_sephyps._str_to_kernel + else: + stk = mc_simple._str_to_kernel # b2 = Two body in use, b3 = Three body in use - b2 = False - b3 = False - many = False - - for s in ['2', 'two', 'Two', 'TWO']: - if (s in name): - b2 = True - for s in ['3', 'three', 'Three', 'THREE']: - if (s in name): - b3 = True - for s in ['mb', 'manybody', 'many', 'Many', 'ManyBody']: - if (s in name): - many = True + str_terms = {'2': ['2', 'two', 'twobody'], + '3': ['3', 'three', 'threebody'], + 'many': ['mb', 'manybody', 'many']} + + if isinstance(kernels, str): + kernels = [kernels] prefix = '' - str_term = {'2': b2, '3': b3, 'many': many} - for term in str_term: - if str_term[term]: - if (len(prefix) > 0): + for term in str_terms: + add = False + for s in str_terms[term]: + for k in kernels: + if s in k.lower(): + add = True + if add: + if len(prefix) > 0: prefix += '+' prefix += term + if len(prefix) == 0: raise RuntimeError( - f"the name has to include at least one number {name}") + f"the name has to include at least one number {kernels}") for suffix in ['', '_grad', '_en', '_force_en']: if prefix+suffix not in stk: @@ -78,53 +84,9 @@ def str_to_kernel_set(name: str, multihyps: bool = False): return stk[prefix], stk[prefix+'_grad'], stk[prefix+'_en'], \ stk[prefix+'_force_en'] -def str_to_mapped_kernel(name: str, multihyps: bool = False, energy=False): - """ - return kernels and kernel gradient function base on a string. - If it contains 'sc', it will use the kernel in sc module; - otherwise, it uses the kernel in mc_simple; - if sc is not included and multihyps is True, - it will use the kernel in mc_sephyps module - otherwise, it will use the kernel in the sc module - Args: - - name (str): name for kernels. example: "2+3mc" - multihyps (bool, optional): True for using multiple hyperparameter groups - energy (bool, optional): True for mapping energy/energy kernel - - :return: mapped kernel function, kernel gradient, energy kernel, - energy_and_force kernel - """ - - if 'sc' in name: - raise NotImplementedError("mapped kernel for single component "\ - "is not implemented") - - # b2 = Two body in use, b3 = Three body in use - b2 = False - many = False - b3 = False - for s in ['3', 'three', 'Three', 'THREE']: - if (s in name): - b3 = True - - if (b3 == True and energy == False): - if (multihyps is True): - tbmfe = map_3b.three_body_mc_en_force_sephyps - tbme = map_3b.three_body_mc_en_sephyps - else: - tbmfe = map_3b.three_body_mc_en_force - tbme = map_3b.three_body_mc_en - else: - raise NotImplementedError("mapped kernel for two-body and manybody kernels "\ - "are not implemented") - - return tbme, tbmfe - - -def from_mask_to_args(hyps, hyps_mask: dict, cutoffs): +def from_mask_to_args(hyps, cutoffs, hyps_mask=None): """ Return the tuple of arguments needed for kernel function. The order of the tuple has to be exactly the same as the one taken by the kernel function. @@ -139,115 +101,72 @@ def from_mask_to_args(hyps, hyps_mask: dict, cutoffs): """ # no special setting - if (hyps_mask is None): - return (hyps, cutoffs) - - if ('map' in hyps_mask): - orig_hyps = hyps_mask['original'] - hm = hyps_mask['map'] - for i, h in enumerate(hyps): - orig_hyps[hm[i]] = h - else: - orig_hyps = hyps + multihyps = True + if hyps_mask is None: + multihyps = False + elif hyps_mask['nspecie'] == 1: + multihyps = False + + if not multihyps: + + cutoffs_array = [0, 0, 0] + cutoffs_array[0] = cutoffs.get('twobody', 0) + cutoffs_array[1] = cutoffs.get('threebody', 0) + cutoffs_array[2] = cutoffs.get('manybody', 0) + return (hyps, cutoffs_array) # setting for mc_sephyps - n2b = hyps_mask.get('nbond', 0) - n3b = hyps_mask.get('ntriplet', 0) - nmb = hyps_mask.get('nmb', 0) + nspecie = hyps_mask['nspecie'] + n2b = hyps_mask.get('ntwobody', 0) + + n3b = hyps_mask.get('nthreebody', 0) + nmanybody = hyps_mask.get('nmanybody', 0) ncut3b = hyps_mask.get('ncut3b', 0) - bond_mask = hyps_mask.get('bond_mask', None) - triplet_mask = hyps_mask.get('triplet_mask', None) - mb_mask = hyps_mask.get('mb_mask', None) + twobody_mask = hyps_mask.get('twobody_mask', None) + threebody_mask = hyps_mask.get('threebody_mask', None) + manybody_mask = hyps_mask.get('manybody_mask', None) cut3b_mask = hyps_mask.get('cut3b_mask', None) - ncutoff = len(cutoffs) - if (ncutoff > 2): - if (cutoffs[2] == 0): - ncutoff = 2 - - if (ncutoff > 0): - cutoff_2b = [cutoffs[0]] - if ('cutoff_2b' in hyps_mask): - cutoff_2b = hyps_mask['cutoff_2b'] - elif (n2b > 0): - cutoff_2b = np.ones(n2b)*cutoffs[0] - - cutoff_3b = None - if (ncutoff > 1): - cutoff_3b = cutoffs[1] - if ('cutoff_3b' in hyps_mask): - cutoff_3b = hyps_mask['cutoff_3b'] - if (ncut3b == 1): - cutoff_3b = cutoff_3b[0] - ncut3b = 0 - cut3b_mask = None - - cutoff_mb = None - if (ncutoff > 2): - cutoff_mb = np.array([cutoffs[2]]) - if ('cutoff_mb' in hyps_mask): - cutoff_mb = hyps_mask['cutoff_mb'] - elif (nmb > 0): - cutoff_mb = np.ones(nmb)*cutoffs[2] - # if the user forget to define nmb - # there has to be a mask, because this is the - # multi hyper parameter mode - if (nmb == 0 and len(orig_hyps)>(n2b*2+n3b*2+1)): - nmb = 1 - nspecie = hyps_mask['nspecie'] - mb_mask = np.zeros(nspecie*nspecie, dtype=int) - - sig2 = None - ls2 = None - sig3 = None - ls3 = None - sigm = None - lsm = None - - if (ncutoff <= 2): - if (n2b != 0): - sig2 = np.array(orig_hyps[:n2b]) - ls2 = np.array(orig_hyps[n2b:n2b * 2]) - if (n3b != 0): - sig3 = np.array(orig_hyps[n2b * 2:n2b * 2 + n3b]) - ls3 = np.array(orig_hyps[n2b * 2 + n3b:n2b * 2 + n3b * 2]) - if (n2b == 0) and (n3b == 0): - raise NameError("Hyperparameter mask missing nbond and/or " - "ntriplet key") - return (cutoff_2b, cutoff_3b, - hyps_mask['nspecie'], hyps_mask['specie_mask'], - n2b, bond_mask, n3b, triplet_mask, - ncut3b, cut3b_mask, - sig2, ls2, sig3, ls3) - - elif (ncutoff == 3): - - if (n2b != 0): - sig2 = np.array(orig_hyps[:n2b]) - ls2 = np.array(orig_hyps[n2b:n2b * 2]) - if (n3b != 0): - start = n2b*2 - sig3 = np.array(orig_hyps[start:start + n3b]) - ls3 = np.array(orig_hyps[start + n3b:start + n3b * 2]) - if (nmb != 0): - start = n2b*2 + n3b*2 - sigm = np.array(orig_hyps[start: start+nmb]) - lsm = np.array(orig_hyps[start+nmb: start+nmb*2]) - - return (cutoff_2b, cutoff_3b, cutoff_mb, - hyps_mask['nspecie'], - np.array(hyps_mask['specie_mask']), - n2b, bond_mask, - n3b, triplet_mask, - ncut3b, cut3b_mask, - nmb, mb_mask, - sig2, ls2, sig3, ls3, sigm, lsm) - else: - raise RuntimeError("only support up to 3 cutoffs") + # TO DO , should instead use the non-sephyps kernel + if (n2b == 1): + twobody_mask = np.zeros(nspecie**2, dtype=int) + if (n3b == 1): + threebody_mask = np.zeros(nspecie**3, dtype=int) + if (nmanybody == 1): + manybody_mask = np.zeros(nspecie**2, dtype=int) + cutoff_2b = cutoffs.get('twobody', 0) + cutoff_3b = cutoffs.get('threebody', 0) + cutoff_mb = cutoffs.get('manybody', 0) -def from_grad_to_mask(grad, hyps_mask): + if 'bond_cutoff_list' in hyps_mask: + cutoff_2b = hyps_mask['bond_cutoff_list'] + else: + cutoff_2b = np.ones(nspecie**2, dtype=float)*cutoff_2b + + if 'threebody_cutoff_list' in hyps_mask: + cutoff_3b = hyps_mask['threebody_cutoff_list'] + if 'manybody_cutoff_list' in hyps_mask: + cutoff_mb = hyps_mask['manybody_cutoff_list'] + + (sig2, ls2) = Parameters.get_component_hyps(hyps_mask, 'twobody', hyps=hyps) + (sig3, ls3) = Parameters.get_component_hyps( + hyps_mask, 'threebody', hyps=hyps) + (sigm, lsm) = Parameters.get_component_hyps( + hyps_mask, 'manybody', hyps=hyps) + + return (cutoff_2b, cutoff_3b, cutoff_mb, + nspecie, + np.array(hyps_mask['specie_mask']), + n2b, twobody_mask, + n3b, threebody_mask, + ncut3b, cut3b_mask, + nmanybody, manybody_mask, + sig2, ls2, sig3, ls3, sigm, lsm) + + +def from_grad_to_mask(grad, hyps_mask=None): """ Return gradient which only includes hyperparameters which are meant to vary @@ -258,23 +177,51 @@ def from_grad_to_mask(grad, hyps_mask): :return: newgrad """ - # no special setting - if (hyps_mask is None): + constrain = True + if hyps_mask is None: + constrain = False + elif 'map' not in hyps_mask: + constrain = False + if not constrain: return grad - # setting for mc_sephyps - # no constrained optimization - if 'map' not in hyps_mask: - return grad + hyp_index = hyps_mask['map'] # setting for mc_sephyps # if the last element is not sigma_noise - if (hyps_mask['map'][-1] == len(grad)): - hm = hyps_mask['map'][:-1] + if hyp_index[-1] == len(grad): + hm = hyp_index[:-1] else: - hm = hyps_mask['map'] + hm = hyp_index newgrad = np.zeros(len(hm), dtype=np.float64) for i, mapid in enumerate(hm): newgrad[i] = grad[mapid] + return newgrad + + +def kernel_str_to_array(kernel_name: str): + """ + Args: + + name (str): name for kernels. example: "2+3mc" + + :return: kernel function, kernel gradient, energy kernel, + energy_and_force kernel + """ + + # kernel name should be replace with kernel array + str_terms = {'twobody': ['2', 'two', 'twobody'], + 'threebody': ['3', 'three', 'threebody'], + 'manybody': ['mb', 'manybody', 'many']} + + array = [] + for term in str_terms: + add = False + for s in str_terms[term]: + if s in kernel_name.lower(): + add = True + if add: + array += [term] + return array diff --git a/flare/mgp/__init__.py b/flare/mgp/__init__.py index 8b1378917..3cec6c94a 100644 --- a/flare/mgp/__init__.py +++ b/flare/mgp/__init__.py @@ -1 +1 @@ - +from flare.mgp.mgp import MappedGaussianProcess diff --git a/flare/mgp/cubic_splines_numba.py b/flare/mgp/cubic_splines_numba.py index b2d9fef32..22109583a 100644 --- a/flare/mgp/cubic_splines_numba.py +++ b/flare/mgp/cubic_splines_numba.py @@ -215,26 +215,6 @@ def vec_eval_cubic_spline_3(a, b, orders, coefs, points, out): out[n] = Phi0_0*(Phi1_0*(Phi2_0*(coefs[i0+0,i1+0,i2+0]) + Phi2_1*(coefs[i0+0,i1+0,i2+1]) + Phi2_2*(coefs[i0+0,i1+0,i2+2]) + Phi2_3*(coefs[i0+0,i1+0,i2+3])) + Phi1_1*(Phi2_0*(coefs[i0+0,i1+1,i2+0]) + Phi2_1*(coefs[i0+0,i1+1,i2+1]) + Phi2_2*(coefs[i0+0,i1+1,i2+2]) + Phi2_3*(coefs[i0+0,i1+1,i2+3])) + Phi1_2*(Phi2_0*(coefs[i0+0,i1+2,i2+0]) + Phi2_1*(coefs[i0+0,i1+2,i2+1]) + Phi2_2*(coefs[i0+0,i1+2,i2+2]) + Phi2_3*(coefs[i0+0,i1+2,i2+3])) + Phi1_3*(Phi2_0*(coefs[i0+0,i1+3,i2+0]) + Phi2_1*(coefs[i0+0,i1+3,i2+1]) + Phi2_2*(coefs[i0+0,i1+3,i2+2]) + Phi2_3*(coefs[i0+0,i1+3,i2+3]))) + Phi0_1*(Phi1_0*(Phi2_0*(coefs[i0+1,i1+0,i2+0]) + Phi2_1*(coefs[i0+1,i1+0,i2+1]) + Phi2_2*(coefs[i0+1,i1+0,i2+2]) + Phi2_3*(coefs[i0+1,i1+0,i2+3])) + Phi1_1*(Phi2_0*(coefs[i0+1,i1+1,i2+0]) + Phi2_1*(coefs[i0+1,i1+1,i2+1]) + Phi2_2*(coefs[i0+1,i1+1,i2+2]) + Phi2_3*(coefs[i0+1,i1+1,i2+3])) + Phi1_2*(Phi2_0*(coefs[i0+1,i1+2,i2+0]) + Phi2_1*(coefs[i0+1,i1+2,i2+1]) + Phi2_2*(coefs[i0+1,i1+2,i2+2]) + Phi2_3*(coefs[i0+1,i1+2,i2+3])) + Phi1_3*(Phi2_0*(coefs[i0+1,i1+3,i2+0]) + Phi2_1*(coefs[i0+1,i1+3,i2+1]) + Phi2_2*(coefs[i0+1,i1+3,i2+2]) + Phi2_3*(coefs[i0+1,i1+3,i2+3]))) + Phi0_2*(Phi1_0*(Phi2_0*(coefs[i0+2,i1+0,i2+0]) + Phi2_1*(coefs[i0+2,i1+0,i2+1]) + Phi2_2*(coefs[i0+2,i1+0,i2+2]) + Phi2_3*(coefs[i0+2,i1+0,i2+3])) + Phi1_1*(Phi2_0*(coefs[i0+2,i1+1,i2+0]) + Phi2_1*(coefs[i0+2,i1+1,i2+1]) + Phi2_2*(coefs[i0+2,i1+1,i2+2]) + Phi2_3*(coefs[i0+2,i1+1,i2+3])) + Phi1_2*(Phi2_0*(coefs[i0+2,i1+2,i2+0]) + Phi2_1*(coefs[i0+2,i1+2,i2+1]) + Phi2_2*(coefs[i0+2,i1+2,i2+2]) + Phi2_3*(coefs[i0+2,i1+2,i2+3])) + Phi1_3*(Phi2_0*(coefs[i0+2,i1+3,i2+0]) + Phi2_1*(coefs[i0+2,i1+3,i2+1]) + Phi2_2*(coefs[i0+2,i1+3,i2+2]) + Phi2_3*(coefs[i0+2,i1+3,i2+3]))) + Phi0_3*(Phi1_0*(Phi2_0*(coefs[i0+3,i1+0,i2+0]) + Phi2_1*(coefs[i0+3,i1+0,i2+1]) + Phi2_2*(coefs[i0+3,i1+0,i2+2]) + Phi2_3*(coefs[i0+3,i1+0,i2+3])) + Phi1_1*(Phi2_0*(coefs[i0+3,i1+1,i2+0]) + Phi2_1*(coefs[i0+3,i1+1,i2+1]) + Phi2_2*(coefs[i0+3,i1+1,i2+2]) + Phi2_3*(coefs[i0+3,i1+1,i2+3])) + Phi1_2*(Phi2_0*(coefs[i0+3,i1+2,i2+0]) + Phi2_1*(coefs[i0+3,i1+2,i2+1]) + Phi2_2*(coefs[i0+3,i1+2,i2+2]) + Phi2_3*(coefs[i0+3,i1+2,i2+3])) + Phi1_3*(Phi2_0*(coefs[i0+3,i1+3,i2+0]) + Phi2_1*(coefs[i0+3,i1+3,i2+1]) + Phi2_2*(coefs[i0+3,i1+3,i2+2]) + Phi2_3*(coefs[i0+3,i1+3,i2+3]))) -Ad = array([ -# t^3 t^2 t 1 - [-1.0/6.0, 3.0/6.0, -3.0/6.0, 1.0/6.0], - [ 3.0/6.0, -6.0/6.0, 0.0/6.0, 4.0/6.0], - [-3.0/6.0, 3.0/6.0, 3.0/6.0, 1.0/6.0], - [ 1.0/6.0, 0.0/6.0, 0.0/6.0, 0.0/6.0] -]) - -dAd = zeros((4,4)) -for i in range(1,4): - Ad_i = Ad[:, i-1] - dAd[:,i] = (4-i) * Ad_i - -d2Ad = zeros((4,4)) -for i in range(1,4): - dAd_i = dAd[:, i-1] - d2Ad[:,i] = (4-i) * dAd_i - - - @njit(cache=True) def vec_eval_cubic_splines_G_1(a, b, orders, coefs, points, vals, dvals): @@ -571,8 +551,6 @@ def filter_coeffs_2d(dinv, data): # First, solve in the X-direction for iy in range(My): - # print(data[:,iy].size) - # print(spline.coefs[:,iy].size) find_coefs_1d(dinv[0], Mx, data[:, iy], coefs[:, iy]) # Now, solve in the Y-direction diff --git a/flare/mgp/grid_kernels_3b.py b/flare/mgp/grid_kernels_3b.py new file mode 100644 index 000000000..3b1915198 --- /dev/null +++ b/flare/mgp/grid_kernels_3b.py @@ -0,0 +1,296 @@ +import numpy as np +from numba import njit +from math import exp, floor +from typing import Callable + +from flare.kernels.cutoffs import quadratic_cutoff + +from time import time + +def grid_kernel_sephyps(kern_type, + data, grids, fj, fdj, + c2, etypes2, + cutoff_2b, cutoff_3b, cutoff_mb, + nspec, spec_mask, + nbond, bond_mask, + ntriplet, triplet_mask, + ncut3b, cut3b_mask, + nmb, mb_mask, + sig2, ls2, sig3, ls3, sigm, lsm, + cutoff_func=quadratic_cutoff): + ''' + Args: + data: a single env of a list of envs + ''' + + bc1 = spec_mask[c2] + bc2 = spec_mask[etypes2[0]] + bc3 = spec_mask[etypes2[1]] + ttype = triplet_mask[nspec * nspec * bc1 + nspec*bc2 + bc3] + ls = ls3[ttype] + sig = sig3[ttype] + cutoffs = [cutoff_2b, cutoff_3b] + + hyps = [sig, ls] + return grid_kernel(kern_type, + data, grids, fj, fdj, + c2, etypes2, + hyps, cutoffs, cutoff_func) + + +def grid_kernel(kern_type, + struc, grids, fj, fdj, + c2, etypes2, + hyps: 'ndarray', cutoffs, + cutoff_func: Callable = quadratic_cutoff): + + r_cut = cutoffs[1] + + if not isinstance(struc, list): + struc = [struc] + + kern = 0 + for env in struc: + kern += grid_kernel_env(kern_type, + env, grids, fj, fdj, + c2, etypes2, + hyps, r_cut, cutoff_func) + + return kern + + +def grid_kernel_env(kern_type, + env1, grids, fj, fdj, + c2, etypes2, + hyps: 'ndarray', r_cut: float, + cutoff_func: Callable = quadratic_cutoff): + + # pre-compute constants that appear in the inner loop + sig = hyps[0] + ls = hyps[1] + derivative = derv_dict[kern_type] + + # collect all the triplets in this training env + triplet_coord_list = get_triplets_for_kern(env1.bond_array_3, env1.ctype, env1.etypes, + env1.cross_bond_inds, env1.cross_bond_dists, env1.triplet_counts, + c2, etypes2) + + if len(triplet_coord_list) == 0: # no triplets + if derivative: + return np.zeros((3, grids.shape[0]), dtype=np.float64) + else: + return np.zeros(grids.shape[0], dtype=np.float64) + + triplet_coord_list = np.array(triplet_coord_list) + triplet_list = triplet_coord_list[:, :3] # (n_triplets, 3) + coord_list = triplet_coord_list[:, 3:] # ((n_triplets, 9) + + # calculate distance difference & exponential part + ls1 = 1 / (2 * ls * ls) + D = 0 + rij_list = [] + for r in range(3): + rj, ri = np.meshgrid(grids[:, r], triplet_list[:, r]) + rij = ri - rj + D += rij * rij # (n_triplets, n_grids) + rij_list.append(rij) + + kern_exp = (sig * sig) * np.exp(- D * ls1) + + # calculate cutoff of the triplets + fi, fdi = triplet_cutoff(triplet_list, r_cut, coord_list, derivative, + cutoff_func) # (n_triplets, 1) + + # calculate the derivative part + kern_func = kern_dict[kern_type] + kern = kern_func(kern_exp, fi, fj, fdi, fdj, + rij_list, coord_list, ls) + + return kern + + +def en_en(kern_exp, fi, fj, *args): + '''energy map + energy block''' + fifj = fi @ fj.T # (n_triplets, n_grids) + kern = np.sum(kern_exp * fifj, axis=0) / 9 # (n_grids,) + return kern + + +def en_force(kern_exp, fi, fj, fdi, fdj, + rij_list, coord_list, ls): + '''energy map + force block''' + fifj = fi @ fj.T # (n_triplets, n_grids) + ls2 = 1 / (ls * ls) + n_trplt, n_grids = kern_exp.shape + kern = np.zeros((3, n_grids), dtype=np.float64) + for d in range(3): + B = 0 + fdij = fdi[:, [d]] @ fj.T + for r in range(3): + rij = rij_list[r] + # column-wise multiplication + # coord_list[:, [r]].shape = (n_triplets, 1) + B += rij * coord_list[:, [3*d+r]] # (n_triplets, n_grids) + + kern[d, :] = - np.sum(kern_exp * (B * ls2 * fifj + fdij), axis=0) / 3 # (n_grids,) + return kern + + +def force_en(kern_exp, fi, fj, fdi, fdj, + rij_list, coord_list, ls): + '''force map + energy block''' + ls2 = 1 / (ls * ls) + fifj = fi @ fj.T # (n_triplets, n_grids) + fdji = fi @ fdj.T + # only r = 0 is non zero, since the grid coords are all (1, 0, 0) + B = rij_list[0] # (n_triplets, n_grids) + kern = np.sum(kern_exp * (B * ls2 * fifj - fdji), axis=0) / 3 # (n_grids,) + return kern + + +def force_force(kern_exp, fi, fj, fdi, fdj, + rij_list, coord_list, ls): + '''force map + force block''' + ls2 = 1 / (ls * ls) + ls3 = ls2 * ls2 + + n_trplt, n_grids = kern_exp.shape + kern = np.zeros((3, n_grids), dtype=np.float64) + + fifj = fi @ fj.T # (n_triplets, n_grids) + fdji = (fi * ls2) @ fdj.T + + B = rij_list[0] * ls3 # B and C both have opposite signs with that in the three_body_helper_1 + + for d in range(3): + fdij = (fdi[:, [d]] * ls2) @ fj.T + I = fdi[:, [d]] @ fdj.T + J = rij_list[0] * fdij # (n_triplets, n_grids) + + A = np.repeat(ls2 * coord_list[:, [3*d]], n_grids, axis=1) + C = 0 + for r in range(3): + rij = rij_list[r] + # column-wise multiplication + # coord_list[:, [r]].shape = (n_triplets, 1) + C += rij * coord_list[:, [3*d+r]] # (n_triplets, n_grids) + + IJKL = I - J + C * fdji + (A - B * C) * fifj + kern[d, :] = np.sum(kern_exp * IJKL, axis=0) + + return kern + + + +def triplet_cutoff(triplets, r_cut, coords, derivative=False, cutoff_func=quadratic_cutoff): + + dfj_list = np.zeros((len(triplets), 3), dtype=np.float64) + + if derivative: + for d in range(3): + s = 3 * d + e = 3 * (d + 1) + f0, df0 = cutoff_func(r_cut, triplets, coords[:, s:e]) + dfj = df0[:, 0] * f0[:, 1] * f0[:, 2] + \ + f0[:, 0] * df0[:, 1] * f0[:, 2] + \ + f0[:, 0] * f0[:, 1] * df0[:, 2] +# dfj = np.expand_dims(dfj, axis=1) + dfj_list[:, d] = dfj + else: + f0, _ = cutoff_func(r_cut, triplets, 0) # (n_grid, 3) + + fj = f0[:, 0] * f0[:, 1] * f0[:, 2] # (n_grid,) + fj = np.expand_dims(fj, axis=1) + + return fj, dfj_list + + +@njit +def get_triplets_for_kern(bond_array_1, c1, etypes1, + cross_bond_inds_1, cross_bond_dists_1, + triplets_1, + c2, etypes2): + + #triplet_list = np.empty((0, 6), dtype=np.float64) + triplet_list = [] + + ej1 = etypes2[0] + ej2 = etypes2[1] + + all_spec = [c2, ej1, ej2] + if c1 in all_spec: + c1_ind = all_spec.index(c1) + ind_list = [0, 1, 2] + ind_list.remove(c1_ind) + all_spec.remove(c1) + + for m in range(bond_array_1.shape[0]): + two_inds = ind_list.copy() + + ri1 = bond_array_1[m, 0] + ci1 = bond_array_1[m, 1:] + ei1 = etypes1[m] + + two_spec = [all_spec[0], all_spec[1]] + if (ei1 in two_spec): + + ei1_ind = ind_list[0] if ei1 == two_spec[0] else ind_list[1] + two_spec.remove(ei1) + two_inds.remove(ei1_ind) + one_spec = two_spec[0] + ei2_ind = two_inds[0] + + for n in range(triplets_1[m]): + ind1 = cross_bond_inds_1[m, m + n + 1] + ei2 = etypes1[ind1] + if (ei2 == one_spec): + + ri2 = bond_array_1[ind1, 0] + ci2 = bond_array_1[ind1, 1:] + + ri3 = cross_bond_dists_1[m, m + n + 1] + ci3 = np.zeros(3) + + # align this triplet to the same species order as r1, r2, r12 + + perms = [] + if (c1 == c2): + if (ei1 == ej1) and (ei2 == ej2): perms.append([0, 1, 2]) + if (ei1 == ej2) and (ei2 == ej1): perms.append([1, 0, 2]) + if (c1 == ej1): + if (ei1 == ej2) and (ei2 == c2): perms.append([1, 2, 0]) + if (ei1 == c2) and (ei2 == ej2): perms.append([0, 2, 1]) + if (c1 == ej2): + if (ei1 == ej1) and (ei2 == c2): perms.append([2, 1, 0]) + if (ei1 == c2) and (ei2 == ej1): perms.append([2, 0, 1]) + + tri = np.array([ri1, ri2, ri3]) + crd1 = np.array([ci1[0], ci2[0], ci3[0]]) + crd2 = np.array([ci1[1], ci2[1], ci3[1]]) + crd3 = np.array([ci1[2], ci2[2], ci3[2]]) + + # append permutations + nperm = len(perms) + for iperm in range(nperm): + perm = perms[iperm] + tricrd = np.take(tri, perm) + crd1_p = np.take(crd1, perm) + crd2_p = np.take(crd2, perm) + crd3_p = np.take(crd3, perm) + tricrd = np.hstack((tricrd, crd1_p, crd2_p, crd3_p)) + triplet_list.append(tricrd) + + return triplet_list + + + +kern_dict = {'energy_energy': en_en, + 'energy_force': en_force, + 'force_energy': force_en, + 'force_force': force_force} + +derv_dict = {'energy_energy': False, + 'energy_force': True, + 'force_energy': False, + 'force_force': True} + diff --git a/flare/mgp/map2b.py b/flare/mgp/map2b.py new file mode 100644 index 000000000..ca8b806ab --- /dev/null +++ b/flare/mgp/map2b.py @@ -0,0 +1,101 @@ +import numpy as np + +from typing import List + +from flare.struc import Structure +from flare.utils.element_coder import Z_to_element + +from flare.mgp.mapxb import MapXbody, SingleMapXbody +from flare.mgp.utils import get_bonds + + +class Map2body(MapXbody): + + def __init__(self, **kwargs,): + ''' + args: the same arguments as MapXbody, to guarantee they have the same + input parameters + ''' + + self.kernel_name = "twobody" + self.singlexbody = SingleMap2body + self.bodies = 2 + super().__init__(**kwargs) + + def build_bond_struc(self, species_list): + ''' + build a bond structure, used in grid generating + ''' + + # 2 body (2 atoms (1 bond) config) + self.spc = [] + for spc1_ind, spc1 in enumerate(species_list): + for spc2 in species_list[spc1_ind:]: + species = [spc1, spc2] + self.spc.append(sorted(species)) + + + def get_arrays(self, atom_env): + + return get_bonds(atom_env.ctype, atom_env.etypes, atom_env.bond_array_2) + + def find_map_index(self, spc): + # use set because of permutational symmetry + return self.spc.index(sorted(spc)) + + + + +class SingleMap2body(SingleMapXbody): + + def __init__(self, **kwargs): + ''' + Build 2-body MGP + + bond_struc: Mock structure used to sample 2-body forces on 2 atoms + ''' + + self.bodies = 2 + self.kernel_name = 'twobody' + + super().__init__(**kwargs) + + # initialize bounds + if self.auto_lower: + self.bounds[0] = np.array([0]) + if self.auto_upper: + self.bounds[1] = np.array([1]) + + spc = self.species + self.species_code = Z_to_element(spc[0]) + '_' + Z_to_element(spc[1]) + + def set_bounds(self, lower_bound, upper_bound): + ''' + lower_bound: scalar or array + upper_bound: scalar or array + ''' + if self.auto_lower: + if isinstance(lower_bound, float): + self.bounds[0] = [lower_bound] + else: + self.bounds[0] = lower_bound + if self.auto_upper: + if isinstance(upper_bound, float): + self.bounds[1] = [upper_bound] + else: + self.bounds[1] = upper_bound + + + def construct_grids(self): + nop = self.grid_num[0] + bond_lengths = np.linspace(self.bounds[0][0], self.bounds[1][0], nop) + return bond_lengths + + + def set_env(self, grid_env, r): + grid_env.bond_array_2 = np.array([[r, 1, 0, 0]]) + return grid_env + + def skip_grid(self, r): + return False + diff --git a/flare/mgp/map3b.py b/flare/mgp/map3b.py new file mode 100644 index 000000000..bf19fb275 --- /dev/null +++ b/flare/mgp/map3b.py @@ -0,0 +1,204 @@ +import numpy as np +from math import floor + +from typing import List + +from flare.struc import Structure +from flare.utils.element_coder import Z_to_element +from flare.gp_algebra import _global_training_data, _global_training_structures +from flare.kernels.utils import from_mask_to_args + +from flare.mgp.mapxb import MapXbody, SingleMapXbody +from flare.mgp.utils import get_triplets, get_kernel_term +from flare.mgp.grid_kernels_3b import triplet_cutoff + + +class Map3body(MapXbody): + + def __init__(self, **kwargs): + + self.kernel_name = "threebody" + self.singlexbody = SingleMap3body + self.bodies = 3 + super().__init__(**kwargs) + + + def build_bond_struc(self, species_list): + ''' + build a bond structure, used in grid generating + ''' + + # 2 body (2 atoms (1 bond) config) + self.spc = [] + N_spc = len(species_list) + for spc1_ind in range(N_spc): + spc1 = species_list[spc1_ind] + for spc2_ind in range(N_spc): # (spc1_ind, N_spc): + spc2 = species_list[spc2_ind] + for spc3_ind in range(N_spc): # (spc2_ind, N_spc): + spc3 = species_list[spc3_ind] + species = [spc1, spc2, spc3] + self.spc.append(species) + + + def get_arrays(self, atom_env): + + spcs, comp_r, comp_xyz = \ + get_triplets(atom_env.ctype, atom_env.etypes, + atom_env.bond_array_3, atom_env.cross_bond_inds, + atom_env.cross_bond_dists, atom_env.triplet_counts) + + return spcs, comp_r, comp_xyz + + def find_map_index(self, spc): + return self.spc.index(spc) + + +class SingleMap3body(SingleMapXbody): + + def __init__(self, **kwargs): + ''' + Build 3-body MGP + + ''' + + self.bodies = 3 + self.kernel_name = 'threebody' + + super().__init__(**kwargs) + + # initialize bounds + self.set_bounds(0, np.ones(3)) + + self.grid_interval = np.min((self.bounds[1]-self.bounds[0])/self.grid_num) + + spc = self.species + self.species_code = Z_to_element(spc[0]) + '_' + \ + Z_to_element(spc[1]) + '_' + Z_to_element(spc[2]) + self.kv3name = f'kv3_{self.species_code}' + + + def set_bounds(self, lower_bound, upper_bound): + if self.auto_lower: + if isinstance(lower_bound, float): + self.bounds[0] = np.ones(3) * lower_bound + else: + self.bounds[0] = lower_bound + if self.auto_upper: + if isinstance(upper_bound, float): + self.bounds[1] = np.ones(3) * upper_bound + else: + self.bounds[1] = upper_bound + + + def construct_grids(self): + ''' + Return: + An array of shape (n_grid, 3) + ''' + # build grids in each dimension + triplets = [] + for d in range(3): + bonds = np.linspace(self.bounds[0][d], self.bounds[1][d], + self.grid_num[d], dtype=np.float64) + triplets.append(bonds) + +# r1 = np.tile(bonds1, (nb12, nb2, 1)) +# r1 = np.moveaxis(r1, -1, 0) +# r2 = np.tile(bonds2, (nb1, nb12, 1)) +# r2 = np.moveaxis(r2, -1, 1) +# r12 = np.tile(bonds12, (nb1, nb2, 1)) + + # concatenate into one array: n_grid x 3 + mesh = np.meshgrid(*triplets, indexing='ij') + del triplets + + mesh_list = [] + n_grid = np.prod(self.grid_num) + for d in range(3): + mesh_list.append(np.reshape(mesh[d], n_grid)) + + mesh_list = np.array(mesh_list).T + + return mesh_list + + + def set_env(self, grid_env, grid_pt): + r1, r2, r12 = grid_pt + dist12 = r12 + grid_env.bond_array_3 = np.array([[r1, 1, 0, 0], [r2, 0, 0, 0]]) + grid_env.cross_bond_dists = np.array([[0, dist12], [dist12, 0]]) + + return grid_env + + + def skip_grid(self, grid_pt): + r1, r2, r12 = grid_pt + + return False + + if not self.map_force: + relaxation = 1/2 * np.max(self.grid_num) * self.grid_interval + if r1 + r2 < r12 - relaxation: + return True + if r1 + r12 < r2 - relaxation: + return True + if r12 + r2 < r1 - relaxation: + return True + + return False + + + def _gengrid_numba(self, name, env12, kernel_info, force_block, s, e): + """ + Loop over different parts of the training set. from element s to element e + + Args: + name: name of the gp instance + s: start index of the training data parition + e: end index of the training data parition + env12: AtomicEnvironment container of the triplet + kernel_info: return value of the get_3b_kernel + """ + + grid_kernel, _, _, cutoffs, hyps, hyps_mask = kernel_info + + args = from_mask_to_args(hyps, cutoffs, hyps_mask) + r_cut = cutoffs['threebody'] + + grids = self.construct_grids() + + if (e-s) == 0: + return np.empty((grids.shape[0], 0), dtype=np.float64) + coords = np.zeros((grids.shape[0], 9), dtype=np.float64) # padding 0 + coords[:, 0] = np.ones_like(coords[:, 0]) + + fj, fdj = triplet_cutoff(grids, r_cut, coords, derivative=True) # TODO: add cutoff func + fdj = fdj[:, [0]] + + if self.map_force: + prefix = 'force' + else: + prefix = 'energy' + + if force_block: + training_data = _global_training_data[name] + kern_type = f'{prefix}_force' + else: + training_data = _global_training_structures[name] + kern_type = f'{prefix}_energy' + + k_v = [] + for m_index in range(s, e): + data = training_data[m_index] + kern_vec = grid_kernel(kern_type, data, grids, fj, fdj, + env12.ctype, env12.etypes, + *args) + k_v.append(kern_vec) + + if len(k_v) > 0: + k_v = np.vstack(k_v).T + else: + k_v = np.zeros((grids.shape[0], 0)) + + return k_v diff --git a/flare/mgp/mapxb.py b/flare/mgp/mapxb.py new file mode 100644 index 000000000..dbbbb727a --- /dev/null +++ b/flare/mgp/mapxb.py @@ -0,0 +1,627 @@ +import json +from flare.utils.element_coder import NumpyEncoder, element_to_Z, Z_to_element + +import os, logging, warnings +import numpy as np +import multiprocessing as mp + +from copy import deepcopy +from math import ceil, floor +from scipy.linalg import solve_triangular +from typing import List + +from flare.env import AtomicEnvironment +from flare.kernels.utils import from_mask_to_args +from flare.gp import GaussianProcess +from flare.gp_algebra import partition_vector, energy_force_vector_unit, \ + force_energy_vector_unit, energy_energy_vector_unit, force_force_vector_unit,\ + _global_training_data, _global_training_structures +from flare.parameters import Parameters +from flare.struc import Structure + +from flare.mgp.utils import get_kernel_term +from flare.mgp.splines_methods import PCASplines, CubicSpline + +global_use_grid_kern = True + +class MapXbody: + def __init__(self, + grid_num: List, + lower_bound: List or str='auto', + upper_bound: List or str='auto', + svd_rank = 'auto', + coded_species: list=[], + map_force: bool=False, + mean_only: bool=True, + container_only: bool=True, + lmp_file_name: str='lmp.mgp', + load_grid: str=None, + lower_bound_relax: float=0.1, + GP: GaussianProcess=None, + n_cpus: int=None, + n_sample: int=100, + hyps_mask: dict=None, + hyps: list=None, + **kwargs): + + # load all arguments as attributes + self.grid_num = np.array(grid_num) + self.lower_bound = lower_bound + self.upper_bound = upper_bound + self.svd_rank = svd_rank + self.coded_species = coded_species + self.map_force = map_force + self.mean_only = mean_only + self.lmp_file_name = lmp_file_name + self.load_grid = load_grid + self.lower_bound_relax = lower_bound_relax + self.n_cpus = n_cpus + self.n_sample = n_sample + + self.spc = [] + self.maps = [] + self.kernel_info = None + self.hyps_mask = hyps_mask + self.hyps = hyps + + self.build_bond_struc(coded_species) + + bounds = [self.lower_bound, self.upper_bound] + self.build_map_container(bounds) + + if (not container_only) and (GP is not None) and \ + (len(GP.training_data) > 0): + self.build_map(GP) + + def build_bond_struc(self, coded_species): + raise NotImplementedError("need to be implemented in child class") + + def get_arrays(self, atom_env): + raise NotImplementedError("need to be implemented in child class") + + def build_map_container(self, bounds): + ''' + construct an empty spline container without coefficients. + ''' + + self.maps = [] + for spc in self.spc: + m = self.singlexbody(bounds=bounds, species=spc, **self.__dict__) + self.maps.append(m) + + + def build_map(self, GP): + ''' + generate/load grids and get spline coefficients + ''' + + self.kernel_info = get_kernel_term(self.kernel_name, + GP.component, GP.hyps_mask, GP.hyps) + self.hyps_mask = GP.hyps_mask + self.hyps = GP.hyps + + for m in self.maps: + m.build_map(GP) + + + def predict(self, atom_env, mean_only): + + assert Parameters.compare_dict(self.hyps_mask, atom_env.cutoffs_mask),\ + 'GP.hyps_mask is not the same as atom_env.cutoffs_mask' + + min_dist = atom_env.bond_array_2[0][0] + lower_bound = np.max(self.maps[0].bounds[0][0]) + if min_dist < lower_bound: + raise ValueError(f'The minimal distance {min_dist:.3f} is below the' + f' mgp lower bound {lower_bound:.3f}') + + if self.mean_only: # if not build mapping for var + mean_only = True + + force_kernel, en_kernel, _, cutoffs, hyps, hyps_mask = self.kernel_info + + args = from_mask_to_args(hyps, cutoffs, hyps_mask) + + kern = 0 + if not mean_only: + if self.map_force: + kern = np.zeros(3) + for d in range(3): + kern[d] = force_kernel(atom_env, atom_env, d+1, d+1, *args) + else: + kern = en_kernel(atom_env, atom_env, *args) + + spcs, comp_r, comp_xyz = self.get_arrays(atom_env) + + # predict for each species + f_spcs = np.zeros(3) + vir_spcs = np.zeros(6) + v_spcs = np.zeros(3) if self.map_force else 0 + e_spcs = 0 + for i, spc in enumerate(spcs): + lengths = np.array(comp_r[i]) + xyzs = np.array(comp_xyz[i]) + map_ind = self.find_map_index(spc) + + f, vir, v, e = self.maps[map_ind].predict(lengths, xyzs, + self.map_force, mean_only) + f_spcs += f + vir_spcs += vir + v_spcs += v + e_spcs += e + + return f_spcs, vir_spcs, kern, v_spcs, e_spcs + + def as_dict(self) -> dict: + """ + Dictionary representation of the MGP model. + """ + + out_dict = deepcopy(dict(vars(self))) + out_dict.pop('kernel_info') + + # Uncertainty mappings currently not serializable; + if not self.mean_only: + out_dict['mean_only'] = True + + # only save the mean coefficients + out_dict['maps'] = [m.mean.__coeffs__ for m in self.maps] + out_dict['bounds'] = [m.bounds for m in self.maps] + + # rm keys since they are built in the __init__ function + key_list = ['singlexbody', 'spc'] + for key in key_list: + if out_dict.get(key) is not None: + del out_dict[key] + + return out_dict + + @staticmethod + def from_dict(dictionary: dict, mapxbody): + """ + Create MGP object from dictionary representation. + """ + + if 'container_only' not in dictionary: + dictionary['container_only'] = True + + new_mgp = mapxbody(**dictionary) + + # Restore kernel_info + new_mgp.kernel_info = get_kernel_term(dictionary['kernel_name'], + 'mc', dictionary['hyps_mask'], dictionary['hyps']) + + # Fill up the model with the saved coeffs + for m in range(len(new_mgp.maps)): + singlexb = new_mgp.maps[m] + bounds = dictionary['bounds'][m] + singlexb.set_bounds(bounds[0], bounds[1]) + singlexb.build_map_container() + singlexb.mean.__coeffs__ = np.array(dictionary['maps'][m]) + + return new_mgp + + + def write(self, f): + for m in self.maps: + m.write(f) + + + +class SingleMapXbody: + def __init__(self, grid_num: int=1, bounds='auto', species: list=[], + map_force=False, svd_rank=0, mean_only: bool=False, + load_grid=None, lower_bound_relax=0.1, + n_cpus: int=None, n_sample: int=100, **kwargs): + + self.grid_num = grid_num + self.bounds = deepcopy(bounds) + self.species = species + self.map_force = map_force + self.svd_rank = svd_rank + self.mean_only = mean_only + self.load_grid = load_grid + self.lower_bound_relax = lower_bound_relax + self.n_cpus = n_cpus + self.n_sample = n_sample + + self.auto_lower = (bounds[0] == 'auto') + if self.auto_lower: + lower_bound = 0 + else: + lower_bound = bounds[0] + + self.auto_upper = (bounds[1] == 'auto') + if self.auto_upper: + upper_bound = 1 + else: + upper_bound = bounds[1] + + self.set_bounds(lower_bound, upper_bound) + + self.hyps_mask = None + self.use_grid_kern = global_use_grid_kern + + if not self.auto_lower and not self.auto_upper: + self.build_map_container() + + def set_bounds(self, lower_bound, upper_bound): + raise NotImplementedError("need to be implemented in child class") + + def construct_grids(self): + raise NotImplementedError("need to be implemented in child class") + + def set_env(self, grid_env, r): + raise NotImplementedError("need to be implemented in child class") + + def skip_grid(self, r): + raise NotImplementedError("need to be implemented in child class") + + def get_grid_env(self, GP): + if isinstance(GP.cutoffs, dict): + max_cut = np.max(list(GP.cutoffs.values())) + else: + max_cut = np.max(GP.cutoffs) + big_cell = np.eye(3) * (2 * max_cut + 1) + positions = [[(i+1)/(self.bodies+1)*0.1, 0, 0] + for i in range(self.bodies)] + grid_struc = Structure(big_cell, self.species, positions) + grid_env = AtomicEnvironment(grid_struc, 0, GP.cutoffs, + cutoffs_mask=GP.hyps_mask) + + return grid_env + + + def GenGrid(self, GP): + ''' + To use GP to predict value on each grid point, we need to generate the + kernel vector kv whose length is the same as the training set size. + + 1. We divide the training set into several batches, corresponding to + different segments of kv + 2. Distribute each batch to a processor, i.e. each processor calculate + the kv segment of one batch for all grids + 3. Collect kv segments and form a complete kv vector for each grid, + and calculate the grid value by multiplying the complete kv vector + with GP.alpha + ''' + + if (self.n_cpus is None): + processes = mp.cpu_count() + else: + processes = self.n_cpus + + # ------ construct grids ------ + n_grid = np.prod(self.grid_num) + grid_mean = np.zeros([n_grid]) + if not self.mean_only: + grid_vars = np.zeros([n_grid, len(GP.alpha)]) + else: + grid_vars = None + + grid_env = self.get_grid_env(GP) + + # -------- get training data info ---------- + n_envs = len(GP.training_data) + n_strucs = len(GP.training_structures) + + if (n_envs == 0) and (n_strucs == 0): + warnings.warn("No training data, will return 0") + return np.zeros([n_grid]), None + + # ------- call gengrid functions --------------- + kernel_info = get_kernel_term(self.kernel_name, GP.component, + GP.hyps_mask, GP.hyps) + if self.use_grid_kern: + try: + kernel_info = get_kernel_term(self.kernel_name, GP.component, + GP.hyps_mask, GP.hyps, grid_kernel=True) + except: + self.use_grid_kern = False + + args = [GP.name, grid_env, kernel_info] + + k12_v_force = self._gengrid_par(args, True, n_envs, processes) + k12_v_energy = self._gengrid_par(args, False, n_strucs, processes) + + k12_v_all = np.hstack([k12_v_force, k12_v_energy]) + del k12_v_force + del k12_v_energy + + # ------- compute bond means and variances --------------- + grid_mean = k12_v_all @ GP.alpha + grid_mean = np.reshape(grid_mean, self.grid_num) + + if not self.mean_only: + grid_vars = solve_triangular(GP.l_mat, k12_v_all.T, lower=True).T + tensor_shape = np.array([*self.grid_num, grid_vars.shape[1]]) + grid_vars = np.reshape(grid_vars, tensor_shape) + + # ------ save mean and var to file ------- + if 'mgp_grids' not in os.listdir('./'): + os.mkdir('mgp_grids') + + grid_path = f'mgp_grids/{self.bodies}_{self.species_code}' + np.save(f'{grid_path}_mean', grid_mean) + np.save(f'{grid_path}_var', grid_vars) + + return grid_mean, grid_vars + + + def _gengrid_par(self, args, force_block, n_envs, processes): + + if n_envs == 0: + n_grid = np.prod(self.grid_num) + return np.empty((n_grid, 0)) + + if self.use_grid_kern: + gengrid_func = self._gengrid_numba + else: + gengrid_func = self._gengrid_inner + + if processes == 1: + return gengrid_func(*args, force_block, 0, n_envs) + + with mp.Pool(processes=processes) as pool: + + block_id, nbatch = \ + partition_vector(self.n_sample, n_envs, processes) + + k12_slice = [] + for ibatch in range(nbatch): + s, e = block_id[ibatch] + k12_slice.append(pool.apply_async(gengrid_func, + args = args + [force_block, s, e])) + k12_matrix = [] + for ibatch in range(nbatch): + k12_matrix += [k12_slice[ibatch].get()] + pool.close() + pool.join() + del k12_slice + k12_v_force = np.hstack(k12_matrix) + del k12_matrix + + return k12_v_force + + + def _gengrid_inner(self, name, grid_env, kern_info, force_block, s, e): + ''' + Calculate kv segments of the given batch of training data for all grids + ''' + + kernel, ek, efk, cutoffs, hyps, hyps_mask = kern_info + if force_block: + size = (e - s) * 3 + force_x_vector_unit = force_force_vector_unit + force_x_kern = kernel + energy_x_vector_unit = energy_force_vector_unit + energy_x_kern = efk + else: + size = e - s + force_x_vector_unit = force_energy_vector_unit + force_x_kern = efk + energy_x_vector_unit = energy_energy_vector_unit + energy_x_kern = ek + + grids = self.construct_grids() + k12_v = np.zeros([len(grids), size]) + + for b in range(grids.shape[0]): + grid_pt = grids[b] + grid_env = self.set_env(grid_env, grid_pt) + + if not self.skip_grid(grid_pt): + if self.map_force: + k12_v[b, :] = force_x_vector_unit(name, s, e, grid_env, + force_x_kern, hyps, cutoffs, hyps_mask, 1) + else: + k12_v[b, :] = energy_x_vector_unit(name, s, e, grid_env, + energy_x_kern, hyps, cutoffs, hyps_mask) + + return k12_v + + + def build_map_container(self): + ''' + build 1-d spline function for mean, 2-d for var + ''' + self.mean = CubicSpline(self.bounds[0], self.bounds[1], + orders=self.grid_num) + + if not self.mean_only: + if self.svd_rank == 'auto': + warnings.warn("The containers for variance are not built because svd_rank='auto'") + if isinstance(self.svd_rank, int): + self.var = PCASplines(self.bounds[0], self.bounds[1], + orders=self.grid_num, + svd_rank=self.svd_rank) + + def update_bounds(self, GP): + rebuild_container = False + + # double check the container and the GP is consistent + if not Parameters.compare_dict(GP.hyps_mask, self.hyps_mask): + rebuild_container = True + + lower_bound = self.bounds[0] + min_dist = self.search_lower_bound(GP) + if min_dist < np.max(lower_bound): # change lower bound + warnings.warn('The minimal distance in training data is lower than \ + the current lower bound, will reset lower bound') + + if self.auto_lower or (min_dist < np.max(lower_bound)): + lower_bound = np.max((min_dist - self.lower_bound_relax, 0)) + rebuild_container = True + + upper_bound = self.bounds[1] + if self.auto_upper: + upper_bound = Parameters.get_cutoff(self.kernel_name, + self.species, GP.hyps_mask) + rebuild_container = True + + if rebuild_container: + self.set_bounds(lower_bound, upper_bound) + self.build_map_container() + + + def build_map(self, GP): + + self.update_bounds(GP) + + if not self.load_grid: + y_mean, y_var = self.GenGrid(GP) + else: + if 'mgp_grids' not in os.listdir(self.load_grid): + raise FileNotFoundError("Please set 'load_grid' as the location of mgp_grids folder") + + grid_path = f'{self.load_grid}/mgp_grids/{self.bodies}_{self.species_code}' + y_mean = np.load(f'{grid_path}_mean.npy') + y_var = np.load(f'{grid_path}_var.npy', allow_pickle=True) + + self.mean.set_values(y_mean) + if not self.mean_only: + if self.svd_rank == 'auto': + self.var = PCASplines(self.bounds[0], self.bounds[1], + orders=self.grid_num, + svd_rank=np.min(y_var.shape)) + self.var.set_values(y_var) + + self.hyps_mask = deepcopy(GP.hyps_mask) + + + def __str__(self): + info = f'''{self.__class__.__name__} + species: {self.species} + lower bound: {self.bounds[0]}, auto_lower = {self.auto_lower} + upper bound: {self.bounds[1]}, auto_upper = {self.auto_upper} + grid num: {self.grid_num} + lower bound relaxation: {self.lower_bound_relax} + load grid from: {self.load_grid}\n''' + + if self.map_force: + info += f' build force mapping\n' + else: + info += f' build energy mapping\n' + + if self.mean_only: + info += f' without variance\n' + else: + info += f' with variance, svd_rank = {self.svd_rank}\n' + + return info + + + def search_lower_bound(self, GP): + ''' + If the lower bound is set to be 'auto', search the minimal interatomic + distances in the training set of GP. + ''' + upper_bound = Parameters.get_cutoff(self.kernel_name, + self.species, GP.hyps_mask) + + lower_bound = np.min(upper_bound) + for env in _global_training_data[GP.name]: + min_dist = env.bond_array_2[0][0] + if min_dist < lower_bound: + lower_bound = min_dist + + for struc in _global_training_structures[GP.name]: + for env in struc: + min_dist = env.bond_array_2[0][0] + if min_dist < lower_bound: + lower_bound = min_dist + + return lower_bound + + + def predict(self, lengths, xyzs, map_force, mean_only): + ''' + predict force and variance contribution of one component + ''' + + assert map_force == self.map_force, f'The mapping is built for'\ + 'map_force={self.map_force}, can not predict for map_force={map_force}' + + lengths = np.array(lengths) + xyzs = np.array(xyzs) + + if self.map_force: + # predict forces and energy + e = 0 + f_0 = self.mean(lengths) + f_d = np.diag(f_0) @ xyzs + f = np.sum(f_d, axis=0) + + # predict var + v = np.zeros(3) + if not mean_only: + v_0 = self.var(lengths) + v_d = v_0 @ xyzs + v = self.var.V @ v_d + + else: + # predict forces and energy + e_0, f_0 = self.mean(lengths, with_derivatives=True) + e = np.sum(e_0) # energy + f_d = np.diag(f_0[:,0,0]) @ xyzs + f = self.bodies * np.sum(f_d, axis=0) + + # predict var + v = 0 + if not mean_only: + v_0 = np.expand_dims(np.sum(self.var(lengths), axis=1), + axis=1) + v = self.var.V @ v_0 + + # predict virial stress + vir = np.zeros(6) + vir_order = ((0,0), (1,1), (2,2), (1,2), (0,2), (0,1)) # match the ASE order + for i in range(6): + vir_i = f_d[:,vir_order[i][0]]\ + * xyzs[:,vir_order[i][1]] * lengths[:,0] + vir[i] = np.sum(vir_i) + + vir *= self.bodies / 2 + return f, vir, v, e + + + def write(self, f): + ''' + Write LAMMPS coefficient file + + This implementation only works for 2b and 3b. User should + implement overload in the actual class if the new kernel + has different coefficient format + + In the future, it should be changed to writing in bin/hex + instead of decimal + ''' + + # write header + elems = self.species_code.split('_') + a = self.bounds[0] + b = self.bounds[1] + order = self.grid_num + + header = ' '.join(elems) + header += ' '+' '.join(map(repr, a)) + header += ' '+' '.join(map(repr, b)) + header += ' '+' '.join(map(str, order)) + f.write(header + '\n') + + # write coefficients + coefs = self.mean.__coeffs__ + self.write_flatten_coeff(f, coefs) + + def write_flatten_coeff(self, f, coefs): + """ + flatten the coefficient and write it as + a block. each line has no more than 5 element. + the accuracy is restricted to .10 + """ + coefs = coefs.reshape([-1]) + for c, coef in enumerate(coefs): + f.write(' '+repr(coef)) + if c % 5 == 4 and c != len(coefs)-1: + f.write('\n') + f.write('\n') diff --git a/flare/mgp/mgp.py b/flare/mgp/mgp.py index 9b6365f3c..f5645c1d1 100644 --- a/flare/mgp/mgp.py +++ b/flare/mgp/mgp.py @@ -10,25 +10,14 @@ import multiprocessing as mp from copy import deepcopy -from math import ceil, floor -from scipy.linalg import solve_triangular from typing import List -from flare.struc import Structure from flare.env import AtomicEnvironment from flare.gp import GaussianProcess -from flare.gp_algebra import partition_vector, energy_force_vector_unit, \ - force_energy_vector_unit, energy_energy_vector_unit, force_force_vector_unit, \ - _global_training_data, _global_training_structures -from flare.kernels.utils import from_mask_to_args, str_to_kernel_set, str_to_mapped_kernel -from flare.kernels.cutoffs import quadratic_cutoff -from flare.utils.element_coder import Z_to_element, NumpyEncoder -from flare.utils.mask_helper import HyperParameterMasking as hpm +from flare.utils.element_coder import NumpyEncoder, element_to_Z, Z_to_element - -from flare.mgp.utils import get_bonds, get_triplets, get_triplets_en, \ - get_2bkernel, get_3bkernel -from flare.mgp.splines_methods import PCASplines, CubicSpline +from flare.mgp.map2b import Map2body +from flare.mgp.map3b import Map3body class MappedGaussianProcess: ''' @@ -36,57 +25,74 @@ class MappedGaussianProcess: and automatically save coefficients for LAMMPS pair style. Args: - struc_params (dict): Parameters for a dummy structure which will be - internally used to probe/store forces associated with different atomic - configurations grid_params (dict): Parameters for the mapping itself, such as - grid size of spline fit, etc. + grid size of spline fit, etc. As described below. + unique_species (dict): List of all the (unique) species included during + the training that need to be mapped map_force (bool): if True, do force mapping; otherwise do energy mapping, default is False - mean_only (bool): if True: only build mapping for mean (force) - container_only (bool): if True: only build splines container - (with no coefficients); if False: Attempt to build map immediately GP (GaussianProcess): None or a GaussianProcess object. If a GP is input, and container_only is False, automatically build a mapping corresponding to the GaussianProcess. + mean_only (bool): if True: only build mapping for mean (force) + container_only (bool): if True: only build splines container + (with no coefficients); if False: Attempt to build map immediately lmp_file_name (str): LAMMPS coefficient file name + n_cpus (int): Default None. Set to the number of cores needed for + parallelization. Used in the construction of the map. + n_sample (int): Default 100. The batch size for building map. Not used now. Examples: - >>> struc_params = {'species': [0, 1], - 'cube_lat': cell, # should input the cell matrix - 'mass_dict': {'0': 27 * unit, '1': 16 * unit}} - >>> grid_params = {'bounds_2': [[1.2], [3.5]], - # [[lower_bound], [upper_bound]] - # These describe the lower and upper - # bounds used to specify the 2-body spline - # fits. - 'bounds_3': [[1.2, 1.2, 1.2], [3.5, 3.5, 3.5]], - # [[lower,lower,lower],[upper,upper,upper]] - # Values describe lower and upper bounds - # for the bondlength-bondlength-bondlength - # grid used to construct and fit 3-body - # kernels; note that for force MGPs - # bondlength-bondlength-costheta - # are the bounds used instead. - 'bodies': [2, 3] # use 2+3 body - 'grid_num_2': 64,# Fidelity of the grid - 'grid_num_3': [16, 16, 16],# Fidelity of the grid - 'svd_rank_2': 64, #Fidelity of uncertainty estimation - 'svd_rank_3': 16**3, - 'update': True, # if True: accelerating grids - # generating by saving intermediate - # coeff when generating grids - 'load_grid': None # Used to load from file - } + >>> # build 2 + 3 body map + >>> grid_params = {'twobody': {'grid_num': [64]}, + ... 'threebody': {'grid_num': [64, 64, 64]}} + + For `grid_params`, the following keys and values are allowed + + Args: + 'two_body' (dict, optional): if 2-body is present, set as a dictionary + of parameters for 2-body mapping. Parameters see below. + 'three_body' (dict, optional): if 3-body is present, set as a dictionary + of parameters for 3-body mapping. Parameters see below. + 'load_grid' (str, optional): Default None. the path to the directory + where the previously generated grids (``grid_*.npy``) are stored. + If no path is specified, MGP will construct grids from scratch. + 'lower_bound_relax' (float, optional): Default 0.1. if 'lower_bound' is + set to 'auto' this value will be used as a relaxation of lower + bound. (see below the description of 'lower_bound') + + For two/three body parameter dictionary, the following keys and values are allowed + + Args: + 'grid_num' (list): a list of integers, the number of grid points for + interpolation. The larger the number, the better the approximation + of MGP is compared with GP. + 'lower_bound' (str or list, optional): Default 'auto', the lower bound + of the spline interpolation will be searched. First, search the + training set of GP and find the minimal interatomic distance r_min. + Then, the ``lower_bound = r_min - lower_bound_relax``. The user + can set their own lower_bound, of the same shape as 'grid_num'. + E.g. for threebody, the customized lower bound can be set as + [1.2, 1.2, 1.2]. + 'upper_bound' (str or list, optional): Default 'auto', the upper bound + of the spline interpolation will be the cutoffs of GP. The user + can set their own upper_bound, of the same shape as 'grid_num'. + E.g. for threebody, the customized lower bound can be set as + [3.5, 3.5, 3.5]. + 'svd_rank' (int, optional): Default 'auto'. If the variance mapping is + needed, it is set as the rank of the mapping. 'auto' uses full + rank, which is the smaller one between the total number of grid + points and training set size. i.e. + ``full_rank = min(np.prod(grid_num), 3 * N_train)`` ''' def __init__(self, grid_params: dict, - struc_params: dict, + unique_species: list=[], map_force: bool=False, GP: GaussianProcess=None, - mean_only: bool=False, + mean_only: bool=True, container_only: bool=True, lmp_file_name: str='lmp.mgp', n_cpus: int=None, @@ -99,356 +105,115 @@ def __init__(self, self.n_cpus = n_cpus self.n_sample = n_sample self.grid_params = grid_params - self.struc_params = struc_params + self.species_labels = [] + self.coded_species = [] + self.hyps_mask = None self.cutoffs = None - # arg_dict = inspect.getargvalues(inspect.currentframe())[3] - # del arg_dict['self'], arg_dict['GP'] - # self.__dict__.update(arg_dict) - self.__dict__.update(grid_params) - - if self.map_force and (3 in self.bodies): - assert (np.abs(self.bounds_3[0][2]) <= 1) or \ - (np.abs(self.bounds_3[1][2]) <= 1), \ - 'The force mapping needs to specify [bond, bond, cos_angle] for \ - 3-body, the 3rd dimension should be in range -1 to 1' - - # if GP exists, the GP setup overrides the grid_params setup - if GP is not None: - - self.cutoffs = GP.cutoffs - self.hyps_mask = GP.hyps_mask - - self.bodies = [] - if "two" in GP.kernel_name: - self.bodies.append(2) - self.kernel2b_info = get_2bkernel(GP) - if "three" in GP.kernel_name: - self.bodies.append(3) - self.kernel3b_info = get_3bkernel(GP) - - self.build_bond_struc(struc_params) - self.maps_2 = [] - self.maps_3 = [] - self.build_map_container(GP) - self.mean_only = mean_only - - if not container_only and (GP is not None) and \ - (len(GP.training_data) > 0): - self.build_map(GP) + for i, ele in enumerate(unique_species): + if isinstance(ele, str): + self.species_labels.append(ele) + self.coded_species.append(element_to_Z(ele)) + elif isinstance(ele, int): + self.coded_species.append(ele) + self.species_labels.append(Z_to_element(ele)) + + self.load_grid = grid_params.get('load_grid', None) + self.update = grid_params.get('update', False) + self.lower_bound_relax = grid_params.get('lower_bound_relax', 0.1) + + self.maps = {} + + optional_xb_params = ['lower_bound', 'upper_bound', 'svd_rank'] + for key in grid_params: + if 'body' in key: + if 'twobody' == key: + mapxbody = Map2body + elif 'threebody' == key: + mapxbody = Map3body + else: + raise KeyError("Only 'twobody' & 'threebody' are allowed") - def build_map_container(self, GP=None): - ''' - construct an empty spline container without coefficients. - ''' + xb_dict = grid_params[key] - if (GP is not None): - self.cutoffs = GP.cutoffs - self.hyps_mask = GP.hyps_mask + # set to 'auto' if the param is not given + args = {} + for oxp in optional_xb_params: + args[oxp] = xb_dict.get(oxp, 'auto') + args['grid_num'] = xb_dict.get('grid_num', None) + for k in xb_dict: + args[k] = xb_dict[k] - if 2 in self.bodies: - for b_struc in self.bond_struc[0]: - if (GP is not None): - self.bounds_2[1][0] = hpm.get_cutoff(b_struc.coded_species, - self.cutoffs, self.hyps_mask) - map_2 = Map2body(self.grid_num_2, self.bounds_2, - b_struc, self.map_force, self.svd_rank_2, - self.mean_only, self.n_cpus, self.n_sample) - self.maps_2.append(map_2) - if 3 in self.bodies: - for b_struc in self.bond_struc[1]: - if (GP is not None): - self.bounds_3[1] = hpm.get_cutoff(b_struc.coded_species, - self.cutoffs, self.hyps_mask) - map_3 = Map3body(self.grid_num_3, self.bounds_3, - b_struc, self.map_force, self.svd_rank_3, - self.mean_only, - self.grid_params['load_grid'], - self.update, self.n_cpus, self.n_sample) - self.maps_3.append(map_3) + xb_maps = mapxbody(**args, **self.__dict__) + self.maps[key] = xb_maps def build_map(self, GP): - ''' - generate/load grids and get spline coefficients - ''' - - # double check the container and the GP is the consistent - if not hpm.compare_dict(GP.hyps_mask, self.hyps_mask): - self.build_map_container(GP) - - if 2 in self.bodies: - self.kernel2b_info = get_2bkernel(GP) - if 3 in self.bodies: - self.kernel3b_info = get_3bkernel(GP) + self.hyps_mask = GP.hyps_mask + self.cutoffs = GP.cutoffs - for map_2 in self.maps_2: - map_2.build_map(GP) - for map_3 in self.maps_3: - map_3.build_map(GP) + for xb in self.maps: + self.maps[xb].build_map(GP) # write to lammps pair style coefficient file self.write_lmp_file(self.lmp_file_name) - def build_bond_struc(self, struc_params): - ''' - build a bond structure, used in grid generating - ''' - - cutoff = 0.1 - cell = struc_params['cube_lat'] - mass_dict = struc_params['mass_dict'] - species_list = struc_params['species'] - N_spc = len(species_list) - - # 2 body (2 atoms (1 bond) config) - bond_struc_2 = [] - spc_2 = [] - spc_2_set = [] - if 2 in self.bodies: - bodies = 2 - for spc1_ind, spc1 in enumerate(species_list): - for spc2 in species_list[spc1_ind:]: - species = [spc1, spc2] - spc_2.append(species) - spc_2_set.append(set(species)) - positions = [[(i+1)/(bodies+1)*cutoff, 0, 0] - for i in range(bodies)] - spc_struc = \ - Structure(cell, species, positions, mass_dict) - spc_struc.coded_species = np.array(species) - bond_struc_2.append(spc_struc) - # 3 body (3 atoms (1 triplet) config) - bond_struc_3 = [] - spc_3 = [] - if 3 in self.bodies: - bodies = 3 - for spc1_ind in range(N_spc): - spc1 = species_list[spc1_ind] - for spc2_ind in range(N_spc): # (spc1_ind, N_spc): - spc2 = species_list[spc2_ind] - for spc3_ind in range(N_spc): # (spc2_ind, N_spc): - spc3 = species_list[spc3_ind] - species = [spc1, spc2, spc3] - spc_3.append(species) - positions = [[(i+1)/(bodies+1)*cutoff, 0, 0] - for i in range(bodies)] - spc_struc = Structure(cell, species, positions, - mass_dict) - spc_struc.coded_species = np.array(species) - bond_struc_3.append(spc_struc) -# if spc1 != spc2: -# species = [spc2, spc3, spc1] -# spc_3.append(species) -# positions = [[(i+1)/(bodies+1)*cutoff, 0, 0] \ -# for i in range(bodies)] -# spc_struc = Structure(cell, species, positions, -# mass_dict) -# spc_struc.coded_species = np.array(species) -# bond_struc_3.append(spc_struc) -# if spc2 != spc3: -# species = [spc3, spc1, spc2] -# spc_3.append(species) -# positions = [[(i+1)/(bodies+1)*cutoff, 0, 0] \ -# for i in range(bodies)] -# spc_struc = Structure(cell, species, positions, -# mass_dict) -# spc_struc.coded_species = np.array(species) -# bond_struc_3.append(spc_struc) - - self.bond_struc = [bond_struc_2, bond_struc_3] - self.spcs = [spc_2, spc_3] - self.spcs_set = [spc_2_set, spc_3] - - def predict(self, atom_env: AtomicEnvironment, mean_only: bool = False)\ - -> (float, 'ndarray', 'ndarray', float): + def predict(self, atom_env: AtomicEnvironment, mean_only: bool = True, + ) -> (float, 'ndarray', 'ndarray', float): ''' predict force, variance, stress and local energy for given - atomic environment + atomic environment + Args: atom_env: atomic environment (with a center atom and its neighbors) mean_only: if True: only predict force (variance is always 0) + Return: force: 3d array of atomic force variance: 3d array of the predictive variance stress: 6d array of the virial stress energy: the local energy (atomic energy) ''' - if self.mean_only: # if not build mapping for var - mean_only = True - # ---------------- predict for two body ------------------- - f2 = vir2 = kern2 = v2 = e2 = 0 - if 2 in self.bodies: + force = virial = kern = v = energy = 0 + for xb in self.maps: + pred = self.maps[xb].predict(atom_env, mean_only) + force += pred[0] + virial += pred[1] + kern += pred[2] + v += pred[3] + energy += pred[4] - f2, vir2, kern2, v2, e2 = \ - self.predict_multicomponent(2, atom_env, self.kernel2b_info, - self.spcs_set[0], - self.maps_2, mean_only) - - # ---------------- predict for three body ------------------- - f3 = vir3 = kern3 = v3 = e3 = 0 - if 3 in self.bodies: - - f3, vir3, kern3, v3, e3 = \ - self.predict_multicomponent(3, atom_env, self.kernel3b_info, - self.spcs[1], self.maps_3, - mean_only) - - force = f2 + f3 - variance = kern2 + kern3 - np.sum((v2 + v3)**2, axis=0) - virial = vir2 + vir3 - energy = e2 + e3 + variance = kern - np.sum(v**2, axis=0) return force, variance, virial, energy - def predict_multicomponent(self, body, atom_env, kernel_info, - spcs_list, mappings, mean_only): - ''' - Add up results from `predict_component` to get the total contribution - of all species - ''' - - kernel, en_kernel, en_force_kernel, cutoffs, hyps, hyps_mask = \ - kernel_info - - args = from_mask_to_args(hyps, hyps_mask, cutoffs) - - kern = np.zeros(3) - for d in range(3): - kern[d] = \ - kernel(atom_env, atom_env, d+1, d+1, *args) - - if (body == 2): - spcs, comp_r, comp_xyz = get_bonds(atom_env.ctype, - atom_env.etypes, atom_env.bond_array_2) - set_spcs = [] - for spc in spcs: - set_spcs += [set(spc)] - spcs = set_spcs - elif (body == 3): - if self.map_force: - get_triplets_func = get_triplets - else: - get_triplets_func = get_triplets_en - - spcs, comp_r, comp_xyz = \ - get_triplets_func(atom_env.ctype, atom_env.etypes, - atom_env.bond_array_3, atom_env.cross_bond_inds, - atom_env.cross_bond_dists, atom_env.triplet_counts) - - # predict for each species - f_spcs = 0 - vir_spcs = 0 - v_spcs = 0 - e_spcs = 0 - for i, spc in enumerate(spcs): - lengths = np.array(comp_r[i]) - xyzs = np.array(comp_xyz[i]) - map_ind = spcs_list.index(spc) - f, vir, v, e = self.predict_component(lengths, xyzs, - mappings[map_ind], mean_only) - f_spcs += f - vir_spcs += vir - v_spcs += v - e_spcs += e - - return f_spcs, vir_spcs, kern, v_spcs, e_spcs - - def predict_component(self, lengths, xyzs, mapping, mean_only): - ''' - predict force and variance contribution of one component - ''' - lengths = np.array(lengths) - xyzs = np.array(xyzs) - - # predict mean - if self.map_force: # force mapping - e = 0 - f_0 = mapping.mean(lengths) - f_d = np.diag(f_0) @ xyzs - f = np.sum(f_d, axis=0) - - # predict stress from force components - vir = np.zeros(6) - vir_order = ((0,0), (1,1), (2,2), (0,1), (0,2), (1,2)) - for i in range(6): - vir_i = f_d[:,vir_order[i][0]]\ - * xyzs[:,vir_order[i][1]] * lengths[:,0] - vir[i] = np.sum(vir_i) - vir *= 0.5 - - else: # energy mapping - e_0, f_0 = mapping.mean(lengths, with_derivatives=True) - e = np.sum(e_0) # energy - - # predict forces and stress - vir = np.zeros(6) - vir_order = ((0,0), (1,1), (2,2), (1,2), (0,2), (0,1)) # match the ASE order - - # two-body - if lengths.shape[-1] == 1: - f_d = np.diag(f_0[:,0,0]) @ xyzs - f = 2 * np.sum(f_d, axis=0) # force: need to check prefactor 2 - - for i in range(6): - vir_i = f_d[:,vir_order[i][0]]\ - * xyzs[:,vir_order[i][1]] * lengths[:,0] - vir[i] = np.sum(vir_i) - - # three-body - if lengths.shape[-1] == 3: - f_d1 = np.diag(f_0[:,0,0]) @ xyzs[:,0,:] - f_d2 = np.diag(f_0[:,1,0]) @ xyzs[:,1,:] - f_d = f_d1 + f_d2 - f = 3 * np.sum(f_d, axis=0) # force: need to check prefactor 3 - - for i in range(6): - vir_i1 = f_d1[:,vir_order[i][0]]\ - * xyzs[:,0,vir_order[i][1]] * lengths[:,0] - vir_i2 = f_d2[:,vir_order[i][0]]\ - * xyzs[:,1,vir_order[i][1]] * lengths[:,1] - vir[i] = np.sum(vir_i1 + vir_i2) - vir *= 1.5 - - - # predict var - # TODO: implement energy var - v = np.zeros(3) - if not mean_only: - v_0 = mapping.var(lengths) - v_d = v_0 @ xyzs - v = mapping.var.V @ v_d - return f, vir, v, e def write_lmp_file(self, lammps_name): ''' write the coefficients to a file that can be used by lammps pair style ''' - # write header f = open(lammps_name, 'w') - header_comment = '''# #2bodyarray #3bodyarray\n# elem1 elem2 a b order - ''' + # write header + header_comment = '''# #2bodyarray #3bodyarray\n# elem1 elem2 a b order\n\n''' f.write(header_comment) + header = '' + xbodies = ['twobody', 'threebody'] + for xb in xbodies: + if xb in self.maps: + num = len(self.maps[xb].maps) + else: + num = 0 + header += f'{num} ' + f.write(header + '\n') - twobodyarray = len(self.spcs[0]) - threebodyarray = len(self.spcs[1]) - header = '\n{} {}\n'.format(twobodyarray, threebodyarray) - f.write(header) - - # write two body - if twobodyarray > 0: - for ind, spc in enumerate(self.spcs[0]): - self.maps_2[ind].write(f, spc) - - # write three body - if threebodyarray > 0: - for ind, spc in enumerate(self.spcs[1]): - self.maps_3[ind].write(f, spc) + # write coefficients + for xb in self.maps: + self.maps[xb].write(f) f.close() @@ -458,6 +223,7 @@ def as_dict(self) -> dict: """ out_dict = deepcopy(dict(vars(self))) + out_dict.pop('maps') # Uncertainty mappings currently not serializable; if not self.mean_only: @@ -466,22 +232,11 @@ def as_dict(self) -> dict: "them.", Warning) out_dict['mean_only'] = True - # Iterate through the mappings for various bodies - for i in self.bodies: - kern_info = f'kernel{i}b_info' - kernel, ek, efk, cutoffs, hyps, hyps_mask = out_dict[kern_info] - out_dict[kern_info] = (kernel.__name__, efk.__name__, - cutoffs, hyps, hyps_mask) - # only save the coefficients - out_dict['maps_2'] = [map_2.mean.__coeffs__ for map_2 in self.maps_2] - out_dict['maps_3'] = [map_3.mean.__coeffs__ for map_3 in self.maps_3] - - # don't need these since they are built in the __init__ function - key_list = ['bond_struc', 'spcs_set', ] - for key in key_list: - if out_dict.get(key) is not None: - del out_dict[key] + maps_dict = {} + for m in self.maps: + maps_dict[m] = self.maps[m].as_dict() + out_dict['maps'] = maps_dict return out_dict @@ -490,42 +245,27 @@ def from_dict(dictionary: dict): """ Create MGP object from dictionary representation. """ - new_mgp = MappedGaussianProcess(grid_params=dictionary['grid_params'], - struc_params=dictionary['struc_params'], - map_force=dictionary['map_force'], - GP=None, - mean_only=dictionary['mean_only'], - container_only=True, - lmp_file_name=dictionary['lmp_file_name'], - n_cpus=dictionary['n_cpus'], - n_sample=dictionary['n_sample']) - # Restore kernel_info - for i in dictionary['bodies']: - kern_info = f'kernel{i}b_info' - hyps_mask = dictionary[kern_info][-1] - if (hyps_mask is None): - multihyps = False - else: - multihyps = True + # Set GP + if dictionary.get('GP'): + GP = GaussianProcess.from_dict(dictionary.get("GP")) + else: + dictionary['GP'] = None - kernel_info = dictionary[kern_info] - kernel_name = kernel_info[0] - kernel, _, ek, efk = str_to_kernel_set(kernel_name, multihyps) - kernel_info[0] = kernel - kernel_info[1] = ek - kernel_info[2] = efk - setattr(new_mgp, kern_info, kernel_info) + dictionary['unique_species'] = list(set(dictionary['species_labels'])) + if 'container_only' not in dictionary: + dictionary['container_only'] = True - # Fill up the model with the saved coeffs - for m, map_2 in enumerate(new_mgp.maps_2): - map_2.mean.__coeffs__ = np.array(dictionary['maps_2'][m]) - for m, map_3 in enumerate(new_mgp.maps_3): - map_3.mean.__coeffs__ = np.array(dictionary['maps_3'][m]) + init_arg_name = ['grid_params', 'unique_species', 'map_force', 'GP', + 'mean_only', 'container_only', 'lmp_file_name', 'n_cpus', 'n_sample'] + kwargs = {key: dictionary[key] for key in init_arg_name} + new_mgp = MappedGaussianProcess(**kwargs) - # Set GP - if dictionary.get('GP'): - new_mgp.GP = GaussianProcess.from_dict(dictionary.get("GP")) + # Fill up the model with the saved coeffs + if 'twobody' in new_mgp.maps: + new_mgp.maps['twobody'] = Map2body.from_dict(dictionary['maps']['twobody'], Map2body) + if 'threebody' in new_mgp.maps: + new_mgp.maps['threebody'] = Map3body.from_dict(dictionary['maps']['threebody'], Map3body) return new_mgp @@ -563,735 +303,3 @@ def from_file(filename: str): raise NotImplementedError -class Map2body: - def __init__(self, grid_num: int, bounds, bond_struc: Structure, - map_force=False, svd_rank=0, mean_only: bool=False, - n_cpus: int=None, n_sample: int=100): - ''' - Build 2-body MGP - - bond_struc: Mock structure used to sample 2-body forces on 2 atoms - ''' - - self.grid_num = grid_num - self.bounds = bounds - self.bond_struc = bond_struc - self.map_force = map_force - self.svd_rank = svd_rank - self.mean_only = mean_only - self.n_cpus = n_cpus - self.n_sample = n_sample - - spc = bond_struc.coded_species - self.species_code = Z_to_element(spc[0]) + '_' + Z_to_element(spc[1]) - -# arg_dict = inspect.getargvalues(inspect.currentframe())[3] -# del arg_dict['self'] -# self.__dict__.update(arg_dict) - - self.build_map_container() - - def GenGrid(self, GP): - ''' - To use GP to predict value on each grid point, we need to generate the - kernel vector kv whose length is the same as the training set size. - - 1. We divide the training set into several batches, corresponding to - different segments of kv - 2. Distribute each batch to a processor, i.e. each processor calculate - the kv segment of one batch for all grids - 3. Collect kv segments and form a complete kv vector for each grid, - and calculate the grid value by multiplying the complete kv vector - with GP.alpha - ''' - - kernel_info = get_2bkernel(GP) - - if (self.n_cpus is None): - processes = mp.cpu_count() - else: - processes = self.n_cpus - - # ------ construct grids ------ - nop = self.grid_num - bond_lengths = np.linspace(self.bounds[0][0], self.bounds[1][0], nop) - env12 = AtomicEnvironment( - self.bond_struc, 0, GP.cutoffs, cutoffs_mask=GP.hyps_mask) - - # --------- calculate force kernels --------------- - n_envs = len(GP.training_data) - n_strucs = len(GP.training_structures) - n_kern = n_envs * 3 + n_strucs - - if (n_envs > 0): - with mp.Pool(processes=processes) as pool: - - block_id, nbatch = \ - partition_vector(self.n_sample, n_envs, processes) - - k12_slice = [] - for ibatch in range(nbatch): - s, e = block_id[ibatch] - k12_slice.append(pool.apply_async( - self._GenGrid_inner, args=(GP.name, s, e, bond_lengths, - env12, kernel_info))) - k12_matrix = [] - for ibatch in range(nbatch): - k12_matrix += [k12_slice[ibatch].get()] - pool.close() - pool.join() - del k12_slice - k12_v_force = np.vstack(k12_matrix) - del k12_matrix - - # --------- calculate energy kernels --------------- - if (n_strucs > 0): - with mp.Pool(processes=processes) as pool: - block_id, nbatch = \ - partition_vector(self.n_sample, n_strucs, processes) - - k12_slice = [] - for ibatch in range(nbatch): - s, e = block_id[ibatch] - k12_slice.append(pool.apply_async( - self._GenGrid_energy, - args=(GP.name, s, e, bond_lengths, env12, kernel_info))) - k12_matrix = [] - for ibatch in range(nbatch): - k12_matrix += [k12_slice[ibatch].get()] - pool.close() - pool.join() - del k12_slice - k12_v_energy = np.vstack(k12_matrix) - del k12_matrix - - if (n_strucs > 0 and n_envs > 0): - k12_v_all = np.vstack([k12_v_force, k12_v_energy]) - k12_v_all = np.moveaxis(k12_v_all, 0, -1) - del k12_v_force - del k12_v_energy - elif (n_strucs > 0): - k12_v_all = np.moveaxis(k12_v_energy, 0, -1) - del k12_v_energy - elif (n_envs > 0): - k12_v_all = np.moveaxis(k12_v_force, 0, -1) - del k12_v_force - else: - return np.zeros([nop]), None - - # ------- compute bond means and variances --------------- - bond_means = np.zeros([nop]) - if not self.mean_only: - bond_vars = np.zeros([nop, len(GP.alpha)]) - else: - bond_vars = None - for b, _ in enumerate(bond_lengths): - k12_v = k12_v_all[b, :] - bond_means[b] = np.matmul(k12_v, GP.alpha) - if not self.mean_only: - bond_vars[b, :] = solve_triangular(GP.l_mat, k12_v, lower=True) - - write_species_name = '' - for x in self.bond_struc.coded_species: - write_species_name += "_" + Z_to_element(x) - # ------ save mean and var to file ------- - np.save('grid2_mean' + write_species_name, bond_means) - np.save('grid2_var' + write_species_name, bond_vars) - - return bond_means, bond_vars - - def _GenGrid_inner(self, name, s, e, bond_lengths, - env12, kernel_info): - ''' - Calculate kv segments of the given batch of training data for all grids - ''' - - kernel, ek, efk, cutoffs, hyps, hyps_mask = kernel_info - size = e - s - k12_v = np.zeros([len(bond_lengths), size*3]) - for b, r in enumerate(bond_lengths): - env12.bond_array_2 = np.array([[r, 1, 0, 0]]) - if self.map_force: - k12_v[b, :] = force_force_vector_unit(name, s, e, env12, kernel, hyps, - cutoffs, hyps_mask, 1) - - else: - k12_v[b, :] = energy_force_vector_unit(name, s, e, - env12, efk, hyps, cutoffs, hyps_mask) - return np.moveaxis(k12_v, 0, -1) - - def _GenGrid_energy(self, name, s, e, bond_lengths, env12, kernel_info): - ''' - Calculate kv segments of the given batch of training data for all grids - ''' - - kernel, ek, efk, cutoffs, hyps, hyps_mask = kernel_info - size = e - s - k12_v = np.zeros([len(bond_lengths), size]) - for b, r in enumerate(bond_lengths): - env12.bond_array_2 = np.array([[r, 1, 0, 0]]) - - if self.map_force: - k12_v[b, :] = force_energy_vector_unit(name, s, e, env12, efk, - hyps, cutoffs, hyps_mask, 1) - else: - k12_v[b, :] = energy_energy_vector_unit(name, s, e, - env12, ek, hyps, cutoffs, hyps_mask) - return np.moveaxis(k12_v, 0, -1) - - - - def build_map_container(self): - ''' - build 1-d spline function for mean, 2-d for var - ''' - self.mean = CubicSpline(self.bounds[0], self.bounds[1], - orders=[self.grid_num]) - - if not self.mean_only: - self.var = PCASplines(self.bounds[0], self.bounds[1], - orders=[self.grid_num], - svd_rank=self.svd_rank) - - def build_map(self, GP): - y_mean, y_var = self.GenGrid(GP) - self.mean.set_values(y_mean) - if not self.mean_only: - self.var.set_values(y_var) - - def write(self, f, spc): - ''' - Write LAMMPS coefficient file - ''' - a = self.bounds[0][0] - b = self.bounds[1][0] - order = self.grid_num - - coefs_2 = self.mean.__coeffs__ - - elem1 = Z_to_element(spc[0]) - elem2 = Z_to_element(spc[1]) - header_2 = '{elem1} {elem2} {a} {b} {order}\n'\ - .format(elem1=elem1, elem2=elem2, a=a, b=b, order=order) - f.write(header_2) - - for c, coef in enumerate(coefs_2): - f.write('{:.10e} '.format(coef)) - if c % 5 == 4 and c != len(coefs_2)-1: - f.write('\n') - - f.write('\n') - - -class Map3body: - - def __init__(self, grid_num, bounds, bond_struc: Structure, - map_force: bool = False, svd_rank: int = 0, mean_only: bool=False, - load_grid: str = '', update: bool = True, n_cpus: int = None, - n_sample: int = 100): - ''' - Build 3-body MGP - - bond_struc: Mock Structure object which contains 3 atoms to get map - from - ''' - self.grid_num = grid_num - self.bounds = bounds - self.bond_struc = bond_struc - self.map_force = map_force - self.svd_rank = svd_rank - self.mean_only = mean_only - self.load_grid = load_grid - self.update = update - self.n_sample = n_sample - - spc = bond_struc.coded_species - self.species_code = Z_to_element(spc[0]) + '_' + \ - Z_to_element(spc[1]) + '_' + Z_to_element(spc[2]) - self.kv3name = f'kv3_{self.species_code}' - - self.build_map_container() - self.n_cpus = n_cpus - self.bounds = bounds - self.mean_only = mean_only - - def GenGrid(self, GP): - ''' - To use GP to predict value on each grid point, we need to generate the - kernel vector kv whose length is the same as the training set size. - - 1. We divide the training set into several batches, corresponding to - different segments of kv - 2. Distribute each batch to a processor, i.e. each processor calculate - the kv segment of one batch for all grids - 3. Collect kv segments and form a complete kv vector for each grid, - and calculate the grid value by multiplying the complete kv vector - with GP.alpha - ''' - - if self.n_cpus is None: - processes = mp.cpu_count() - else: - processes = self.n_cpus - - # ------ get 3body kernel info ------ - kernel_info = get_3bkernel(GP) - - # ------ construct grids ------ - n1, n2, n12 = self.grid_num - bonds1 = np.linspace(self.bounds[0][0], self.bounds[1][0], n1) - bonds2 = np.linspace(self.bounds[0][1], self.bounds[1][1], n2) - bonds12 = np.linspace(self.bounds[0][2], self.bounds[1][2], n12) - grid_means = np.zeros([n1, n2, n12]) - - if not self.mean_only: - grid_vars = np.zeros([n1, n2, n12, len(GP.alpha)]) - else: - grid_vars = None - - env12 = AtomicEnvironment(self.bond_struc, 0, GP.cutoffs, - cutoffs_mask=GP.hyps_mask) - n_envs = len(GP.training_data) - n_strucs = len(GP.training_structures) - n_kern = n_envs * 3 + n_strucs - - mapk = str_to_mapped_kernel('3', GP.multihyps) - mapped_kernel_info = (kernel_info[0], mapk[0], mapk[1], - kernel_info[3], kernel_info[4], kernel_info[5]) - - if processes == 1: - if self.update: - raise NotImplementedError("the update function is " - "not yet implemented") - else: - if (n_envs > 0): - k12_v_force = \ - self._GenGrid_numba(GP.name, 0, n_envs, self.bounds, - n1, n2, n12, env12, mapped_kernel_info) - if (n_strucs > 0): - k12_v_energy = \ - self._GenGrid_energy(GP.name, 0, n_strucs, bonds1, bonds2, - bonds12, env12, kernel_info) - else: - - # ------------ force kernels ------------- - if (n_envs > 0): - if self.update: - - self.UpdateGrid() - - - - else: - block_id, nbatch = \ - partition_vector(self.n_sample, n_envs, processes) - - k12_slice = [] - with mp.Pool(processes=processes) as pool: - for ibatch in range(nbatch): - s, e = block_id[ibatch] - k12_slice.append(pool.apply_async( - self._GenGrid_inner, - args=(GP.name, s, e, bonds1, bonds2, bonds12, - env12, kernel_info))) - k12_matrix = [] - for ibatch in range(nbatch): - k12_matrix += [k12_slice[ibatch].get()] - pool.close() - pool.join() - - del k12_slice - k12_v_force = np.vstack(k12_matrix) - del k12_matrix - - # set OMB_NUM_THREADS mkl threads number to # of logical cores, per_atom_par=False - # ------------ force kernels ------------- - if (n_strucs > 0): - if self.update: - - self.UpdateGrid() - - - - else: - block_id, nbatch = \ - partition_vector(self.n_sample, n_strucs, processes) - - k12_slice = [] - with mp.Pool(processes=processes) as pool: - for ibatch in range(nbatch): - s, e = block_id[ibatch] - k12_slice.append(pool.apply_async( - self._GenGrid_energy, - args=(GP.name, s, e, bonds1, bonds2, bonds12, - env12, kernel_info))) - k12_matrix = [] - for ibatch in range(nbatch): - k12_matrix += [k12_slice[ibatch].get()] - pool.close() - pool.join() - - del k12_slice - k12_v_energy = np.vstack(k12_matrix) - del k12_matrix - - if (n_envs > 0 and n_strucs > 0): - k12_v_all = np.vstack([k12_v_force, k12_v_energy]) - k12_v_all = np.moveaxis(k12_v_all, 0, -1) - del k12_v_force - del k12_v_energy - elif (n_envs > 0): - k12_v_all = np.moveaxis(k12_v_force, 0, -1) - del k12_v_force - elif (n_strucs > 0): - k12_v_all = np.moveaxis(k12_v_energy, 0, -1) - del k12_v_energy - else: - return np.zeros(n1, n2, n12), None - - for b12 in range(len(bonds12)): - for b1 in range(len(bonds1)): - for b2 in range(len(bonds2)): - k12_v = k12_v_all[b1, b2, b12, :] - grid_means[b1, b2, b12] = np.matmul(k12_v, GP.alpha) - if not self.mean_only: - grid_vars[b1, b2, b12, :] = solve_triangular(GP.l_mat, - k12_v, lower=True) - - - # Construct file names according to current mapping - - # ------ save mean and var to file ------- - np.save('grid3_mean_'+self.species_code, grid_means) - np.save('grid3_var_'+self.species_code, grid_vars) - - return grid_means, grid_vars - - def UpdateGrid(self): - raise NotImplementedError("the update function is " - "not yet implemented") - - if self.kv3name in os.listdir(): - subprocess.run(['rm', '-rf', self.kv3name]) - - os.mkdir(self.kv3name) - - # get the size of saved kv vector - kv_filename = f'{self.kv3name}/{0}' - if kv_filename in os.listdir(self.kv3name): - old_kv_file = np.load(kv_filename+'.npy') - last_size = int(old_kv_file[0,0]) - new_kv_file[i, :, :last_size] = old_kv_file - - k12_v_all = np.zeros([len(bonds1), len(bonds2), len(bonds12), - size * 3]) - - for i in range(n12): - if f'{self.kv3name}/{i}.npy' in os.listdir(self.kv3name): - old_kv_file = np.load(f'{self.kv3name}/{i}.npy') - last_size = int(old_kv_file[0,0]) - #TODO k12_v_all[] - else: - last_size = 0 - - # parallelize based on grids, since usually the number of - # the added training points are small - ngrids = int(ceil(n12 / processes)) - nbatch = int(ceil(n12 / ngrids)) - - block_id = [] - for ibatch in range(nbatch): - s = int(ibatch * processes) - e = int(np.min(((ibatch+1)*processes, n12))) - block_id += [(s, e)] - - k12_slice = [] - for ibatch in range(nbatch): - k12_slice.append(pool.apply_async(self._GenGrid_inner, - args=(GP.name, last_size, size, - bonds1, bonds2, bonds12[s:e], - env12, kernel_info))) - - for ibatch in range(nbatch): - s, e = block_id[ibatch] - k12_v_all[:, :, s:e, :] = k12_slice[ibatch].get() - - - def _GenGrid_inner(self, name, s, e, bonds1, bonds2, bonds12, env12, kernel_info): - - ''' - Calculate kv segments of the given batch of training data for all grids - ''' - - kernel, ek, efk, cutoffs, hyps, hyps_mask = kernel_info - - # open saved k vector file, and write to new file - size = (e - s) * 3 - k12_v = np.zeros([len(bonds1), len(bonds2), len(bonds12), size]) - for b12, r12 in enumerate(bonds12): - for b1, r1 in enumerate(bonds1): - for b2, r2 in enumerate(bonds2): - - if self.map_force: - cos_angle12 = r12 - x2 = r2 * cos_angle12 - y2 = r2 * np.sqrt(1-cos_angle12**2) - dist12 = np.linalg.norm(np.array([x2-r1, y2, 0])) - else: - dist12 = r12 - - env12.bond_array_3 = np.array([[r1, 1, 0, 0], - [r2, 0, 0, 0]]) - env12.cross_bond_dists = np.array([[0, dist12], [dist12, 0]]) - - if self.map_force: - k12_v[b1, b2, b12, :] = \ - force_force_vector_unit(name, s, e, env12, kernel, hyps, - cutoffs, hyps_mask, 1) - else: - k12_v[b1, b2, b12, :] = energy_force_vector_unit(name, s, e, - env12, efk, - hyps, cutoffs, hyps_mask) - - # open saved k vector file, and write to new file - if self.update: - self.UpdateGrid_inner() - - return np.moveaxis(k12_v, -1, 0) - - def _GenGrid_energy(self, name, s, e, bonds1, bonds2, bonds12, env12, kernel_info): - - ''' - Calculate kv segments of the given batch of training data for all grids - ''' - - kernel, ek, efk, cutoffs, hyps, hyps_mask = kernel_info - - # open saved k vector file, and write to new file - size = e - s - k12_v = np.zeros([len(bonds1), len(bonds2), len(bonds12), size]) - for b12, r12 in enumerate(bonds12): - for b1, r1 in enumerate(bonds1): - for b2, r2 in enumerate(bonds2): - - if self.map_force: - cos_angle12 = r12 - x2 = r2 * cos_angle12 - y2 = r2 * np.sqrt(1-cos_angle12**2) - dist12 = np.linalg.norm(np.array([x2-r1, y2, 0])) - else: - dist12 = r12 - - env12.bond_array_3 = np.array([[r1, 1, 0, 0], - [r2, 0, 0, 0]]) - env12.cross_bond_dists = np.array([[0, dist12], [dist12, 0]]) - - if self.map_force: - k12_v[b1, b2, b12, :] = \ - force_energy_vector_unit(name, s, e, env12, efk, hyps, - cutoffs, hyps_mask, 1) - else: - k12_v[b1, b2, b12, :] = energy_energy_vector_unit(name, s, e, - env12, ek, - hyps, cutoffs, hyps_mask) - - # open saved k vector file, and write to new file - if self.update: - self.UpdateGrid_inner() - - return np.moveaxis(k12_v, -1, 0) - - - def _GenGrid_numba(self, name, s, e, bounds, nb1, nb2, nb12, env12, kernel_info): - """ - Loop over different parts of the training set. from element s to element e - - Args: - name: name of the gp instance - s: start index of the training data parition - e: end index of the training data parition - bonds1: list of bond to consider for edge center-1 - bonds2: list of bond to consider for edge center-2 - bonds12: list of bond to consider for edge 1-2 - env12: AtomicEnvironment container of the triplet - kernel_info: return value of the get_3b_kernel - """ - - kernel, en_kernel, en_force_kernel, cutoffs, hyps, hyps_mask = \ - kernel_info - - training_data = _global_training_data[name] - - ds = [1, 2, 3] - size = (e-s) * 3 - - bonds1 = np.linspace(bounds[0][0], bounds[1][0], nb1) - bonds2 = np.linspace(bounds[0][0], bounds[1][0], nb2) - bonds12 = np.linspace(bounds[0][2], bounds[1][2], nb12) - - r1 = np.ones([nb1, nb2, nb12], dtype=np.float64) - r2 = np.ones([nb1, nb2, nb12], dtype=np.float64) - r12 = np.ones([nb1, nb2, nb12], dtype=np.float64) - for b12 in range(nb12): - for b1 in range(nb1): - for b2 in range(nb2): - r1[b1, b2, b12] = bonds1[b1] - r2[b1, b2, b12] = bonds2[b2] - r12[b1, b2, b12] = bonds12[b12] - del bonds1 - del bonds2 - del bonds12 - - args = from_mask_to_args(hyps, hyps_mask, cutoffs) - - k_v = [] - for m_index in range(size): - x_2 = training_data[int(floor(m_index / 3))+s] - d_2 = ds[m_index % 3] - k_v += [[en_force_kernel(x_2, r1, r2, r12, - env12.ctype, env12.etypes, - d_2, *args)]] - - return np.vstack(k_v) - - def _GenGrid_energy_numba(self, name, s, e, bounds, nb1, nb2, nb12, env12, kernel_info): - """ - Loop over different parts of the training set. from element s to element e - - Args: - name: name of the gp instance - s: start index of the training data parition - e: end index of the training data parition - bonds1: list of bond to consider for edge center-1 - bonds2: list of bond to consider for edge center-2 - bonds12: list of bond to consider for edge 1-2 - env12: AtomicEnvironment container of the triplet - kernel_info: return value of the get_3b_kernel - """ - - kernel, en_kernel, en_force_kernel, cutoffs, hyps, hyps_mask = \ - kernel_info - - training_data = _global_training_structures[name] - - ds = [1, 2, 3] - size = (e-s) * 3 - - bonds1 = np.linspace(bounds[0][0], bounds[1][0], nb1) - bonds2 = np.linspace(bounds[0][0], bounds[1][0], nb2) - bonds12 = np.linspace(bounds[0][2], bounds[1][2], nb12) - - r1 = np.ones([nb1, nb2, nb12], dtype=np.float64) - r2 = np.ones([nb1, nb2, nb12], dtype=np.float64) - r12 = np.ones([nb1, nb2, nb12], dtype=np.float64) - for b12 in range(nb12): - for b1 in range(nb1): - for b2 in range(nb2): - r1[b1, b2, b12] = bonds1[b1] - r2[b1, b2, b12] = bonds2[b2] - r12[b1, b2, b12] = bonds12[b12] - del bonds1 - del bonds2 - del bonds12 - - args = from_mask_to_args(hyps, hyps_mask, cutoffs) - - k_v = [] - for m_index in range(size): - structure = training_structures[m_index + s] - kern_curr = 0 - for environment in structure: - kern_curr += en_kernel(x, environment, *args) - kv += [kern_curr] - - return np.hstack(k_v) - - - def UpdateGrid_inner(self): - raise NotImplementedError("the update function is not yet"\ - "implemented") - - s, e = block - chunk = e - s - new_kv_file = np.zeros((chunk, - self.grid_num[0]*self.grid_num[1]+1, - total_size)) - new_kv_file[:,0,0] = np.ones(chunk) * total_size - for i in range(s, e): - kv_filename = f'{self.kv3name}/{i}' - if kv_filename in os.listdir(self.kv3name): - old_kv_file = np.load(kv_filename+'.npy') - last_size = int(old_kv_file[0,0]) - new_kv_file[i, :, :last_size] = old_kv_file - else: - last_size = 0 - ds = [1, 2, 3] - nop = self.grid_num[0] - - k12_v = new_kv_file[:,1:,:] - for i in range(s, e): - np.save(f'{self.kv3name}/{i}', new_kv_file[i,:,:]) - - - - def build_map_container(self): - ''' - build 3-d spline function for mean, - 3-d for the low rank approximation of L^{-1}k* - ''' - - # create spline interpolation class object - self.mean = CubicSpline(self.bounds[0], self.bounds[1], - orders=self.grid_num) - - if not self.mean_only: - self.var = PCASplines(self.bounds[0], self.bounds[1], - orders=self.grid_num, - svd_rank=self.svd_rank) - - def build_map(self, GP): - # Load grid or generate grid values - # If load grid was not specified, will be none - if not self.load_grid: - y_mean, y_var = self.GenGrid(GP) - # If load grid is blank string '' or pre-fix, load in - else: - y_mean = np.load(self.load_grid+'grid3_mean_' + - self.species_code+'.npy') - y_var = np.load(self.load_grid+'grid3_var_' + - self.species_code+'.npy') - - self.mean.set_values(y_mean) - if not self.mean_only: - self.var.set_values(y_var) - - def write(self, f, spc): - a = self.bounds[0] - b = self.bounds[1] - order = self.grid_num - - coefs_3 = self.mean.__coeffs__ - - elem1 = Z_to_element(spc[0]) - elem2 = Z_to_element(spc[1]) - elem3 = Z_to_element(spc[2]) - - header_3 = '{elem1} {elem2} {elem3} {a1} {a2} {a3} {b1}'\ - ' {b2} {b3:.10e} {order1} {order2} {order3}\n'\ - .format(elem1=elem1, elem2=elem2, elem3=elem3, - a1=a[0], a2=a[1], a3=a[2], - b1=b[0], b2=b[1], b3=b[2], - order1=order[0], order2=order[1], order3=order[2]) - f.write(header_3) - - n = 0 - for i in range(coefs_3.shape[0]): - for j in range(coefs_3.shape[1]): - for k in range(coefs_3.shape[2]): - coef = coefs_3[i, j, k] - f.write('{:.10e} '.format(coef)) - if n % 5 == 4: - f.write('\n') - n += 1 - - f.write('\n') diff --git a/flare/mgp/splines_methods.py b/flare/mgp/splines_methods.py index eecd24ab3..fe2fe8720 100644 --- a/flare/mgp/splines_methods.py +++ b/flare/mgp/splines_methods.py @@ -56,7 +56,8 @@ def set_values(self, y): def __call__(self, x): y_pred = [] - for r in range(self.svd_rank): + rank = self.svd_rank + for r in range(rank): y_pred.append(self.models[r](x)) return np.array(y_pred) diff --git a/flare/mgp/utils.py b/flare/mgp/utils.py index 33fb90024..ca69d3d65 100644 --- a/flare/mgp/utils.py +++ b/flare/mgp/utils.py @@ -1,78 +1,83 @@ -import io, os, sys, time, random, math -import multiprocessing as mp +import warnings import numpy as np from numpy import array from numba import njit +from math import exp, floor +from typing import Callable from flare.env import AtomicEnvironment from flare.kernels.cutoffs import quadratic_cutoff -from flare.kernels.kernels import three_body_helper_1, \ - three_body_helper_2, force_helper -from flare.kernels.utils import str_to_kernel_set as stks -from flare.utils.mask_helper import HyperParameterMasking - - -def get_2bkernel(GP): - if 'mc' in GP.kernel_name: - kernel, _, ek, efk = stks('2', GP.multihyps) +from flare.kernels.utils import str_to_kernel_set +from flare.parameters import Parameters + +from flare.mgp.grid_kernels_3b import grid_kernel, grid_kernel_sephyps + + +def str_to_mapped_kernel(name: str, component: str = "mc", + hyps_mask: dict = None): + """ + Return kernels and kernel gradient function based on a string. + If it contains 'sc', it will use the kernel in sc module; + otherwise, it uses the kernel in mc_simple; + if sc is not included and multihyps is True, + it will use the kernel in mc_sephyps module. + Otherwise, it will use the kernel in the sc module. + + Args: + + name (str): name for kernels. example: "2+3mc" + multihyps (bool, optional): True for using multiple hyperparameter groups + + :return: mapped kernel function, kernel gradient, energy kernel, + energy_and_force kernel + + """ + + multihyps = True + if hyps_mask is None: + multihyps = False + elif hyps_mask['nspecie'] == 1: + multihyps = False + + # b2 = Two body in use, b3 = Three body in use + b2 = False + many = False + b3 = False + for s in ['3', 'three']: + if s in name.lower() or s == name.lower(): + b3 = True + + if b3: + if multihyps: + return grid_kernel_sephyps, None, None, None + else: + return grid_kernel, None, None, None else: - kernel, _, ek, efk = stks('2sc', GP.multihyps) - - cutoffs = [GP.cutoffs[0]] - - hyps, hyps_mask = HyperParameterMasking.get_2b_hyps(GP.hyps, GP.hyps_mask, GP.multihyps) - - return (kernel, ek, efk, cutoffs, hyps, hyps_mask) - - -def get_3bkernel(GP): - - if 'mc' in GP.kernel_name: - kernel, _, ek, efk = stks('3', GP.multihyps) + warnings.Warn(NotImplemented("mapped kernel for two-body and manybody kernels " + "are not implemented")) + return None + +def get_kernel_term(kernel_name, component, hyps_mask, hyps, grid_kernel=False): + """ + Args + term (str): 'twobody' or 'threebody' + """ + if grid_kernel: + stks = str_to_mapped_kernel + kernel_name_list = kernel_name else: - kernel, _, ek, efk = stks('3sc', GP.multihyps) + stks = str_to_kernel_set + kernel_name_list = [kernel_name] - base = 0 - for t in ['two', '2']: - if t in GP.kernel_name: - base = 2 + kernel, _, ek, efk = stks(kernel_name_list, component, hyps_mask) - cutoffs = np.copy(GP.cutoffs) - - hyps, hyps_mask = \ - HyperParameterMasking.get_3b_hyps(\ - GP.hyps, GP.hyps_mask, GP.multihyps) + # hyps_mask is modified here + hyps, cutoffs, hyps_mask = Parameters.get_component_mask(hyps_mask, kernel_name, hyps=hyps) return (kernel, ek, efk, cutoffs, hyps, hyps_mask) -def get_l_bound(curr_l_bound, structure, two_d=False): - positions = structure.positions - if two_d: - cell = structure.cell[:2] - else: - cell = structure.cell - - min_dist = curr_l_bound - for ind1, pos01 in enumerate(positions): - for i1 in range(2): - for vec1 in cell: - pos1 = pos01 + i1 * vec1 - - for ind2, pos02 in enumerate(positions): - for i2 in range(2): - for vec2 in cell: - pos2 = pos02 + i2 * vec2 - - if np.all(pos1 == pos2): - continue - dist12 = np.linalg.norm(pos1-pos2) - if dist12 < min_dist: - min_dist = dist12 - min_atoms = (ind1, ind2) - return min_dist - @njit def get_bonds(ctype, etypes, bond_array): @@ -98,6 +103,7 @@ def get_bonds(ctype, etypes, bond_array): bond_dirs.append([b_dir]) return exist_species, bond_lengths, bond_dirs + @njit def get_triplets(ctype, etypes, bond_array, cross_bond_inds, cross_bond_dists, triplets): @@ -114,27 +120,16 @@ def get_triplets(ctype, etypes, bond_array, cross_bond_inds, ind1 = cross_bond_inds[m, m+n+1] r2 = bond_array[ind1, 0] c2 = bond_array[ind1, 1:] - c12 = np.sum(c1*c2) - if c12 > 1: # to prevent numerical error - c12 = 1 - elif c12 < -1: - c12 = -1 spc2 = etypes[ind1] -# if spc1 == spc2: -# spcs_list = [[ctype, spc1, spc2], [ctype, spc1, spc2]] -# elif ctype == spc1: # spc1 != spc2 -# spcs_list = [[ctype, spc1, spc2], [spc2, ctype, spc1]] -# elif ctype == spc2: # spc1 != spc2 -# spcs_list = [[spc1, spc2, ctype], [spc2, ctype, spc1]] -# else: # all different -# spcs_list = [[ctype, spc1, spc2], [ctype, spc2, spc1]] + c12 = np.sum(c1*c2) + r12 = np.sqrt(r1**2 + r2**2 - 2*r1*r2*c12) spcs_list = [[ctype, spc1, spc2], [ctype, spc2, spc1]] for i in range(2): spcs = spcs_list[i] - triplet = array([r2, r1, c12]) if i else array([r1, r2, c12]) - coord = c2 if i else c1 + triplet = array([r2, r1, r12]) if i else array([r1, r2, r12]) + coord = c2 if i else c1 # TODO: figure out what's wrong. why not [c1, c2] for force map if spcs not in exist_species: exist_species.append(spcs) tris.append([triplet]) @@ -145,54 +140,3 @@ def get_triplets(ctype, etypes, bond_array, cross_bond_inds, tri_dir[k].append(coord) return exist_species, tris, tri_dir - -@njit -def get_triplets_en(ctype, etypes, bond_array, cross_bond_inds, - cross_bond_dists, triplets): - exist_species = [] - tris = [] - tri_dir = [] - - for m in range(bond_array.shape[0]): - r1 = bond_array[m, 0] - c1 = bond_array[m, 1:] - spc1 = etypes[m] - - for n in range(triplets[m]): - ind1 = cross_bond_inds[m, m+n+1] - r2 = bond_array[ind1, 0] - c2 = bond_array[ind1, 1:] - c12 = np.sum(c1*c2) - r12 = np.sqrt(r1**2 + r2**2 - 2*r1*r2*c12) - - spc2 = etypes[ind1] - -# if spc1 == spc2: -# spcs_list = [[ctype, spc1, spc2], [ctype, spc1, spc2]] -# elif ctype == spc1: # spc1 != spc2 -# spcs_list = [[ctype, spc1, spc2], [spc2, ctype, spc1]] -# elif ctype == spc2: # spc1 != spc2 -# spcs_list = [[spc1, spc2, ctype], [spc2, ctype, spc1]] -# else: # all different -# spcs_list = [[ctype, spc1, spc2], [ctype, spc2, spc1]] - - if spc1 <= spc2: - spcs = [ctype, spc1, spc2] - triplet = array([r1, r2, r12]) - coord = [c1, c2] - else: - spcs = [ctype, spc2, spc1] - triplet = array([r2, r1, r12]) - coord = [c2, c1] - - if spcs not in exist_species: - exist_species.append(spcs) - tris.append([triplet]) - tri_dir.append([coord]) - else: - k = exist_species.index(spcs) - tris[k].append(triplet) - tri_dir[k].append(coord) - - return exist_species, tris, tri_dir - diff --git a/flare/otf.py b/flare/otf.py index 47d702120..a4628bdae 100644 --- a/flare/otf.py +++ b/flare/otf.py @@ -1,12 +1,11 @@ -import sys +import logging import numpy as np import time -import copy -import multiprocessing as mp -import subprocess + +from copy import deepcopy +from datetime import datetime from shutil import copyfile from typing import List, Tuple, Union -from datetime import datetime import flare.predict as predict from flare import struc, gp, env, md @@ -20,28 +19,33 @@ class OTF: molecular dynamics. Args: - dft_input (str): Input file. dt (float): MD timestep. number_of_steps (int): Number of timesteps in the training simulation. + prev_pos_init ([type], optional): Previous positions. Defaults + to None. + rescale_steps (List[int], optional): List of frames for which the + velocities of the atoms are rescaled. Defaults to []. + rescale_temps (List[int], optional): List of rescaled temperatures. + Defaults to []. + gp (gp.GaussianProcess): Initial GP model. - dft_loc (str): Location of DFT executable. + calculate_energy (bool, optional): If True, the energy of each + frame is calculated with the GP. Defaults to False. + write_model (int, optional): If 0, write never. If 1, write at + end of run. If 2, write after each training and end of run. + If 3, write after each time atoms are added and end of run. + std_tolerance_factor (float, optional): Threshold that determines when DFT is called. Specifies a multiple of the current noise hyperparameter. If the epistemic uncertainty on a force component exceeds this value, DFT is called. Defaults to 1. - prev_pos_init ([type], optional): Previous positions. Defaults - to None. - par (bool, optional): If True, force predictions are made in - parallel. Defaults to False. skip (int, optional): Number of frames that are skipped when dumping to the output file. Defaults to 0. init_atoms (List[int], optional): List of atoms from the input structure whose local environments and force components are used to train the initial GP model. If None is specified, all atoms are used to train the initial GP. Defaults to None. - calculate_energy (bool, optional): If True, the energy of each - frame is calculated with the GP. Defaults to False. output_name (str, optional): Name of the output file. Defaults to 'otf_run'. max_atoms_added (int, optional): Number of atoms added each time @@ -50,10 +54,7 @@ class OTF: hyperparameters of the GP are optimized. After this many updates to the GP, the hyperparameters are frozen. Defaults to 10. - rescale_steps (List[int], optional): List of frames for which the - velocities of the atoms are rescaled. Defaults to []. - rescale_temps (List[int], optional): List of rescaled temperatures. - Defaults to []. + force_source (Union[str, object], optional): DFT code used to calculate ab initio forces during training. A custom module can be used here in place of the DFT modules available in the FLARE package. The @@ -62,12 +63,13 @@ class OTF: species, cell, and masses of a structure of atoms; and run_dft_par, which takes a number of DFT related inputs and returns the forces on all atoms. Defaults to "qe". - n_cpus (int, optional): Number of cpus used during training. - Defaults to 1. npool (int, optional): Number of k-point pools for DFT calculations. Defaults to None. mpi (str, optional): Determines how mpi is called. Defaults to "srun". + dft_loc (str): Location of DFT executable. + dft_input (str): Input file. + dft_output (str): Output file. dft_kwargs ([type], optional): Additional arguments which are passed when DFT is called; keyword arguments vary based on the program (e.g. ESPRESSO vs. VASP). Defaults to None. @@ -80,24 +82,33 @@ class OTF: single file name, or a list of several. Copied files will be prepended with the date and time with the format 'Year.Month.Day:Hour:Minute:Second:'. - write_model (int, optional): If 0, write never. If 1, write at - end of run. If 2, write after each training and end of run. - If 3, write after each time atoms are added and end of run. + + n_cpus (int, optional): Number of cpus used during training. + Defaults to 1. """ - def __init__(self, dft_input: str, dt: float, number_of_steps: int, - gp: gp.GaussianProcess, dft_loc: str, + def __init__(self, + # md args + dt: float, number_of_steps: int, + prev_pos_init: 'ndarray' = None, + rescale_steps: List[int] = [], rescale_temps: List[int] = [], + # flare args + gp: gp.GaussianProcess = None, + calculate_energy: bool = False, + write_model: int = 0, + # otf args std_tolerance_factor: float = 1, - prev_pos_init: 'ndarray' = None, par: bool = False, skip: int = 0, init_atoms: List[int] = None, - calculate_energy: bool = False, output_name: str = 'otf_run', + output_name: str = 'otf_run', max_atoms_added: int = 1, freeze_hyps: int = 10, - rescale_steps: List[int] = [], rescale_temps: List[int] = [], + # dft args force_source: str = "qe", - n_cpus: int = 1, npool: int = None, mpi: str = "srun", - dft_kwargs=None, dft_output='dft.out', - store_dft_output: Tuple[Union[str, List[str]],str] = None, - write_model: int = 0): + npool: int = None, mpi: str = "srun", dft_loc: str = None, + dft_input: str = None, dft_output='dft.out', dft_kwargs=None, + store_dft_output: Tuple[Union[str, List[str]], str] = None, + # par args + n_cpus: int = 1, + ): self.dft_input = dft_input self.dft_output = dft_output @@ -146,22 +157,23 @@ def __init__(self, dft_input: str, dt: float, number_of_steps: int, self.dft_count = 0 # set pred function - if (par and gp.per_atom_par and gp.par) and not calculate_energy: + if (n_cpus > 1 and gp.per_atom_par and gp.parallel) and not calculate_energy: self.pred_func = predict.predict_on_structure_par elif not calculate_energy: self.pred_func = predict.predict_on_structure - elif (par and gp.per_atom_par and gp.par): + elif (n_cpus > 1 and gp.per_atom_par and gp.parallel): self.pred_func = predict.predict_on_structure_par_en else: self.pred_func = predict.predict_on_structure_en - self.par = par # set rescale attributes self.rescale_steps = rescale_steps self.rescale_temps = rescale_temps + # set logger self.output = Output(output_name, always_flush=True) self.output_name = output_name + # set number of cpus and npool for DFT runs self.n_cpus = n_cpus self.npool = npool @@ -180,8 +192,7 @@ def run(self): 'Year.Month.Day:Hour:Minute:Second:'. """ - self.output.write_header(self.gp.cutoffs, self.gp.kernel_name, - self.gp.hyps, self.gp.opt_algorithm, + self.output.write_header(str(self.gp), self.dt, self.number_of_steps, self.structure, self.std_tolerance) @@ -190,29 +201,19 @@ def run(self): while self.curr_step < self.number_of_steps: # run DFT and train initial model if first step and DFT is on - if self.curr_step == 0 and self.std_tolerance != 0 and len(self.gp.training_data)==0: - # call dft and update positions - self.run_dft() - dft_frcs = copy.deepcopy(self.structure.forces) - new_pos = md.update_positions(self.dt, self.noa, - self.structure) + if self.curr_step == 0 and self.std_tolerance != 0 and len(self.gp.training_data) == 0: + + self.initialize_train() + new_pos = self.md_step() self.update_temperature(new_pos) self.record_state() - # make initial gp model and predict forces - self.update_gp(self.init_atoms, dft_frcs) - if (self.dft_count-1) < self.freeze_hyps: - self.train_gp() - if self.write_model >= 2: - self.gp.write_model(self.output_name+"_model") - # after step 1, try predicting with GP model else: - self.gp.check_L_alpha() - self.pred_func(self.structure, self.gp, self.n_cpus) + # compute forces and stds with GP self.dft_step = False - new_pos = md.update_positions(self.dt, self.noa, - self.structure) + self.compute_properties() + new_pos = self.md_step() # get max uncertainty atoms std_in_bound, target_atoms = \ @@ -224,36 +225,22 @@ def run(self): # record GP forces self.update_temperature(new_pos) self.record_state() - gp_frcs = copy.deepcopy(self.structure.forces) + gp_frcs = deepcopy(self.structure.forces) # run DFT and record forces self.dft_step = True self.run_dft() - dft_frcs = copy.deepcopy(self.structure.forces) - new_pos = md.update_positions(self.dt, self.noa, - self.structure) - self.update_temperature(new_pos) + dft_frcs = deepcopy(self.structure.forces) + + # run MD step & record the state self.record_state() # compute mae and write to output - mae = np.mean(np.abs(gp_frcs - dft_frcs)) - mac = np.mean(np.abs(dft_frcs)) - - self.output.write_to_log('\nmean absolute error:' - ' %.4f eV/A \n' % mae) - self.output.write_to_log('mean absolute dft component:' - ' %.4f eV/A \n' % mac) + self.compute_mae(gp_frcs, dft_frcs) # add max uncertainty atoms to training set self.update_gp(target_atoms, dft_frcs) - if (self.dft_count-1) < self.freeze_hyps: - self.train_gp() - if self.write_model == 2: - self.gp.write_model(self.output_name+"_model") - if self.write_model == 3: - self.gp.write_model(self.output_name+'_model') - # write gp forces if counter >= self.skip and not self.dft_step: self.update_temperature(new_pos) @@ -269,6 +256,27 @@ def run(self): if self.write_model >= 1: self.gp.write_model(self.output_name+"_model") + def initialize_train(self): + # call dft and update positions + self.run_dft() + dft_frcs = deepcopy(self.structure.forces) + + # make initial gp model and predict forces + self.update_gp(self.init_atoms, dft_frcs) + + def compute_properties(self): + ''' + In ASE-OTF, it will be replaced by subclass method + ''' + self.gp.check_L_alpha() + self.pred_func(self.structure, self.gp, self.n_cpus) + + def md_step(self): + ''' + In ASE-OTF, it will be replaced by subclass method + ''' + return md.update_positions(self.dt, self.noa, self.structure) + def run_dft(self): """Calculates DFT forces on atoms in the current structure. @@ -278,24 +286,20 @@ def run_dft(self): Calculates DFT forces on atoms in the current structure.""" - self.output.write_to_log('\nCalling DFT...\n') + f = logging.getLogger(self.output.basename+'log') + f.info('\nCalling DFT...\n') # calculate DFT forces - forces = self.dft_module.run_dft_par(self.dft_input, self.structure, - self.dft_loc, - n_cpus=self.n_cpus, - dft_out=self.dft_output, - npool=self.npool, - mpi=self.mpi, - dft_kwargs=self.dft_kwargs) + forces = self.dft_module.run_dft_par( + self.dft_input, self.structure, self.dft_loc, n_cpus=self.n_cpus, + dft_out=self.dft_output, npool=self.npool, mpi=self.mpi, + dft_kwargs=self.dft_kwargs) + self.structure.forces = forces # write wall time of DFT calculation self.dft_count += 1 - self.output.write_to_log('DFT run complete.\n') - time_curr = time.time() - self.start_time - self.output.write_to_log('number of DFT calls: %i \n' % self.dft_count) - self.output.write_to_log('wall time from start: %.2f s \n' % time_curr) + self.output.conclude_dft(self.dft_count, self.start_time) # Store DFT outputs in another folder if desired # specified in self.store_dft_output @@ -321,10 +325,7 @@ def update_gp(self, train_atoms: List[int], dft_frcs: 'ndarray'): will be added to the training set. dft_frcs (np.ndarray): DFT forces on all atoms in the structure. """ - self.output.write_to_log('\nAdding atom {} to the training set.\n' - .format(train_atoms)) - self.output.write_to_log('Uncertainty: {}.\n' - .format(self.structure.stds[train_atoms[0]])) + self.output.add_atom_info(train_atoms, self.structure.stds) # update gp model self.gp.update_db(self.structure, dft_frcs, @@ -332,15 +333,31 @@ def update_gp(self, train_atoms: List[int], dft_frcs: 'ndarray'): self.gp.set_L_alpha() + # write model + if (self.dft_count-1) < self.freeze_hyps: + self.train_gp() + if self.write_model == 2: + self.gp.write_model(self.output_name+"_model") + if self.write_model == 3: + self.gp.write_model(self.output_name+'_model') + def train_gp(self): """Optimizes the hyperparameters of the current GP model.""" - self.gp.train(self.output) + self.gp.train() self.output.write_hyps(self.gp.hyp_labels, self.gp.hyps, self.start_time, self.gp.likelihood, self.gp.likelihood_gradient, hyps_mask=self.gp.hyps_mask) + def compute_mae(self, gp_frcs, dft_frcs): + mae = np.mean(np.abs(gp_frcs - dft_frcs)) + mac = np.mean(np.abs(dft_frcs)) + + f = logging.getLogger(self.output.basename+'log') + f.info(f'mean absolute error: {mae:.4f} eV/A') + f.info(f'mean absolute dft component: {mac:.4f} eV/A') + def update_positions(self, new_pos: 'ndarray'): """Performs a Verlet update of the atomic positions. @@ -372,8 +389,7 @@ def update_temperature(self, new_pos: 'ndarray'): self.velocities = velocities def record_state(self): - self.output.write_md_config(self.dt, self.curr_step, self.structure, - self.temperature, self.KE, - self.local_energies, self.start_time, - self.dft_step, - self.velocities) + self.output.write_md_config( + self.dt, self.curr_step, self.structure, self.temperature, + self.KE, self.local_energies, self.start_time, self.dft_step, + self.velocities) diff --git a/flare/otf_parser.py b/flare/otf_parser.py index 522573271..9e96cb56d 100644 --- a/flare/otf_parser.py +++ b/flare/otf_parser.py @@ -1,7 +1,10 @@ import sys import numpy as np + +from copy import deepcopy from typing import List, Tuple from flare import gp, env, struc, otf +from flare.gp import GaussianProcess class OtfAnalysis: @@ -42,48 +45,38 @@ def __init__(self, filename, calculate_energy=False): self.gp_species_list = gp_species_list self.gp_atom_count = gp_atom_count - def make_gp(self, cell=None, kernel_name=None, algo=None, - call_no=None, cutoffs=None, hyps=None, init_gp=None, - hyp_no=None, par=True, kernel=None): + def make_gp(self, cell=None, call_no=None, hyps=None, init_gp=None, + hyp_no=None, **kwargs,): + + if call_no is None: + call_no = len(self.gp_position_list) + if hyp_no is None: + hyp_no = call_no + if hyps is None: + # check out the last non-empty element from the list + for icall in reversed(range(hyp_no)): + if len(self.gp_hyp_list[icall]) > 0: + hyps = self.gp_hyp_list[icall][-1] + break + if cell is None: + cell = self.header['cell'] + if init_gp is None: # Use run's values as extracted from header # TODO Allow for kernel gradient in header - if cell is None: - cell = self.header['cell'] - if kernel_name is None: - kernel_name = self.header['kernel_name'] - if algo is None: - algo = self.header['algo'] - if cutoffs is None: - cutoffs = self.header['cutoffs'] - if call_no is None: - call_no = len(self.gp_position_list) - if hyp_no is None: - hyp_no = call_no - if hyps is None: - # check out the last non-empty element from the list - for icall in reversed(range(hyp_no)): - if len(self.gp_hyp_list[icall]) > 0: - gp_hyps = self.gp_hyp_list[icall][-1] - break - else: - gp_hyps = hyps - if (kernel is not None) and (kernel_name is None): - DeprecationWarning("kernel replaced with kernel_name") - kernel_name = kernel.__name__ + dictionary = deepcopy(self.header) + dictionary['hyps'] = hyps + for k in kwargs: + if kwargs[k] is not None: + dictionary[k] = kwargs[k] gp_model = \ - gp.GaussianProcess(kernel_name=kernel_name, - hyps=gp_hyps, - cutoffs=cutoffs, opt_algorithm=algo, - par=par) + GaussianProcess.from_dict(dictionary) else: gp_model = init_gp - call_no = len(self.gp_position_list) - gp_hyps = self.gp_hyp_list[hyp_no-1][-1] - gp_model.hyps = gp_hyps + gp_model.hyps = hyps for (positions, forces, atoms, _, species) in \ zip(self.gp_position_list[:call_no], @@ -350,9 +343,12 @@ def parse_header_information(outfile: str = 'otf_run.out') -> dict: if stopreading is None: raise Exception("OTF File is malformed") + cutoffs_dict = {} for i, line in enumerate(lines[:stopreading]): - # TODO Update this in full - if 'cutoffs' in line: + line_lower = line.lower() + + # gp related + if 'cutoffs' in line_lower: line = line.split(':')[1].strip() line = line.strip('[').strip(']') line = line.split() @@ -363,28 +359,43 @@ def parse_header_information(outfile: str = 'otf_run.out') -> dict: except: cutoffs.append(float(val[:-1])) header_info['cutoffs'] = cutoffs - if 'frames' in line: - header_info['frames'] = int(line.split(':')[1]) - if 'kernel_name' in line: + elif 'cutoff' in line_lower: + line = line.split(':') + name = line[0][7:] + value = float(line[1]) + cutoffs_dict[name] = value + header_info['cutoffs'] = cutoffs_dict + + if 'kernel_name' in line_lower: header_info['kernel_name'] = line.split(':')[1].strip() - if 'kernel' in line: + elif 'kernels' in line_lower: + line = line.split(':')[1].strip() + line = line.strip('[').strip(']') + line = line.split() + header_info['kernels'] = line + elif 'kernel' in line_lower: header_info['kernel_name'] = line.split(':')[1].strip() - if 'number of hyperparameters:' in line: + + if 'number of hyperparameters:' in line_lower: header_info['n_hyps'] = int(line.split(':')[1]) - if 'optimization algorithm' in line: + if 'optimization algorithm' in line_lower: header_info['algo'] = line.split(':')[1].strip() - if 'number of atoms' in line: + + # otf related + if 'frames' in line_lower: + header_info['frames'] = int(line.split(':')[1]) + if 'number of atoms' in line_lower: header_info['atoms'] = int(line.split(':')[1]) - if 'timestep' in line: + if 'timestep' in line_lower: header_info['dt'] = float(line.split(':')[1]) - if 'system species' in line: + if 'system species' in line_lower: line = line.split(':')[1] line = line.split("'") species = [item for item in line if item.isalpha()] header_info['species_set'] = set(species) - if 'periodic cell' in line: + if 'periodic cell' in line_lower: vectors = [] for cell_line in lines[i+1:i+4]: cell_line = \ @@ -393,7 +404,7 @@ def parse_header_information(outfile: str = 'otf_run.out') -> dict: vector = [float(vec[0]), float(vec[1]), float(vec[2])] vectors.append(vector) header_info['cell'] = np.array(vectors) - if 'previous positions' in line: + if 'previous positions' in line_lower: struc_spec = [] prev_positions = [] for pos_line in lines[i+1:i+1+header_info.get('atoms', 0)]: diff --git a/flare/output.py b/flare/output.py index 7c9928ce5..159fe1677 100644 --- a/flare/output.py +++ b/flare/output.py @@ -4,12 +4,13 @@ or running an MD simulation updated on-the-fly. """ import datetime -import os -import shutil +import logging import time -import multiprocessing import numpy as np +from logging import FileHandler, StreamHandler, Logger +from os.path import isfile +from shutil import move as movefile from typing import Union from flare.struc import Structure @@ -22,26 +23,30 @@ class Output: class. It is also used in get_neg_like_grad and get_neg_likelihood in gp_algebra to print intermediate results. - It opens and prints files with the basename prefix and different + It opens and print files with the basename prefix and different suffixes corresponding to different kinds of output data. :param basename: Base output file name, suffixes will be added :type basename: str, optional + :param verbose: print level. The same as logging level. It can be + CRITICAL, ERROR, WARNING, INFO, DEBUG, NOTSET + :type verbose: str, optional :param always_flush: Always write to file instantly :type always_flus: bool, optional """ def __init__(self, basename: str = 'otf_run', + verbose: str = 'INFO', always_flush: bool = False): """ Construction. Open files. """ self.basename = f"{basename}" - self.outfiles = {} filesuffix = {'log': '.out', 'hyps': '-hyps.dat'} + self.logger = [] - for filetype in filesuffix.keys(): - self.open_new_log(filetype, filesuffix[filetype]) + for filetype in filesuffix: + self.open_new_log(filetype, filesuffix[filetype], verbose) self.always_flush = always_flush @@ -50,33 +55,27 @@ def conclude_run(self): destruction function that closes all files """ - print('-' * 20, file=self.outfiles['log']) - print('Run complete.', file=self.outfiles['log']) - for (k, v) in self.outfiles.items(): - v.close() - del self.outfiles - self.outfiles = {} + logger = logging.getLogger(self.basename+'log') + logger.info('-' * 20) + logger.info('Run complete.') + logging.shutdown() + self.logger = [] - def open_new_log(self, filetype: str, suffix: str): + def open_new_log(self, filetype: str, suffix: str, verbose='info'): """ Open files. If files with the same name are exist, they are backed up with a suffix "-bak". - :param filetype: the key name in self.outfiles + :param filetype: the key name for logging :param suffix: the suffix of the file to be opened + :param verbose: the verbose level for the logger """ - filename = self.basename + suffix - - # if the file exists, back up - if os.path.isfile(filename): - shutil.copy(filename, filename + "-bak") - - if filetype in self.outfiles.keys(): - if self.outfiles[filetype].closed: - self.outfiles[filetype] = open(filename, "w+") - else: - self.outfiles[filetype] = open(filename, "w+") + if filetype not in self.logger: + set_logger(self.basename+filetype, stream=False, + fileout_name=self.basename+suffix, + verbose=verbose) + self.logger += [filetype] def write_to_log(self, logstring: str, name: str = "log", flush: bool = False): @@ -84,27 +83,27 @@ def write_to_log(self, logstring: str, name: str = "log", Write any string to logfile :param logstring: the string to write - :param name: the key name of the file to print + :param name: the key name of the file to logger named 'log' :param flush: whether it should be flushed """ - self.outfiles[name].write(logstring) + logger = logging.getLogger(self.basename+name) + logger.info(logstring) if flush or self.always_flush: - self.outfiles[name].flush() + logger.handlers[0].flush() - def write_header(self, cutoffs, kernel_name: str, - hyps, algo: str, dt: float = None, - Nsteps: int = None, structure: Structure= None, + def write_header(self, gp_str: str, + dt: float = None, + Nsteps: int = None, structure: Structure = None, std_tolerance: Union[float, int] = None, optional: dict = None): """ + TO DO: this should be replace by the string method of GP and OTF, GPFA + Write header to the log function. Designed for Trajectory Trainer and OTF runs and can take flexible input for both. - :param cutoffs: GP cutoffs - :param kernel_name: Kernel names - :param hyps: list of hyper-parameters - :param algo: algorithm for hyper parameter optimization + :param gp_str: string representation of the GP :param dt: timestep for OTF MD :param Nsteps: total number of steps for OTF MD :param structure: initial structure @@ -112,8 +111,8 @@ def write_header(self, cutoffs, kernel_name: str, :param optional: a dictionary of all the other parameters """ - f = self.outfiles['log'] - f.write(f'{datetime.datetime.now()} \n') + f = logging.getLogger(self.basename+'log') + f.info(f'{datetime.datetime.now()}') if isinstance(std_tolerance, tuple): std_string = 'relative uncertainty tolerance: ' \ @@ -126,18 +125,13 @@ def write_header(self, cutoffs, kernel_name: str, elif std_tolerance > 0: std_string = \ f'uncertainty tolerance: {np.abs(std_tolerance)} ' \ - 'times noise hyperparameter \n' + 'times noise hyperparameter \n' else: std_string = '' - headerstring = '' - headerstring += \ - f'number of cpu cores: {multiprocessing.cpu_count()}\n' - headerstring += f'cutoffs: {cutoffs}\n' - headerstring += f'kernel_name: {kernel_name}\n' - headerstring += f'number of hyperparameters: {len(hyps)}\n' - headerstring += f'hyperparameters: {str(hyps)}\n' - headerstring += f'hyperparameter optimization algorithm: {algo}\n' + headerstring = '\n' + headerstring += gp_str + headerstring += '\n' headerstring += std_string if dt is not None: headerstring += f'timestep (ps): {dt}\n' @@ -162,10 +156,10 @@ def write_header(self, cutoffs, kernel_name: str, headerstring += '\n' headerstring += '-' * 80 + '\n' - f.write(headerstring) + f.info(headerstring) if self.always_flush: - f.flush() + f.handlers[0].flush() def write_md_config(self, dt, curr_step, structure, temperature, KE, local_energies, @@ -240,13 +234,12 @@ def write_md_config(self, dt, curr_step, structure, f'potential energy: {pot_en:.6f} eV \n' string += f'total energy: {tot_en:.6f} eV \n' - string += 'wall time from start: ' - string += f'{(time.time() - start_time):.2f} s \n' - - self.outfiles['log'].write(string) + logger = logging.getLogger(self.basename+'log') + logger.info(string) + self.write_wall_time(start_time) if self.always_flush: - self.outfiles['log'].flush() + logger.handlers[0].flush() def write_xyz(self, curr_step: int, pos: np.array, species: list, filename: str, @@ -258,8 +251,8 @@ def write_xyz(self, curr_step: int, pos: np.array, species: list, :param curr_step: Int, number of frames to note in the comment line :param pos: nx3 matrix of forces, positions, or nything :param species: n element list of symbols - :param filename: file to print - :param header: header printed in comments + :param filename: key of logger + :param header: header in comments :param forces: list of forces on atoms predicted by GP :param stds: uncertainties predicted by GP :param forces_2: true forces from ab initio source @@ -289,10 +282,11 @@ def write_xyz(self, curr_step: int, pos: np.array, species: list, else: string += '\n' - self.outfiles[filename].write(string) + logger = logging.getLogger(self.basename+filename) + logger.info(string) if self.always_flush: - self.outfiles[filename].flush() + logger.handlers[0].flush() def write_xyz_config(self, curr_step, structure, dft_step, forces: np.array = None, stds: np.array = None, @@ -302,8 +296,8 @@ def write_xyz_config(self, curr_step, structure, dft_step, :param curr_step: Int, number of frames to note in the comment line :param structure: Structure, contain positions and forces :param dft_step: Boolean, whether this is a DFT call. - :param forces: Optional list of forces to print in xyz file - :param stds: Optional list of uncertanties to print in xyz file + :param forces: Optional list of forces to xyz file + :param stds: Optional list of uncertanties to xyz file :param forces_2: Optional second list of forces (e.g. DFT forces) :return: @@ -332,35 +326,46 @@ def write_hyps(self, hyp_labels, hyps, start_time, like, like_grad, :return: """ - f = self.outfiles[name] - f.write('\nGP hyperparameters: \n') + f = logging.getLogger(self.basename+name) - if hyps_mask is not None: - if 'map' in hyps_mask: - hyps = hyps_mask['original'] - if len(hyp_labels)!=len(hyps): - hyp_labels = None + f.info('\nGP hyperparameters: ') if hyp_labels is not None: for i, label in enumerate(hyp_labels): - f.write(f'Hyp{i} : {label} = {hyps[i]:.4f}\n') + f.info(f'Hyp{i} : {label} = {hyps[i]:.4f}') else: for i, hyp in enumerate(hyps): - f.write(f'Hyp{i} : {hyp:.4f}\n') + f.info(f'Hyp{i} : {hyp:.4f}') + + f.info(f'likelihood: {like:.4f}') + f.info(f'likelihood gradient: {like_grad}') - f.write(f'likelihood: {like:.4f}\n') - f.write(f'likelihood gradient: {like_grad}\n') if start_time: - time_curr = time.time() - start_time - f.write(f'wall time from start: {time_curr:.2f} s \n') + self.write_wall_time(start_time) if self.always_flush: - f.flush() + f.handlers[0].flush() + + def write_wall_time(self, start_time): + time_curr = time.time() - start_time + f = logging.getLogger(self.basename+'log') + f.info(f'wall time from start: {time_curr:.2f} s') + + def conclude_dft(self, dft_count, start_time): + f = logging.getLogger(self.basename+'log') + f.info('DFT run complete.') + f.info(f'number of DFT calls: {dft_count}') + self.write_wall_time(start_time) + + def add_atom_info(self, train_atoms, stds): + f = logging.getLogger(self.basename+'log') + f.info(f'Adding atom {train_atoms} to the training set.') + f.info(f'Uncertainty: {stds[train_atoms[0]]}') def write_gp_dft_comparison(self, curr_step, frame, start_time, dft_forces, error, local_energies=None, KE=None, - mgp= False): + mgp=False): """ write the comparison to logfile :param curr_step: current timestep @@ -446,12 +451,88 @@ def write_gp_dft_comparison(self, curr_step, frame, string += f'total energy: {tot_en:10.6} eV \n' stat += f' {pot_en:10.6} {tot_en:10.6}' - dt = time.time() - start_time - string += f'wall time from start: {dt:10.2}\n' - stat += f' {dt}\n' + f = logging.getLogger(self.basename+'log') + f.info(string) + self.write_wall_time(start_time) - self.outfiles['log'].write(string) - # self.outfiles['stat'].write(stat) + # stat += f' {dt}\n' + # logging.getLogger('stat').write(stat) if self.always_flush: - self.outfiles['log'].flush() + f.handlers[0].flush() + + +def add_stream(logger: Logger, verbose: str = "info"): + ''' + set up screen sctream handler to the logger with handlers + + :param logger: the logger + :param verbose: verbose level + :type verbose: str + ''' + + stream_defined = False + for handler in logger.handlers: + if isinstance(handler, StreamHandler): + stream_defined = True + + if not stream_defined: + ch = StreamHandler() + ch.setLevel(logging.DEBUG) + # formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') + # ch.setFormatter(formatter) + logger.addHandler(ch) + + +def add_file(logger: Logger, filename: str, verbose: str = "info"): + ''' + set up file handler to the logger with handlers + + :param logger: the logger + :param filename: name of the logfile + :type filename: str + :param verbose: verbose level + :type verbose: str + ''' + + file_defined = False + for handler in logger.handlers: + if isinstance(handler, FileHandler): + file_defined = True + + if not file_defined: + + # back up + if isfile(filename): + movefile(filename, filename+"-bak") + + fh = FileHandler(filename) + verbose = getattr(logging, verbose.upper()) + logger.setLevel(verbose) + fh.setLevel(logging.DEBUG) + logger.addHandler(fh) + + +def set_logger(name: str, stream: bool, fileout_name: str = None, + verbose: str = "info"): + ''' + set up a logger with handlers + + :param name: unique name of the logger in logging module + :type name: str + :param stream: if True, set up a screen output + :type stream: bool + :param fileout_name: name for log file + :type fileout_name: str + :param verbose: verbose level + :type verbose: str + ''' + logger = logging.getLogger(name) + logger.propagate = False + logger.handlers = [] + logger.setLevel(getattr(logging, verbose.upper())) + if stream: + add_stream(logger, verbose) + if fileout_name is not None: + add_file(logger, fileout_name, verbose) + return logger diff --git a/flare/parameters.py b/flare/parameters.py new file mode 100644 index 000000000..63b39350a --- /dev/null +++ b/flare/parameters.py @@ -0,0 +1,481 @@ +import inspect +import json +import logging +import math +import numpy as np +import pickle +import time + +from copy import deepcopy +from itertools import combinations_with_replacement, permutations +from numpy.random import random +from numpy import array as nparray +from numpy import max as npmax +from typing import List, Callable, Union +from warnings import warn +from sys import stdout +from os import devnull + +from flare.output import set_logger +from flare.utils.element_coder import element_to_Z, Z_to_element + +class Parameters(): + ''' + ''' + + all_kernel_types = ['twobody', 'threebody', 'manybody'] + ndim = {'twobody': 2, 'threebody': 3, 'manybody': 2, 'cut3b': 2} + n_kernel_parameters = {'twobody': 2, 'threebody': 2, 'manybody': 2, 'cut3b': 0} + + logger = set_logger("Parameters", stream=True, + fileout_name=None, verbose="info") + + def __init__(self): + ''' + Enumerate all keys and their default values that hyps_mask should store + ''' + + self.param_dict = {'nspecie': 1, + 'ntwobody': 0, + 'nthreebody': 0, + 'ncut3b': 0, + 'nmanybody': 0, + 'specie_mask': None, + 'twobody_mask': None, + 'threebody_mask': None, + 'cut3b_mask': None, + 'manybody_mask': None, + 'twobody_cutoff_list': None, + 'threebody_cutoff_list': None, + 'manybody_cutoff_list': None, + 'train_noise': True, + 'energy_noise': 0, + 'map': None, + 'original_hyps': [], + 'original_labels': [] + } + self.hyps = None + self.hyp_labels = None + self.cutoffs = {} + self.kernels = [] + + @staticmethod + def cutoff_array_to_dict(cutoffs): + ''' + Convert old cutoffs array to the new dictionary format + ''' + + if isinstance(cutoffs, dict): + return cutoffs + + if (cutoffs is not None) and not isinstance(cutoffs, dict): + DeprecationWarning("cutoffs is replace by dictionary") + newcutoffs = {'twobody': cutoffs[0]} + if len(cutoffs) > 1: + newcutoffs['threebody'] = cutoffs[1] + if len(cutoffs) > 2: + newcutoffs['manybody'] = cutoffs[2] + Parameters.logger.debug("Convert cutoffs array to cutoffs dict") + Parameters.logger.debug("Original", cutoffs) + Parameters.logger.debug("Now", newcutoffs) + return newcutoffs + else: + raise TypeError("cannot handle cutoffs with {type(cutoffs)} type") + + @staticmethod + def backward(kernels, param_dict): + + if param_dict is None: + param_dict = {} + + # update old keys to new keys. for example nspec to nspecies + replace_list = {'spec': 'specie', 'bond': 'twobody', + 'triplet': 'threebody', 'mb': 'manybody'} + keys = list(param_dict.keys()) + for key in keys: + for original in replace_list: + if original in key and replace_list[original] not in key: + newkey = key.replace(original, replace_list[original]) + param_dict[newkey] = param_dict[key] + DeprecationWarning( + "{key} is being replaced with {newkey}") + + # add a couple new keys that was not there previously + if 'train_noise' not in param_dict: + param_dict['train_noise'] = True + DeprecationWarning("train_noise has to be in hyps_mask, set to True") + if 'nspecie' not in param_dict: + param_dict['nspecie'] = 1 + + # sort the kernels dictionary again. but this can result in + # wrong results... + if set(kernels) != set(param_dict.get("kernels", [])): + + start = 0 + for k in Parameters.all_kernel_types: + if k in kernels: + if k+'_start' not in param_dict: + param_dict[k+'_start'] = start + if 'n'+k not in param_dict: + Parameters.logger.debug("add in hyper parameter separators"\ + "for", k) + param_dict['n'+k] = 1 + start += Parameters.n_kernel_parameters[k] + else: + start += param_dict['n'+k] * Parameters.n_kernel_parameters[k] + else: + Warning("inconsistency between input kernel and kernel list"\ + "stored in hyps_mask") + + Parameters.logger.debug("Replace kernel array in param_dict") + param_dict['kernels'] = deepcopy(kernels) + + return param_dict + + @staticmethod + def check_instantiation(hyps, cutoffs, kernels, param_dict): + """ + Runs a series of checks to ensure that the user has not supplied + contradictory arguments which will result in undefined behavior + with multiple hyperparameters. + + :return: + """ + + assert isinstance(param_dict, dict) + assert isinstance(cutoffs, dict) + assert isinstance(kernels, (list)) + + param_dict['cutoffs'] = cutoffs + + # double check nspecie is there + nspecie = param_dict['nspecie'] + if nspecie > 1: + assert 'specie_mask' in param_dict, "specie_mask key " \ + "missing " \ + "in param_dict dictionary" + param_dict['specie_mask'] = nparray( + param_dict['specie_mask'], dtype=np.int) + + # for each kernel, check whether it is defined + # and the length of corresponding hyper-parameters + hyps_length = 0 + used_parameters = np.zeros_like(hyps, dtype=bool) + for kernel in kernels+['cut3b']: + + n = param_dict.get(f'n{kernel}', 0) + assert isinstance(n, int) + + if kernel != 'cut3b': + hyps_length += Parameters.n_kernel_parameters[kernel]*n + assert n > 0, f"{kernel} has n 0" + + # check all corresponding keys exist + assert kernel in cutoffs + assert kernel+"_start" in param_dict + + # check the partition of hyperparameters are not used + start = param_dict[kernel+"_start"] + length = Parameters.n_kernel_parameters[kernel]*n + assert not used_parameters[start:start+length].any() + used_parameters[start:start+length] = True + + if n > 1: + + assert f'{kernel}_mask' in param_dict, f"{kernel}_mask key " \ + "missing " \ + "in param_dict dictionary" + + # check mask has the right dimension and values + mask = param_dict[f'{kernel}_mask'] + param_dict[f'{kernel}_mask'] = nparray(mask, dtype=np.int) + + assert (npmax(mask) < n) + dim = Parameters.ndim[kernel] + assert len(mask) == nspecie ** dim, \ + f"wrong dimension of {kernel}_mask: " \ + f" {len(mask)} != nspec ^ {dim} {nspecie**dim}" + + # check whether the mask array is symmetrical + # enumerate all possible combinations + all_comb = list(combinations_with_replacement( + np.arange(nspecie), dim)) + for comb in all_comb: + mask_value = None + perm = list(permutations(comb)) + for ele_list in perm: + mask_id = 0 + for ele in ele_list: + mask_id += ele + mask_id *= nspecie + mask_id = mask_id // nspecie + if mask_value == None: + mask_value = mask[mask_id] + else: + assert mask[mask_id] == mask_value, \ + f'{kernel}_mask has to be symmetrical' + + if kernel != 'cut3b': + if kernel+'_cutoff_list' in param_dict: + cutoff_list = param_dict[kernel+'_cutoff_list'] + assert len(cutoff_list) == n, \ + f'number of cutoffs should be the same as n {n}' + assert npmax(cutoff_list) <= cutoffs[kernel] + else: + assert f'{kernel}_mask' not in param_dict + assert f'{kernel}_cutof_list' not in param_dict + + if 'map' in param_dict: + + assert ('original_hyps' in param_dict), \ + "original hyper parameters have to be defined" + + # Ensure typed correctly as numpy array + param_dict['original_hyps'] = nparray( + param_dict['original_hyps'], dtype=np.float) + if (len(param_dict['original_hyps']) - 1) not in param_dict['map']: + assert param_dict['train_noise'] is False, \ + "train_noise should be False when noise is not in hyps" + + assert len(param_dict['map']) == len(hyps), \ + "the hyperparmeter length is inconsistent with the mask" + assert npmax(param_dict['map']) < len(param_dict['original_hyps']) + + else: + assert param_dict['train_noise'] is True, \ + "train_noise should be True when map is not used" + + hyps = Parameters.get_hyps(param_dict, hyps) + + hyps_length += 1 + assert hyps_length == len(hyps), \ + "the hyperparmeter length is inconsistent with the mask" + + return param_dict + + @staticmethod + def get_component_hyps(param_dict, kernel_name, hyps=None, constraint=False, noise=False): + ''' + return the hyper-parameters correspond to the kernel specified by kernel_name + + Args: + + param_dict (dict): the hyps_mask dictionary used/stored in GaussianProcess + kernel_name (str): the name of the kernel. + hyps (np.array): if hyps is None, use the one stored in param_dict + constraint (bool): if True, return one additional list that shows whether the + hyper-parmaeters can be trained + noise (bool): if True, the last element of returned hyper-parameters is + the noise variance. + + return: hyper-parameters, and whether they can be optimized + ''' + + if kernel_name not in param_dict['kernels']: + if constraint: + if noise: + return [None, None, None], [None, None] + else: + return [None, None], [None, None] + else: + if noise: + return [None, None, None] + else: + return [None, None] + + hyps, opt = Parameters.get_hyps(param_dict, hyps=hyps, constraint=True) + s = param_dict[kernel_name+'_start'] + n = param_dict[f'n{kernel_name}'] + + newhyps = [hyps[s:s+n], hyps[s+n:s+2*n]] + newopt = [opt[s:s+n], opt[s+n:s+2*n]] + + if noise: + newhyps += [hyps[-1]] + + if constraint: + return newhyps, newopt + else: + return newhyps + + @staticmethod + def get_component_mask(param_dict, kernel_name, hyps=None): + ''' + return the hyper-parameter masking correspond to the kernel specified by kernel_name + + Args: + + param_dict (dict): the hyps_mask dictionary used/stored in GaussianProcess + kernel_name (str): the name of the kernel. + hyps (np.array): if hyps is None, use the one stored in param_dict + + return: hyper-parameters, cutoffs, and new hyps_mask + ''' + + if kernel_name in param_dict['kernels']: + new_dict = {} + new_dict['kernels'] = [kernel_name] + + new_dict[kernel_name+'_start'] = 0 + + name_list = ['nspecie', 'specie_mask', + 'n'+kernel_name, kernel_name+'_mask', + kernel_name+'_cutoff_list'] + + if kernel_name == 'threebody': + name_list += ['ncut3b', 'cut3b_mask'] + + for name in name_list: + if name in param_dict: + new_dict[name] = deepcopy(param_dict[name]) + + hyps = np.hstack(Parameters.get_component_hyps( + param_dict, kernel_name, hyps=hyps, noise=True)) + + cutoffs = {} + if 'twobody' in param_dict['cutoffs']: + cutoffs['twobody'] = param_dict['cutoffs']['twobody'] + cutoffs[kernel_name] = param_dict['cutoffs'][kernel_name] + + return hyps, cutoffs, new_dict + else: + return [], {}, {} + + @staticmethod + def get_noise(param_dict, hyps=None, constraint=False): + ''' + get the noise parameters + + Args: + + constraint (bool): if True, return one additional list that shows whether the + hyper-parmaeters can be trained + noise (bool): if True, the last element of returned hyper-parameters is + the noise variance. + ''' + hyps = Parameters.get_hyps(param_dict, hyps=hyps) + if constraint: + return hyps[-1], param_dict['train_noise'] + else: + return hyps[-1] + + @staticmethod + def get_cutoff(kernel_name, coded_species, param_dict): + ''' + get the cutoff + + Args: + + kernel_name (str): name of the kernel + coded_species (list): list of element names + param_dict (dict): hyps_mask + + ''' + + cutoffs = param_dict['cutoffs'] + universal_cutoff = cutoffs[kernel_name] + + if f'{kernel_name}_cutoff_list' in param_dict: + + specie_mask = param_dict['species_mask'] + cutoff_list = param_dict[f'{kernel_name}_cutoff_list'] + + if kernel_name != 'threebody': + mask_id = 0 + for ele in coded_specie: + mask_id += specie_mask[ele] + mask_id *= nspecie + mask_id = mask_id // nspecie + mask_id = param_dict[kernel_name+'_mask'][mask_id] + return cutoff_list[mask_id] + else: + cut3b_mask = param_dict['cut3b_mask'] + ele1 = species_mask[coded_species[0]] + ele2 = species_mask[coded_species[1]] + ele3 = species_mask[coded_species[2]] + twobody1 = cut3b_mask[param_dict['nspecie']*ele1 + ele2] + twobody2 = cut3b_mask[param_dict['nspecie']*ele1 + ele3] + twobody12 = cut3b_mask[param_dict['nspecie']*ele2 + ele3] + return np.array([cutoff_list[twobody1], + cutoff_list[twobody2], + cutoff_list[twobody12]]) + else: + if kernel_name != 'threebody': + return [universal_cutoff] + else: + return [universal_cutoff]*3 + + @staticmethod + def get_hyps(param_dict, hyps=None, constraint=False): + ''' + get the cutoff + + Args: + + kernel_name (str): name of the kernel + coded_species (list): list of element names + param_dict (dict): hyps_mask + + ''' + + if hyps is None: + hyps = param_dict['hyps'] + + if 'map' in param_dict: + newhyps = np.copy(param_dict['original_hyps']) + opt = np.zeros_like(newhyps, dtype=bool) + for i, ori in enumerate(param_dict['map']): + newhyps[ori] = hyps[i] + opt[ori] = True + else: + newhyps = np.copy(hyps) + opt = np.zeros_like(hyps, dtype=bool) + + if constraint: + return newhyps, opt + else: + return newhyps + + @staticmethod + def compare_dict(dict1, dict2): + ''' + compare whether two hyps_masks are the same + ''' + + if type(dict1) != type(dict2): + return False + + if dict1 is None: + return True + + list_of_names = ['nspecie', 'specie_mask', 'map', 'original_hyps'] + for k in Parameters.all_kernel_types: + list_of_names += ['n'+k] + list_of_names += [k+'_mask'] + list_of_names += ['cutoff_'+k] + list_of_names += [k+'_cutoff_list'] + list_of_names += ['ncut3b'] + list_of_names += ['cut3b_mask'] + + for k in list_of_names: + if (k in dict1) != (k in dict2): + return False + elif k in dict1: + if not (np.isclose(dict1[k], dict2[k]).all()): + return False + + for k in ['hyp_labels', 'original_labels']: + if (k in dict1) != (k in dict2): + return False + elif k in dict1: + if not (dict1[k] == dict2[k]): + return False + + for k in ['train_noise']: + if (k in dict1) != (k in dict2): + return False + elif k in dict1: + if dict1[k] != dict2[k]: + return False + + return True diff --git a/flare/predict.py b/flare/predict.py index f9d48ddc6..8f9711f38 100644 --- a/flare/predict.py +++ b/flare/predict.py @@ -267,7 +267,7 @@ def predict_on_structure_en(structure: Structure, gp: GaussianProcess, forces[n][i] = float(force) stds[n][i] = np.sqrt(np.abs(var)) - if write_to_structure: + if write_to_structure and structure.forces is not None: structure.forces[n][i] = float(force) structure.stds[n][i] = np.sqrt(np.abs(var)) @@ -314,20 +314,13 @@ def predict_on_structure_par_en(structure: Structure, gp: GaussianProcess, else: selective_atoms = [] - # Work in serial if the number of cpus is 1 - if n_cpus is 1: - return predict_on_structure_en(structure, gp, - write_to_structure=write_to_structure, - selective_atoms=selective_atoms, - skipped_atom_value=skipped_atom_value) - if n_cpus is None: pool = mp.Pool(processes=mp.cpu_count()) else: pool = mp.Pool(processes=n_cpus) - results = [] # Parallelize over atoms in structure + results = [] for atom_i in range(structure.nat): if atom_i not in selective_atoms and selective_atoms: @@ -357,9 +350,10 @@ def predict_on_structure_par_en(structure: Structure, gp: GaussianProcess, return forces, stds, local_energies -def predict_on_atom_mgp(atom: int, structure, cutoffs, mgp, +def predict_on_atom_mgp(atom: int, structure, mgp, write_to_structure=False): - chemenv = AtomicEnvironment(structure, atom, cutoffs) + chemenv = AtomicEnvironment(structure, atom, mgp.cutoffs, + cutoffs_mask=mgp.hyps_mask) # predict force components and standard deviations force, var, virial, local_energy = mgp.predict(chemenv) comps = force @@ -398,7 +392,7 @@ def predict_on_structure_mgp(structure, mgp, output=None, continue forces[n, :], stds[n, :], _ = \ - predict_on_atom_mgp(n, structure, mgp.cutoffs, mgp, + predict_on_atom_mgp(n, structure, mgp, write_to_structure) return forces, stds diff --git a/flare/struc.py b/flare/struc.py index e3a2d8857..2abcfa010 100644 --- a/flare/struc.py +++ b/flare/struc.py @@ -265,7 +265,7 @@ def from_dict(dictionary: dict) -> 'flare.struc.Structure': return struc @staticmethod - def from_ase_atoms(atoms: 'ase.Atoms') -> 'flare.struc.Structure': + def from_ase_atoms(atoms: 'ase.Atoms', cell=None) -> 'flare.struc.Structure': """ From an ASE Atoms object, return a FLARE structure @@ -273,7 +273,10 @@ def from_ase_atoms(atoms: 'ase.Atoms') -> 'flare.struc.Structure': :type atoms: ASE Atoms object :return: A FLARE structure from an ASE atoms object """ - struc = Structure(cell=np.array(atoms.cell), + + if cell is None: + cell = np.array(atoms.cell) + struc = Structure(cell=cell, positions=atoms.positions, species=atoms.get_chemical_symbols()) return struc diff --git a/flare/utils/element_coder.py b/flare/utils/element_coder.py index 1847628cf..0fc220d45 100644 --- a/flare/utils/element_coder.py +++ b/flare/utils/element_coder.py @@ -208,3 +208,114 @@ def Z_to_element(Z: int) -> str: raise ValueError("Input Z is not a number. It should be an " "integer") return _Z_to_element[Z] + +_Z_to_mass = {1:1.0079, + 2:4.0026, + 3:6.941, + 4:9.0122, + 5:10.811, + 6:12.0107, + 7:14.0067, + 8:15.9994, + 9:18.9984, + 10:20.1797, + 11:22.9897, + 12:24.305, + 13:26.9815, + 14:28.0855, + 15:30.9738, + 16:32.065, + 17:35.453, + 19:39.0983, + 18:39.948, + 20:40.078, + 21:44.9559, + 22:47.867, + 23:50.9415, + 24:51.9961, + 25:54.938, + 26:55.845, + 28:58.6934, + 27:58.9332, + 29:63.546, + 30:65.39, + 31:69.723, + 32:72.64, + 33:74.9216, + 34:78.96, + 35:79.904, + 36:83.8, + 37:85.4678, + 38:87.62, + 39:88.9059, + 40:91.224, + 41:92.9064, + 42:95.94, + 43:98, + 44:101.07, + 45:102.9055, + 46:106.42, + 47:107.8682, + 48:112.411, + 49:114.818, + 50:118.71, + 51:121.76, + 53:126.9045, + 52:127.6, + 54:131.293, + 55:132.9055, + 56:137.327, + 57:138.9055, + 58:140.116, + 59:140.9077, + 60:144.24, + 61:145, + 62:150.36, + 63:151.964, + 64:157.25, + 65:158.9253, + 66:162.5, + 67:164.9303, + 68:167.259, + 69:168.9342, + 70:173.04, + 71:174.967, + 72:178.49, + 73:180.9479, + 74:183.84, + 75:186.207, + 76:190.23, + 77:192.217, + 78:195.078, + 79:196.9665, + 80:200.59, + 81:204.3833, + 82:207.2, + 83:208.9804, + 84:209, + 85:210, + 86:222, + 87:223, + 88:226, + 89:227, + 91:231.0359, + 90:232.0381, + 93:237, + 92:238.0289, + 95:243, + 94:244, + 96:247, + 97:247, + 98:251, + 99:252, + 100:257, + 101:258, + 102:259, + 104:261, + 103:262, + 105:262, + 107:264, + 106:266, + 109:268, + 111:272, + 108:277} diff --git a/flare/utils/env_getarray.py b/flare/utils/env_getarray.py new file mode 100644 index 000000000..0ba8383da --- /dev/null +++ b/flare/utils/env_getarray.py @@ -0,0 +1,479 @@ +from math import sqrt +from numba import njit +import numpy as np +import flare.kernels.cutoffs as cf +from flare.kernels.kernels import coordination_number, q_value_mc + +@njit +def get_2_body_arrays(positions, atom: int, cell, r_cut, cutoff_2, species, sweep, + nspecie, specie_mask, twobody_mask): + """Returns distances, coordinates, species of atoms, and indices of neighbors + in the 2-body local environment. This method is implemented outside + the AtomicEnvironment class to allow for njit acceleration with Numba. + + :param positions: Positions of atoms in the structure. + :type positions: np.ndarray + :param atom: Index of the central atom of the local environment. + :type atom: int + :param cell: 3x3 array whose rows are the Bravais lattice vectors of the + cell. + :type cell: np.ndarray + :param cutoff_2: 2-body cutoff radius. + :type cutoff_2: np.ndarray + :param species: Numpy array of species represented by their atomic numbers. + :type species: np.ndarray + :param nspecie: number of atom types to define bonds + :type: int + :param specie_mask: mapping from atomic number to atom types + :type: np.ndarray + :param twobody_mask: mapping from the types of end atoms to bond types + :type: np.ndarray + :return: Tuple of arrays describing pairs of atoms in the 2-body local + environment. + + bond_array_2: Array containing the distances and relative + coordinates of atoms in the 2-body local environment. First column + contains distances, remaining columns contain Cartesian coordinates + divided by the distance (with the origin defined as the position of the + central atom). The rows are sorted by distance from the central atom. + + bond_positions_2: Coordinates of atoms in the 2-body local environment. + + etypes: Species of atoms in the 2-body local environment represented by + their atomic number. + + bond_indices: Structure indices of atoms in the local environment. + + :rtype: np.ndarray, np.ndarray, np.ndarray, np.ndarray + """ + noa = len(positions) + pos_atom = positions[atom] + super_count = sweep.shape[0]**3 + coords = np.zeros((noa, 3, super_count), dtype=np.float64) + dists = np.zeros((noa, super_count), dtype=np.float64) + cutoff_count = 0 + + vec1 = cell[0] + vec2 = cell[1] + vec3 = cell[2] + + sepcut = False + bcn = 0 + if nspecie > 1 and cutoff_2 is not None: + sepcut = True + bc = specie_mask[species[atom]] + bcn = nspecie * bc + + # record distances and positions of images + for n in range(noa): + diff_curr = positions[n] - pos_atom + im_count = 0 + if sepcut and (specie_mask is not None) and (cutoff_2 is not None): + bn = specie_mask[species[n]] + r_cut = cutoff_2[twobody_mask[bn+bcn]] + + for s1 in sweep: + for s2 in sweep: + for s3 in sweep: + im = diff_curr + s1 * vec1 + s2 * vec2 + s3 * vec3 + dist = sqrt(im[0] * im[0] + im[1] * im[1] + im[2] * im[2]) + if (dist < r_cut) and (dist != 0): + dists[n, im_count] = dist + coords[n, :, im_count] = im + cutoff_count += 1 + im_count += 1 + + # create 2-body bond array + bond_indices = np.zeros(cutoff_count, dtype=np.int8) + bond_array_2 = np.zeros((cutoff_count, 4), dtype=np.float64) + bond_positions_2 = np.zeros((cutoff_count, 3), dtype=np.float64) + etypes = np.zeros(cutoff_count, dtype=np.int8) + bond_count = 0 + + for m in range(noa): + spec_curr = species[m] + if sepcut and (specie_mask is not None) and (cutoff_2 is not None): + bm = specie_mask[species[m]] + r_cut = cutoff_2[twobody_mask[bm+bcn]] + for im_count in range(super_count): + dist_curr = dists[m, im_count] + if (dist_curr < r_cut) and (dist_curr != 0): + coord = coords[m, :, im_count] + bond_array_2[bond_count, 0] = dist_curr + bond_array_2[bond_count, 1:4] = coord / dist_curr + bond_positions_2[bond_count, :] = coord + etypes[bond_count] = spec_curr + bond_indices[bond_count] = m + bond_count += 1 + + # sort by distance + sort_inds = bond_array_2[:, 0].argsort() + bond_array_2 = bond_array_2[sort_inds] + bond_positions_2 = bond_positions_2[sort_inds] + bond_indices = bond_indices[sort_inds] + etypes = etypes[sort_inds] + + return bond_array_2, bond_positions_2, etypes, bond_indices + + +@njit +def get_3_body_arrays(bond_array_2, bond_positions_2, ctype, + etypes, r_cut, cutoff_3, + nspecie, specie_mask, cut3b_mask): + """Returns distances and coordinates of triplets of atoms in the + 3-body local environment. + + :param bond_array_2: 2-body bond array. + :type bond_array_2: np.ndarray + :param bond_positions_2: Coordinates of atoms in the 2-body local + environment. + :type bond_positions_2: np.ndarray + :param ctype: atomic number of the center atom + :type: int + :param cutoff_3: 3-body cutoff radius. + :type cutoff_3: np.ndarray + :param nspecie: number of atom types to define bonds + :type: int + :param specie_mask: mapping from atomic number to atom types + :type: np.ndarray + :param cut3b_mask: mapping from the types of end atoms to bond types + :type: np.ndarray + :return: Tuple of 4 arrays describing triplets of atoms in the 3-body local + environment. + + bond_array_3: Array containing the distances and relative + coordinates of atoms in the 3-body local environment. First column + contains distances, remaining columns contain Cartesian coordinates + divided by the distance (with the origin defined as the position of the + central atom). The rows are sorted by distance from the central atom. + + cross_bond_inds: Two dimensional array whose row m contains the indices + of atoms n > m that are within a distance cutoff_3 of both atom n and the + central atom. + + cross_bond_dists: Two dimensional array whose row m contains the + distances from atom m of atoms n > m that are within a distance cutoff_3 + of both atom n and the central atom. + + triplet_counts: One dimensional array of integers whose entry m is the + number of atoms that are within a distance cutoff_3 of atom m. + + :rtype: (np.ndarray, np.ndarray, np.ndarray, np.ndarray) + """ + + sepcut = False + if nspecie > 1 and cutoff_3 is not None: + bc = specie_mask[ctype] + bcn = nspecie * bc + r_cut = np.max(cutoff_3) + sepcut = True + + # get 3-body bond array + ind_3_l = np.where(bond_array_2[:, 0] > r_cut)[0] + if (ind_3_l.shape[0] > 0): + ind_3 = ind_3_l[0] + else: + ind_3 = bond_array_2.shape[0] + + bond_array_3 = bond_array_2[0:ind_3, :] + bond_positions_3 = bond_positions_2[0:ind_3, :] + + cut_m = r_cut + cut_n = r_cut + cut_mn = r_cut + + # get cross bond array + cross_bond_inds = np.zeros((ind_3, ind_3), dtype=np.int8) - 1 + cross_bond_dists = np.zeros((ind_3, ind_3), dtype=np.float64) + triplet_counts = np.zeros(ind_3, dtype=np.int8) + for m in range(ind_3): + pos1 = bond_positions_3[m] + count = m + 1 + trips = 0 + + if sepcut and (specie_mask is not None) and (cut3b_mask is not None) and (cutoff_3 is not None): + # choose bond dependent bond + bm = specie_mask[etypes[m]] + btype_m = cut3b_mask[bm + bcn] # (m, c) + cut_m = cutoff_3[btype_m] + bmn = nspecie * bm # for cross_dist usage + + for n in range(m + 1, ind_3): + + if sepcut and (specie_mask is not None) and (cut3b_mask is not None) and (cutoff_3 is not None): + bn = specie_mask[etypes[n]] + btype_n = cut3b_mask[bn + bcn] # (n, c) + cut_n = cutoff_3[btype_n] + + # for cross_dist (m,n) pair + btype_mn = cut3b_mask[bn + bmn] + cut_mn = cutoff_3[btype_mn] + + pos2 = bond_positions_3[n] + diff = pos2 - pos1 + dist_curr = sqrt( + diff[0] * diff[0] + diff[1] * diff[1] + diff[2] * diff[2]) + + if dist_curr < cut_mn \ + and bond_array_2[m, 0] < cut_m \ + and bond_array_2[n, 0] < cut_n: + cross_bond_inds[m, count] = n + cross_bond_dists[m, count] = dist_curr + count += 1 + trips += 1 + + triplet_counts[m] = trips + + return bond_array_3, cross_bond_inds, cross_bond_dists, triplet_counts + +@njit +def get_m2_body_arrays(positions, atom: int, cell, r_cut, manybody_cutoff_list, + species, sweep: np.ndarray, nspec, spec_mask, manybody_mask, + cutoff_func=cf.quadratic_cutoff): + # TODO: + # 1. need to deal with the conflict of cutoff functions if other funcs are used + # 2. complete the docs of "Return" + # TODO: this can be probably improved using stored arrays, redundant calls to get_2_body_arrays + # Get distances, positions, species and indices of neighbouring atoms + """ + Args: + positions (np.ndarray): Positions of atoms in the structure. + atom (int): Index of the central atom of the local environment. + cell (np.ndarray): 3x3 array whose rows are the Bravais lattice vectors of the + cell. + manybody_cutoff_list (float): 2-body cutoff radius. + species (np.ndarray): Numpy array of species represented by their atomic numbers. + + Return: + Tuple of arrays describing pairs of atoms in the 2-body local + environment. + """ + # Get distances, positions, species and indexes of neighbouring atoms + bond_array_mb, _, etypes, bond_inds = get_2_body_arrays( + positions, atom, cell, r_cut, manybody_cutoff_list, species, sweep, + nspec, spec_mask, manybody_mask) + + sepcut = False + if nspec > 1 and manybody_cutoff_list is not None: + bc = spec_mask[species[atom]] + bcn = bc * nspec + sepcut = True + + species_list = np.array(list(set(species)), dtype=np.int8) + n_bonds = len(bond_inds) + n_specs = len(species_list) + qs = np.zeros(n_specs, dtype=np.float64) + qs_neigh = np.zeros((n_bonds, n_specs), dtype=np.float64) + q_neigh_grads = np.zeros((n_bonds, 3), dtype=np.float64) + + # get coordination number of center atom for each species + for s in range(n_specs): + if sepcut and (spec_mask is not None) and (manybody_mask is not None) and (manybody_cutoff_list is not None): + bs = spec_mask[species_list[s]] + mbtype = manybody_mask[bcn + bs] + r_cut = manybody_cutoff_list[mbtype] + + qs[s] = q_value_mc(bond_array_mb[:, 0], r_cut, species_list[s], + etypes, cutoff_func) + + # get coordination number of all neighbor atoms for each species + for i in range(n_bonds): + if sepcut and (spec_mask is not None) and (manybody_mask is not None) and (manybody_cutoff_list is not None): + be = spec_mask[etypes[i]] + ben = be * nspec + + neigh_bond_array, __, neigh_etypes, ___ = \ + get_2_body_arrays(positions, bond_inds[i], cell, r_cut, + manybody_cutoff_list, species, sweep, nspec, spec_mask, manybody_mask) + for s in range(n_specs): + if sepcut and (spec_mask is not None) and (manybody_mask is not None) and (manybody_cutoff_list is not None): + bs = spec_mask[species_list[s]] + mbtype = manybody_mask[bs + ben] + r_cut = manybody_cutoff_list[mbtype] + + qs_neigh[i, s] = q_value_mc(neigh_bond_array[:, 0], r_cut, + species_list[s], neigh_etypes, cutoff_func) + + # get grad from each neighbor atom + for i in range(n_bonds): + if sepcut and (spec_mask is not None) and (manybody_mask is not None) and (manybody_cutoff_list is not None): + be = spec_mask[etypes[i]] + mbtype = manybody_mask[bcn + be] + r_cut = manybody_cutoff_list[mbtype] + + ri = bond_array_mb[i, 0] + for d in range(3): + ci = bond_array_mb[i, d+1] + + ____, q_neigh_grads[i, d] = coordination_number(ri, ci, r_cut, + cutoff_func) + + # get grads of the center atom + q_grads = q2_grads_mc(q_neigh_grads, species_list, etypes) + + return qs, qs_neigh, q_grads, q_neigh_grads, species_list, etypes + +@njit +def q2_grads_mc(neigh_grads, species_list, etypes): + n_specs = len(species_list) + n_neigh = neigh_grads.shape[0] + grads = np.zeros((n_specs, 3)) + for i in range(n_neigh): + si = np.where(species_list==etypes[i])[0][0] + grads[si, :] += neigh_grads[i, :] + + return grads + + +@njit +def get_m3_body_arrays(positions, atom: int, cell, cutoff: float, species, + sweep, cutoff_func=cf.quadratic_cutoff): + """ + Note: here we assume the cutoff is not too large, + i.e., 2 * cutoff < cell_size + """ + species_list = np.array(list(set(species)), dtype=np.int8) + + q_func = coordination_number + + bond_array, bond_positions, etypes, bond_inds = \ + get_2_body_arrays(positions, atom, cell, cutoff, species, sweep) + + bond_array_m3b, cross_bond_inds, cross_bond_dists, triplets = \ + get_3_body_arrays(bond_array, bond_positions, cutoff) + + # get descriptor of center atom for each species + m3b_array = q3_value_mc(bond_array_m3b[:, 0], cross_bond_inds, + cross_bond_dists, triplets, cutoff, species_list, etypes, + cutoff_func, q_func) + + + # get descriptor of all neighbor atoms for each species + n_bonds = len(bond_array_m3b) + n_specs = len(species_list) + m3b_neigh_array = np.zeros((n_bonds, n_specs, n_specs)) + for i in range(n_bonds): + neigh_bond_array, neigh_positions, neigh_etypes, _ = \ + get_2_body_arrays(positions, bond_inds[i], cell, cutoff, species, sweep) + + neigh_array_m3b, neigh_cross_inds, neigh_cross_dists, neigh_triplets = \ + get_3_body_arrays(neigh_bond_array, neigh_positions, cutoff) + + m3b_neigh_array[i, :, :] = q3_value_mc(neigh_array_m3b[:, 0], + neigh_cross_inds, neigh_cross_dists, neigh_triplets, + cutoff, species_list, neigh_etypes, cutoff_func, q_func) + + # get grad from each neighbor atom, assume the cutoff is not too large + # such that 2 * cutoff < cell_size + m3b_neigh_grads = q3_neigh_grads_mc(bond_array_m3b, cross_bond_inds, + cross_bond_dists, triplets, cutoff, species_list, etypes, + cutoff_func, q_func) + + # get grads of the center atom + m3b_grads = q3_grads_mc(m3b_neigh_grads, species_list, etypes) + + return m3b_array, m3b_neigh_array, m3b_grads, m3b_neigh_grads, species_list, etypes + +@njit +def q3_grads_mc(neigh_grads, species_list, etypes): + n_specs = len(species_list) + n_neigh = neigh_grads.shape[0] + grads = np.zeros((n_specs, n_specs, 3)) + for i in range(n_neigh): + si = np.where(species_list==etypes[i])[0][0] + for spec_j in species_list: + sj = np.where(species_list==spec_j)[0][0] + if si == sj: + grads[si, sj, :] += neigh_grads[i, sj, :] / 2 + else: + grads[si, sj, :] += neigh_grads[i, sj, :] + + return grads + +@njit +def q3_neigh_grads_mc(bond_array_m3b, cross_bond_inds, cross_bond_dists, + triplets, r_cut, species_list, etypes, cutoff_func, + q_func=coordination_number): + + n_bonds = len(bond_array_m3b) + n_specs = len(species_list) + m3b_grads = np.zeros((n_bonds, n_specs, 3)) + + # get grad from each neighbor atom + for i in range(n_bonds): + + # get grad of q_func + ri = bond_array_m3b[i, 0] + si = np.where(species_list==etypes[i])[0][0] + qi, _ = q_func(ri, 0, r_cut, cutoff_func) + + qi_grads = np.zeros(3) + for d in range(3): + ci = bond_array_m3b[i, d + 1] + _, qi_grads[d] = q_func(ri, ci, r_cut, cutoff_func) + + # go through all triplets with "atom" and "i" + for ind in range(triplets[i]): + j = cross_bond_inds[i, i + ind + 1] + rj = bond_array_m3b[j, 0] + sj = np.where(species_list==etypes[j])[0][0] + qj, _ = q_func(rj, 0, r_cut, cutoff_func) + + qj_grads = np.zeros(3) + for d in range(3): + cj = bond_array_m3b[j, d + 1] + _, qj_grads[d] = q_func(rj, cj, r_cut, cutoff_func) + + rij = cross_bond_dists[i, i + ind + 1] + qij, _ = q_func(rij, 0, r_cut, cutoff_func) + + q_grad = (qi_grads * qj + qi * qj_grads) * qij + + # remove duplicant + # if si == sj: + # q_grad /= 2 + m3b_grads[i, sj, :] += q_grad + m3b_grads[j, si, :] += q_grad + + return m3b_grads + + +@njit +def q3_value_mc(distances, cross_bond_inds, cross_bond_dists, triplets, + r_cut, species_list, etypes, cutoff_func, q_func=coordination_number): + """Compute value of many-body many components descriptor based + on distances of atoms in the local many-body environment. + + Args: + distances (np.ndarray): distances between atoms i and j + r_cut (float): cutoff hyperparameter + ref_species (int): species to consider to compute the contribution + etypes (np.ndarray): atomic species of neighbours + cutoff_func (callable): cutoff function + q_func (callable): many-body pairwise descrptor function + + Return: + float: the value of the many-body descriptor + """ + n_specs = len(species_list) + mb3_array = np.zeros((n_specs, n_specs)) + n_bonds = len(distances) + + for m in range(n_bonds): + q1, _ = q_func(distances[m], 0, r_cut, cutoff_func) + s1 = np.where(species_list==etypes[m])[0][0] + + for n in range(triplets[m]): + ind = cross_bond_inds[m, m + n + 1] + s2 = np.where(species_list==etypes[ind])[0][0] + q2, _ = q_func(distances[ind], 0, r_cut, cutoff_func) + + r3 = cross_bond_dists[m, m + n + 1] + q3, _ = q_func(r3, 0, r_cut, cutoff_func) + + mb3_array[s1, s2] += q1 * q2 * q3 + if s1 != s2: + mb3_array[s2, s1] += q1 * q2 * q3 + + return mb3_array + diff --git a/flare/utils/mask_helper.py b/flare/utils/mask_helper.py deleted file mode 100644 index 0a3df43b0..000000000 --- a/flare/utils/mask_helper.py +++ /dev/null @@ -1,1361 +0,0 @@ -import time -import math -import pickle -import inspect -import json - -import numpy as np -from copy import deepcopy -from numpy.random import random -from numpy import array as nparray -from numpy import max as npmax -from typing import List, Callable, Union -from warnings import warn -from sys import stdout -from os import devnull - -from flare.utils.element_coder import element_to_Z, Z_to_element - - -class HyperParameterMasking(): - """ - A helper class to construct the hyps_mask dictionary for AtomicEnvironment - , GaussianProcess and MappedGaussianProcess - - Examples: - - pm = HyperParameterMasking(species=['C', 'H', 'O'], - bonds=[['*', '*'], ['O','O']], - triplets=[['*', '*', '*'], - ['O','O', 'O']], - parameters={'bond0':[1, 0.5, 1], 'bond1':[2, 0.2, 2], - 'triplet0':[1, 0.5], 'triplet1':[2, 0.2], - 'cutoff3b':1}, - constraints={'bond0':[False, True]}) - hm = pm.hyps_mask - hyps = hm['hyps'] - cutoffs = hm['cutoffs'] - - In this example, four atomic species are involved. There are many kinds - of bonds and triplets. But we only want to use eight different sigmas - and lengthscales. - - In order to do so, we first define all the bonds to be group "bond0", by - listing "*-*" as the first element in the bond argument. The second - element O-O is then defined to be group "bond1". Note that the order - matters here. The later element overrides the ealier one. If - bonds=[['O', 'O'], ['*', '*']], then all bonds belong to group "bond1". - - Similarly, O-O-O is defined as triplet1, while all remaining ones - are left as triplet0. - - The hyperpameters for each group is listed in the order of - [sig, ls, cutoff] in the parameters argument. So in this example, - O-O interaction will use [2, 0.2, 2] as its sigma, length scale, and - cutoff. - - For triplet, the parameter arrays only come with two elements. So there - is no cutoff associated with triplet0 or triplet1; instead, a universal - cutoff is used, which is defined as 'cutoff3b'. - - The constraints argument define which hyper-parameters will be optimized. - True for optimized and false for being fixed. - - See more examples in tests/test_mask_helper.py - - """ - - def __init__(self, hyps_mask=None, species=None, bonds=None, - triplets=None, cut3b=None, mb=None, parameters=None, - constraints={}, allseparate=False, random=False, verbose=False): - """ Initialization function - - :param hyps_mask: Not implemented yet - :type hyps_mask: dict - :param species: list or dictionary that define specie groups - :type species: [dict, list] - :param bonds: list or dictionary that define bond groups - :type bonds: [dict, list, bool] - :param triplets: list or dictionary that define triplet groups - :type triplets: [dict, list, bool] - :param cut3b: list or dictionary that define 3b-cutoff groups - :type cut3b: [dict, list] - :param mb: list or dictionary that define many-body groups - :type mb: [dict, list, bool] - :param parameters: dictionary of parameters - :type parameters: dict - :param constraints: whether the hyperparmeters are optimized (True) or not (False) - :constraints: dict - :param random: if True, define each single bond type into a separate group and randomized initial parameters - :type random: bool - :param verbose: print the process to screen or to null - :type verbose: bool - - See format of species, bonds, triplets, cut3b, mb in list_groups() function. - - See format of parameters and constraints in list_parameters() function. - - """ - - if (verbose): - self.fout = stdout - else: - self.fout = open(devnull, 'w') - - self.n = {} - self.groups = {} - self.all_members = {} - self.all_group_names = {} - self.all_names = [] - self.all_types = ['specie', 'bond', 'triplet', 'mb', 'cut3b'] - - for group_type in self.all_types: - self.n[group_type] = 0 - self.groups[group_type] = [] - self.all_members[group_type] = [] - self.all_group_names[group_type] = [] - self.sigma = {} - self.ls = {} - self.all_cutoff = {} - self.hyps_sig = {} - self.hyps_ls = {} - self.hyps_opt = {} - self.opt = {'noise': True} - self.mask = {} - self.cutoff_list = {} - self.noise = 0.05 - self.universal = {} - - self.cutoffs_array = np.array([0, 0, 0], dtype=np.float) - self.hyps = None - - if (species is not None): - self.list_groups('specie', species) - if (not allseparate): - if (bonds is not None): - self.list_groups('bond', bonds) - if (triplets is not None): - self.list_groups('triplet', triplets) - if (cut3b is not None): - self.list_groups('cut3b', cut3b) - if (mb is not None): - self.list_groups('mb', mb) - if (parameters is not None): - self.list_parameters(parameters, constraints) - try: - self.hyps_mask = self.generate_dict() - except: - print("more parameters needed to generate the hypsmask", - file=self.fout) - else: - if (parameters is not None): - self.list_parameters(parameters, constraints) - if (not random): - assert 'lengthscale' in self.universal - assert 'sigma' in self.universal - - # sort the types - if (bonds is not None): - assert isinstance(bonds, bool) - else: - bonds = False - if (triplets is not None): - assert isinstance(triplets, bool) - else: - triplets = False - if (mb is not None): - assert isinstance(mb, bool) - else: - mb = False - - if bonds: - self.fill_in_parameters('bond', random) - if triplets: - self.fill_in_parameters('triplet', random) - if mb: - self.fill_in_parameters('mb', random) - self.hyps_mask = self.generate_dict() - - def list_parameters(self, parameter_dict, constraints={}): - """Define many groups of parameters - - :param parameter_dict: dictionary of all parameters - :type parameter_dict: dict - :param constraints: dictionary of all constraints - :type constraints: dict - - example: parameter_dict={"name":[sig, ls, cutoffs], ...} - constraints={"name":[True, False, False], ...} - - The name of parameters can be the group name previously defined in - define_group or list_groups function. Aside from the group name, - "noise", "cutoff2b", "cutoff3b", and "cutoffmb" are reserved for - noise parmater and universal cutoffs. - - For non-reserved keys, the value should be a list of 2-3 elements, - correspond to the sigma, lengthscale (and cutoff if the third one - is defined). For reserved keys, the value should be a scalar. - - The parameter_dict and constraints should uses the same set of keys. - The keys in constraints but not in parameter_dict will be ignored. - - The value in the constraints can be either a single bool, which apply - to all parameters, or list of bools that apply to each parameter. - """ - - for name in parameter_dict: - self.set_parameters( - name, parameter_dict[name], constraints.get(name, True)) - - def list_groups(self, group_type, definition_list): - """define groups in batches. - - Args: - - group_type (str): "specie", "bond", "triplet", "cut3b", "mb" - definition_list (list, dict): list of elements - - This function runs define_group in batch. Please first read - the manual of define_group. - - If the definition_list is a list, it is equivalent to - executing define_group through the definition_list. - - | for all terms in the list: - | define_group(group_type, group_type+'n', the nth term in the list) - - So the first bond defined will be group bond0, second one will be - group bond1. For specie, it will define all the listed elements as - groups with only one element with their original name. - - If the definition_list is a dictionary, it is equivalent to - - | for k, v in the dict: - | define_group(group_type, k, v) - - It is not recommended to use the dictionary mode, especially when - the group definitions are conflicting with each other. There is no - guarantee that the looping order is the same as you want. - - Unlike define_group, it can only be called once for each - group_type, and not after any define_group calls. - - """ - if (group_type == 'specie'): - if (len(self.all_group_names['specie']) > 0): - raise RuntimeError("this function has to be run " - "before any define_group") - if (isinstance(definition_list, list)): - for ele in definition_list: - if isinstance(ele, list): - self.define_group('specie', ele, ele) - else: - self.define_group('specie', ele, [ele]) - elif (isinstance(elemnt_list, dict)): - for ele in definition_list: - self.define_group('specie', ele, definition_list[ele]) - else: - raise RuntimeError("type unknown") - else: - if (len(self.all_group_names['specie']) == 0): - raise RuntimeError("this function has to be run " - "before any define_group") - if (isinstance(definition_list, list)): - ngroup = len(definition_list) - for idg in range(ngroup): - self.define_group(group_type, f"{group_type}{idg}", - definition_list[idg]) - elif (isinstance(definition_list, dict)): - for name in definition_list: - if (isinstance(definition_list[name][0], list)): - for ele in definition_list[name]: - self.define_group(group_type, name, ele) - else: - self.define_group(group_type, name, - definition_list[name]) - - def fill_in_parameters(self, group_type, random=False): - """Separate all possible types of bonds, triplets, mb. - One type per group. And fill in either universal ls and sigma from - pre-defined parameters from set_parameters("sigma", ..) and set_parameters("ls", ..) - or random parameters if random is True. - - Args: - - group_type (str): "specie", "bond", "triplet", "cut3b", "mb" - definition_list (list, dict): list of elements - - """ - nspec = len(self.all_group_names['specie']) - if (nspec < 1): - raise RuntimeError("the specie group has to be defined in advance") - if (group_type in ['bond', 'mb']): - tid = 0 - for i in range(nspec): - ele1 = self.all_group_names['specie'][i] - for j in range(i, nspec): - ele2 = self.all_group_names['specie'][j] - if (random): - self.define_group(group_type, f'{group_type}{tid}', - [ele1, ele2], parameters=np.random.random(2)) - else: - self.define_group(group_type, f'{group_type}{tid}', - [ele1, ele2], - parameters=[self.universal['sigma'], - self.universal['lengthscale']]) - tid += 1 - elif (group_type == 'triplet'): - tid = 0 - for i in range(nspec): - ele1 = self.all_group_names['specie'][i] - for j in range(i, nspec): - ele2 = self.all_group_names['specie'][j] - for k in range(j, nspec): - ele3 = self.all_group_names['specie'][k] - if (random): - self.define_group(group_type, f'{group_type}{tid}', - [ele1, ele2, ele3], parameters=np.random.random(2)) - else: - self.define_group(group_type, f'{group_type}{tid}', - [ele1, ele2, ele3], - parameters=[self.universal['sigma'], - self.universal['lengthscale']]) - tid += 1 - else: - print(group_type, "will be ignored", file=self.fout) - - def define_group(self, group_type, name, element_list, parameters=None, atomic_str=False): - """Define specie/bond/triplet/3b cutoff/manybody group - - Args: - group_type (str): "specie", "bond", "triplet", "cut3b", "mb" - name (str): the name use for indexing. can be anything but "*" - element_list (list): list of elements - parameters (list): corresponding parameters for this group - atomic_str (bool): whether the element in element_list is - group name or periodic table element name. - - The function is helped to define different groups for specie/bond/triplet - /3b cutoff/manybody terms. This function can be used for many times. - The later one always overrides the former one. - - The name of the group has to be unique string (but not "*"), that - define a group of species or bonds, etc. If the same name is used, - in two function calls, the definitions of the group will be merged. - Both calls will be effective. - - element_list has to be a list of atomic elements, or a list of - specie group names (which should be defined in previous calls), or "*". - "*" will loop the function over all previously defined species. - It has to be two elements for bond/3b cutoff/manybody term, or - three elements for triplet. For specie group definition, it can be - as many elements as you want. - - If multiple define_group calls have conflict with element, the later one - has higher priority. For example, bond 1-2 are defined as group1 in - the first call, and as group2 in the second call. In the end, the bond - will be left as group2. - - Example 1: - - define_group('specie', 'water', ['H', 'O']) - define_group('specie', 'salt', ['Cl', 'Na']) - - They define H and O to be group water, and Na and Cl to be group salt. - - Example 2.1: - - define_group('bond', 'in-water', ['H', 'H'], atomic_str=True) - define_group('bond', 'in-water', ['H', 'O'], atomic_str=True) - define_group('bond', 'in-water', ['O', 'O'], atomic_str=True) - - Example 2.2: - define_group('bond', 'in-water', ['water', 'water']) - - The 2.1 is equivalent to 2.2. - - Example 3.1: - - define_group('specie', '1', ['H']) - define_group('specie', '2', ['O']) - define_group('bond', 'Hgroup', ['H', 'H'], atomic_str=True) - define_group('bond', 'Hgroup', ['H', 'O'], atomic_str=True) - define_group('bond', 'OO', ['O', 'O'], atomic_str=True) - - Example 3.2: - - define_group('specie', '1', ['H']) - define_group('specie', '2', ['O']) - define_group('bond', 'Hgroup', ['H', '*'], atomic_str=True) - define_group('bond', 'OO', ['O', 'O'], atomic_str=True) - - Example 3.3: - - list_groups('specie', ['H', 'O']) - define_group('bond', 'Hgroup', ['H', '*']) - define_group('bond', 'OO', ['O', 'O']) - - Example 3.4: - - list_groups('specie', ['H', 'O']) - define_group('bond', 'OO', ['*', '*']) - define_group('bond', 'Hgroup', ['H', '*']) - - 3.1 to 3.4 are all equivalent. - """ - - if (name == '*'): - raise ValueError("* is reserved for substitution, cannot be used " - "as a group name") - - if (group_type != 'specie'): - - # Check all the other group_type to - exclude_list = deepcopy(self.all_types) - ide = exclude_list.index(group_type) - exclude_list.pop(ide) - - for gt in exclude_list: - if (name in self.all_group_names[gt]): - raise ValueError("group name has to be unique across all types. " - f"{name} is found in type {gt}") - - if (name in self.all_group_names[group_type]): - groupid = self.all_group_names[group_type].index(name) - else: - groupid = self.n[group_type] - self.all_group_names[group_type].append(name) - self.groups[group_type].append([]) - self.n[group_type] += 1 - - if (group_type == 'specie'): - for ele in element_list: - assert ele not in self.all_members['specie'], \ - "The element has already been defined" - self.groups['specie'][groupid].append(ele) - self.all_members['specie'].append(ele) - print( - f"Element {ele} will be defined as group {name}", file=self.fout) - else: - if (len(self.all_group_names['specie']) == 0): - raise RuntimeError("The atomic species have to be" - "defined in advance") - if ("*" not in element_list): - gid = [] - for ele_name in element_list: - if (atomic_str): - for idx in range(self.n['specie']): - if (ele_name in self.groups['specie'][idx]): - gid += [idx] - print(f"Warning: Element {ele_name} is used for " - f"definition, but the whole group " - f"{self.all_group_names[idx]} is affected", file=self.fout) - else: - gid += [self.all_group_names['specie'].index(ele_name)] - - for ele in self.all_members[group_type]: - if set(gid) == set(ele): - print( - f"Warning: the definition of {group_type} {ele} will be overriden", file=self.fout) - self.groups[group_type][groupid].append(gid) - self.all_members[group_type].append(gid) - print( - f"{group_type} {gid} will be defined as group {name}", file=self.fout) - if (parameters is not None): - self.set_parameters(name, parameters) - else: - one_star_less = deepcopy(element_list) - idstar = element_list.index('*') - one_star_less.pop(idstar) - for sub in self.all_group_names['specie']: - self.define_group(group_type, name, - one_star_less + [sub], parameters=parameters, atomic_str=atomic_str) - - def find_group(self, group_type, element_list, atomic_str=False): - - # remember the later command override the earlier ones - if (group_type == 'specie'): - if (not isinstance(element_list, str)): - print("for element, it has to be a string", file=self.fout) - return None - name = None - for igroup in range(self.n['specie']): - gname = self.all_group_names[group_type][igroup] - allspec = self.groups[group_type][igroup] - if (element_list in allspec): - name = gname - return name - print("cannot find the group", file=self.fout) - return None - else: - if ("*" in element_list): - print("* cannot be used for find", file=self.fout) - return None - gid = [] - for ele_name in element_list: - gid += [self.all_group_names['specie'].index(ele_name)] - setlist = set(gid) - name = None - for igroup in range(self.n[group_type]): - gname = self.all_group_names[group_type][igroup] - for ele in self.groups[group_type][igroup]: - if set(gid) == set(ele): - name = gname - return name - # self.groups[group_type][groupid].append(gid) - # self.all_members[group_type].append(gid) - # print( - # f"{group_type} {gid} will be defined as group {name}", file=self.fout) - # if (parameters is not None): - # self.set_parameters(name, parameters) - # else: - # one_star_less = deepcopy(element_list) - # idstar = element_list.index('*') - # one_star_less.pop(idstar) - # for sub in self.all_group_names['specie']: - # self.define_group(group_type, name, - # one_star_less + [sub], parameters=parameters, atomic_str=atomic_str) - - def set_parameters(self, name, parameters, opt=True): - """Set the parameters for certain group - - :param name: name of the patermeters - :type name: str - :param parameters: the sigma, lengthscale, and cutoff of each group. - :type parameters: list - :param opt: whether to optimize the parameter or not - :type opt: bool, list - - The name of parameters can be the group name previously defined in - define_group or list_groups function. Aside from the group name, - "noise", "cutoff2b", "cutoff3b", and "cutoffmb" are reserved for - noise parmater and universal cutoffs. - - The parameter should be a list of 2-3 elements, for sigma, - lengthscale (and cutoff if the third one is defined). - - The optimization flag can be a single bool, which apply to all - parameters, or list of bools that apply to each parameter. - """ - - if (name == 'noise'): - self.noise = parameters - self.opt['noise'] = opt - return - - if (name in ['cutoff2b', 'cutoff3b', 'cutoffmb']): - cutstr2index = {'cutoff2b': 0, 'cutoff3b': 1, 'cutoffmb': 2} - self.cutoffs_array[cutstr2index[name]] = parameters - return - - if (name in ['sigma', 'lengthscale']): - self.universal[name] = parameters - return - - if (isinstance(opt, bool)): - opt = [opt, opt, opt] - - if ('cut3b' not in name): - if (name in self.sigma): - print( - f"Warning, the sig, ls of group {name} is overriden", file=self.fout) - self.sigma[name] = parameters[0] - self.ls[name] = parameters[1] - self.opt[name+'sig'] = opt[0] - self.opt[name+'ls'] = opt[1] - print(f"Parameters for group {name} will be set as " - f"sig={parameters[0]} ({opt[0]}) " - f"ls={parameters[1]} ({opt[1]})", file=self.fout) - if (len(parameters) > 2): - if (name in self.all_cutoff): - print( - f"Warning, the cutoff of group {name} is overriden", file=self.fout) - self.all_cutoff[name] = parameters[2] - print(f"Cutoff for group {name} will be set as " - f"{parameters[2]}", file=self.fout) - else: - self.all_cutoff[name] = parameters - - def set_constraints(self, name, opt): - """Set the parameters for certain group - - :param name: name of the patermeters - :type name: str - :param opt: whether to optimize the parameter or not - :type opt: bool, list - - The name of parameters can be the group name previously defined in - define_group or list_groups function. Aside from the group name, - "noise", "cutoff2b", "cutoff3b", and "cutoffmb" are reserved for - noise parmater and universal cutoffs. - - The optimization flag can be a single bool, which apply to all - parameters under that name, or list of bools that apply to each - parameter. - """ - - if (name == 'noise'): - self.opt['noise'] = opt - return - - if (name in ['cutoff2b', 'cutoff3b', 'cutoffmb']): - cutstr2index = {'cutoff2b': 0, 'cutoff3b': 1, 'cutoffmb': 2} - return - - if (isinstance(opt, bool)): - opt = [opt, opt, opt] - - if ('cut3b' not in name): - if (name in self.sigma): - print( - f"Warning, the sig, ls of group {name} is overriden", file=self.fout) - self.opt[name+'sig'] = opt[0] - self.opt[name+'ls'] = opt[1] - print(f"Parameters for group {name} will be set as " - f"sig {opt[0]} " - f"ls {opt[1]}", file=self.fout) - - def summarize_group(self, group_type): - """Sort and combine all the previous definition to internal varialbes - - Args: - - group_type (str): species, bond, triplet, cut3b, mb - """ - aeg = self.all_group_names[group_type] - nspecie = self.n['specie'] - if (group_type == "specie"): - self.nspecie = nspecie - self.specie_mask = np.ones(118, dtype=np.int)*(nspecie-1) - for idt in range(self.nspecie): - for ele in self.groups['specie'][idt]: - atom_n = element_to_Z(ele) - self.specie_mask[atom_n] = idt - print(f"elemtn {ele} is defined as type {idt} with name " - f"{aeg[idt]}", file=self.fout) - print( - f"All the remaining elements are left as type {idt}", file=self.fout) - elif (group_type in ['bond', 'cut3b', 'mb']): - if (self.n[group_type] == 0): - print(group_type, "is not defined. Skipped", file=self.fout) - return - self.mask[group_type] = np.ones( - nspecie**2, dtype=np.int)*(self.n[group_type]-1) - self.hyps_sig[group_type] = [] - self.hyps_ls[group_type] = [] - self.hyps_opt[group_type] = [] - for idt in range(self.n[group_type]): - name = aeg[idt] - for bond in self.groups[group_type][idt]: - g1 = bond[0] - g2 = bond[1] - self.mask[group_type][g1+g2*nspecie] = idt - self.mask[group_type][g2+g1*nspecie] = idt - s1 = self.groups['specie'][g1] - s2 = self.groups['specie'][g2] - print(f"{group_type} {s1} - {s2} is defined as type {idt} " - f"with name {name}", file=self.fout) - if (group_type != 'cut3b'): - sig = self.sigma[name] - ls = self.ls[name] - self.hyps_sig[group_type] += [sig] - self.hyps_ls[group_type] += [ls] - self.hyps_opt[group_type] += [self.opt[name+'sig']] - self.hyps_opt[group_type] += [self.opt[name+'ls']] - print(f" using hyper-parameters of {sig:6.2g} "\ - f"{ls:6.2g}", file=self.fout) - print( - f"All the remaining elements are left as type {idt}", file=self.fout) - - cutstr2index = {'bond': 0, 'cut3b': 1, 'mb': 2} - - # sort out the cutoffs - allcut = [] - alldefine = True - for idt in range(self.n[group_type]): - if (aeg[idt] in self.all_cutoff): - allcut += [self.all_cutoff[aeg[idt]]] - else: - alldefine = False - print(f"Warning, {aeg[idt]} cutoff is not define. "\ - "it's going to use the universal cutoff.") - - if len(allcut)>0: - universal_cutoff = self.cutoffs_array[cutstr2index[group_type]] - if (universal_cutoff <= 0): - universal_cutoff = np.max(allcut) - print(f"Warning, universal cutoffs {cutstr2index[group_type]}for " - f"{group_type} is defined as zero! reset it to {universal_cutoff}") - - self.cutoff_list[group_type] = [] - for idt in range(self.n[group_type]): - self.cutoff_list[group_type] += [ - self.all_cutoff.get(aeg[idt], universal_cutoff)] - - max_cutoff = np.max(self.cutoff_list[group_type]) - # update the universal cutoff to make it higher than - if (alldefine): - self.cutoffs_array[cutstr2index[group_type]] = \ - max_cutoff - elif (not np.any(self.cutoff_list[group_type]-max_cutoff)): - # if not all the cutoffs are defined separately - # and they are all the same value. so - del self.cutoff_list[group_type] - if (group_type == 'cut3b'): - self.cutoffs_array[cutstr2index[group_type]] = max_cutoff - self.n['cut3b'] = 0 - - if (self.cutoffs_array[cutstr2index[group_type]] <= 0): - raise RuntimeError( - f"cutoffs for {group_type} is undefined") - - elif (group_type == "triplet"): - self.ntriplet = self.n['triplet'] - if (self.ntriplet == 0): - print(group_type, "is not defined. Skipped", file=self.fout) - return - self.mask[group_type] = np.ones( - nspecie**3, dtype=np.int)*(self.ntriplet-1) - self.hyps_sig[group_type] = [] - self.hyps_ls[group_type] = [] - self.hyps_opt[group_type] = [] - for idt in range(self.n['triplet']): - name = aeg[idt] - for triplet in self.groups['triplet'][idt]: - g1 = triplet[0] - g2 = triplet[1] - g3 = triplet[2] - self.mask[group_type][g1+g2*nspecie+g3*nspecie**2] = idt - self.mask[group_type][g1+g3*nspecie+g2*nspecie**2] = idt - self.mask[group_type][g2+g1*nspecie+g3*nspecie**2] = idt - self.mask[group_type][g2+g3*nspecie+g1*nspecie**2] = idt - self.mask[group_type][g3+g1*nspecie+g2*nspecie**2] = idt - self.mask[group_type][g3+g2*nspecie+g1*nspecie**2] = idt - s1 = self.groups['specie'][g1] - s2 = self.groups['specie'][g2] - s3 = self.groups['specie'][g3] - print(f"triplet {s1} - {s2} - {s3} is defined as type {idt} with name " - f"{name}", file=self.fout) - sig = self.sigma[name] - ls = self.ls[name] - self.hyps_sig[group_type] += [sig] - self.hyps_ls[group_type] += [ls] - self.hyps_opt[group_type] += [self.opt[name+'sig']] - self.hyps_opt[group_type] += [self.opt[name+'ls']] - print( - f" using hyper-parameters of {sig} {ls}", file=self.fout) - print( - f"all the remaining elements are left as type {idt}", file=self.fout) - if (self.cutoffs_array[1] == 0): - allcut = [] - for idt in range(self.n[group_type]): - allcut += [self.all_cutoff.get(aeg[idt], 0)] - aeg_ = self.all_group_names['cut3b'] - for idt in range(self.n['cut3b']): - allcut += [self.all_cutoff.get(aeg_[idt], 0)] - if (len(allcut)>0): - self.cutoffs_array[1] = np.max(allcut) - else: - raise RuntimeError( - f"cutoffs for {group_type} is undefined") - else: - pass - - def generate_dict(self): - """Dictionary representation of the mask. The output can be used for AtomicEnvironment - or the GaussianProcess - """ - - # sort out all the definitions and resolve conflicts - # cut3b has to be summarize before triplet - # because the universal triplet cutoff is checked - # at the end of triplet search - self.summarize_group('specie') - self.summarize_group('bond') - self.summarize_group('cut3b') - self.summarize_group('triplet') - self.summarize_group('mb') - - hyps_mask = {} - hyps_mask['nspecie'] = self.n['specie'] - hyps_mask['specie_mask'] = self.specie_mask - - hyps = [] - hyps_label = [] - opt = [] - for group in ['bond', 'triplet', 'mb']: - if (self.n[group] >= 1): - # copy the mask - hyps_mask['n'+group] = self.n[group] - hyps_mask[group+'_mask'] = self.mask[group] - hyps += [self.hyps_sig[group]] - hyps += [self.hyps_ls[group]] - # check parameters - opt += [self.hyps_opt[group]] - aeg = self.all_group_names[group] - for idt in range(self.n[group]): - hyps_label += ['Signal_Var._'+aeg[idt]] - for idt in range(self.n[group]): - hyps_label += ['Length_Scale_'+group] - opt += [self.opt['noise']] - hyps_label += ['Noise_Var.'] - hyps_mask['hyps_label'] = hyps_label - hyps += [self.noise] - - # handle partial optimization if any constraints are defined - hyps_mask['original'] = np.hstack(hyps) - - opt = np.hstack(opt) - hyps_mask['train_noise'] = self.opt['noise'] - if (not opt.all()): - nhyps = len(hyps_mask['original']) - hyps_mask['original_labels'] = hyps_mask['hyps_label'] - mapping = [] - hyps_mask['hyps_label'] = [] - for i in range(nhyps): - if (opt[i]): - mapping += [i] - hyps_mask['hyps_label'] += [hyps_label[i]] - newhyps = hyps_mask['original'][mapping] - hyps_mask['map'] = np.array(mapping, dtype=np.int) - elif (opt.any()): - newhyps = hyps_mask['original'] - else: - raise RuntimeError("hyps has length zero." - "at least one component of the hyper-parameters" - "should be allowed to be optimized. \n") - hyps_mask['hyps'] = newhyps - - # checkout universal cutoffs and seperate cutoffs - nbond = hyps_mask.get('nbond', 0) - ntriplet = hyps_mask.get('ntriplet', 0) - nmb = hyps_mask.get('nmb', 0) - if len(self.cutoff_list.get('bond', [])) > 0 \ - and nbond > 0: - hyps_mask['cutoff_2b'] = np.array( - self.cutoff_list['bond'], dtype=np.float) - if len(self.cutoff_list.get('cut3b', [])) > 0 \ - and ntriplet > 0: - hyps_mask['cutoff_3b'] = np.array( - self.cutoff_list['cut3b'], dtype=np.float) - hyps_mask['ncut3b'] = self.n['cut3b'] - hyps_mask['cut3b_mask'] = self.mask['cut3b'] - if len(self.cutoff_list.get('mb', [])) > 0 \ - and nmb > 0: - hyps_mask['cutoff_mb'] = np.array( - self.cutoff_list['mb'], dtype=np.float) - - self.hyps_mask = hyps_mask - if (self.cutoffs_array[2] > 0) and nmb > 0: - hyps_mask['cutoffs'] = self.cutoffs_array - else: - hyps_mask['cutoffs'] = self.cutoffs_array[:2] - - if self.n['specie'] < 2: - print("only one type of elements was defined. Please use multihyps=False", - file=self.fout) - - return hyps_mask - - @staticmethod - def from_dict(hyps_mask, verbose=False, init_spec=[]): - """ convert dictionary mask to HM instance - This function is not tested yet - """ - - HyperParameterMasking.check_instantiation(hyps_mask) - - pm = HyperParameterMasking(verbose=verbose) - - hyps = hyps_mask['hyps'] - - if 'map' in hyps_mask: - ihyps = hyps - hyps = hyps_mask['original'] - constraints = np.zeros(len(hyps), dtype=bool) - for ele in hyps_mask['map']: - constraints[ele] = True - for i, ori in enumerate(hyps_mask['map']): - hyps[ori] = ihyps[i] - else: - constraints = np.ones(len(hyps), dtype=bool) - - pm.nspecie = hyps_mask['nspecie'] - nele = len(hyps_mask['specie_mask']) - max_species = np.max(hyps_mask['specie_mask']) - specie_mask = hyps_mask['specie_mask'] - for i in range(max_species+1): - elelist= np.where(specie_mask == i)[0] - if len(elelist) > 0: - for ele in elelist: - if (ele != 0): - elename = Z_to_element(ele) - if (len(init_spec) > 0): - if elename in init_spec: - pm.define_group( - "specie", i, [elename]) - else: - pm.define_group("specie", i, [elename]) - - nbond = hyps_mask.get('nbond', 0) - ntriplet = hyps_mask.get('ntriplet', 0) - nmb = hyps_mask.get('nmb', 0) - for t in ['bond', 'mb']: - if (f'n{t}' in hyps_mask): - if (t == 'bond'): - cutoffname = 'cutoff_2b' - sig = hyps[:nbond] - ls = hyps[nbond:2*nbond] - csig = constraints[:nbond] - cls = constraints[nbond:2*nbond] - else: - cutoffname = 'cutoff_mb' - sig = hyps[nbond*2+ntriplet*2:nbond*2+ntriplet*2+nmb] - ls = hyps[nbond*2+ntriplet*2+nmb:nbond*2+ntriplet*2+nmb*2] - csig = constraints[nbond*2+ntriplet*2:] - cls = constraints[nbond*2+ntriplet*2+nmb:] - for i in range(pm.nspecie): - for j in range(i, pm.nspecie): - ttype = hyps_mask[f'{t}_mask'][i+j*pm.nspecie] - pm.define_group(f"{t}", f"{t}{ttype}", [i, j]) - for i in range(hyps_mask[f'n{t}']): - if (cutoffname in hyps_mask): - pm.set_parameters(f"{t}{i}", [sig[i], ls[i], hyps_mask[cutoffname][i]], - opt=[csig[i], cls[i]]) - else: - pm.set_parameters(f"{t}{i}", [sig[i], ls[i]], - opt=[csig[i], cls[i]]) - if ('ntriplet' in hyps_mask): - sig = hyps[nbond*2:nbond*2+ntriplet] - ls = hyps[nbond*2+ntriplet:nbond*2+ntriplet*2] - csig = constraints[nbond*2:nbond*2+ntriplet] - cls = constraints[nbond*2+ntriplet:nbond*2+ntriplet*2] - for i in range(pm.nspecie): - for j in range(i, pm.nspecie): - for k in range(j, pm.nspecie): - triplettype = hyps_mask[f'triplet_mask'][i + - j*pm.nspecie+k*pm.nspecie*pm.nspecie] - pm.define_group( - f"triplet", f"triplet{triplettype}", [i, j, k]) - for i in range(hyps_mask[f'ntriplet']): - pm.set_parameters(f"triplet{i}", [sig[i], ls[i]], - opt=[csig[i], cls[i]]) - if (f'ncut3b' in hyps_mask): - for i in range(pm.nspecie): - for j in range(i, pm.nspecie): - ttype = hyps_mask[f'cut3b_mask'][i+j*pm.nspecie] - pm.define_group("cut3b", f"cut3b{ttype}", [i, j]) - for i in range(hyps_mask['ncut3b']): - pm.set_parameters( - f"cut3b{i}", [0, 0, hyps_mask['cutoff_3b'][i]]) - - pm.set_parameters('noise', hyps[-1]) - if 'cutoffs' in hyps_mask: - cut = hyps_mask['cutoffs'] - pm.set_parameters(f"cutoff2b", cut[0]) - try: - pm.set_parameters(f"cutoff3b", cut[1]) - except: - pass - try: - pm.set_parameters(f"cutoffmb", cut[2]) - except: - pass - - return pm - - @staticmethod - def check_instantiation(hyps_mask): - """ - Runs a series of checks to ensure that the user has not supplied - contradictory arguments which will result in undefined behavior - with multiple hyperparameters. - :return: - """ - - # backward compatability - if ('nspec' in hyps_mask): - hyps_mask['nspecie'] = hyps_mask['nspec'] - if ('spec_mask' in hyps_mask): - hyps_mask['specie_mask'] = hyps_mask['spec_mask'] - if ('train_noise' not in hyps_mask): - hyps_mask['train_noise'] = True - - assert isinstance(hyps_mask, dict) - - assert 'nspecie' in hyps_mask, "nspecie key missing in " \ - "hyps_mask dictionary" - assert 'specie_mask' in hyps_mask, "specie_mask key " \ - "missing " \ - "in hyps_mask dicticnary" - - nspecie = hyps_mask['nspecie'] - hyps_mask['specie_mask'] = nparray( - hyps_mask['specie_mask'], dtype=np.int) - - if 'nbond' in hyps_mask: - n2b = hyps_mask['nbond'] - assert n2b > 0 - assert isinstance(n2b, int) - hyps_mask['bond_mask'] = nparray( - hyps_mask['bond_mask'], dtype=np.int) - if n2b > 0: - bmask = hyps_mask['bond_mask'] - assert (npmax(bmask) < n2b) - assert len(bmask) == nspecie ** 2, \ - f"wrong dimension of bond_mask: " \ - f" {len(bmask)} != nspecie^2 {nspecie**2}" - for t2b in range(nspecie): - for t2b_2 in range(t2b, nspecie): - assert bmask[t2b*nspecie+t2b_2] == bmask[t2b_2*nspecie+t2b], \ - 'bond_mask has to be symmetric' - else: - n2b = 0 - - if 'ntriplet' in hyps_mask: - n3b = hyps_mask['ntriplet'] - assert n3b > 0 - assert isinstance(n3b, int) - hyps_mask['triplet_mask'] = nparray( - hyps_mask['triplet_mask'], dtype=np.int) - if n3b > 0: - tmask = hyps_mask['triplet_mask'] - assert (npmax(tmask) < n3b) - assert len(tmask) == nspecie ** 3, \ - f"wrong dimension of bond_mask: " \ - f" {len(tmask)} != nspecie^3 {nspecie**3}" - - for t3b in range(nspecie): - for t3b_2 in range(t3b, nspecie): - for t3b_3 in range(t3b_2, nspecie): - assert tmask[t3b*nspecie*nspecie+t3b_2*nspecie+t3b_3] \ - == tmask[t3b*nspecie*nspecie+t3b_3*nspecie+t3b_2], \ - 'bond_mask has to be symmetric' - assert tmask[t3b*nspecie*nspecie+t3b_2*nspecie+t3b_3] \ - == tmask[t3b_2*nspecie*nspecie+t3b*nspecie+t3b_3], \ - 'bond_mask has to be symmetric' - assert tmask[t3b*nspecie*nspecie+t3b_2*nspecie+t3b_3] \ - == tmask[t3b_2*nspecie*nspecie+t3b_3*nspecie+t3b], \ - 'bond_mask has to be symmetric' - assert tmask[t3b*nspecie*nspecie+t3b_2*nspecie+t3b_3] \ - == tmask[t3b_3*nspecie*nspecie+t3b*nspecie+t3b_2], \ - 'bond_mask has to be symmetric' - assert tmask[t3b*nspecie*nspecie+t3b_2*nspecie+t3b_3] \ - == tmask[t3b_3*nspecie*nspecie+t3b_2*nspecie+t3b], \ - 'bond_mask has to be symmetric' - else: - n3b = 0 - - if 'nmb' in hyps_mask: - nmb = hyps_mask['nmb'] - assert nmb > 0 - assert isinstance(nmb, int) - hyps_mask['mb_mask'] = nparray(hyps_mask['mb_mask'], dtype=np.int) - if nmb > 0: - bmask = hyps_mask['mb_mask'] - assert (npmax(bmask) < nmb) - assert len(bmask) == nspecie ** 2, \ - f"wrong dimension of mb_mask: " \ - f" {len(bmask)} != nspecie^2 {nspecie**2}" - for tmb in range(nspecie): - for tmb_2 in range(tmb, nspecie): - assert bmask[tmb*nspecie+tmb_2] == bmask[tmb_2*nspecie+tmb], \ - 'mb_mask has to be symmetric' - # else: - # nmb = 1 - # hyps_mask['mb_mask'] = np.zeros(nspecie**2, dtype=np.int) - - if 'map' in hyps_mask: - assert ('original' in hyps_mask), \ - "original hyper parameters have to be defined" - # Ensure typed correctly as numpy array - hyps_mask['original'] = nparray( - hyps_mask['original'], dtype=np.float) - - if (len(hyps_mask['original']) - 1) not in hyps_mask['map']: - assert hyps_mask['train_noise'] is False, \ - "train_noise should be False when noise is not in hyps" - else: - assert hyps_mask['train_noise'] is True, \ - "train_noise should be True when map is not used" - - if 'cutoff_2b' in hyps_mask: - c2b = hyps_mask['cutoff_2b'] - assert len(c2b) == n2b, \ - f'number of 2b cutoff should be the same as n2b {n2b}' - - if 'cutoff_3b' in hyps_mask: - c3b = hyps_mask['cutoff_3b'] - assert nc3b > 0 - assert isinstance(nc3b, int) - hyps_mask['cut3b_mask'] = nparray( - hyps_mask['cut3b_mask'], dtype=int) - assert len(c3b) == hyps_mask['ncut3b'], \ - f'number of 3b cutoff should be the same as ncut3b {ncut3b}' - assert len(hyps_mask['cut3b_mask']) == nspecie ** 2, \ - f"wrong dimension of cut3b_mask: " \ - f" {len(bmask)} != nspecie^2 {nspecie**2}" - assert npmax(hyps_mask['cut3b_mask']) < hyps_mask['ncut3b'], \ - f"wrong dimension of cut3b_mask: " \ - f" {len(bmask)} != nspecie^2 {nspecie**2}" - - if 'cutoff_mb' in hyps_mask: - cmb = hyps_mask['cutoff_mb'] - assert len(cmb) == nmb, \ - f'number of mb cutoff should be the same as nmb {nmb}' - - return hyps_mask - - @staticmethod - def check_matching(hyps_mask, hyps, cutoffs): - """ - check whether hyps_mask, hyps and cutoffs are compatible - used in GaussianProcess - """ - - n2b = hyps_mask.get('nbond', 0) - n3b = hyps_mask.get('ntriplet', 0) - nmb = hyps_mask.get('nmb', 0) - - if (len(cutoffs) <= 2): - assert ((n2b + n3b) > 0) - else: - assert ((n2b + n3b + nmb) > 0) - - if 'map' in hyps_mask: - if (len(cutoffs) <= 2): - assert (n2b * 2 + n3b * 2 + 1) == len(hyps_mask['original']), \ - "the hyperparmeter length is inconsistent with the mask" - else: - if (nmb == 0): - nmb = 1 - hyps_mask['mb_mask'] = np.zeros(hyps_mask['nspecie']**2, dtype=np.int) - assert (n2b * 2 + n3b * 2 + nmb * 2 + 1) == len(hyps_mask['original']), \ - "the hyperparmeter length is inconsistent with the mask" - assert len(hyps_mask['map']) == len(hyps), \ - "the hyperparmeter length is inconsistent with the mask" - else: - if (len(cutoffs) <= 2): - assert (n2b * 2 + n3b * 2 + 1) == len(hyps), \ - "the hyperparmeter length is inconsistent with the mask" - else: - if (nmb == 0): - nmb = 1 - hyps_mask['mb_mask'] = np.zeros(hyps_mask['nspecie']**2, dtype=np.int) - assert (n2b * 2 + n3b * 2 + nmb*2 + 1) == len(hyps), \ - "the hyperparmeter length is inconsistent with the mask" - - if 'cutoff_2b' in hyps_mask: - assert cutoffs[0] >= npmax(hyps_mask['cutoff_2b']), \ - 'general cutoff should be larger than all cutoffs listed in hyps_mask' - - if 'cutoff_3b' in hyps_mask: - assert cutoffs[0] >= npmax(hyps_mask['cutoff_3b']), \ - 'general cutoff should be larger than all cutoffs listed in hyps_mask' - - if 'cutoff_mb' in hyps_mask: - assert cutoffs[0] >= npmax(hyps_mask['cutoff_mb']), \ - 'general cutoff should be larger than all cutoffs listed in hyps_mask' - - @staticmethod - def mask2cutoff(cutoffs, cutoffs_mask): - """use in flare.env AtomicEnvironment to resolve what cutoff to use""" - - ncutoffs = len(cutoffs) - scalar_cutoff_2 = cutoffs[0] - scalar_cutoff_3 = 0 - scalar_cutoff_mb = 0 - if (ncutoffs > 1): - scalar_cutoff_3 = cutoffs[1] - if (ncutoffs > 2): - scalar_cutoff_mb = cutoffs[2] - - if (scalar_cutoff_2 == 0): - scalar_cutoff_2 = np.max([scalar_cutoff_3, scalar_cutoff_mb]) - - if (cutoffs_mask is None): - return scalar_cutoff_2, scalar_cutoff_3, scalar_cutoff_mb, \ - None, None, None, \ - 1, 1, 1, 1, None, None, None, None - - nspecie = cutoffs_mask.get('nspecie', 1) - nspecie = nspecie - if (nspecie == 1): - return scalar_cutoff_2, scalar_cutoff_3, scalar_cutoff_mb, \ - None, None, None, \ - 1, 1, 1, 1, None, None, None, None - - n2b = cutoffs_mask.get('nbond', 1) - n3b = cutoffs_mask.get('ncut3b', 1) - nmb = cutoffs_mask.get('nmb', 1) - specie_mask = cutoffs_mask.get('specie_mask', None) - bond_mask = cutoffs_mask.get('bond_mask', None) - cut3b_mask = cutoffs_mask.get('cut3b_mask', None) - mb_mask = cutoffs_mask.get('mb_mask', None) - cutoff_2b = cutoffs_mask.get('cutoff_2b', None) - cutoff_3b = cutoffs_mask.get('cutoff_3b', None) - cutoff_mb = cutoffs_mask.get('cutoff_mb', None) - - if cutoff_2b is not None: - scalar_cutoff_2 = np.max(cutoff_2b) - else: - n2b = 1 - - if cutoff_3b is not None: - scalar_cutoff_3 = np.max(cutoff_3b) - else: - n3b = 1 - - if cutoff_mb is not None: - scalar_cutoff_mb = np.max(cutoff_mb) - else: - nmb = 1 - - return scalar_cutoff_2, scalar_cutoff_3, scalar_cutoff_mb, \ - cutoff_2b, cutoff_3b, cutoff_mb, \ - nspecie, n2b, n3b, nmb, specie_mask, bond_mask, cut3b_mask, mb_mask - - @staticmethod - def get_2b_hyps(hyps, hyps_mask, multihyps=False): - - original_hyps = np.copy(hyps) - if (multihyps is True): - new_hyps = HyperParameterMasking.get_hyps(hyps_mask, hyps) - n2b = hyps_mask['nbond'] - new_hyps = np.hstack([new_hyps[:n2b*2], new_hyps[-1]]) - new_hyps_mask = {'nbond': n2b, 'ntriplet': 0, - 'nspecie': hyps_mask['nspecie'], - 'specie_mask': hyps_mask['specie_mask'], - 'bond_mask': hyps_mask['bond_mask']} - if ('cutoff_2b' in hyps_mask): - new_hyps_mask['cutoff_2b'] = hyps_mask['cutoff_2b'] - else: - new_hyps = [hyps[0], hyps[1], hyps[-1]] - new_hyps_mask = None - - return new_hyps, new_hyps_mask - - @staticmethod - def get_3b_hyps(hyps, hyps_mask, multihyps=False): - - if (multihyps is True): - new_hyps = HyperParameterMasking.get_hyps(hyps_mask, hyps) - n2b = hyps_mask.get('nbond', 0) - n3b = hyps_mask['ntriplet'] - new_hyps = np.hstack([new_hyps[n2b*2:n2b*2+n3b*2], new_hyps[-1]]) - new_hyps_mask = {'ntriplet': n3b, 'nbond': 0, - 'nspecie': hyps_mask['nspecie'], - 'specie_mask': hyps_mask['specie_mask'], - 'triplet_mask': hyps_mask['triplet_mask']} - ncut3b = hyps_mask.get('ncut3b', 0) - if (ncut3b > 0): - new_hyps_mask['ncut3b'] = hyps_mask['cut3b_mask'] - new_hyps_mask['cut3b_mask'] = hyps_mask['cut3b_mask'] - new_hyps_mask['cutoff_3b'] = hyps_mask['cutoff_3b'] - else: - # kind of assuming that 2-body is there - base = 2 - new_hyps = np.hstack([hyps[0+base], hyps[1+base], hyps[-1]]) - new_hyps_mask = None - - return hyps, hyps_mask - - @staticmethod - def get_mb_hyps(hyps, hyps_mask, multihyps=False): - - if (multihyps is True): - new_hyps = HyperParameterMasking.get_hyps(hyps_mask, hyps) - n2b = hyps_mask.get('n2b', 0) - n3b = hyps_mask.get('n3b', 0) - n23b2 = (n2b+n3b)*2 - nmb = hyps_mask['nmb'] - - new_hyps = np.hstack([new_hyps[n23b2:n23b2+nmb*2], new_hyps[-1]]) - - new_hyps_mask = {'nmb': nmb, 'nbond': 0, 'ntriplet':0, - 'nspecie': hyps_mask['nspecie'], - 'specie_mask': hyps_mask['specie_mask'], - 'mb_mask': hyps_mask['mb_mask']} - - if ('cutoff_mb' in hyps_mask): - new_hyps_mask['cutoff_mb'] = hyps_mask['cutoff_mb'] - else: - # kind of assuming that 2+3 are there - base = 4 - new_hyps = np.hstack([hyps[0+base], hyps[1+base], hyps[-1]]) - new_hyps_mask = None - - return new_hyps, new_hyps_mask - - @staticmethod - def get_cutoff(coded_species, cutoff, hyps_mask): - - if (len(coded_species)==2): - if (hyps_mask is None): - return cutoff[0] - elif ('cutoff_2b' not in hyps_mask): - return cutoff[0] - - ele1 = hyps_mask['species_mask'][coded_species[0]] - ele2 = hyps_mask['species_mask'][coded_species[1]] - bond_type = hyps_mask['bond_mask'][ \ - hyps_mask['nspecie']*ele1 + ele2] - return hyps_mask['cutoff_2b'][bond_type] - - elif (len(coded_species)==3): - if (hyps_mask is None): - return np.ones(3)*cutoff[1] - elif ('cutoff_3b' not in hyps_mask): - return np.ones(3)*cutoff[1] - - ele1 = hyps_mask['species_mask'][coded_species[0]] - ele2 = hyps_mask['species_mask'][coded_species[1]] - ele3 = hyps_mask['species_mask'][coded_species[2]] - bond1 = hyps_mask['cut3b_mask'][ \ - hyps_mask['nspecie']*ele1 + ele2] - bond2 = hyps_mask['cut3b_mask'][ \ - hyps_mask['nspecie']*ele1 + ele3] - bond12 = hyps_mask['cut3b_mask'][ \ - hyps_mask['nspecie']*ele2 + ele3] - return np.array([hyps_mask['cutoff_3b'][bond1], - hyps_mask['cutoff_3b'][bond2], - hyps_mask['cutoff_3b'][bond12]]) - else: - raise NotImplementedError - - @staticmethod - def get_hyps(hyps_mask, hyps): - if 'map' in hyps_mask: - newhyps = np.copy(hyps_mask['original']) - for i, ori in enumerate(hyps_mask['map']): - newhyps[ori] = hyps[i] - return newhyps - else: - return hyps - - @staticmethod - def compare_dict(dict1, dict2): - - if type(dict1) != type(dict2): - return False - - if dict1 is None: - return True - - for k in ['nspecie', 'specie_mask', 'nbond', 'bond_mask', - 'cutoff_2b', 'ntriplet', 'triplet_mask', - 'n3b', 'cut3b_mask', 'nmb', 'mb_mask', - 'cutoff_mb', 'map']: #, 'train_noise']: - if (k in dict1) != (k in dict2): - return False - elif (k in dict1): - if not (np.isclose(dict1[k], dict2[k]).all()): - return False - - for k in ['train_noise']: - if (k in dict1) != (k in dict2): - return False - elif (k in dict1): - if dict1[k] !=dict2[k]: - return False - return True diff --git a/flare/utils/parameter_helper.py b/flare/utils/parameter_helper.py new file mode 100644 index 000000000..b8c1414f4 --- /dev/null +++ b/flare/utils/parameter_helper.py @@ -0,0 +1,1107 @@ +""" +For multi-component systems, the configurational space can be highly complicated. +One may want to use different hyper-parameters and cutoffs for different interactions, +or do constraint optimisation for hyper-parameters. + +To use more hyper-parameters, we need special kernel function that can differentiate different +pairs, triplets and other descriptors and determine which number to use for what interaction. + +This kernel can be enabled by using the ``hyps_mask`` argument of the GaussianProcess class. +It contains multiple arrays to describe how to break down the array of hyper-parameters and +apply them when computing the kernel. Detail descriptions of this argument can be seen in +kernel/mc_sephyps.py. + +The ParameterHelper class is to generate the hyps_mask with a more human readable interface. + +Example: + +>>> pm = ParameterHelper(species=['C', 'H', 'O'], +... kernels={'twobody':[['*', '*'], ['O','O']], +... 'threebody':[['*', '*', '*'], +... ['O','O', 'O']]}, +... parameters={'twobody0':[1, 0.5, 1], 'twobody1':[2, 0.2, 2], +... 'threebody0':[1, 0.5], 'threebody1':[2, 0.2], +... 'cutoff_threebody':1}, +... constraints={'twobody0':[False, True]}) +>>> hm = pm.hyps_mask +>>> hyps = hm['hyps'] +>>> cutoffs = hm['cutoffs'] +>>> kernels = hm['kernels'] +>>> gp_model = GaussianProcess(kernels=kernels, cutoffs=cutoffs, +... hyps=hyps, hyps_mask=hm) + +In this example, four atomic species are involved. There are many kinds +of twobodys and threebodys. But we only want to use eight different signal variance +and length-scales. + +In order to do so, we first define all the twobodys to be group "twobody0", by +listing "*-*" as the first element in the twobody argument. The second +element O-O is then defined to be group "twobody1". Note that the order +matters here. The later element overrides the ealier one. If +twobodys=[['O', 'O'], ['*', '*']], then all twobodys belong to group "twobody1". + +Similarly, O-O-O is defined as threebody1, while all remaining ones +are left as threebody0. + +The hyperpameters for each group is listed in the order of +[sig, ls, cutoff] in the parameters argument. So in this example, +O-O interaction will use [2, 0.2, 2] as its sigma, length scale, and +cutoff. + +For threebody, the parameter arrays only come with two elements. So there +is no cutoff associated with threebody0 or threebody1; instead, a universal +cutoff is used, which is defined as 'cutoff_threebody'. + +The constraints argument define which hyper-parameters will be optimized. +True for optimized and false for being fixed. + +Here are a couple more simple examples. + +Define a 5-parameter 2+3 kernel (1, 0.5, 1, 0.5, 0.05) + +>>> pm = ParameterHelper(kernels=['twobody', 'threebody'], +... parameters={'sigma': 1, +... 'lengthscale': 0.5, +... 'cutoff_twobody': 2, +... 'cutoff_threebody': 1, +... 'noise': 0.05}) + +Define a 5-parameter 2+3 kernel (1, 1, 1, 1, 0.05) + +>>> pm = ParameterHelper(kernels=['twobody', 'threebody'], +... parameters={'cutoff_twobody': 2, +... 'cutoff_threebody': 1, +... 'noise': 0.05}, +... ones=ones, +... random=not ones) + +Define a 9-parameter 2+3 kernel + +>>> pm = ParameterHelper() +>>> pm.define_group('specie', 'O', ['O']) +>>> pm.define_group('specie', 'rest', ['C', 'H']) +>>> pm.define_group('twobody', '**', ['*', '*']) +>>> pm.define_group('twobody', 'OO', ['O', 'O']) +>>> pm.define_group('threebody', '***', ['*', '*', '*']) +>>> pm.define_group('threebody', 'Oall', ['O', 'O', 'O']) +>>> pm.set_parameters('**', [1, 0.5]) +>>> pm.set_parameters('OO', [1, 0.5]) +>>> pm.set_parameters('Oall', [1, 0.5]) +>>> pm.set_parameters('***', [1, 0.5]) +>>> pm.set_parameters('cutoff_twobody', 5) +>>> pm.set_parameters('cutoff_threebody', 4) + +See more examples in functions ``ParameterHelper.define_group`` , ``ParameterHelper.set_parameters``, +and in the tests ``tests/test_parameters.py`` +""" + +import inspect +import json +import logging +import math +import numpy as np +import pickle +import time + +from copy import deepcopy +from itertools import combinations_with_replacement, permutations +from numpy import array as nparray +from numpy import max as npmax +from typing import List, Callable, Union + +from flare.output import set_logger +from flare.parameters import Parameters +from flare.utils.element_coder import element_to_Z, Z_to_element + + +class ParameterHelper(): + """ + A helper class to construct the hyps_mask dictionary for AtomicEnvironment + , GaussianProcess and MappedGaussianProcess + + Args: + hyps_mask (dict): Not implemented yet + species (dict, list): Define specie groups + kernels (dict, list): Define kernels and groups for the kernels + cutoff_groups (dict): Define different cutoffs for different species + parameters (dict): Define signal variance, length scales, and cutoffs + constraints (dict): If listed as False, the cooresponding hyperparmeters + will not be trained + allseparate (bool): If True, define each type pair/triplet into a + separate group. + random (bool): If True, randomized all signal variances and lengthscales + one (bool): If True, set all signal variances and lengthscales to one + verbose (str): Level to print with "ERROR", "WARNING", "INFO", "DEBUG" + + * the ``species`` is an optional input. It can be left as None if the user only wants + to set up one group of hyper-parameters for each kernel. + * the ``kernels`` can be defined along with or without groups. But the later mode + is not compatible with the ``allseparate`` flag. + + >>> kernels=['twobody', 'threebody'], + + or + + >>> kernels={'twobody':[['*', '*'], ['O','O']], + ... 'threebody':[['*', '*', '*'], + ... ['O','O', 'O']]}, + + Current options for the kernels are twobody, threebody and manybody (based on coordination number). + * See format of ``species``, ``kernels`` (dict), and ``cutoff_groups`` in ``list_groups()`` function. + * See format of ``parameters`` and ``constraints`` in ``list_parameters()`` function. + """ + + # TO DO, sync it to kernel class + # need to be synced with kernel class + + # name of the kernels + all_kernel_types = ['twobody', 'threebody', 'manybody'] + additional_groups = ['cut3b'] + + # dimension of the kernels + ndim = {'twobody': 2, 'threebody': 3, 'manybody': 2, 'cut3b': 2} + n_kernel_parameters = {'twobody': 2, + 'threebody': 2, 'manybody': 2, 'cut3b': 0} + + def __init__(self, hyps_mask=None, species=None, kernels={}, + cutoff_groups={}, parameters=None, + constraints={}, allseparate=False, random=False, ones=False, + verbose="WARNING"): + + self.logger = set_logger("ParameterHelper", stream=True, + fileout_name=None, verbose="info") + + self.all_types = ['specie'] + \ + ParameterHelper.all_kernel_types + ParameterHelper.additional_groups + + self.all_group_types = ParameterHelper.all_kernel_types + \ + ParameterHelper.additional_groups + + # number of groups {'twobody': 1, 'threebody': 2} + self.n = {} + # definition of groups {'specie': [['C', 'H'], ['O']], 'twobody': [[['*', '*']], [[ele1, ele2]]]} + self.groups = {} + # joint values of the groups {'specie': ['C', 'H', 'O'], 'twobody': [['*', '*'], [ele1, ele2]]} + self.all_members = {} + # names of each group {'specie': ['group1', 'group2'], 'twobody': ['twobody0', 'twobody1']} + self.all_group_names = {} + # joint list of all the keys in self.all_group_names + self.all_names = [] + + # set up empty container + for group_type in self.all_types: + self.n[group_type] = 0 + self.groups[group_type] = [] + self.all_members[group_type] = [] + self.all_group_names[group_type] = [] + + # store parameters, key should be the one used in + # all_group_names or kernel_name + self.sigma = {} + self.ls = {} + self.noise = 0.05 + self.energy_noise = 0.1 + self.opt = {'noise': True} + + # key should be sigma, lengthscale + # cutoff_kernel_name + self.universal = {} + + # key should be in all_group_names + self.all_cutoff = {} + + # used for as_dict + self.hyps_sig = {} + self.hyps_ls = {} + self.hyps_opt = {} + self.cutoff_list = {} + self.mask = {} + + self.hyps = None + + if isinstance(kernels, dict): + self.kernel_dict = kernels + self.kernels = list(kernels.keys()) + assert (not allseparate) + elif isinstance(kernels, list): + self.kernels = kernels + # by default, there is only one group of hyperparameters + # for each type of the kernel + # unless allseparate is defined + self.kernel_dict = {} + for ktype in kernels: + self.kernel_dict[ktype] = [['*']*ParameterHelper.ndim[ktype]] + + if species is not None: + self.list_groups('specie', species) + + # define groups + if allseparate: + for ktype in self.kernels: + self.all_separate_groups(ktype) + else: + for ktype in self.kernels: + self.list_groups(ktype, self.kernel_dict[ktype]) + + # check for cut3b + for group in cutoff_groups: + self.list_groups(group, cutoff_groups[group]) + + # define parameters + if parameters is not None: + self.list_parameters(parameters, constraints) + + if 'lengthscale' in self.universal and 'sigma' in self.universal: + universal = True + else: + universal = False + + if (random+ones+universal) > 1: + raise RuntimeError( + "random and ones cannot be simultaneously True") + elif random or ones or universal: + for ktype in self.kernels: + self.fill_in_parameters( + ktype, random=random, ones=ones, universal=universal) + + elif len(self.kernels) > 0: + self.list_groups('specie', ['*']) + + # define groups + for ktype in self.kernels: + self.list_groups(ktype, self.kernel_dict[ktype]) + + # check for cut3b + for group in cutoff_groups: + self.list_groups(group, cutoff_groups[group]) + + # define parameters + if parameters is not None: + self.list_parameters(parameters, constraints) + + if 'lengthscale' in self.universal and 'sigma' in self.universal: + universal = True + else: + universal = False + + if (random+ones+universal) > 1: + raise RuntimeError( + "random and ones cannot be simultaneously True") + elif random or ones or universal: + for ktype in self.kernels: + self.fill_in_parameters( + ktype, random=random, ones=ones, universal=universal) + + def list_parameters(self, parameter_dict:dict, constraints:dict={}): + """Define many groups of parameters + + Args: + parameter_dict (dict): dictionary of all parameters + constraints (dict): dictionary of all constraints + + Example: + + >>> parameter_dict={"group_name":[sig, ls, cutoffs], ...} + >>> constraints={"group_name":[True, False, False], ...} + + The name of parameters can be the group name previously defined in + define_group or list_groups function. Aside from the group name, + ``noise``, ``cutoff_twobody``, ``cutoff_threebody``, and + ``cutoff_manybody`` are reserved for noise parmater + and universal cutoffs, while ``sigma`` and ``lengthscale`` are + reserved for universal signal variances and length scales. + + For non-reserved keys, the value should be a list of 2 to 3 elements, + corresponding to the sigma, lengthscale and cutoff (if the third one + is defined). For reserved keys, the value should be a float number. + + The parameter_dict and constraints should use the same set of keys. + If a key in constraints is not used in parameter_dict, it will be ignored. + + The value in the constraints can be either a single bool, which apply + to all parameters, or list of bools that apply to each parameter. + """ + + for name in parameter_dict: + self.set_parameters( + name, parameter_dict[name], constraints.get(name, True)) + + def list_groups(self, group_type, definition_list): + """define groups in batches. + + Args: + group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody" + definition_list (list, dict): list of elements + + This function runs define_group in batch. Please first read + the manual of define_group. + + If the definition_list is a list, it is equivalent to + executing define_group through the definition_list. + + >>> for all terms in the list: + >>> define_group(group_type, group_type+'n', the nth term in the list) + + So the first twobody defined will be group twobody0, second one will be + group twobody1. For specie, it will define all the listed elements as + groups with only one element with their original name. + + If the definition_list is a dictionary, it is equivalent to + + >>> for k, v in the dict: + >>> define_group(group_type, k, v) + + It is not recommended to use the dictionary mode, especially when + the group definitions are conflicting with each other. There is no + guarantee that the priority order is the same as you want. + + Unlike ParameterHelper.define_group(), it can only be called once for each + group_type, and not after any ParameterHelper.define_group() calls. + + """ + + if group_type == 'specie': + if len(self.all_group_names['specie']) > 0: + raise RuntimeError("this function has to be run " + "before any define_group") + if isinstance(definition_list, list): + for ele in definition_list: + if isinstance(ele, list): + self.define_group('specie', ele, ele) + else: + self.define_group('specie', ele, [ele]) + elif isinstance(elemnt_list, dict): + for ele in definition_list: + self.define_group('specie', ele, definition_list[ele]) + else: + raise RuntimeError("type unknown") + else: + if self.n['specie'] == 0: + raise RuntimeError("this function has to be run " + "before any define_group") + if isinstance(definition_list, list): + ngroup = len(definition_list) + for idg in range(ngroup): + self.define_group(group_type, f"{group_type}{idg}", + definition_list[idg]) + elif isinstance(definition_list, dict): + for name in definition_list: + if isinstance(definition_list[name][0], list): + for ele in definition_list[name]: + self.define_group(group_type, name, ele) + else: + self.define_group(group_type, name, + definition_list[name]) + + def all_separate_groups(self, group_type): + """Separate all possible types of twobodys, threebodys, manybody. + One type per group. + + Args: + group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody" + + """ + nspec = len(self.all_group_names['specie']) + if nspec < 1: + raise RuntimeError("the specie group has to be defined in advance") + if group_type in self.all_group_types: + # TO DO: the two blocks below can be replace by some upper triangle operation + + # generate all possible combination of group + ele_grid = self.all_group_names['specie'] + grid = np.meshgrid(*[ele_grid]*ParameterHelper.ndim[group_type]) + grid = np.array(grid).T.reshape(-1, + ParameterHelper.ndim[group_type]) + + # remove the redundant groups + allgroup = [] + for group in grid: + exist = False + set_list_group = set(list(group)) + for prev_group in allgroup: + if set(prev_group) == set_list_group: + exist = True + if not exist: + allgroup += [list(group)] + + # define the group + tid = 0 + for group in allgroup: + self.define_group(group_type, f'{group_type}{tid}', + group) + tid += 1 + else: + logger.info(f"{group_type} will be ignored") + + def fill_in_parameters(self, group_type, random=False, ones=False, universal=False): + """Separate all possible types of twobodys, threebodys, manybody. + One type per group. And fill in either universal ls and sigma from + pre-defined parameters from set_parameters("sigma", ..) and set_parameters("ls", ..) + or random parameters if random is True. + + Args: + group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody" + definition_list (list, dict): list of elements + + """ + nspec = len(self.all_group_names['specie']) + if nspec < 1: + raise RuntimeError("the specie group has to be defined in advance") + if random: + for group_name in self.all_group_names[group_type]: + self.set_parameters(group_name, parameters=np.random.random(2)) + elif ones: + for group_name in self.all_group_names[group_type]: + self.set_parameters(group_name, parameters=np.ones(2)) + elif universal: + for group_name in self.all_group_names[group_type]: + self.set_parameters(group_name, + parameters=[self.universal['sigma'], + self.universal['lengthscale']]) + + def define_group(self, group_type, name, element_list, parameters=None, atomic_str=False): + """Define specie/twobody/threebody/3b cutoff/manybody group + + Args: + group_type (str): "specie", "twobody", "threebody", "cut3b", "manybody" + name (str): the name use for indexing. can be anything but "*" + element_list (list): list of elements + parameters (list): corresponding parameters for this group + atomic_str (bool): whether the elements in element_list are + specified by group names or periodic table element names. + + The function is helped to define different groups for specie/twobody/threebody + /3b cutoff/manybody terms. This function can be used for many times. + The later one always overrides the former one. + + The name of the group has to be unique string (but not "*"), that + define a group of species or twobodys, etc. If the same name is used, + in two function calls, the definitions of the group will be merged. + Both calls will be effective. + + element_list has to be a list of atomic elements, or a list of + specie group names (which should be defined in previous calls), or "*". + "*" will loop the function over all previously defined species. + It has to be two elements for twobody/3b cutoff/manybody term, or + three elements for threebody. For specie group definition, it can be + as many elements as you want. + + If multiple define_group calls have conflict with element, the later one + has higher priority. For example, twobody 1-2 are defined as group1 in + the first call, and as group2 in the second call. In the end, the twobody + will be left as group2. + + Example 1: + + >>> define_group('specie', 'water', ['H', 'O']) + >>> define_group('specie', 'salt', ['Cl', 'Na']) + + They define H and O to be group water, and Na and Cl to be group salt. + + Example 2.1: + + >>> define_group('twobody', 'in-water', ['H', 'H'], atomic_str=True) + >>> define_group('twobody', 'in-water', ['H', 'O'], atomic_str=True) + >>> define_group('twobody', 'in-water', ['O', 'O'], atomic_str=True) + + Example 2.2: + + >>> define_group('twobody', 'in-water', ['water', 'water']) + + The 2.1 is equivalent to 2.2. + + Example 3.1: + + >>> define_group('specie', '1', ['H']) + >>> define_group('specie', '2', ['O']) + >>> define_group('twobody', 'Hgroup', ['H', 'H'], atomic_str=True) + >>> define_group('twobody', 'Hgroup', ['H', 'O'], atomic_str=True) + >>> define_group('twobody', 'OO', ['O', 'O'], atomic_str=True) + + Example 3.2: + + >>> define_group('specie', '1', ['H']) + >>> define_group('specie', '2', ['O']) + >>> define_group('twobody', 'Hgroup', ['H', '*'], atomic_str=True) + >>> define_group('twobody', 'OO', ['O', 'O'], atomic_str=True) + + Example 3.3: + + >>> list_groups('specie', ['H', 'O']) + >>> define_group('twobody', 'Hgroup', ['H', '*']) + >>> define_group('twobody', 'OO', ['O', 'O']) + + Example 3.4: + + >>> list_groups('specie', ['H', 'O']) + >>> define_group('twobody', 'OO', ['*', '*']) + >>> define_group('twobody', 'Hgroup', ['H', '*']) + + 3.1 to 3.4 are all equivalent. + """ + + if name == '*' and group_type == 'specie': + name = 'allspecie' + element_list = ['H'] + elif name == '*': + raise ValueError("* is reserved for substitution, cannot be used " + "as a group name") + + if group_type != 'specie': + + assert len(element_list) == ParameterHelper.ndim[group_type] + + # Check all the other group_type to + exclude_list = deepcopy(self.all_types) + ide = exclude_list.index(group_type) + exclude_list.pop(ide) + + for gt in exclude_list: + if name in self.all_group_names[gt]: + raise ValueError("group name has to be unique across all types. " + f"{name} is found in type {gt}") + + if name in self.all_group_names[group_type]: + groupid = self.all_group_names[group_type].index(name) + else: + groupid = self.n[group_type] + self.all_group_names[group_type].append(name) + self.groups[group_type].append([]) + self.n[group_type] += 1 + + if group_type == 'specie': + for ele in element_list: + assert ele not in self.all_members['specie'], \ + "The element has already been defined" + self.groups['specie'][groupid].append(ele) + self.all_members['specie'].append(ele) + self.logger.debug( + f"Element {ele} will be defined as group {name}") + else: + if len(self.all_group_names['specie']) == 0: + raise RuntimeError("The atomic species have to be" + "defined in advance") + + # first translate element/group name to group name + group_name_list = [] + if atomic_str: + for ele_name in element_list: + if ele_name == "*": + gid += ["*"] + else: + for idx in range(self.n['specie']): + group_name = self.all_group_names['species'][idx] + if ele_name in self.groups['specie'][idx]: + group_name_list += [group_name] + self.logger.debug(f"Element {ele_name} is used for " + f"definition, but the whole group " + f"{group_name} is affected") + else: + group_name_list = element_list + + if "*" not in group_name_list: + + gid = [] + for ele_name in group_name_list: + gid += [self.all_group_names['specie'].index(ele_name)] + + for ele in self.all_members[group_type]: + if set(gid) == set(ele): + self.logger.debug( + f"the definition of {group_type} {ele} will be overriden") + self.groups[group_type][groupid].append(gid) + self.all_members[group_type].append(gid) + self.logger.debug( + f"{group_type} {gid} will be defined as group {name}") + if parameters is not None: + self.set_parameters(name, parameters) + else: + one_star_less = deepcopy(group_name_list) + idstar = group_name_list.index('*') + one_star_less.pop(idstar) + for sub in self.all_group_names['specie']: + self.logger.debug(f"{sub}, {one_star_less}") + self.define_group(group_type, name, + one_star_less + [sub], parameters=parameters, + atomic_str=False) + + def find_group(self, group_type, element_list, atomic_str=False): + """ find the group that contains the input pair + + Args: + group_type (str): species, twobody, threebody, cut3b, manybody + element_list (list): list of elements for a pair/triplet/coordination-pair + atomic_str (bool): whether the elements in element_list are + specified by group names or periodic table element names. + + Return: + name (str): + """ + + # remember the later command override the earlier ones + if group_type == 'specie': + if not isinstance(element_list, str): + self.logger.debug("for element, it has to be a string") + return None + name = None + for igroup in range(self.n['specie']): + gname = self.all_group_names[group_type][igroup] + allspec = self.groups[group_type][igroup] + if element_list in allspec: + name = gname + return name + self.logger.debug("cannot find the group") + return None + else: + if "*" in element_list: + self.logger.debug("* cannot be used for find") + return None + gid = [] + for ele_name in element_list: + gid += [self.all_group_names['specie'].index(ele_name)] + setlist = set(gid) + name = None + for igroup in range(self.n[group_type]): + gname = self.all_group_names[group_type][igroup] + for ele in self.groups[group_type][igroup]: + if set(gid) == set(ele): + name = gname + self.logger.debug(f"find the group {name}") + return name + + def set_parameters(self, name, parameters, opt=True): + """Set the parameters for certain group + + Args: + name (str): name of the patermeters + parameters (list): the sigma, lengthscale, and cutoff of each group. + opt (bool, list): whether to optimize the parameter or not + + The name of parameters can be the group name previously defined in + define_group or list_groups function. Aside from the group name, + ``noise``, ``cutoff_twobody``, ``cutoff_threebody``, and + ``cutoff_manybody`` are reserved for noise parmater + and universal cutoffs, while ``sigma`` and ``lengthscale`` are + reserved for universal signal variances and length scales. + + The parameter should be a list of 2-3 elements, for sigma, + lengthscale (and cutoff if the third one is defined). + + The optimization flag can be a single bool, which apply to all + parameters, or list of bools that apply to each parameter. + """ + + if name == 'noise': + self.noise = parameters + self.opt['noise'] = opt + return + elif name == 'energy_noise': + self.energy_noise = parameters + return + elif 'cutoff' in name: + self.universal[name] = parameters + return + elif name in ['sigma', 'lengthscale']: + self.universal[name] = parameters + self.opt[name] = opt + return + + if isinstance(opt, bool): + opt = [opt]*2 + + if 'cut3b' not in name: + if name in self.sigma: + self.logger.debug( + f"the sig, ls of group {name} is overriden") + self.sigma[name] = parameters[0] + self.ls[name] = parameters[1] + self.opt[name+'sig'] = opt[0] + self.opt[name+'ls'] = opt[1] + self.logger.debug(f"ParameterHelper for group {name} will be set as " + f"sig={parameters[0]} ({opt[0]}) " + f"ls={parameters[1]} ({opt[1]})") + if len(parameters) > 2: + if name in self.all_cutoff: + self.logger.debug( + f"the cutoff of group {name} is overriden") + self.all_cutoff[name] = parameters[2] + self.logger.debug(f"Cutoff for group {name} will be set as " + f"{parameters[2]}") + else: + self.all_cutoff[name] = parameters + + def set_constraints(self, name, opt): + """Set the parameters for certain group + + Args: + name (str): name of the patermeters + opt (bool, list): whether to optimize the parameter or not + + The name of parameters can be the group name previously defined in + define_group or list_groups function. Aside from the group name, + ``noise``, ``cutoff_twobody``, ``cutoff_threebody``, and + ``cutoff_manybody`` are reserved for noise parmater + and universal cutoffs, while ``sigma`` and ``lengthscale`` are + reserved for universal signal variances and length scales. + + The optimization flag can be a single bool, which apply to all + parameters under that name, or list of bools that apply to each + parameter. + """ + + if name == 'noise': + self.opt['noise'] = opt + return + + if isinstance(opt, bool): + opt = [opt, opt, opt] + + if 'cut3b' not in name: + if name in self.sigma: + self.logger.debug( + f"the sig, ls of group {name} is overriden") + self.opt[name+'sig'] = opt[0] + self.opt[name+'ls'] = opt[1] + self.logger.debug(f"ParameterHelper for group {name} will be set as " + f"sig {opt[0]} " + f"ls {opt[1]}") + + def summarize_group(self, group_type): + """Sort and combine all the previous definition to internal varialbes + + Args: + group_type (str): species, twobody, threebody, cut3b, manybody + """ + + aeg = self.all_group_names[group_type] + nspecie = self.n['specie'] + if group_type == "specie": + self.nspecie = nspecie + self.specie_mask = np.ones(118, dtype=np.int)*(nspecie-1) + for idt in range(self.nspecie): + for ele in self.groups['specie'][idt]: + atom_n = element_to_Z(ele) + self.specie_mask[atom_n] = idt + self.logger.debug(f"elemtn {ele} is defined as type {idt} with name " + f"{aeg[idt]}") + self.logger.debug( + f"All the remaining elements are left as type {idt}") + + elif group_type in self.all_group_types: + + if self.n[group_type] == 0: + self.logger.debug(f"{group_type} is not defined. Skipped") + return + + if group_type not in self.kernels and group_type in ParameterHelper.all_kernel_types: + self.kernels.append(group_type) + + self.mask[group_type] = np.ones( + nspecie**ParameterHelper.ndim[group_type], dtype=np.int)*(self.n[group_type]-1) + + self.hyps_sig[group_type] = [] + self.hyps_ls[group_type] = [] + self.hyps_opt[group_type] = [] + + for idt in range(self.n[group_type]): + name = aeg[idt] + for ele_list in self.groups[group_type][idt]: + # generate all possible permutation + perms = list(permutations(ele_list)) + for ele_list in perms: + mask_id = 0 + for ele in ele_list: + mask_id += ele + mask_id *= nspecie + mask_id = mask_id // nspecie + self.mask[group_type][mask_id] = idt + def_str = "-".join(map(str, self.groups['specie'])) + self.logger.debug(f"{group_type} {def_str} is defined as type {idt} " + f"with name {name}") + + if group_type != 'cut3b': + sig = self.sigma.get(name, -1) + opt_sig = self.opt.get(name+'sig', True) + if sig == -1: + sig = self.sigma.get(group_type, -1) + opt_sig = self.opt.get(group_type+'sig', True) + if sig == -1: + sig = self.universal.get('sigma', -1) + opt_sig = self.opt.get('sigma', True) + + ls = self.ls.get(name, -1) + opt_ls = self.opt.get(name+'ls', True) + if ls == -1: + ls = self.ls.get(group_type, -1) + opt_ls = self.opt.get(group_type+'ls', True) + if ls == -1: + ls = self.universal.get('lengthscale', -1) + opt_ls = self.opt.get('lengthscale', True) + + if sig < 0 or ls < 0: + self.logger.error(f"hyper parameters for group {name}" + " is not defined") + raise RuntimeError + self.hyps_sig[group_type] += [sig] + self.hyps_ls[group_type] += [ls] + self.hyps_opt[group_type] += [opt_sig] + self.hyps_opt[group_type] += [opt_ls] + self.logger.debug(f" using hyper-parameters of {sig:6.2g} " + f"{ls:6.2g}") + self.logger.debug( + f"All the remaining elements are left as type {idt}") + + # sort out the cutoffs + if group_type == 'cut3b': + universal_cutoff = self.universal.get('cutoff_threebody', 0) + else: + universal_cutoff = self.universal.get('cutoff_'+group_type, 0) + + allcut = [] + alldefine = True + for idt in range(self.n[group_type]): + if aeg[idt] in self.all_cutoff: + allcut += [self.all_cutoff[aeg[idt]]] + else: + alldefine = False + self.logger.info(f"{aeg[idt]} cutoff is not defined. " + "it's going to use the universal cutoff.") + + if group_type != 'threebody': + + if len(allcut) > 0: + if universal_cutoff <= 0: + universal_cutoff = np.max(allcut) + self.logger.info(f"universal cutoffs {cutstr2index[group_type]}for " + f"{group_type} is defined as zero! reset it to {universal_cutoff}") + + self.cutoff_list[group_type] = [] + for idt in range(self.n[group_type]): + self.cutoff_list[group_type] += [ + self.all_cutoff.get(aeg[idt], universal_cutoff)] + self.cutoff_list[group_type] = np.array( + self.cutoff_list[group_type], dtype=float) + + max_cutoff = np.max(self.cutoff_list[group_type]) + + # update the universal cutoff to make it higher than + if alldefine: + universal_cutoff = max_cutoff + self.logger.info(f"universal cutoff is updated to " + f"{universal_cutoff}") + elif not np.any(self.cutoff_list[group_type]-max_cutoff): + # if not all the cutoffs are defined separately + # and they are all the same value + del self.cutoff_list[group_type] + universal_cutoff = max_cutoff + if group_type == 'cut3b': + self.n['cut3b'] = 0 + self.logger.info(f"universal cutoff is updated to " + f"{universal_cutoff}") + + else: + if universal_cutoff <= 0 and len(allcut) > 0: + universal_cutoff = np.max(allcut) + self.logger.info(f"threebody universal cutoff is updated to" + f"{universal_cutoff}, but the separate definitions will" + "be ignored") + + if universal_cutoff > 0: + if group_type == 'cut_3b': + self.universal['cutoff_threebody'] = universal_cutoff + else: + self.universal['cutoff_'+group_type] = universal_cutoff + else: + self.logger.error(f"cutoffs for {group_type} is undefined") + raise RuntimeError + + else: + pass + + def as_dict(self): + """Dictionary representation of the mask. The output can be used for AtomicEnvironment + or the GaussianProcess + """ + + # sort out all the definitions and resolve conflicts + # cut3b has to be summarize before threebody + # because the universal threebody cutoff is checked + # at the end of threebody search + + self.summarize_group('specie') + for ktype in ParameterHelper.additional_groups: + self.summarize_group(ktype) + for ktype in ParameterHelper.all_kernel_types: + self.summarize_group(ktype) + + hyps_mask = {} + cutoff_dict = {} + + nspecie = self.n['specie'] + hyps_mask['nspecie'] = self.n['specie'] + if self.n['specie'] > 1: + hyps_mask['specie_mask'] = self.specie_mask + + hyps = [] + hyp_labels = [] + opt = [] + for group in self.kernels: + + hyps_mask['n'+group] = self.n[group] + hyps_mask[group+'_start'] = len(hyps) + hyps += [self.hyps_sig[group]] + hyps += [self.hyps_ls[group]] + hyps = list(np.hstack(hyps)) + opt += [self.hyps_opt[group]] + cutoff_dict[group] = self.universal['cutoff_'+group] + + if self.n[group] > 1: + hyps_mask[group+'_mask'] = self.mask[group] + # check parameters + aeg = self.all_group_names[group] + for idt in range(self.n[group]): + hyp_labels += ['Signal Var. '+aeg[idt]] + for idt in range(self.n[group]): + hyp_labels += ['Length '+group] + else: + hyp_labels += ['Signal Var. '+group] + hyp_labels += ['Length '+group] + + if group in self.cutoff_list: + hyps_mask[group+'_cutoff_list'] = self.cutoff_list[group] + + if self.n['cut3b'] >= 1: + hyps_mask['ncut3b'] = self.n[group] + hyps_mask['cut3b_mask'] = self.mask[group] + hyps_mask['threebody_cutoff_list'] = self.cutoff_list['cut3b'] + + hyps_mask['train_noise'] = self.opt['noise'] + hyps_mask['energy_noise'] = self.energy_noise + + opt += [self.opt['noise']] + hyp_labels += ['Noise Var.'] + hyps += [self.noise] + hyps = np.hstack(hyps) + opt = np.hstack(opt) + + # handle partial optimization if any constraints are defined + if not opt.all(): + nhyps = len(hyps) + hyps_mask['original_hyps'] = hyps + hyps_mask['original_labels'] = hyp_labels + mapping = [] + new_labels = [] + for i in range(nhyps): + if opt[i]: + mapping += [i] + new_labels += [hyp_labels[i]] + newhyps = hyps[mapping] + hyps_mask['map'] = np.array(mapping, dtype=np.int) + elif opt.any(): + newhyps = hyps + new_labels = hyp_labels + else: + raise RuntimeError("hyps has length zero." + "at least one component of the hyper-parameters" + "should be allowed to be optimized. \n") + + if self.n['specie'] < 2: + self.logger.debug( + "only one type of elements was defined. Please use multihyps=False") + + hyps_mask['kernels'] = self.kernels + hyps_mask['kernel_name'] = "+".join(hyps_mask['kernels']) + hyps_mask['cutoffs'] = cutoff_dict + hyps_mask['hyps'] = newhyps + hyps_mask['hyp_labels'] = new_labels + + logging.debug(str(hyps_mask)) + + return hyps_mask + + @staticmethod + def from_dict(hyps_mask, verbose=False, init_spec=[]): + """ convert dictionary mask to HM instance + This function is not tested yet + """ + + Parameters.check_instantiation(hyps_mask['hyps'], + hyps_mask['cutoffs'], + hyps_mask['kernels'], + hyps_mask) + + pm = ParameterHelper(verbose=verbose) + + nspecie = hyps_mask['nspecie'] + if nspecie > 1: + nele = len(hyps_mask['specie_mask']) + max_species = np.max(hyps_mask['specie_mask']) + specie_mask = hyps_mask['specie_mask'] + for i in range(max_species+1): + elelist = np.where(specie_mask == i)[0] + if len(elelist) > 0: + for ele in elelist: + if ele != 0: + elename = Z_to_element(ele) + if len(init_spec) > 0: + if elename in init_spec: + pm.define_group( + "specie", i, [elename]) + else: + pm.define_group("specie", i, [elename]) + else: + pm.define_group("specie", i, ['*']) + + for kernel in hyps_mask['kernels']+['cut3b']: + n = hyps_mask.get('n'+kernel, 0) + if n >= 0: + if kernel != 'cut3b': + chyps, copt = Parameters.get_component_hyps(hyps_mask, kernel, + constraint=True, noise=False) + sig = chyps[0] + ls = chyps[1] + csig = copt[0] + cls = copt[1] + cutoff = hyps_mask['cutoffs'][kernel] + pm.set_parameters('cutoff_'+kernel, cutoff) + cutoff_list = hyps_mask.get( + f'{kernel}_cutoff_list', np.ones(len(sig))*cutoff) + elif kernel == 'cut3b' and n > 1: + cutoff_list = hyps_mask['threebody_cutoff_list'] + + if n > 1: + all_specie = np.arange(nspecie) + all_comb = combinations_with_replacement( + all_specie, ParameterHelper.ndim[kernel]) + for comb in all_comb: + mask_id = 0 + for ele in comb: + mask_id += ele + mask_id *= nspecie + mask_id = mask_id // nspecie + ttype = hyps_mask[f'{kernel}_mask'][mask_id] + pm.define_group(f"{kernel}", f"{kernel}{ttype}", comb) + if kernel != 'cut3b' and kernel != 'threebody': + pm.set_parameters(f"{kernel}{ttype}", [sig[ttype], ls[ttype], cutoff_list[ttype]], + opt=[csig[ttype], cls[ttype]]) + elif kernel == 'threebody': + pm.set_parameters(f"{kernel}{ttype}", [sig[ttype], ls[ttype]], + opt=[csig[ttype], cls[ttype]]) + else: + pm.set_parameters( + f"{kernel}{ttype}", cutoff_list[ttype]) + else: + pm.define_group(kernel, kernel, [ + '*']*ParameterHelper.ndim[kernel]) + pm.set_parameters(kernel, parameters=np.hstack( + [sig, ls, cutoff]), opt=copt) + + hyps = Parameters.get_hyps(hyps_mask) + pm.set_parameters('noise', hyps[-1]) + + if 'cutoffs' in hyps_mask: + cutoffs = hyps_mask['cutoffs'] + for k in cutoffs: + pm.set_parameters(f"cutoff_{k}", cutoffs[k]) + + return pm diff --git a/requirements.txt b/requirements.txt index 17463ca91..cd3ae5a20 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,10 @@ -numpy +numpy>=1.16.0 scipy memory_profiler numba ase pymatgen nptyping +nbsphinx +IPython +pytest>=4.6 diff --git a/tests/fake_gp.py b/tests/fake_gp.py index e75514403..93e4066cd 100644 --- a/tests/fake_gp.py +++ b/tests/fake_gp.py @@ -6,164 +6,124 @@ from flare.gp import GaussianProcess from flare.env import AtomicEnvironment from flare.struc import Structure -from flare.utils.mask_helper import HyperParameterMasking +from flare.utils.parameter_helper import ParameterHelper -def get_random_structure(cell, unique_species, noa): +def get_random_structure(cell, unique_species, noa, set_seed:int = None): """Create a random test structure """ - np.random.seed(0) + if set_seed: + np.random.seed(set_seed) - positions = [] - forces = [] - species = [] - - for n in range(noa): - positions.append(np.random.uniform(-1, 1, 3)) - forces.append(np.random.uniform(-1, 1, 3)) - species.append(unique_species[np.random.randint(0, - len(unique_species))]) + forces = (np.random.random([noa, 3])-0.5)*2 + positions = np.random.random([noa, 3]) + species = [unique_species[np.random.randint(0, len(unique_species))] \ + for i in range(noa)] test_structure = Structure(cell, species, positions) return test_structure, forces -def generate_hm(nbond, ntriplet, nmb=1, constraint=False, multihyps=True): +def generate_hm(ntwobody, nthreebody, nmanybody=1, constraint=False, multihyps=True): + cutoff = 0.8 if (multihyps is False): - hyps_label = [] - if (nbond > 0): - nbond = 1 - hyps_label += ['Length', 'Signal Var.'] - if (ntriplet > 0): - ntriplet = 1 - hyps_label += ['Length', 'Signal Var.'] - hyps_label += ['Length', 'Signal Var.'] - hyps_label += ['Noise Var.'] - return random((nbond+ntriplet+1)*2+1), {'hyps_label': hyps_label}, np.ones(3, dtype=np.float)*0.8 - - pm = HyperParameterMasking(species=['H', 'He'], parameters={'cutoff2b': 0.8, - 'cutoff3b': 0.8, 'cutoffmb': 0.8, 'noise':0.05}) - pm.define_group('bond', 'b1', ['*', '*'], parameters=random(2)) - pm.define_group('triplet', 't1', ['*', '*', '*'], parameters=random(2)) - if (nmb>0): - pm.define_group('mb', 'mb1', ['*', '*'], parameters=random(2)) - if (nbond > 1): - pm.define_group('bond', 'b2', ['H', 'H'], parameters=random(2)) - if (ntriplet > 1): - pm.define_group('triplet', 't2', ['H', 'H', 'H'], parameters=random(2)) - - hm = pm.generate_dict() - hyps = hm['hyps'] - cut = hm['cutoffs'] + kernels = [] + parameters = {} + if (ntwobody > 0): + kernels += ['twobody'] + parameters['cutoff_twobody'] = cutoff + if (nthreebody > 0): + kernels += ['threebody'] + parameters['cutoff_threebody'] = cutoff + if (nmanybody > 0): + kernels += ['manybody'] + parameters['cutoff_manybody'] = cutoff + pm = ParameterHelper(kernels=kernels, random=True, + parameters=parameters) + hm = pm.as_dict() + hyps = hm['hyps'] + cut = hm['cutoffs'] + return hyps, hm, cut + + pm = ParameterHelper(species=['H', 'He'], parameters={'noise':0.05}) + if (ntwobody > 0): + pm.define_group('twobody', 'b1', ['*', '*'], parameters=random(2)) + pm.set_parameters('cutoff_twobody', cutoff) + if (nthreebody > 0): + pm.define_group('threebody', 't1', ['*', '*', '*'], parameters=random(2)) + pm.set_parameters('cutoff_threebody', cutoff) + if (nmanybody > 0): + pm.define_group('manybody', 'manybody1', ['*', '*'], parameters=random(2)) + pm.set_parameters('cutoff_manybody', cutoff) + if (ntwobody > 1): + pm.define_group('twobody', 'b2', ['H', 'H'], parameters=random(2)) + if (nthreebody > 1): + pm.define_group('threebody', 't2', ['H', 'H', 'H'], parameters=random(2)) if (constraint is False): - print(hyps) - print(hm) - print(cut) + hm = pm.as_dict() + hyps = hm['hyps'] + cut = hm['cutoffs'] return hyps, hm, cut - pm.set_parameters('b1', parameters=random(2), opt=[True, False]) - pm.set_parameters('t1', parameters=random(2), opt=[False, True]) - hm = pm.generate_dict() + pm.set_constraints('b1', opt=[True, False]) + pm.set_constraints('t1', opt=[False, True]) + hm = pm.as_dict() hyps = hm['hyps'] cut = hm['cutoffs'] return hyps, hm, cut -def get_gp(bodies, kernel_type='mc', multihyps=True, cellabc=[1, 1, 1.5]) -> GaussianProcess: - """Returns a GP instance with a two-body numba-based kernel""" - print("\nSetting up...\n") - - # params - cell = np.diag(cellabc) - unique_species = [2, 1] - cutoffs = np.array([0.8, 0.8]) - noa = 5 - - nbond = 0 - ntriplet = 0 - prefix = bodies - if ('2' in bodies or 'two' in bodies): - nbond = 1 - if ('3' in bodies or 'three' in bodies): - ntriplet = 1 - - hyps, hm, _ = generate_hm(nbond, ntriplet, nmb=0, multihyps=multihyps) - - # create test structure - test_structure, forces = get_random_structure(cell, unique_species, - noa) - energy = 3.14 - - hl = hm['hyps_label'] - if (multihyps is False): - hm = None - - # test update_db - gaussian = \ - GaussianProcess(kernel_name=f'{prefix}{kernel_type}', - hyps=hyps, - hyp_labels=hl, - cutoffs=cutoffs, multihyps=multihyps, hyps_mask=hm, - parallel=False, n_cpus=1) - gaussian.update_db(test_structure, forces, energy=energy) - gaussian.check_L_alpha() - - print('alpha:') - print(gaussian.alpha) - - return gaussian - - -def get_force_gp(bodies, kernel_type='mc', multihyps=True, cellabc=[1,1,1.5]) -> GaussianProcess: +def get_gp(bodies, kernel_type='mc', multihyps=True, cellabc=[1, 1, 1.5], + force_only=False, noa=5) -> GaussianProcess: """Returns a GP instance with a two-body numba-based kernel""" print("\nSetting up...\n") # params cell = np.diag(cellabc) - unique_species = [2, 1] - cutoffs = np.array([0.8, 0.8]) - noa = 5 + unique_species = [2, 1, 3] - nbond = 0 - ntriplet = 0 + ntwobody = 0 + nthreebody = 0 prefix = bodies if ('2' in bodies or 'two' in bodies): - nbond = 1 + ntwobody = 1 if ('3' in bodies or 'three' in bodies): - ntriplet = 1 + nthreebody = 1 - hyps, hm, _ = generate_hm(nbond, ntriplet, multihyps=multihyps) + hyps, hm, _ = generate_hm(ntwobody, nthreebody, nmanybody=0, multihyps=multihyps) + cutoffs = hm['cutoffs'] + kernels = hm['kernels'] + hl = hm['hyp_labels'] # create test structure test_structure, forces = get_random_structure(cell, unique_species, noa) energy = 3.14 - hl = hm['hyps_label'] - if (multihyps is False): - hm = None - # test update_db gaussian = \ - GaussianProcess(kernel_name=f'{prefix}{kernel_type}', + GaussianProcess(kernels=kernels, + component=kernel_type, hyps=hyps, hyp_labels=hl, - cutoffs=cutoffs, multihyps=multihyps, hyps_mask=hm, + cutoffs=cutoffs, hyps_mask=hm, parallel=False, n_cpus=1) - gaussian.update_db(test_structure, forces) + if force_only: + gaussian.update_db(test_structure, forces) + else: + gaussian.update_db(test_structure, forces, energy=energy) gaussian.check_L_alpha() - print('alpha:') - print(gaussian.alpha) + #print(gaussian.alpha) return gaussian def get_params(): parameters = {'unique_species': [2, 1], - 'cutoff': 0.8, + 'cutoffs': {'twobody': 0.8}, 'noa': 5, 'cell': np.eye(3), 'db_pts': 30} @@ -175,7 +135,7 @@ def get_tstp(hm=None) -> AtomicEnvironment: # params cell = np.eye(3) unique_species = [2, 1] - cutoffs = np.ones(3)*0.8 + cutoffs = {'twobody':0.8, 'threebody': 0.8, 'manybody': 0.8} noa = 10 test_structure_2, _ = get_random_structure(cell, unique_species, @@ -193,9 +153,9 @@ def generate_mb_envs(cutoffs, cell, delt, d1, mask=None, kern_type='mc'): [0.0, 0., 0.3], [1., 1., 0.]]) positions0[1:] += 0.1*np.random.random([4, 3]) - triplet = [1, 1, 2, 1] - np.random.shuffle(triplet) - species_1 = np.hstack([triplet, randint(1, 2)]) + threebody = [1, 1, 2, 1] + np.random.shuffle(threebody) + species_1 = np.hstack([threebody, randint(1, 2)]) if kern_type == 'sc': species_1 = np.ones(species_1.shape) return generate_mb_envs_pos(positions0, species_1, cutoffs, cell, delt, d1, mask) @@ -256,7 +216,7 @@ def generate_envs(cutoffs, delta): """ # create env 1 # perturb the x direction of atom 0 for +- delta - cell = np.eye(3)*np.max(cutoffs+0.1) + cell = np.eye(3)*(np.max(list(cutoffs.values()))+0.1) atom_1 = 0 pos_1 = np.vstack([[0, 0, 0], random([3, 3])]) pos_2 = deepcopy(pos_1) diff --git a/tests/test_OTF_vasp.py b/tests/test_OTF_vasp.py index 633efec78..823cdb3dc 100644 --- a/tests/test_OTF_vasp.py +++ b/tests/test_OTF_vasp.py @@ -1,6 +1,5 @@ import pytest import os -import sys import numpy as np from flare.otf import OTF from flare.gp import GaussianProcess @@ -13,6 +12,11 @@ # test otf runs # ------------------------------------------------------ +@pytest.mark.skipif(not os.environ.get('VASP_COMMAND', + False), reason='VASP_COMMAND not found ' + 'in environment: Please install VASP ' + ' and set the VASP_COMMAND env. ' + 'variable to point to cp2k.popt') def test_otf_h2(): """ :return: @@ -22,7 +26,7 @@ def test_otf_h2(): vasp_input = './POSCAR' dt = 0.0001 number_of_steps = 5 - cutoffs = np.array([5]) + cutoffs = {'twobody':5} dft_loc = 'cp ./test_files/test_vasprun_h2.xml vasprun.xml' std_tolerance_factor = -0.1 @@ -36,12 +40,16 @@ def test_otf_h2(): hyp_labels=hyp_labels, maxiter=50) - otf = OTF(vasp_input, dt, number_of_steps, gp, dft_loc, - std_tolerance_factor, init_atoms=[0], - calculate_energy=True, max_atoms_added=1, - n_cpus=1, force_source='vasp', + otf = OTF(dt=dt, number_of_steps=number_of_steps, + gp=gp, calculate_energy=True, + std_tolerance_factor=std_tolerance_factor, + init_atoms=[0], + output_name='h2_otf_vasp', + max_atoms_added=1, + force_source='vasp', + dft_input=vasp_input, dft_loc=dft_loc, dft_output="vasprun.xml", - output_name='h2_otf_vasp') + n_cpus=1) otf.run() diff --git a/tests/test_ase_otf.py b/tests/test_ase_otf.py index 0ae47f6fd..865122e2a 100644 --- a/tests/test_ase_otf.py +++ b/tests/test_ase_otf.py @@ -5,28 +5,44 @@ from flare import otf, kernels from flare.gp import GaussianProcess -from flare.mgp.mgp import MappedGaussianProcess +from flare.mgp import MappedGaussianProcess from flare.ase.calculator import FLARE_Calculator -from flare.ase.otf_md import otf_md -from flare.ase.logger import OTFLogger +from flare.ase.otf import ASE_OTF +from flare.utils.parameter_helper import ParameterHelper +# from flare.ase.logger import OTFLogger from ase import units from ase.md.velocitydistribution import (MaxwellBoltzmannDistribution, Stationary, ZeroRotation) -from ase.spacegroup import crystal -from ase.calculators.espresso import Espresso md_list = ['VelocityVerlet', 'NVTBerendsen', 'NPTBerendsen', 'NPT', 'Langevin'] +@pytest.fixture(scope='module') +def md_params(): + + md_dict = {'temperature': 500} + for md_engine in md_list: + if md_engine == 'VelocityVerlet': + md_dict[md_engine] = {} + else: + md_dict[md_engine] = {'temperature': md_dict['temperature']} + + md_dict['NVTBerendsen'].update({'taut': 0.5e3 * units.fs}) + md_dict['NPT'].update({'externalstress': 0, 'ttime': 25, 'pfactor': 3375}) + md_dict['Langevin'].update({'friction': 0.02}) + + yield md_dict + del md_dict + + @pytest.fixture(scope='module') def super_cell(): - # create primitive cell based on materials project - # url: https://materialsproject.org/materials/mp-22915/ + from ase.spacegroup import crystal a = 3.855 alpha = 90 - atoms = crystal(['H', 'He'], # Ag, I + atoms = crystal(['H', 'He'], basis=[(0, 0, 0), (0.5, 0.5, 0.5)], size=(2, 1, 1), cellpar=[a, a, a, alpha, alpha, alpha]) @@ -46,122 +62,76 @@ def flare_calc(): for md_engine in md_list: # ---------- create gaussian process model ------------------- - gp_model = GaussianProcess(kernel_name='2+3_mc', - hyps=[0.1, 1., 0.001, 1, 0.06], - cutoffs=(5.0, 5.0), - hyp_labels=['sig2', 'ls2', 'sig3', - 'ls3', 'noise'], - opt_algorithm='BFGS', - par=False) + + # set up GP hyperparameters + kernels = ['twobody', 'threebody'] # use 2+3 body kernel + parameters = {'cutoff_twobody': 5.0, + 'cutoff_threebody': 3.5} + pm = ParameterHelper( + kernels = kernels, + random = True, + parameters=parameters + ) + + hm = pm.as_dict() + hyps = hm['hyps'] + cut = hm['cutoffs'] + print('hyps', hyps) + + gp_model = GaussianProcess( + kernels = kernels, + component = 'sc', # single-component. For multi-comp, use 'mc' + hyps = hyps, + cutoffs = cut, + hyp_labels = ['sig2','ls2','sig3','ls3','noise'], + opt_algorithm = 'L-BFGS-B', + n_cpus = 1 + ) # ----------- create mapped gaussian process ------------------ - struc_params = {'species': [1, 2], - 'cube_lat': np.eye(3) * 100, - 'mass_dict': {'0': 2, '1': 4}} - - # grid parameters - lower_cut = 2.5 - two_cut, three_cut = gp_model.cutoffs - grid_num_2 = 8 - grid_num_3 = 8 - grid_params = {'bounds_2': [[lower_cut], [two_cut]], - 'bounds_3': [[lower_cut, lower_cut, -1], - [three_cut, three_cut, 1]], - 'grid_num_2': grid_num_2, - 'grid_num_3': [grid_num_3, grid_num_3, grid_num_3], - 'svd_rank_2': 0, - 'svd_rank_3': 0, - 'bodies': [2, 3], - 'load_grid': None, - 'update': False} - - mgp_model = MappedGaussianProcess(grid_params, - struc_params, - map_force=True, - GP=gp_model, - mean_only=False, - container_only=False, - lmp_file_name='lmp.mgp', - n_cpus=1) + grid_params = {'twobody': {'grid_num': [64]}, + 'threebody': {'grid_num': [16, 16, 16]}} + + mgp_model = MappedGaussianProcess(grid_params = grid_params, + unique_species = [1, 2], + n_cpus = 1, + map_force = False, + mean_only = False) # ------------ create ASE's flare calculator ----------------------- - flare_calculator = FLARE_Calculator(gp_model, mgp_model, + flare_calculator = FLARE_Calculator(gp_model, mgp_model=mgp_model, par=True, use_mapping=True) - flare_calc_dict[md_engine] = flare_calculator print(md_engine) yield flare_calc_dict del flare_calc_dict -@pytest.mark.skipif(not os.environ.get('PWSCF_COMMAND', - False), reason='PWSCF_COMMAND not found ' - 'in environment: Please install Quantum ' - 'ESPRESSO and set the PWSCF_COMMAND env. ' - 'variable to point to pw.x.') @pytest.fixture(scope='module') def qe_calc(): - # set up executable - label = 'scf' - input_file = label+'.pwi' - output_file = label+'.pwo' - no_cpus = 1 - pw = os.environ.get('PWSCF_COMMAND') - os.environ['ASE_ESPRESSO_COMMAND'] = f'{pw} < {input_file} > {output_file}' - - # set up input parameters - input_data = {'control': {'prefix': label, - 'pseudo_dir': 'test_files/pseudos/', - 'outdir': './out', - 'calculation': 'scf'}, - 'system': {'ibrav': 0, - 'ecutwfc': 20, - 'ecutrho': 40, - 'smearing': 'gauss', - 'degauss': 0.02, - 'occupations': 'smearing'}, - 'electrons': {'conv_thr': 1.0e-02, - 'electron_maxstep': 100, - 'mixing_beta': 0.7}} - - # pseudo-potentials - ion_pseudo = {'H': 'H.pbe-kjpaw.UPF', - 'He': 'He.pbe-kjpaw_psl.1.0.0.UPF'} - - # create ASE calculator - dft_calculator = Espresso(pseudopotentials=ion_pseudo, label=label, - tstress=True, tprnfor=True, nosym=True, - input_data=input_data, kpts=(1,1,1)) + + from ase.calculators.lj import LennardJones + dft_calculator = LennardJones() yield dft_calculator del dft_calculator -@pytest.mark.skipif(not os.environ.get('PWSCF_COMMAND', - False), reason='PWSCF_COMMAND not found ' - 'in environment: Please install Quantum ' - 'ESPRESSO and set the PWSCF_COMMAND env. ' - 'variable to point to pw.x.') @pytest.mark.parametrize('md_engine', md_list) -def test_otf_md(md_engine, super_cell, flare_calc, qe_calc): +def test_otf_md(md_engine, md_params, super_cell, flare_calc, qe_calc): np.random.seed(12345) flare_calculator = flare_calc[md_engine] # set up OTF MD engine - md_params = {'timestep': 1 * units.fs, 'trajectory': None, 'dt': 1* - units.fs, - 'externalstress': 0, 'ttime': 25, 'pfactor': 3375, - 'mask': None, 'temperature': 500, 'taut': 1, 'taup': 1, - 'pressure': 0, 'compressibility': 0, 'fixcm': 1, - 'friction': 0.02} - - otf_params = {'dft_calc': qe_calc, - 'init_atoms': [0, 1, 2, 3], + otf_params = {'init_atoms': [0, 1, 2, 3], + 'output_name': md_engine, 'std_tolerance_factor': 2, 'max_atoms_added' : len(super_cell.positions), - 'freeze_hyps': 10, - 'use_mapping': flare_calculator.use_mapping} + 'freeze_hyps': 10} +# 'use_mapping': flare_calculator.use_mapping} + + md_kwargs = md_params[md_engine] # intialize velocity temperature = md_params['temperature'] @@ -170,16 +140,23 @@ def test_otf_md(md_engine, super_cell, flare_calc, qe_calc): ZeroRotation(super_cell) # zero angular momentum super_cell.set_calculator(flare_calculator) - test_otf = otf_md(md_engine, super_cell, md_params, otf_params) + test_otf = ASE_OTF(super_cell, + timestep = 1 * units.fs, + number_of_steps = 3, + dft_calc = qe_calc, + md_engine = md_engine, + md_kwargs = md_kwargs, + **otf_params) + + # TODO: test if mgp matches gp + # TODO: see if there's difference between MD timestep & OTF timestep # set up logger - otf_logger = OTFLogger(test_otf, super_cell, - logfile=md_engine+'.log', mode="w", data_in_logfile=True) - test_otf.attach(otf_logger, interval=1) +# otf_logger = OTFLogger(test_otf, super_cell, +# logfile=md_engine+'.log', mode="w", data_in_logfile=True) +# test_otf.attach(otf_logger, interval=1) - # run otf - number_of_steps = 3 - test_otf.otf_run(number_of_steps) + test_otf.run() for f in glob.glob("scf.pw*"): os.remove(f) @@ -189,7 +166,7 @@ def test_otf_md(md_engine, super_cell, flare_calc, qe_calc): shutil.rmtree(f) for f in os.listdir("./"): - if f in [f'{md_engine}.log', 'lmp.mgp']: + if f in [f'{md_engine}.out', f'{md_engine}-hyps.dat', 'lmp.mgp']: os.remove(f) if f in ['out', 'otf_data']: shutil.rmtree(f) diff --git a/tests/test_cp2k_util.py b/tests/test_cp2k_util.py index 661656fe4..a484272dd 100644 --- a/tests/test_cp2k_util.py +++ b/tests/test_cp2k_util.py @@ -118,8 +118,8 @@ def test_cp2k_input_edit(): positions, species, cell, masses = parse_dft_input(newfilename) - assert np.equal(positions[0], structure.positions[0]).all() - assert np.equal(structure.vec1, cell[0, :]).all() + assert np.isclose(positions[0], structure.positions[0]).all() + assert np.isclose(structure.vec1, cell[0, :]).all() remove(newfilename) cleanup() diff --git a/tests/test_env.py b/tests/test_env.py index 96a7f1a8c..801533e86 100644 --- a/tests/test_env.py +++ b/tests/test_env.py @@ -6,14 +6,14 @@ np.random.seed(0) -cutoff_mask_list = [(True, np.array([1]), [10]), - (False, np.array([1]), [16]), - (False, np.array([1, 0.05]), [16, 0]), - (False, np.array([1, 0.8]), [16, 1]), - (False, np.array([1, 0.9]), [16, 21]), - (True, np.array([1, 0.8]), [16, 9]), - (True, np.array([1, 0.05, 0.4]), [16, 0]), - (False, np.array([1, 0.05, 0.4]), [16, 0])] +cutoff_mask_list = [# (True, np.array([1]), [10]), + (False, {'twobody':1}, [16]), + (False, {'twobody':1, 'threebody':0.05}, [16, 0]), + (False, {'twobody':1, 'threebody':0.8}, [16, 1]), + (False, {'twobody':1, 'threebody':0.9}, [16, 21]), + (True, {'twobody':1, 'threebody':0.8}, [16, 9]), + (True, {'twobody':1, 'threebody':0.05, 'manybody':0.4}, [16, 0]), + (False, {'twobody':1, 'threebody':0.05, 'manybody':0.4}, [16, 0])] @pytest.fixture(scope='module') @@ -22,7 +22,7 @@ def structure() -> Structure: Returns a GP instance with a two-body numba-based kernel """ - # list of all bonds and triplets can be found in test_files/test_env_list + # list of all twobodys and threebodys can be found in test_files/test_env_list cell = np.eye(3) species = [1, 2, 3, 1] positions = np.array([[0, 0, 0], [0.5, 0.5, 0.5], @@ -93,10 +93,10 @@ def generate_mask(cutoff): # (1, 1) uses 0.5 cutoff, (1, 2) (1, 3) (2, 3) use 0.9 cutoff mask = {'nspecie': 2, 'specie_mask': np.ones(118, dtype=int)} mask['specie_mask'][1] = 0 - mask['cutoff_2b'] = np.array([0.5, 0.9]) - mask['nbond'] = 2 - mask['bond_mask'] = np.ones(4, dtype=int) - mask['bond_mask'][0] = 0 + mask['twobody_cutoff_list'] = np.array([0.5, 0.9]) + mask['ntwobody'] = 2 + mask['twobody_mask'] = np.ones(4, dtype=int) + mask['twobody_mask'][0] = 0 elif (ncutoff == 2): # the 3b mask is the same structure as 2b @@ -120,7 +120,7 @@ def generate_mask(cutoff): mask = {'nspecie': nspecie, 'specie_mask': specie_mask, - 'cutoff_3b': np.array([0.5, 0.9, 0.8, 0.9, 0.05]), + 'threebody_cutoff_list': np.array([0.5, 0.9, 0.8, 0.9, 0.05]), 'ncut3b': ncut3b, 'cut3b_mask': tmask} @@ -128,10 +128,11 @@ def generate_mask(cutoff): # (1, 1) uses 0.5 cutoff, (1, 2) (1, 3) (2, 3) use 0.9 cutoff mask = {'nspecie': 2, 'specie_mask': np.ones(118, dtype=int)} mask['specie_mask'][1] = 0 - mask['cutoff_mb'] = np.array([0.5, 0.9]) - mask['nmb'] = 2 - mask['mb_mask'] = np.ones(4, dtype=int) - mask['mb_mask'][0] = 0 + mask['manybody_cutoff_list'] = np.array([0.5, 0.9]) + mask['nmanybody'] = 2 + mask['manybody_mask'] = np.ones(4, dtype=int) + mask['manybody_mask'][0] = 0 + mask['cutoffs'] = cutoff return mask diff --git a/tests/test_files/VelocityVerlet.log b/tests/test_files/VelocityVerlet.log deleted file mode 100644 index e3c84fb56..000000000 --- a/tests/test_files/VelocityVerlet.log +++ /dev/null @@ -1,89 +0,0 @@ -2019-10-16 22:17:58.842938 -number of cpu cores: -cutoffs: [5. 5.] -kernel: two_plus_three_body_mc -number of hyperparameters: 5 -hyperparameters: [0.1 1. 0.001 1. 0.06 ] -hyperparameter optimization algorithm: BFGS -uncertainty tolerance: 1 times noise -timestep (ps): 0.001 -number of frames: 0 -number of atoms: 4 -system species: ['Ag', 'I', 'Ag', 'I'] -periodic cell: -[[7.71 0. 0. ] - [0. 3.855 0. ] - [0. 0. 3.855]] -previous positions (A): -Ag 0.489004 0.419311 -0.107638 -I 1.989324 1.740464 1.813266 -Ag 4.344619 -0.112703 -0.166401 -I 5.740768 1.822571 2.234367 --------------------------------------------------- -*-Frame: 0 -Simulation time: 0.0 ps -Type Positions DFT Forces Uncertainties Velocities -Ag 0.485901 0.410331 -0.141231 -0.936668 -0.620960 0.498037 0.000000 0.000000 0.000000 -0.003103 -0.008980 -0.033593 -I 1.965587 1.753340 1.848765 0.723168 0.308194 -0.394375 0.000000 0.000000 0.000000 -0.023737 0.012876 0.035499 -Ag 4.351446 -0.114926 -0.154143 -0.438520 0.736847 0.025200 0.000000 0.000000 0.000000 0.006827 -0.002223 0.012258 -I 5.761340 1.819218 2.217003 0.652020 -0.424080 -0.128863 0.000000 0.000000 0.000000 0.020572 -0.003354 -0.017364 - -temperature: 482.6435569534418 K -kinetic energy: 0.24954593792383903 eV -wall time from start: 6.3175413608551025 s - -Adding atom [0] to the training set. -Uncertainty: [0. 0. 0.]. -Adding atom [1] to the training set. -Uncertainty: [0. 0. 0.]. -Adding atom [2] to the training set. -Uncertainty: [0. 0. 0.]. -Adding atom [3] to the training set. -Uncertainty: [0. 0. 0.]. - -GP hyperparameters: -Hyp0 : sig2 = -0.19398990737605573 -Hyp1 : ls2 = 0.908822798289524 -Hyp2 : sig3 = 0.9575224391705098 -Hyp3 : ls3 = 21.169915974370866 -Hyp4 : noise = 0.4253621583432521 -likelihood: -19.202850165167334 -likelihood gradient: [ 4.09655043 1.51271249 -4.29716492 0.10055218 -5.18645278] -wall time from start: 9.89 s --------------------------------------------------- --Frame: 1 -Simulation time: 0.001 ps -Type Positions GP Forces Uncertainties Velocities -Ag 0.478457 0.398473 -0.172516 -0.735715 -0.969873 0.343320 0.384351 0.313210 0.167310 -0.010855 -0.016354 -0.029693 -I 1.944699 1.767430 1.882710 1.120653 0.373864 -0.478936 0.561959 0.213593 0.298612 -0.016472 0.015563 0.032058 -Ag 4.356240 -0.113734 -0.141769 -0.274217 0.462411 0.317277 0.432477 0.311834 0.194696 0.003523 0.003336 0.013846 -I 5.784480 1.814193 2.199132 -0.110721 0.133598 -0.181661 0.491638 0.225157 0.257373 0.022704 -0.004498 -0.018587 - -temperature: 451.9032725218858 K -kinetic energy: 0.23365198678743504 eV -wall time from start: 11.301556587219238 s --------------------------------------------------- --Frame: 2 -Simulation time: 0.002 ps -Type Positions GP Forces Uncertainties Velocities -Ag 0.464192 0.377624 -0.200617 -0.756198 -1.040380 0.490954 0.596210 0.322811 0.206975 -0.017770 -0.025672 -0.025826 -I 1.932642 1.784466 1.912881 1.523414 0.429261 -0.512766 0.917508 0.222169 0.289644 -0.006055 0.018727 0.028151 -Ag 4.358492 -0.108254 -0.126452 -0.215417 0.447744 0.206183 0.657219 0.300057 0.188372 0.001254 0.007555 0.016272 -I 5.806748 1.810221 2.179828 -0.551798 0.163375 -0.184371 0.850820 0.206953 0.227729 0.020094 -0.003328 -0.020029 - -temperature: 449.9356800301974 K -kinetic energy: 0.23263466311924058 eV -wall time from start: 11.358051061630249 s --------------------------------------------------- --Frame: 3 -Simulation time: 0.003 ps -Type Positions GP Forces Uncertainties Velocities -Ag 0.442916 0.347130 -0.224168 -0.786557 -1.034843 0.597918 0.905371 0.324346 0.223627 -0.024921 -0.035291 -0.020779 -I 1.932589 1.804885 1.939011 1.702739 0.452557 -0.534727 1.107930 0.223327 0.269765 0.006656 0.022202 0.024024 -Ag 4.358747 -0.098624 -0.109225 -0.140954 0.418403 0.116154 0.963443 0.281925 0.190087 -0.000398 0.011570 0.017766 -I 5.824668 1.807537 2.159073 -0.775229 0.163882 -0.179345 1.044492 0.188329 0.205767 0.014865 -0.002039 -0.021463 - -temperature: 507.5620559823112 K -kinetic energy: 0.2624297941822027 eV -wall time from start: 11.41139554977417 s -Run complete. diff --git a/tests/test_gp.py b/tests/test_gp.py index 4c3ed8cae..97da5baf4 100644 --- a/tests/test_gp.py +++ b/tests/test_gp.py @@ -9,6 +9,7 @@ from scipy.optimize import OptimizeResult import flare +from flare.predict import predict_on_structure from flare.gp import GaussianProcess from flare.env import AtomicEnvironment from flare.struc import Structure @@ -18,7 +19,6 @@ from .fake_gp import generate_hm, get_tstp, get_random_structure from copy import deepcopy - multihyps_list = [True, False] @@ -28,20 +28,18 @@ def all_gps() -> GaussianProcess: gp_dict = {True: None, False: None} for multihyps in multihyps_list: + hyps, hm, cutoffs = generate_hm(1, 1, multihyps=multihyps) - hl = hm['hyps_label'] - if (multihyps is False): - hm = None + hl = hm['hyp_labels'] # test update_db - gpname = '2+3+mb_mc' gp_dict[multihyps] = \ - GaussianProcess(kernel_name=gpname, + GaussianProcess(kernels=hm['kernels'], hyps=hyps, hyp_labels=hl, cutoffs=cutoffs, - multihyps=multihyps, hyps_mask=hm, + hyps_mask=hm, parallel=False, n_cpus=1) test_structure, forces = \ @@ -67,6 +65,7 @@ def params(): @pytest.fixture(scope='module') def validation_env() -> AtomicEnvironment: + np.random.seed(0) test_pt = get_tstp(None) yield test_pt del test_pt @@ -111,13 +110,14 @@ def test_train(self, all_gps, params, par, n_cpus, multihyps): # train gp test_gp.hyps = np.ones(len(test_gp.hyps)) - hyp = list(test_gp.hyps) + hyps = tuple(test_gp.hyps) + test_gp.train() - hyp_post = list(test_gp.hyps) + hyp_post = tuple(test_gp.hyps) # check if hyperparams have been updated - assert (hyp != hyp_post) + assert (hyps != hyp_post) def test_train_failure(self, all_gps, params, mocker): """ @@ -171,8 +171,12 @@ def test_constrained_optimization_simple(self, all_gps): hyps, hm, cutoffs = generate_hm(1, 1, constraint=True, multihyps=True) test_gp.hyps_mask = hm - test_gp.hyp_labels = hm['hyps_label'] + test_gp.hyp_labels = hm['hyp_labels'] test_gp.hyps = hyps + test_gp.update_kernel(hm['kernel_name'], "mc", hm) + test_gp.set_L_alpha() + + hyp = list(test_gp.hyps) # Check that the hyperparameters were updated test_gp.maxiter = 1 @@ -183,52 +187,52 @@ def test_constrained_optimization_simple(self, all_gps): class TestAlgebra(): - @pytest.mark.parametrize('par, per_atom_par, n_cpus', - [(False, False, 1), - (True, True, 2), - (True, False, 2)]) - @pytest.mark.parametrize('multihyps', multihyps_list) - def test_predict(self, all_gps, validation_env, - par, per_atom_par, n_cpus, multihyps): - test_gp = all_gps[multihyps] - test_gp.parallel = par - test_gp.per_atom_par = per_atom_par - pred = test_gp.predict(x_t=validation_env, d=1) - assert (len(pred) == 2) - assert (isinstance(pred[0], float)) - assert (isinstance(pred[1], float)) - - @pytest.mark.parametrize('par, n_cpus', [(True, 2), - (False, 1)]) - @pytest.mark.parametrize('multihyps', multihyps_list) - def test_set_L_alpha(self, all_gps, params, par, n_cpus, multihyps): - test_gp = all_gps[multihyps] - test_gp.parallel = par - test_gp.n_cpus = n_cpus - test_gp.set_L_alpha() - - @pytest.mark.parametrize('par, n_cpus', [(True, 2), - (False, 1)]) - @pytest.mark.parametrize('multihyps', multihyps_list) - def test_update_L_alpha(self, all_gps, params, par, n_cpus, multihyps): - # set up gp model - test_gp = all_gps[multihyps] - test_gp.parallel = par - test_gp.n_cpus = n_cpus - - test_structure, forces = \ - get_random_structure(params['cell'], params['unique_species'], 2) - energy = 3.14 - test_gp.check_L_alpha() - test_gp.update_db(test_structure, forces, energy=energy) - test_gp.update_L_alpha() - - # compare results with set_L_alpha - ky_mat_from_update = np.copy(test_gp.ky_mat) - test_gp.set_L_alpha() - ky_mat_from_set = np.copy(test_gp.ky_mat) - - assert (np.all(np.absolute(ky_mat_from_update - ky_mat_from_set)) < 1e-6) + @pytest.mark.parametrize('par, per_atom_par, n_cpus', + [(False, False, 1), + (True, True, 2), + (True, False, 2)]) + @pytest.mark.parametrize('multihyps', multihyps_list) + def test_predict(self, all_gps, validation_env, + par, per_atom_par, n_cpus, multihyps): + test_gp = all_gps[multihyps] + test_gp.parallel = par + test_gp.per_atom_par = per_atom_par + pred = test_gp.predict(x_t=validation_env, d=1) + assert (len(pred) == 2) + assert (isinstance(pred[0], float)) + assert (isinstance(pred[1], float)) + + @pytest.mark.parametrize('par, n_cpus', [(True, 2), + (False, 1)]) + @pytest.mark.parametrize('multihyps', multihyps_list) + def test_set_L_alpha(self, all_gps, params, par, n_cpus, multihyps): + test_gp = all_gps[multihyps] + test_gp.parallel = par + test_gp.n_cpus = n_cpus + test_gp.set_L_alpha() + + @pytest.mark.parametrize('par, n_cpus', [(True, 2), + (False, 1)]) + @pytest.mark.parametrize('multihyps', multihyps_list) + def test_update_L_alpha(self, all_gps, params, par, n_cpus, multihyps): + # set up gp model + test_gp = all_gps[multihyps] + test_gp.parallel = par + test_gp.n_cpus = n_cpus + + test_structure, forces = \ + get_random_structure(params['cell'], params['unique_species'], 2) + energy = 3.14 + test_gp.check_L_alpha() + test_gp.update_db(test_structure, forces, energy=energy) + test_gp.update_L_alpha() + + # compare results with set_L_alpha + ky_mat_from_update = np.copy(test_gp.ky_mat) + test_gp.set_L_alpha() + ky_mat_from_set = np.copy(test_gp.ky_mat) + + assert (np.all(np.absolute(ky_mat_from_update - ky_mat_from_set)) < 1e-6) class TestIO(): @@ -236,17 +240,15 @@ class TestIO(): def test_representation_method(self, all_gps, multihyps): test_gp = all_gps[multihyps] the_str = str(test_gp) + print(the_str) assert 'GaussianProcess Object' in the_str - if (multihyps): - assert 'Kernel: two_three_many_body_mc' in the_str - else: - assert 'Kernel: two_plus_three_plus_many_body_mc' in the_str - assert 'Cutoffs: [0.8 0.8 0.8]' in the_str + assert 'Kernel: [\'twobody\', \'threebody\', \'manybody\']' in the_str + assert 'Cutoffs: {\'twobody\': 0.8, \'threebody\': 0.8, \'manybody\': 0.8}' in the_str assert 'Model Likelihood: ' in the_str if not multihyps: - assert 'Length: ' in the_str - assert 'Signal Var.: ' in the_str - assert "Noise Var.: " in the_str + assert 'Length ' in the_str + assert 'Signal Var. ' in the_str + assert "Noise Var." in the_str @pytest.mark.parametrize('multihyps', multihyps_list) def test_serialization_method(self, all_gps, validation_env, multihyps): @@ -268,6 +270,7 @@ def test_serialization_method(self, all_gps, validation_env, multihyps): for d in [1, 2, 3]: assert np.all(test_gp.predict(x_t=validation_env, d=d) == new_gp.predict(x_t=validation_env, d=d)) + assert new_gp.training_data is not test_gp.training_data @pytest.mark.parametrize('multihyps', multihyps_list) def test_load_and_reload(self, all_gps, validation_env, multihyps): @@ -281,7 +284,11 @@ def test_load_and_reload(self, all_gps, validation_env, multihyps): for d in [1, 2, 3]: assert np.all(test_gp.predict(x_t=validation_env, d=d) == new_gp.predict(x_t=validation_env, d=d)) - os.remove('test_gp_write.pickle') + + try: + os.remove('test_gp_write.pickle') + except: + pass test_gp.write_model('test_gp_write', 'json') @@ -315,6 +322,7 @@ def test_load_reload_huge(self, all_gps): new_gp = GaussianProcess.from_file('test_gp_write.json') assert np.array_equal(prev_ky_mat, new_gp.ky_mat) assert np.array_equal(prev_l_mat, new_gp.l_mat) + assert new_gp.training_data is not test_gp.training_data os.remove('test_gp_write.json') @@ -332,9 +340,7 @@ def dumpcompare(obj1, obj2): assert k1 == k2, f"key {k1} is not the same as {k2}" - print(k1) - - if (k1 != "name"): + if 'name' not in k1: if (obj1[k1] is None): continue else: @@ -389,6 +395,76 @@ def test_training_statistics(): test_structure.coded_species)) +def test_remove_force_data(): + """ + Train a GP on one fake structure. Store forces from prediction. + Add a new fake structure and ensure predictions change; then remove + the structure and ensure predictions go back to normal. + :return: + """ + + test_structure, forces = get_random_structure(5.0*np.eye(3), + ['H', 'Be'], + 5) + + test_structure_2, forces_2 = get_random_structure(5.0*np.eye(3), + ['H', 'Be'], + 5) + + gp = GaussianProcess(kernels=['twobody'], cutoffs={'twobody':0.8}) + + gp.update_db(test_structure, forces) + + with raises(ValueError): + gp.remove_force_data(1000000) + + init_forces, init_stds = predict_on_structure(test_structure, gp, + write_to_structure=False) + init_forces_2, init_stds_2 = predict_on_structure(test_structure_2, gp, + write_to_structure=False) + + # Alternate adding in the entire structure and adding in only one atom. + for custom_range in [None, [0]]: + + # Add in data and ensure the predictions change in reponse + gp.update_db(test_structure_2, forces_2, custom_range=custom_range) + + new_forces, new_stds = predict_on_structure(test_structure, gp, + write_to_structure=False) + + new_forces_2, new_stds_2 = predict_on_structure(test_structure_2, gp, + write_to_structure=False) + + assert not np.array_equal(init_forces, new_forces) + assert not np.array_equal(init_forces_2, new_forces_2) + assert not np.array_equal(init_stds, new_stds) + assert not np.array_equal(init_stds_2, new_stds_2) + + # Remove that data and test to see that the predictions revert to + # what they were previously + if custom_range == [0]: + popped_strucs, popped_forces = gp.remove_force_data(5) + else: + popped_strucs, popped_forces = gp.remove_force_data([5, 6, 7, 8, + 9]) + + for i in range(len(popped_forces)): + assert np.array_equal(popped_forces[i],forces_2[i]) + assert np.array_equal(popped_strucs[i].structure.positions, + test_structure_2.positions) + + final_forces, final_stds = predict_on_structure(test_structure, gp, + write_to_structure=False) + final_forces_2, final_stds_2 = predict_on_structure(test_structure_2, gp, + write_to_structure=False) + + assert np.array_equal(init_forces, final_forces) + assert np.array_equal(init_stds, final_stds) + + assert np.array_equal(init_forces_2, final_forces_2) + assert np.array_equal(init_stds_2, final_stds_2) + + class TestHelper(): def test_adjust_cutoffs(self, all_gps): @@ -399,11 +475,15 @@ def test_adjust_cutoffs(self, all_gps): # testing on the predictions made, just that the cutoffs in the # atomic environments are correctly re-created - old_cutoffs = np.copy(test_gp.cutoffs) - - test_gp.adjust_cutoffs(np.array(test_gp.cutoffs) + .5, train=False) + old_cutoffs = {} + new_cutoffs = {} + for k in test_gp.cutoffs: + old_cutoffs[k] = test_gp.cutoffs[k] + new_cutoffs[k] = 0.5+old_cutoffs[k] + test_gp.hyps_mask['cutoffs']=new_cutoffs + test_gp.adjust_cutoffs(new_cutoffs, train=False, new_hyps_mask=test_gp.hyps_mask) - assert np.array_equal(test_gp.cutoffs, old_cutoffs + .5) + assert np.array_equal(list(test_gp.cutoffs.values()), np.array(list(old_cutoffs.values()), dtype=float) + .5) for env in test_gp.training_data: - assert np.array_equal(env.cutoffs, test_gp.cutoffs) + assert env.cutoffs == test_gp.cutoffs diff --git a/tests/test_gp_algebra.py b/tests/test_gp_algebra.py index d87510c42..c0432347c 100644 --- a/tests/test_gp_algebra.py +++ b/tests/test_gp_algebra.py @@ -4,15 +4,9 @@ import flare.gp_algebra from flare.env import AtomicEnvironment +from flare.kernels.utils import str_to_kernel_set from flare.struc import Structure -from flare.kernels import mc_sephyps -from flare.kernels.mc_simple import two_plus_three_body_mc, \ - two_plus_three_body_mc_grad, two_plus_three_mc_en,\ - two_plus_three_mc_force_en -from flare.kernels.mc_sephyps import two_plus_three_body_mc \ - as two_plus_three_body_mc_multi -from flare.kernels.mc_sephyps import two_plus_three_body_mc_grad \ - as two_plus_three_body_mc_grad_multi +from flare.utils.parameter_helper import ParameterHelper from flare.gp_algebra import get_like_grad_from_mats, \ get_force_block, get_force_energy_block, \ @@ -22,7 +16,7 @@ get_Ky_mat, update_force_block, update_energy_block, \ update_force_energy_block -from .fake_gp import get_tstp +from tests.fake_gp import get_tstp @pytest.fixture(scope='module') @@ -37,7 +31,7 @@ def params(): def get_random_training_set(nenv, nstruc): """Create a random training_set array with parameters And generate four different kinds of hyperparameter sets: - * multi hypper parameters with two bond type and two triplet type + * multi hypper parameters with two twobody type and two threebody type * constrained optimization, with noise parameter optimized * constrained optimization, without noise parameter optimized * simple hyper parameters without multihyps set up @@ -45,61 +39,49 @@ def get_random_training_set(nenv, nstruc): np.random.seed(0) - cutoffs = np.array([0.8, 0.8]) - hyps = np.ones(5, dtype=float) - kernel = (two_plus_three_body_mc, two_plus_three_body_mc_grad, - two_plus_three_mc_en, two_plus_three_mc_force_en) - kernel_m = \ - (two_plus_three_body_mc_multi, two_plus_three_body_mc_grad_multi, - mc_sephyps.two_plus_three_mc_en, - mc_sephyps.two_plus_three_mc_force_en) + cutoffs = {'twobody':0.8, 'threebody':0.8} + parameters={'cutoff_twobody': 0.8, + 'cutoff_threebody': 0.8, + 'noise': 0.05} # 9 different hyper-parameters - hyps_mask1 = {'nspecie': 2, - 'specie_mask': np.zeros(118, dtype=int), - 'nbond': 2, - 'bond_mask': np.array([0, 1, 1, 1]), - 'triplet_mask': np.array([0, 1, 1, 1, 1, 1, 1, 1]), - 'ntriplet': 2} - hyps_mask1['specie_mask'][2] = 1 - hyps1 = np.ones(9, dtype=float) + pm = ParameterHelper(species=['H', 'He'], + kernels={'twobody':[['*', '*'], ['H', 'H']], + 'threebody':[['*', '*', '*'], + ['H', 'H', 'H']]}, + parameters=parameters, + ones=True, random=False, + verbose="DEBUG") + hyps_mask1 = pm.as_dict() # 9 different hyper-parameters, onlye train the 0, 2, 4, 6, 8 - hyps_mask2 = {'nspecie': 2, - 'specie_mask': np.zeros(118, dtype=int), - 'nbond': 2, - 'bond_mask': np.array([0, 1, 1, 1]), - 'ntriplet': 2, - 'triplet_mask': np.array([0, 1, 1, 1, 1, 1, 1, 1]), - 'train_noise':True, - 'map':[0,2,4,6,8], - 'original':np.array([1, 1, 1, 1, 1, 1, 1, 1, 1])} - hyps_mask2['specie_mask'][2] = 1 - hyps2 = np.ones(5, dtype=float) + pm.set_constraints('twobody0', [True, True]) + pm.set_constraints('twobody1', [False, False]) + pm.set_constraints('threebody0', [True, True]) + pm.set_constraints('threebody1', [False, False]) + hyps_mask2 = pm.as_dict() # 9 different hyper-parameters, only train the 0, 2, 4, 6 - hyps_mask3 = {'nspecie': 2, - 'specie_mask': np.zeros(118, dtype=int), - 'nbond': 2, - 'bond_mask': np.array([0, 1, 1, 1]), - 'ntriplet': 2, - 'triplet_mask': np.array([0, 1, 1, 1, 1, 1, 1, 1]), - 'train_noise':False, - 'map':[0,2,4,6], - 'original':np.array([1, 1, 1, 1, 1, 1, 1, 1, 1])} - hyps_mask3['specie_mask'][2] = 1 - hyps3 = np.ones(4, dtype=float) + pm.set_constraints('noise', False) + hyps_mask3 = pm.as_dict() # 5 different hyper-parameters, equivalent to no multihyps - hyps_mask4 = {'nspecie': 1, - 'specie_mask': np.zeros(118, dtype=int), - 'nbond': 1, - 'bond_mask': np.array([0]), - 'ntriplet': 1, - 'triplet_mask': np.array([0])} - hyps4 = np.ones(5, dtype=float) - hyps_list = [hyps1, hyps2, hyps3, hyps4, hyps] - hyps_mask_list = [hyps_mask1, hyps_mask2, hyps_mask3, hyps_mask4, None] + pm = ParameterHelper(species=['H', 'He'], + kernels={'twobody':[['*', '*']], + 'threebody':[['*', '*', '*']]}, + parameters=parameters, + ones=True, random=False, + verbose="DEBUG") + hyps_mask4 = pm.as_dict() + + # 5 different hyper-parameters, no multihyps + pm = ParameterHelper(kernels=['twobody', 'threebody'], + parameters=parameters, + ones=True, random=False, + verbose="DEBUG") + hyps_mask5 = pm.as_dict() + + hyps_mask_list = [hyps_mask1, hyps_mask2, hyps_mask3, hyps_mask4, hyps_mask5] # create training environments and forces cell = np.eye(3) @@ -138,15 +120,18 @@ def get_random_training_set(nenv, nstruc): energy_noise = 0.01 - return hyps, name, kernel, cutoffs, kernel_m, hyps_list, hyps_mask_list, \ - energy_noise + return name, cutoffs, hyps_mask_list, energy_noise -def test_ky_mat(params): - hyps, name, kernel, cutoffs, kernel_m, hyps_list, hyps_mask_list, \ - energy_noise = params +@pytest.fixture(scope='module') +def ky_mat_ref(params): + name, cutoffs, hyps_mask_list, energy_noise = params # get the reference without multi hyps + hyps_mask = hyps_mask_list[-1] + hyps = hyps_mask['hyps'] + kernel = str_to_kernel_set(hyps_mask['kernels'], 'mc', hyps_mask) + time0 = time.time() ky_mat0 = get_Ky_mat(hyps, name, kernel[0], kernel[2], kernel[3], energy_noise, cutoffs) @@ -159,54 +144,49 @@ def test_ky_mat(params): energy_noise, cutoffs, n_cpus=2, n_sample=5) print("compute ky_mat parallel", time.time()-time0) - print(ky_mat) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff == 0), "parallel implementation is wrong" - - # compute the ky_mat with different parameters - for i in range(len(hyps_list)): - - hyps = hyps_list[i] - hyps_mask = hyps_mask_list[i] - - if hyps_mask is None: - ker1 = kernel[0] - ker2 = kernel[2] - ker3 = kernel[3] - else: - ker1 = kernel_m[0] - ker2 = kernel_m[2] - ker3 = kernel_m[3] - - # serial implementation - time0 = time.time() - ky_mat = get_Ky_mat(hyps, name, ker1, ker2, ker3, - energy_noise, cutoffs, hyps_mask) - print(f"compute ky_mat with multihyps, test {i}, n_cpus=1", - time.time()-time0) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff == 0), "multi hyps implementation is wrong"\ - f"with case {i}" - - # parallel implementation - time0 = time.time() - ky_mat = get_Ky_mat(hyps, name, ker1, ker2, ker3, - energy_noise, cutoffs, hyps_mask, n_cpus=2, - n_sample=20) - print(f"compute ky_mat with multihyps, test {i}, n_cpus=2", - time.time()-time0) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff == 0), "multi hyps parallel "\ - "implementation is wrong with case {i}" - - -def test_ky_mat_update(params): - """ - check ky_mat_update function - """ + assert np.isclose(ky_mat, ky_mat0, rtol=1e-3).all(), \ + "parallel implementation is wrong" + + yield ky_mat0 + del ky_mat0 + + +@pytest.mark.parametrize('ihyps', [0, 1]) +def test_ky_mat(params, ihyps, ky_mat_ref): + + name, cutoffs, hyps_mask_list, energy_noise = params + hyps_mask = hyps_mask_list[ihyps] + hyps = hyps_mask['hyps'] + kernel = str_to_kernel_set(hyps_mask['kernels'], 'mc', hyps_mask) + + # serial implementation + time0 = time.time() + ky_mat = get_Ky_mat(hyps, name, kernel[0], kernel[2], kernel[3], + energy_noise, cutoffs, hyps_mask) + print(f"compute ky_mat with multihyps, test {ihyps}, n_cpus=1", + time.time()-time0) + assert np.isclose(ky_mat, ky_mat_ref, rtol=1e-3).all(), \ + "multi hyps implementation is wrong"\ + f"with case {ihyps}" + + # parallel implementation + time0 = time.time() + ky_mat = get_Ky_mat(hyps, name, kernel[0], kernel[2], kernel[3], + energy_noise, cutoffs, hyps_mask, n_cpus=2, + n_sample=20) + print(f"compute ky_mat with multihyps, test {ihyps}, n_cpus=2", + time.time()-time0) + assert np.isclose(ky_mat, ky_mat_ref, rtol=1e-3).all(), \ + f"multi hyps parallel implementation is wrong with case {ihyps}" + - hyps, name, kernel, cutoffs, \ - kernel_m, hyps_list, hyps_mask_list, energy_noise = params +@pytest.mark.parametrize('ihyps', [0, 1, -1]) +def test_ky_mat_update(params, ihyps): + + name, cutoffs, hyps_mask_list, energy_noise = params + hyps_mask = hyps_mask_list[ihyps] + hyps = hyps_mask['hyps'] + kernel = str_to_kernel_set(hyps_mask['kernels'], 'mc', hyps_mask) # prepare old data set as the starting point n = 5 @@ -218,222 +198,135 @@ def test_ky_mat_update(params): training_structures[:s] func = [get_Ky_mat, get_ky_mat_update] + # get the reference - ky_mat0 = func[0](hyps, name, kernel[0], kernel[2], kernel[3], - energy_noise, cutoffs) - ky_mat_old = func[0](hyps, 'old', kernel[0], kernel[2], kernel[3], - energy_noise, cutoffs) + ky_mat0 = func[0](hyps, name, kernel[0], kernel[2], + kernel[3], energy_noise, cutoffs, hyps_mask) + ky_mat_old = func[0](hyps, 'old', kernel[0], kernel[2], + kernel[3], energy_noise, cutoffs, hyps_mask) # update ky_mat = func[1](ky_mat_old, n, hyps, name, kernel[0], kernel[2], - kernel[3], energy_noise, cutoffs) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - - assert (diff <= 1e-15), "update function is wrong" + kernel[3], energy_noise, cutoffs, hyps_mask) + assert np.isclose(ky_mat, ky_mat0, rtol=1e-10).all(), \ + "update function is wrong" # parallel version ky_mat = func[1](ky_mat_old, n, hyps, name, kernel[0], kernel[2], - kernel[3], energy_noise, cutoffs, n_cpus=2, n_sample=20) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - - assert (diff == 0), "parallel implementation is wrong" - - # check multi hyps implementation - for i in range(len(hyps_list)): - - hyps = hyps_list[i] - hyps_mask = hyps_mask_list[i] + kernel[3], energy_noise, cutoffs, hyps_mask, n_cpus=2, + n_sample=20) + assert np.isclose(ky_mat, ky_mat0, rtol=1e-10).all(), \ + "update function is wrong" - if hyps_mask is None: - ker1 = kernel[0] - ker2 = kernel[2] - ker3 = kernel[3] - else: - ker1 = kernel_m[0] - ker2 = kernel_m[2] - ker3 = kernel_m[3] - # serial implementation - ky_mat = func[1](ky_mat_old, n, hyps, name, ker1, ker2, ker3, - energy_noise, cutoffs, hyps_mask) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff < 1e-12), "multi hyps parameter implementation is wrong" +@pytest.mark.parametrize('ihyps', [0, 1, -1]) +def test_kernel_vector(params, ihyps): - # parallel implementation - ky_mat = func[1](ky_mat_old, n, hyps, name, ker1, ker2, ker3, - energy_noise, cutoffs, hyps_mask, n_cpus=2, - n_sample=20) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff < 1e-12), "multi hyps parameter parallel "\ - "implementation is wrong" - - -def test_kernel_vector(params): - - hyps, name, kernel, cutoffs, _, _, _, _ = params + name, cutoffs, hyps_mask_list, _ = params + np.random.seed(10) test_point = get_tstp() size1 = len(flare.gp_algebra._global_training_data[name]) size2 = len(flare.gp_algebra._global_training_structures[name]) + hyps_mask = hyps_mask_list[ihyps] + hyps = hyps_mask['hyps'] + kernel = str_to_kernel_set(hyps_mask['kernels'], 'mc', hyps_mask) + # test the parallel implementation for multihyps vec = get_kernel_vector(name, kernel[0], kernel[3], test_point, 1, - hyps, cutoffs) + hyps, cutoffs, hyps_mask) vec_par = \ get_kernel_vector(name, kernel[0], kernel[3], test_point, 1, hyps, - cutoffs, n_cpus=2, n_sample=100) + cutoffs, hyps_mask, n_cpus=2, n_sample=100) - assert (all(np.equal(vec, vec_par))), "parallel implementation is wrong" + assert (np.isclose(vec, vec_par, rtol=1e-4).all()), "parallel implementation is wrong" assert (vec.shape[0] == size1 * 3 + size2) -def test_en_kern_vec(params): +@pytest.mark.parametrize('ihyps', [0, 1, -1]) +def test_en_kern_vec(params, ihyps): - hyps, name, kernel, cutoffs, _, _, _, _ = params + name, cutoffs, hyps_mask_list, _ = params + hyps_mask = hyps_mask_list[ihyps] + hyps = hyps_mask['hyps'] + kernel = str_to_kernel_set(hyps_mask['kernels'], 'mc', hyps_mask) + np.random.seed(10) test_point = get_tstp() size1 = len(flare.gp_algebra._global_training_data[name]) size2 = len(flare.gp_algebra._global_training_structures[name]) # test the parallel implementation for multihyps - vec = en_kern_vec(name, kernel[3], kernel[2], test_point, hyps, cutoffs) + vec = en_kern_vec(name, kernel[3], kernel[2], test_point, hyps, cutoffs, hyps_mask) vec_par = \ - en_kern_vec(name, kernel[3], kernel[2], test_point, hyps, cutoffs, + en_kern_vec(name, kernel[3], kernel[2], test_point, hyps, cutoffs, hyps_mask, n_cpus=2, n_sample=100) assert (all(np.equal(vec, vec_par))), "parallel implementation is wrong" assert (vec.shape[0] == size1 * 3 + size2) -def test_ky_and_hyp(params): +@pytest.mark.parametrize('ihyps', [0, 1, 2, 3, 4]) +def test_ky_and_hyp(params, ihyps, ky_mat_ref): - hyps, name, kernel, cutoffs, \ - kernel_m, hyps_list, hyps_mask_list, _ = params + name, cutoffs, hyps_mask_list, _ = params + hyps_mask = hyps_mask_list[ihyps] + hyps = hyps_mask['hyps'] + kernel = str_to_kernel_set(hyps_mask['kernels'], 'mc', hyps_mask) - hypmat_0, ky_mat0 = get_ky_and_hyp(hyps, name, kernel[1], cutoffs) - - # parallel version - hypmat, ky_mat = get_ky_and_hyp(hyps, name, kernel[1], cutoffs, n_cpus=2) - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff == 0), "parallel implementation is wrong" - - # check all cases - for i in range(len(hyps_list)): - hyps = hyps_list[i] - hyps_mask = hyps_mask_list[i] - - if hyps_mask is None: - ker = kernel[1] - else: - ker = kernel_m[1] - - # serial implementation - hypmat, ky_mat = get_ky_and_hyp(hyps, name, ker, cutoffs, hyps_mask) - - if (i == 0): - hypmat9 = hypmat - diff = (np.max(np.abs(ky_mat-ky_mat0))) - assert (diff == 0), "multi hyps parameter implementation is wrong" - - # compare to no hyps_mask version - diff = 0 - if (i == 1): - diff = (np.max(np.abs(hypmat-hypmat9[[0, 2, 4, 6, 8], :, :]))) - elif (i == 2): - diff = (np.max(np.abs(hypmat-hypmat9[[0, 2, 4, 6], :, :]))) - elif (i == 3): - diff = (np.max(np.abs(hypmat-hypmat_0))) - elif (i == 4): - diff = (np.max(np.abs(hypmat-hypmat_0))) - assert (diff == 0), "multi hyps implementation is wrong"\ - f"in case {i}" - - # parallel implementation - hypmat_par, ky_mat_par = \ - get_ky_and_hyp(hyps, name, ker, cutoffs, hyps_mask, - n_cpus=2, n_sample=2) - - # compare to serial implementation - diff = (np.max(np.abs(ky_mat-ky_mat_par))) - assert (diff == 0), f"multi hyps parallel "\ - f"implementation is wrong in case {i}" - - diff = (np.max(np.abs(hypmat_par-hypmat))) - assert (diff == 0), f"multi hyps parallel implementation is wrong"\ - f" in case{i}" - - -def test_grad(params): - hyps, name, kernel, cutoffs, \ - kernel_m, hyps_list, hyps_mask_list, _ = params - - # obtain reference - func = get_ky_and_hyp - hyp_mat, ky_mat = func(hyps, name, kernel[1], cutoffs) - like0, like_grad0 = \ - get_like_grad_from_mats(ky_mat, hyp_mat, name) - - # serial implementation func = get_ky_and_hyp - hyp_mat, ky_mat = func(hyps, name, kernel[1], cutoffs) - like, like_grad = \ - get_like_grad_from_mats(ky_mat, hyp_mat, name) - assert (like == like0), "wrong likelihood" - assert np.max(np.abs(like_grad-like_grad0)) == 0, "wrong likelihood" + # serial version + hypmat_ser, ky_mat_ser = func(hyps, name, kernel[1], cutoffs, + hyps_mask) + # parallel version + hypmat_par, ky_mat_par = func(hyps, name, kernel[1], cutoffs, + hyps_mask, n_cpus=2) - func = get_ky_and_hyp - for i in range(len(hyps_list)): - hyps = hyps_list[i] - hyps_mask = hyps_mask_list[i] - - if hyps_mask is None: - ker = kernel[1] - else: - ker = kernel_m[1] - - hyp_mat, ky_mat = func(hyps, name, ker, cutoffs, hyps_mask) - like, like_grad = get_like_grad_from_mats(ky_mat, hyp_mat, name) - assert (like == like0), "wrong likelihood" - - if (i == 0): - like_grad9 = like_grad - - diff = 0 - if (i == 1): - diff = (np.max(np.abs(like_grad-like_grad9[[0, 2, 4, 6, 8]]))) - elif (i == 2): - diff = (np.max(np.abs(like_grad-like_grad9[[0, 2, 4, 6]]))) - elif (i == 3): - diff = (np.max(np.abs(like_grad-like_grad0))) - elif (i == 4): - diff = (np.max(np.abs(like_grad-like_grad0))) - assert (diff == 0), "multi hyps implementation is wrong"\ - f"in case {i}" - - -def test_ky_hyp_grad(params): - hyps, name, kernel, cutoffs, _, _, _, _ = params + ref = ky_mat_ref[:ky_mat_ser.shape[0], :ky_mat_ser.shape[1]] - func = get_ky_and_hyp + assert np.isclose(ky_mat_ser, ref, rtol=1e-5).all(), \ + "serial implementation is not consistent with get_Ky_mat" + assert np.isclose(ky_mat_par, ref, rtol=1e-5).all(), \ + "parallel implementation is not consistent with get_Ky_mat" + assert np.isclose(hypmat_ser, hypmat_par, rtol=1e-5).all(), \ + "serial implementation is not consistent with parallel implementation" - hyp_mat, ky_mat = func(hyps, name, kernel[1], cutoffs) + # analytical form + hyp_mat, ky_mat = func(hyps, name, kernel[1], cutoffs, hyps_mask) _, like_grad = get_like_grad_from_mats(ky_mat, hyp_mat, name) + delta = 0.001 for i in range(len(hyps)): + newhyps = np.copy(hyps) + newhyps[i] += delta - hyp_mat_p, ky_mat_p = func(newhyps, name, kernel[1], cutoffs) + hyp_mat_p, ky_mat_p = func(newhyps, name, kernel[1], + cutoffs, hyps_mask) like_p, _ = \ get_like_grad_from_mats(ky_mat_p, hyp_mat_p, name) + newhyps[i] -= 2*delta - hyp_mat_m, ky_mat_m = func(newhyps, name, kernel[1], cutoffs) + hyp_mat_m, ky_mat_m = func(newhyps, name, kernel[1], + cutoffs, hyps_mask) like_m, _ = \ get_like_grad_from_mats(ky_mat_m, hyp_mat_m, name) - diff = np.abs(like_grad[i]-(like_p-like_m)/2./delta) - assert (diff < 1e-3), "wrong calculation of hyp_mat" + + # numerical form + numeric = (like_p-like_m)/2./delta + assert np.isclose(like_grad[i], numeric, rtol=1e-3 ), \ + f"wrong calculation of hyp_mat {i}" + +if __name__ == "__main__": + + import cProfile + import re + params = get_random_training_set(10, 2) + cProfile.run('test_ky_and_hyp(params)') diff --git a/tests/test_gp_from_aimd.py b/tests/test_gp_from_aimd.py index a642413d2..1c7b7ed41 100644 --- a/tests/test_gp_from_aimd.py +++ b/tests/test_gp_from_aimd.py @@ -10,12 +10,12 @@ from flare.env import AtomicEnvironment from flare.struc import Structure from flare.gp import GaussianProcess -from flare.mgp.mgp import MappedGaussianProcess +from flare.mgp import MappedGaussianProcess from flare.gp_from_aimd import TrajectoryTrainer,\ parse_trajectory_trainer_output from flare.utils.learner import subset_of_frame_by_element -from .test_mgp_unit import all_mgp, all_gp, get_random_structure +from tests.test_mgp import all_mgp, all_gp, get_random_structure from .fake_gp import get_gp @pytest.fixture @@ -44,7 +44,7 @@ def methanol_gp(): @pytest.fixture def fake_gp(): return GaussianProcess(kernel_name="2+3", - hyps=np.array([1]), + hyps=np.array([1, 1, 1, 1, 1]), cutoffs=np.array([4, 3])) @@ -210,44 +210,28 @@ def test_mgp_gpfa(all_mgp, all_gp): :return: ''' + np.random.seed(10) gp_model = get_gp('3', 'mc', False) gp_model.set_L_alpha() grid_num_2 = 5 grid_num_3 = 3 lower_cut = 0.01 - two_cut = gp_model.cutoffs[0] - three_cut = gp_model.cutoffs[1] - # set struc params. cell and masses arbitrary? - mapped_cell = np.eye(3) * 2 - struc_params = {'species': [1, 2], - 'cube_lat': mapped_cell, - 'mass_dict': {'0': 27, '1': 16}} - - # grid parameters - train_size = len(gp_model.training_data) - grid_params = {'bodies': [2], - 'cutoffs': gp_model.cutoffs, - 'bounds_2': [[lower_cut], [two_cut]], - 'bounds_3': [[lower_cut, lower_cut, lower_cut], - [three_cut, three_cut, three_cut]], - 'grid_num_2': grid_num_2, - 'grid_num_3': [grid_num_3, grid_num_3, grid_num_3], - 'svd_rank_2': np.min((grid_num_2, 3 * train_size)), - 'svd_rank_3': np.min((grid_num_3 ** 3, 3 * train_size)), - 'load_grid': None, + grid_params_3b = {'lower_bound': [lower_cut for d in range(3)], + 'grid_num': [grid_num_3 for d in range(3)], + 'svd_rank': 'auto'} + grid_params = {'load_grid': None, 'update': False} + grid_params['threebody'] = grid_params_3b + unique_species = gp_model.training_statistics['species'] - struc_params = {'species': [1, 2], - 'cube_lat': np.eye(3) * 2, - 'mass_dict': {'0': 27, '1': 16}} - - mgp_model = MappedGaussianProcess(grid_params, struc_params) + mgp_model = MappedGaussianProcess(grid_params=grid_params, unique_species=unique_species, n_cpus=1, + map_force=False) mgp_model.build_map(gp_model) + nenv = 10 cell = np.eye(3) - unique_species = gp_model.training_data[0].species struc, f = get_random_structure(cell, unique_species, nenv) struc.forces = np.array(f) diff --git a/tests/test_grid_kernel.py b/tests/test_grid_kernel.py new file mode 100644 index 000000000..bb192b0d3 --- /dev/null +++ b/tests/test_grid_kernel.py @@ -0,0 +1,138 @@ +import numpy as np +import pytest +import sys + +from copy import deepcopy +from itertools import combinations_with_replacement, permutations +from numpy import isclose +from numpy.random import random, randint + +from flare.env import AtomicEnvironment +from flare.kernels.utils import from_mask_to_args, str_to_kernel_set +from flare.kernels.cutoffs import quadratic_cutoff_bound +from flare.parameters import Parameters +from flare.struc import Structure +from flare.utils.parameter_helper import ParameterHelper + +from tests.fake_gp import generate_mb_envs, generate_mb_twin_envs +from tests.test_mc_sephyps import generate_same_hm, generate_diff_hm +from flare.mgp.utils import get_triplets, get_kernel_term +from flare.mgp.grid_kernels_3b import triplet_cutoff, grid_kernel_sephyps, \ + grid_kernel, get_triplets_for_kern + +# multi_cut = [False, True] +hyps_list = [True, False] + +@pytest.mark.parametrize('same_hyps', hyps_list) +@pytest.mark.parametrize('prefix', ['energy']) #, 'force']) +def test_start(parameter, same_hyps, prefix): + + env1, env2, hm1, hm2 = parameter + + # get environments, hyps and arguments for training data + env = env1 if same_hyps else env2 + hm = hm1 if same_hyps else hm2 + kernel = grid_kernel if same_hyps else grid_kernel_sephyps + args = from_mask_to_args(hm['hyps'], hm['cutoffs'], None if same_hyps else hm) + # # debug + # for k in env.__dict__: + # print(k, env.__dict__[k]) + + # get all possible triplet grids + list_of_triplet = list(combinations_with_replacement([1, 2], 3)) + # # debug + # print(list_of_triplet) + + for comb in list_of_triplet: + for species in set(permutations(comb)): + + grid_env, grid = get_grid_env(species, parameter, same_hyps) + + # # debug + # print(species) + # for k in grid_env.__dict__: + # print(k, grid_env.__dict__[k]) + # print(grid) + + reference = get_reference(grid_env, species, parameter, same_hyps) + + + coords = np.zeros((1, 9), dtype=np.float64) + coords[:, 0] += 1.0 + + fj, fdj = triplet_cutoff(grid, hm['cutoffs']['threebody'], + coords, derivative=True) + fdj = fdj[:, [0]] + + kern_type = f'{prefix}_force' + kern_vec = kernel(kern_type, env, grid, fj, fdj, + grid_env.ctype, grid_env.etypes, + *args) + kern_vec = np.hstack(kern_vec) + print('species, reference, kern_vec, reference-kern_vec') + print(species, reference, kern_vec, reference-kern_vec) + +@pytest.fixture(scope='module') +def parameter(): + + np.random.seed(10) + kernels = ['threebody'] + + delta = 1e-8 + cutoffs, hyps1, hyps2, hm1, hm2 = generate_same_hm( + kernels, multi_cutoff=False) + cutoffs, hyps2, hm2, = generate_diff_hm( + kernels, diff_cutoff=False, constraint=False) + + cell = 1e7 * np.eye(3) + env1 = generate_mb_envs(hm1['cutoffs'], cell, delta, 1) + env2 = generate_mb_envs(hm2['cutoffs'], cell, delta, 2) + env1 = env1[0][0] + env2 = env2[0][0] + + all_list = (env1, env2, hm1, hm2) + + yield all_list + del all_list + +def get_grid_env(species, parameter, same_hyps): + '''generate a single triplet environment''' + + env1, env2, hm1, hm2 = parameter + + big_cell = np.eye(3) * 100 + r1 = 0.5 + r2 = 0.5 + positions = [[0, 0, 0], [r1, 0, 0], [0, r2, 0]] + grid_struc = Structure(big_cell, species, positions) + if same_hyps: + env = AtomicEnvironment(grid_struc, 0, hm1['cutoffs'], hm1) + else: + env = AtomicEnvironment(grid_struc, 0, hm2['cutoffs'], hm2) + + env.bond_array_3 = np.array([[r1, 1, 0, 0], [r2, 0, 0, 0]]) + + grid = np.array([[r1, r2, np.sqrt(r1**2+r2**2)]]) + return env, grid + +def get_reference(grid_env, species, parameter, same_hyps): + + env1, env2, hm1, hm2 = parameter + env = env1 if same_hyps else env2 + hm = hm1 if same_hyps else hm2 + + kernel, kg, en_kernel, force_en_kernel = str_to_kernel_set( + hm['kernels'], "mc", None if same_hyps else hm) + args = from_mask_to_args(hm['hyps'], hm['cutoffs'], None if same_hyps else hm) + + energy_force = np.zeros(3, dtype=np.float) + # force_force = np.zeros(3, dtype=np.float) + # force_energy = np.zeros(3, dtype=np.float) + # energy_energy = np.zeros(3, dtype=np.float) + for i in range(3): + energy_force[i] = force_en_kernel(env, grid_env, i+1, *args) + # force_energy[i] = force_en_kernel(env, grid_env, i, *args) + # force_force[i] = kernel(grid_env, env, 0, i, *args) +# result = funcs[1][i](env1, env2, d1, *args1) + return energy_force # , force_energy, force_force, energy_energy + diff --git a/tests/test_kernel.py b/tests/test_kernel.py index ce1565841..1630dec7b 100644 --- a/tests/test_kernel.py +++ b/tests/test_kernel.py @@ -10,21 +10,23 @@ from .fake_gp import generate_mb_envs -list_to_test = ['2', '3', '2+3', 'mb', '2+3+mb'] +list_to_test = [['2'], ['3'], + ['2', '3'], + ['2', '3', 'many']] list_type = ['sc', 'mc'] -def generate_hm(kernel_name): +def generate_hm(kernels): hyps = [] - for term in ['2', '3', 'mb']: - if (term in kernel_name): + for term in ['2', '3', 'many']: + if (term in kernels): hyps += [random(2)] hyps += [random()] return np.hstack(hyps) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('kernel_type', list_type) -def test_force_en(kernel_name, kernel_type): +def test_force_en(kernels, kernel_type): """Check that the analytical force/en kernel matches finite difference of energy kernel.""" @@ -40,20 +42,20 @@ def test_force_en(kernel_name, kernel_type): env1 = generate_mb_envs(cutoffs, cell, delta, d1, kern_type=kernel_type) env2 = generate_mb_envs(cutoffs, cell, delta, d1, kern_type=kernel_type) - hyps = generate_hm(kernel_name) + hyps = generate_hm(kernels) _, __, en_kernel, force_en_kernel = \ - str_to_kernel_set(kernel_name+kernel_type) + str_to_kernel_set(kernels, kernel_type) print(force_en_kernel.__name__) nterm = 0 - for term in ['2', '3', 'mb']: - if (term in kernel_name): + for term in ['2', '3', 'many']: + if (term in kernels): nterm += 1 kern_finite_diff = 0 - if ('mb' in kernel_name): - _, __, enm_kernel, ___ = str_to_kernel_set('mb'+kernel_type) + if ('many' in kernels): + _, __, enm_kernel, ___ = str_to_kernel_set(['many'], kernel_type) mhyps = hyps[(nterm-1)*2:] calc = 0 nat = len(env1[0]) @@ -63,21 +65,21 @@ def test_force_en(kernel_name, kernel_type): mb_diff = calc / (2 * delta) kern_finite_diff += mb_diff - if ('2' in kernel_name): - nbond = 1 - _, __, en2_kernel, ___ = str_to_kernel_set('2'+kernel_type) - calc1 = en2_kernel(env1[2][0], env2[0][0], hyps[0:nbond * 2], cutoffs) - calc2 = en2_kernel(env1[1][0], env2[0][0], hyps[0:nbond * 2], cutoffs) + if ('2' in kernels): + ntwobody = 1 + _, __, en2_kernel, ___ = str_to_kernel_set('2', kernel_type) + calc1 = en2_kernel(env1[2][0], env2[0][0], hyps[0:ntwobody * 2], cutoffs) + calc2 = en2_kernel(env1[1][0], env2[0][0], hyps[0:ntwobody * 2], cutoffs) diff2b = 4 * (calc1 - calc2) / 2.0 / 2.0 / delta kern_finite_diff += diff2b else: - nbond = 0 + ntwobody = 0 - if ('3' in kernel_name): - _, __, en3_kernel, ___ = str_to_kernel_set('3'+kernel_type) - calc1 = en3_kernel(env1[2][0], env2[0][0], hyps[nbond * 2:], cutoffs) - calc2 = en3_kernel(env1[1][0], env2[0][0], hyps[nbond * 2:], cutoffs) + if ('3' in kernels): + _, __, en3_kernel, ___ = str_to_kernel_set('3', kernel_type) + calc1 = en3_kernel(env1[2][0], env2[0][0], hyps[ntwobody * 2:], cutoffs) + calc2 = en3_kernel(env1[1][0], env2[0][0], hyps[ntwobody * 2:], cutoffs) diff3b = 9 * (calc1 - calc2) / 2.0 / 3.0 / delta kern_finite_diff += diff3b @@ -85,14 +87,14 @@ def test_force_en(kernel_name, kernel_type): kern_analytical = \ force_en_kernel(env1[0][0], env2[0][0], d1, hyps, cutoffs) - print("\nforce_en", kernel_name, kern_finite_diff, kern_analytical) + print("\nforce_en", kernels, kern_finite_diff, kern_analytical) assert (isclose(kern_finite_diff, kern_analytical, rtol=tol)) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('kernel_type', list_type) -def test_force(kernel_name, kernel_type): +def test_force(kernels, kernel_type): """Check that the analytical force kernel matches finite difference of energy kernel.""" @@ -105,14 +107,14 @@ def test_force(kernel_name, kernel_type): np.random.seed(10) - hyps = generate_hm(kernel_name) + hyps = generate_hm(kernels) kernel, kg, en_kernel, fek = \ - str_to_kernel_set(kernel_name+kernel_type, False) + str_to_kernel_set(kernels, kernel_type) args = (hyps, cutoffs) nterm = 0 - for term in ['2', '3', 'mb']: - if (term in kernel_name): + for term in ['2', '3', 'many']: + if (term in kernels): nterm += 1 env1 = generate_mb_envs(cutoffs, cell, delta, d1, kern_type=kernel_type) @@ -120,8 +122,8 @@ def test_force(kernel_name, kernel_type): # check force kernel kern_finite_diff = 0 - if ('mb' == kernel_name): - _, __, enm_kernel, ___ = str_to_kernel_set('mb'+kernel_type) + if ('many' == kernels): + _, __, enm_kernel, ___ = str_to_kernel_set('many', kernel_type) mhyps = hyps[(nterm-1)*2:] print(hyps) print(mhyps) @@ -137,26 +139,26 @@ def test_force(kernel_name, kernel_type): # TODO: Establish why 2+3+MB fails (numerical error?) return - if ('2' in kernel_name): - nbond = 1 - _, __, en2_kernel, ___ = str_to_kernel_set('2'+kernel_type) - print(hyps[0:nbond * 2]) + if ('2' in kernels): + ntwobody = 1 + _, __, en2_kernel, ___ = str_to_kernel_set(['2'], kernel_type) + print(hyps[0:ntwobody * 2]) - calc1 = en2_kernel(env1[1][0], env2[1][0], hyps[0:nbond * 2], cutoffs) - calc2 = en2_kernel(env1[2][0], env2[2][0], hyps[0:nbond * 2], cutoffs) - calc3 = en2_kernel(env1[1][0], env2[2][0], hyps[0:nbond * 2], cutoffs) - calc4 = en2_kernel(env1[2][0], env2[1][0], hyps[0:nbond * 2], cutoffs) + calc1 = en2_kernel(env1[1][0], env2[1][0], hyps[0:ntwobody * 2], cutoffs) + calc2 = en2_kernel(env1[2][0], env2[2][0], hyps[0:ntwobody * 2], cutoffs) + calc3 = en2_kernel(env1[1][0], env2[2][0], hyps[0:ntwobody * 2], cutoffs) + calc4 = en2_kernel(env1[2][0], env2[1][0], hyps[0:ntwobody * 2], cutoffs) kern_finite_diff += 4 * (calc1 + calc2 - calc3 - calc4) / (4*delta**2) else: - nbond = 0 - - if ('3' in kernel_name): - _, __, en3_kernel, ___ = str_to_kernel_set('3'+kernel_type) - print(hyps[nbond * 2:]) - calc1 = en3_kernel(env1[1][0], env2[1][0], hyps[nbond * 2:], cutoffs) - calc2 = en3_kernel(env1[2][0], env2[2][0], hyps[nbond * 2:], cutoffs) - calc3 = en3_kernel(env1[1][0], env2[2][0], hyps[nbond * 2:], cutoffs) - calc4 = en3_kernel(env1[2][0], env2[1][0], hyps[nbond * 2:], cutoffs) + ntwobody = 0 + + if ('3' in kernels): + _, __, en3_kernel, ___ = str_to_kernel_set(['3'], kernel_type) + print(hyps[ntwobody * 2:]) + calc1 = en3_kernel(env1[1][0], env2[1][0], hyps[ntwobody * 2:], cutoffs) + calc2 = en3_kernel(env1[2][0], env2[2][0], hyps[ntwobody * 2:], cutoffs) + calc3 = en3_kernel(env1[1][0], env2[2][0], hyps[ntwobody * 2:], cutoffs) + calc4 = en3_kernel(env1[2][0], env2[1][0], hyps[ntwobody * 2:], cutoffs) kern_finite_diff += 9 * (calc1 + calc2 - calc3 - calc4) / (4*delta**2) kern_analytical = kernel(env1[0][0], env2[0][0], d1, d2, *args) @@ -164,9 +166,9 @@ def test_force(kernel_name, kernel_type): assert(isclose(kern_finite_diff, kern_analytical, rtol=tol)) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('kernel_type', list_type) -def test_hyps_grad(kernel_name, kernel_type): +def test_hyps_grad(kernels, kernel_type): d1 = randint(1, 3) d2 = randint(1, 3) @@ -176,11 +178,11 @@ def test_hyps_grad(kernel_name, kernel_type): cutoffs = np.ones(3)*1.2 np.random.seed(10) - hyps = generate_hm(kernel_name) + hyps = generate_hm(kernels) env1 = generate_mb_envs(cutoffs, cell, 0, d1, kern_type=kernel_type)[0][0] env2 = generate_mb_envs(cutoffs, cell, 0, d2, kern_type=kernel_type)[0][0] - kernel, kernel_grad, _, _ = str_to_kernel_set(kernel_name+kernel_type, False) + kernel, kernel_grad, _, _ = str_to_kernel_set(kernels, kernel_type) grad_test = kernel_grad(env1, env2, d1, d2, hyps, cutoffs) diff --git a/tests/test_lmp.py b/tests/test_lmp.py index d6b2eaeca..f51aec02b 100644 --- a/tests/test_lmp.py +++ b/tests/test_lmp.py @@ -5,16 +5,16 @@ from flare import struc, env, gp from flare import otf_parser -from flare.mgp.mgp import MappedGaussianProcess -from flare.lammps import lammps_calculator from flare.ase.calculator import FLARE_Calculator +from flare.mgp import MappedGaussianProcess +from flare.lammps import lammps_calculator +from flare.utils.element_coder import _Z_to_mass, _element_to_Z from ase.calculators.lammpsrun import LAMMPS from tests.fake_gp import get_gp, get_random_structure -body_list = ['2', '3'] -multi_list = [False, True] curr_path = os.getcwd() +force_block_only = True def clean(): for f in os.listdir("./"): @@ -31,153 +31,80 @@ def clean(): 'in environment: Please install LAMMPS ' 'and set the $lmp env. ' 'variable to point to the executatble.') - @pytest.fixture(scope='module') -def all_gp(): - - allgp_dict = {} +def gp_model(): np.random.seed(0) - for bodies in ['2', '3', '2+3']: - for multihyps in [False, True]: - gp_model = get_gp(bodies, 'mc', multihyps) - gp_model.parallel = True - gp_model.n_cpus = 2 - allgp_dict[f'{bodies}{multihyps}'] = gp_model + # TO DO, should be 2+3 eventually + gauss = get_gp('2', 'mc', False, cellabc=[1, 1, 1], + force_only=force_block_only, noa=5) + gauss.parallel = True + gauss.n_cpus = 2 + yield gauss + del gauss - yield allgp_dict - del allgp_dict @pytest.fixture(scope='module') -def all_mgp(): +def mgp_model(gp_model): + """ + test the init function + """ - allmgp_dict = {} - for bodies in ['2', '3', '2+3']: - for multihyps in [False, True]: - allmgp_dict[f'{bodies}{multihyps}'] = None + grid_params={} + if 'twobody' in gp_model.kernels: + grid_params['twobody']={'grid_num': [64], + 'lower_bound':[0.1], + 'svd_rank': 14} + if 'threebody' in gp_model.kernels: + grid_params['threebody']={'grid_num': [16]*3, + 'lower_bound':[0.1]*3, + 'svd_rank': 14} + species_list = [1, 2, 3] + lammps_location = f'tmp_lmp.mgp' + mapped_model = MappedGaussianProcess(grid_params=grid_params, unique_species=species_list, n_cpus=1, + map_force=False, lmp_file_name=lammps_location, mean_only=True) - yield allmgp_dict - del allmgp_dict + # import flare.mgp.mapxb + # flare.mgp.mapxb.global_use_grid_kern = False -@pytest.fixture(scope='module') -def all_ase_calc(): + mapped_model.build_map(gp_model) - all_ase_calc_dict = {} - for bodies in ['2', '3', '2+3']: - for multihyps in [False, True]: - all_ase_calc_dict[f'{bodies}{multihyps}'] = None + yield mapped_model + del mapped_model - yield all_ase_calc_dict - del all_ase_calc_dict @pytest.fixture(scope='module') -def all_lmp_calc(): - - if 'tmp' not in os.listdir("./"): - os.mkdir('tmp') - - all_lmp_calc_dict = {} - for bodies in ['2', '3', '2+3']: - for multihyps in [False, True]: - all_lmp_calc_dict[f'{bodies}{multihyps}'] = None - - yield all_lmp_calc_dict - del all_lmp_calc_dict - - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -def test_init(bodies, multihyps, all_mgp, all_gp): - """ - test the init function - """ - - gp_model = all_gp[f'{bodies}{multihyps}'] - - grid_num_2 = 64 - grid_num_3 = 16 - lower_cut = 0.01 - two_cut = gp_model.cutoffs[0] - three_cut = gp_model.cutoffs[1] - lammps_location = f'{bodies}{multihyps}.mgp' - - # set struc params. cell and masses arbitrary? - mapped_cell = np.eye(3) * 20 - struc_params = {'species': [1, 2], - 'cube_lat': mapped_cell, - 'mass_dict': {'0': 2, '1': 4}} - - # grid parameters - blist = [] - if ('2' in bodies): - blist+= [2] - if ('3' in bodies): - blist+= [3] - train_size = len(gp_model.training_data) - grid_params = {'bodies': blist, - 'cutoffs':gp_model.cutoffs, - 'bounds_2': [[lower_cut], [two_cut]], - 'bounds_3': [[lower_cut, lower_cut, lower_cut], - [three_cut, three_cut, three_cut]], - 'grid_num_2': grid_num_2, - 'grid_num_3': [grid_num_3, grid_num_3, grid_num_3], - 'svd_rank_2': 14, - 'svd_rank_3': 14, - 'load_grid': None, - 'update': False} - - struc_params = {'species': [1, 2], - 'cube_lat': np.eye(3)*2, - 'mass_dict': {'0': 27, '1': 16}} - mgp_model = MappedGaussianProcess(grid_params, struc_params, n_cpus=4, - mean_only=True, lmp_file_name=lammps_location) - all_mgp[f'{bodies}{multihyps}'] = mgp_model - - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -def test_build_map(all_gp, all_mgp, all_ase_calc, bodies, multihyps): +def ase_calculator(gp_model, mgp_model): """ test the mapping for mc_simple kernel """ + cal = FLARE_Calculator(gp_model, mgp_model, par=False, use_mapping=True) + yield cal + del cal - # multihyps = False - gp_model = all_gp[f'{bodies}{multihyps}'] - mgp_model = all_mgp[f'{bodies}{multihyps}'] - mgp_model.build_map(gp_model) - - all_ase_calc[f'{bodies}{multihyps}'] = FLARE_Calculator(gp_model, - mgp_model, par=False, use_mapping=True) - - clean() +@pytest.fixture(scope='module') +def lmp_calculator(gp_model, mgp_model): -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -def test_lmp_calc(bodies, multihyps, all_lmp_calc): + species = gp_model.training_statistics['species'] + specie_symbol_list = " ".join(species) + masses=[f"{i} {_Z_to_mass[_element_to_Z[species[i]]]}" for i in range(len(species))] - label = f'{bodies}{multihyps}' # set up input params - - by = 'no' - ty = 'no' - if '2' in bodies: - by = 'yes' - if '3' in bodies: - ty = 'yes' - + label = 'tmp_lmp' + by = 'yes' if 'twobody' in gp_model.kernels else 'no' + ty = 'yes' if 'threebody' in gp_model.kernels else 'no' parameters = {'command': os.environ.get('lmp'), # set up executable for ASE 'newton': 'off', 'pair_style': 'mgp', - 'pair_coeff': [f'* * {label}.mgp H He {by} {ty}'], - 'mass': ['1 2', '2 4']} + 'pair_coeff': [f'* * {label}.mgp {specie_symbol_list} {by} {ty}'], + 'mass': masses} files = [f'{label}.mgp'] - # create ASE calc - lmp_calc = LAMMPS(label=f'tmp{label}', keep_tmp_files=True, tmp_dir='./tmp/', - parameters=parameters, files=files) - - all_lmp_calc[label] = lmp_calc + lmp_calc = LAMMPS(label=f'tmp{label}', keep_tmp_files=True, tmp_dir='./tmp/', + parameters=parameters, files=files, specorder=species) + yield lmp_calc + del lmp_calc @pytest.mark.skipif(not os.environ.get('lmp', @@ -185,14 +112,14 @@ def test_lmp_calc(bodies, multihyps, all_lmp_calc): 'in environment: Please install LAMMPS ' 'and set the $lmp env. ' 'variable to point to the executatble.') -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -def test_lmp_predict(all_ase_calc, all_lmp_calc, bodies, multihyps): +def test_lmp_predict(gp_model, mgp_model, ase_calculator, lmp_calculator): """ test the lammps implementation """ - label = f'{bodies}{multihyps}' + currdir = os.getcwd() + + label = 'tmp_lmp' for f in os.listdir("./"): if label in f: @@ -201,45 +128,51 @@ def test_lmp_predict(all_ase_calc, all_lmp_calc, bodies, multihyps): os.remove(f) clean() - flare_calc = all_ase_calc[label] - lmp_calc = all_lmp_calc[label] - - gp_model = flare_calc.gp_model - mgp_model = flare_calc.mgp_model lammps_location = mgp_model.lmp_file_name # lmp file is automatically written now every time MGP is constructed mgp_model.write_lmp_file(lammps_location) # create test structure - cell = np.diag(np.array([1, 1, 1.5])) * 4 + np.random.seed(1) + cell = np.diag(np.array([1, 1, 1])) * 4 nenv = 10 - unique_species = gp_model.training_data[0].species + unique_species = gp_model.training_statistics['species'] cutoffs = gp_model.cutoffs struc_test, f = get_random_structure(cell, unique_species, nenv) - struc_test.positions *= 4 # build ase atom from struc ase_atoms_flare = struc_test.to_ase_atoms() - ase_atoms_flare.set_calculator(flare_calc) - ase_atoms_lmp = struc_test.to_ase_atoms() - ase_atoms_lmp.set_calculator(lmp_calc) - - lmp_en = ase_atoms_lmp.get_potential_energy() - flare_en = ase_atoms_flare.get_potential_energy() - - lmp_stress = ase_atoms_lmp.get_stress() - flare_stress = ase_atoms_flare.get_stress() + ase_atoms_flare.set_calculator(ase_calculator) - lmp_forces = ase_atoms_lmp.get_forces() - flare_forces = ase_atoms_flare.get_forces() - - # check that lammps agrees with gp to within 1 meV/A - assert np.all(np.abs(lmp_en - flare_en) < 1e-4) - assert np.all(np.abs(lmp_forces - flare_forces) < 1e-4) - assert np.all(np.abs(lmp_stress - flare_stress) < 1e-3) - - for f in os.listdir('./'): - if (label in f) or (f in ['log.lammps']): - os.remove(f) + ase_atoms_lmp = struc_test.to_ase_atoms() + ase_atoms_lmp.set_calculator(lmp_calculator) + + try: + lmp_en = ase_atoms_lmp.get_potential_energy() + flare_en = ase_atoms_flare.get_potential_energy() + + lmp_stress = ase_atoms_lmp.get_stress() + flare_stress = ase_atoms_flare.get_stress() + + lmp_forces = ase_atoms_lmp.get_forces() + flare_forces = ase_atoms_flare.get_forces() + except Exception as e: + os.chdir(currdir) + print(e) + raise e + + os.chdir(currdir) + + # check that lammps agrees with mgp to within 1 meV/A + print("energy", lmp_en, flare_en) + assert np.isclose(lmp_en, flare_en, atol=1e-3).all() + print("force", lmp_forces, flare_forces) + assert np.isclose(lmp_forces, flare_forces, atol=1e-3).all() + print("stress", lmp_stress, flare_stress) + assert np.isclose(lmp_stress, flare_stress, atol=1e-3).all() + + # for f in os.listdir('./'): + # if (label in f) or (f in ['log.lammps']): + # os.remove(f) diff --git a/tests/test_mask_helper.py b/tests/test_mask_helper.py deleted file mode 100644 index 6f4555237..000000000 --- a/tests/test_mask_helper.py +++ /dev/null @@ -1,137 +0,0 @@ -import pytest -import numpy as np -from flare.struc import Structure -from flare.utils.mask_helper import HyperParameterMasking -from .test_gp import dumpcompare - - -def test_generate_by_line(): - - pm = HyperParameterMasking() - pm.define_group('specie', 'O', ['O']) - pm.define_group('specie', 'C', ['C']) - pm.define_group('specie', 'H', ['H']) - pm.define_group('bond', '**', ['C', 'H']) - pm.define_group('bond', 'OO', ['O', 'O']) - pm.define_group('triplet', '***', ['O', 'O', 'C']) - pm.define_group('triplet', 'OOO', ['O', 'O', 'O']) - pm.define_group('mb', '1.5', ['C', 'H']) - pm.define_group('mb', '1.5', ['C', 'O']) - pm.define_group('mb', '1.5', ['O', 'H']) - pm.define_group('mb', '2', ['O', 'O']) - pm.define_group('mb', '2', ['H', 'O']) - pm.define_group('mb', '2.8', ['O', 'O']) - pm.set_parameters('**', [1, 0.5]) - pm.set_parameters('OO', [1, 0.5]) - pm.set_parameters('***', [1, 0.5]) - pm.set_parameters('OOO', [1, 0.5]) - pm.set_parameters('1.5', [1, 0.5, 1.5]) - pm.set_parameters('2', [1, 0.5, 2]) - pm.set_parameters('2.8', [1, 0.5, 2.8]) - pm.set_parameters('cutoff2b', 5) - pm.set_parameters('cutoff3b', 4) - pm.set_parameters('cutoffmb', 3) - hm = pm.generate_dict() - print(hm) - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - -def test_generate_by_line2(): - - pm = HyperParameterMasking() - pm.define_group('specie', 'O', ['O']) - pm.define_group('specie', 'rest', ['C', 'H']) - pm.define_group('bond', '**', ['*', '*']) - pm.define_group('bond', 'OO', ['O', 'O']) - pm.define_group('triplet', '***', ['*', '*', '*']) - pm.define_group('triplet', 'Oall', ['O', 'O', 'O']) - pm.set_parameters('**', [1, 0.5]) - pm.set_parameters('OO', [1, 0.5]) - pm.set_parameters('Oall', [1, 0.5]) - pm.set_parameters('***', [1, 0.5]) - pm.set_parameters('cutoff2b', 5) - pm.set_parameters('cutoff3b', 4) - hm = pm.generate_dict() - print(hm) - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - -def test_generate_by_list(): - - pm = HyperParameterMasking() - pm.list_groups('specie', ['O', 'C', 'H']) - pm.list_groups('bond', [['*', '*'], ['O','O']]) - pm.list_groups('triplet', [['*', '*', '*'], ['O','O', 'O']]) - pm.list_parameters({'bond0':[1, 0.5], 'bond1':[2, 0.2], - 'triplet0':[1, 0.5], 'triplet1':[2, 0.2], - 'cutoff2b':2, 'cutoff3b':1}) - hm = pm.generate_dict() - print(hm) - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - -def test_initialization(): - pm = HyperParameterMasking(species=['O', 'C', 'H'], - bonds=[['*', '*'], ['O','O']], - triplets=[['*', '*', '*'], ['O','O', 'O']], - parameters={'bond0':[1, 0.5], 'bond1':[2, 0.2], - 'triplet0':[1, 0.5], 'triplet1':[2, 0.2], - 'cutoff2b':2, 'cutoff3b':1}) - hm = pm.hyps_mask - print(hm) - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - -def test_opt(): - pm = HyperParameterMasking(species=['O', 'C', 'H'], - bonds=[['*', '*'], ['O','O']], - triplets=[['*', '*', '*'], ['O','O', 'O']], - parameters={'bond0':[1, 0.5, 1], 'bond1':[2, 0.2, 2], - 'triplet0':[1, 0.5], 'triplet1':[2, 0.2], - 'cutoff2b':2, 'cutoff3b':1}, - constraints={'bond0':[False, True]}) - hm = pm.hyps_mask - print(hm) - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - -def test_randomization(): - pm = HyperParameterMasking(species=['O', 'C', 'H'], - bonds=True, triplets=True, - mb=False, allseparate=True, - random=True, - parameters={'cutoff2b': 7, - 'cutoff3b': 4.5, - 'cutoffmb': 3}, - verbose=True) - hm = pm.hyps_mask - print(hm) - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - name = pm.find_group('specie', 'O') - print("find group name for O", name) - name = pm.find_group('bond', ['O', 'C']) - print("find group name for O-C", name) - -def test_from_dict(): - pm = HyperParameterMasking(species=['O', 'C', 'H'], - bonds=True, triplets=True, - mb=False, allseparate=True, - random=True, - parameters={'cutoff2b': 7, - 'cutoff3b': 4.5, - 'cutoffmb': 3}, - verbose=True) - hm = pm.hyps_mask - HyperParameterMasking.check_instantiation(hm) - HyperParameterMasking.check_matching(hm, hm['hyps'], hm['cutoffs']) - print(hm['hyps']) - print("obtain test hm", hm) - - pm1 = HyperParameterMasking.from_dict(hm, verbose=True) - print("from_dict") - hm1 = pm1.generate_dict() - print(hm['hyps']) - print(hm1['hyps'][:33], hm1['hyps'][33:]) - - dumpcompare(hm, hm1) diff --git a/tests/test_mc_sephyps.py b/tests/test_mc_sephyps.py index ac83a6bd2..f3858d76d 100644 --- a/tests/test_mc_sephyps.py +++ b/tests/test_mc_sephyps.py @@ -6,20 +6,21 @@ from numpy.random import random, randint from numpy import isclose -from flare.kernels.mc_sephyps import _str_to_kernel as stk from flare.kernels.utils import from_mask_to_args, str_to_kernel_set from flare.kernels.cutoffs import quadratic_cutoff_bound -from flare.utils.mask_helper import HyperParameterMasking +from flare.parameters import Parameters +from flare.utils.parameter_helper import ParameterHelper from .fake_gp import generate_mb_envs, generate_mb_twin_envs -list_to_test = ['2', '3', 'mb', '2+3', '2+3+mb'] +list_to_test = [['twobody'], ['threebody'], + ['twobody', 'threebody'], + ['twobody', 'threebody', 'manybody']] multi_cut = [False, True] - -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('multi_cutoff', multi_cut) -def test_force_en_multi_vs_simple(kernel_name, multi_cutoff): +def test_force_en_multi_vs_simple(kernels, multi_cutoff): """Check that the analytical kernel matches the one implemented in mc_simple.py""" @@ -30,7 +31,7 @@ def test_force_en_multi_vs_simple(kernel_name, multi_cutoff): # set hyperparameters cutoffs, hyps1, hyps2, hm1, hm2 = generate_same_hm( - kernel_name, multi_cutoff) + kernels, multi_cutoff) delta = 1e-8 env1 = generate_mb_envs(cutoffs, cell, delta, d1) @@ -40,20 +41,17 @@ def test_force_en_multi_vs_simple(kernel_name, multi_cutoff): # mc_simple kernel0, kg0, en_kernel0, force_en_kernel0 = str_to_kernel_set( - kernel_name, False) - args0 = (hyps1, cutoffs) - print("args0", args0) + kernels, "mc", None) + args0 = from_mask_to_args(hyps1, cutoffs) # mc_sephyps # args1 and args 2 use 1 and 2 groups of hyper-parameters # if (diff_cutoff), 1 or 2 group of cutoffs # but same value as in args0 kernel, kg, en_kernel, force_en_kernel = str_to_kernel_set( - kernel_name, True) - args1 = from_mask_to_args(hyps1, hm1, cutoffs) - args2 = from_mask_to_args(hyps2, hm2, cutoffs) - print("args1", args1) - print("args2", args2) + kernels, "mc", hm2) + args1 = from_mask_to_args(hyps1, cutoffs, hm1) + args2 = from_mask_to_args(hyps2, cutoffs, hm2) funcs = [[kernel0, kg0, en_kernel0, force_en_kernel0], [kernel, kg, en_kernel, force_en_kernel]] @@ -63,39 +61,39 @@ def test_force_en_multi_vs_simple(kernel_name, multi_cutoff): i = 2 reference = funcs[0][i](env1, env2, *args0) result = funcs[1][i](env1, env2, *args1) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) assert(isclose(reference, result, rtol=tol)) result = funcs[1][i](env1, env2, *args2) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) assert(isclose(reference, result, rtol=tol)) i = 3 reference = funcs[0][i](env1, env2, d1, *args0) result = funcs[1][i](env1, env2, d1, *args1) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) assert(isclose(reference, result, rtol=tol)) result = funcs[1][i](env1, env2, d1, *args2) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) assert(isclose(reference, result, rtol=tol)) i = 0 reference = funcs[0][i](env1, env2, d1, d2, *args0) result = funcs[1][i](env1, env2, d1, d2, *args1) assert(isclose(reference, result, rtol=tol)) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) result = funcs[1][i](env1, env2, d1, d2, *args2) assert(isclose(reference, result, rtol=tol)) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) i = 1 reference = funcs[0][i](env1, env2, d1, d2, *args0) result = funcs[1][i](env1, env2, d1, d2, *args1) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) assert(isclose(reference[0], result[0], rtol=tol)) assert(isclose(reference[1], result[1], rtol=tol).all()) result = funcs[1][i](env1, env2, d1, d2, *args2) - print(kernel_name, i, reference, result) + print(kernels, i, reference, result) assert(isclose(reference[0], result[0], rtol=tol)) joint_grad = np.zeros(len(result[1])//2) for i in range(joint_grad.shape[0]): @@ -103,9 +101,9 @@ def test_force_en_multi_vs_simple(kernel_name, multi_cutoff): assert(isclose(reference[1], joint_grad, rtol=tol).all()) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('diff_cutoff', multi_cut) -def test_check_sig_scale(kernel_name, diff_cutoff): +def test_check_sig_scale(kernels, diff_cutoff): """Check whether the grouping is properly assign with four environments @@ -125,9 +123,7 @@ def test_check_sig_scale(kernel_name, diff_cutoff): tol = 1e-4 scale = 2 - cutoffs, hyps0, hm = generate_diff_hm(kernel_name, diff_cutoff) - print(cutoffs) - print(hm) + cutoffs, hyps0, hm = generate_diff_hm(kernels, diff_cutoff) delta = 1e-8 env1, env1_t = generate_mb_twin_envs(cutoffs, np.eye(3)*100, delta, d1, hm) @@ -144,10 +140,10 @@ def test_check_sig_scale(kernel_name, diff_cutoff): hyps1[1::4] *= scale kernel, kg, en_kernel, force_en_kernel = str_to_kernel_set( - kernel_name, True) + kernels, 'mc', hm) - args0 = from_mask_to_args(hyps0, hm, cutoffs) - args1 = from_mask_to_args(hyps1, hm, cutoffs) + args0 = from_mask_to_args(hyps0, cutoffs, hm) + args1 = from_mask_to_args(hyps1, cutoffs, hm) reference = en_kernel(env1, env2, *args0) result = en_kernel(env1_t, env2_t, *args1) @@ -182,9 +178,9 @@ def test_check_sig_scale(kernel_name, diff_cutoff): [idx], scale**2, rtol=tol) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('diff_cutoff', multi_cut) -def test_force_bound_cutoff_compare(kernel_name, diff_cutoff): +def test_force_bound_cutoff_compare(kernels, diff_cutoff): """Check that the analytical kernel matches the one implemented in mc_simple.py""" @@ -194,10 +190,10 @@ def test_force_bound_cutoff_compare(kernel_name, diff_cutoff): cell = 1e7 * np.eye(3) delta = 1e-8 - cutoffs, hyps, hm = generate_diff_hm(kernel_name, diff_cutoff) + cutoffs, hyps, hm = generate_diff_hm(kernels, diff_cutoff) kernel, kg, en_kernel, force_en_kernel = str_to_kernel_set( - kernel_name, True) - args = from_mask_to_args(hyps, hm, cutoffs) + kernels, "mc", hm) + args = from_mask_to_args(hyps, cutoffs, hm) np.random.seed(10) env1 = generate_mb_envs(cutoffs, cell, delta, d1, hm) @@ -224,13 +220,13 @@ def test_force_bound_cutoff_compare(kernel_name, diff_cutoff): assert(isclose(reference, result, rtol=tol)) -@pytest.mark.parametrize('kernel_name', ['2+3']) +@pytest.mark.parametrize('kernels', [['twobody', 'threebody']]) @pytest.mark.parametrize('diff_cutoff', multi_cut) -def test_constraint(kernel_name, diff_cutoff): +def test_constraint(kernels, diff_cutoff): """Check that the analytical force/en kernel matches finite difference of energy kernel.""" - if ('mb' in kernel_name): + if ('manybody' in kernels): return d1 = 1 @@ -239,11 +235,11 @@ def test_constraint(kernel_name, diff_cutoff): delta = 1e-8 cutoffs, hyps, hm = generate_diff_hm( - kernel_name, diff_cutoff=diff_cutoff, constraint=True) + kernels, diff_cutoff=diff_cutoff, constraint=True) - _, __, en_kernel, force_en_kernel = str_to_kernel_set(kernel_name, True) + _, __, en_kernel, force_en_kernel = str_to_kernel_set(kernels, "mc", hm) - args0 = from_mask_to_args(hyps, hm, cutoffs) + args0 = from_mask_to_args(hyps, cutoffs, hm) np.random.seed(10) env1 = generate_mb_envs(cutoffs, cell, delta, d1, hm) @@ -251,14 +247,14 @@ def test_constraint(kernel_name, diff_cutoff): kern_finite_diff = 0 - if ('2' in kernel_name): - _, __, en2_kernel, fek2 = str_to_kernel_set('2', True) + if ('twobody' in kernels): + _, __, en2_kernel, fek2 = str_to_kernel_set(['twobody'], "mc", hm) calc1 = en2_kernel(env1[1][0], env2[0][0], *args0) calc2 = en2_kernel(env1[0][0], env2[0][0], *args0) kern_finite_diff += 4*(calc1 - calc2) / 2.0 / delta - if ('3' in kernel_name): - _, __, en3_kernel, fek3 = str_to_kernel_set('3', True) + if ('threebody' in kernels): + _, __, en3_kernel, fek3 = str_to_kernel_set(['threebody'], "mc", hm) calc1 = en3_kernel(env1[1][0], env2[0][0], *args0) calc2 = en3_kernel(env1[0][0], env2[0][0], *args0) kern_finite_diff += 9*(calc1 - calc2) / 3.0 / delta @@ -266,12 +262,13 @@ def test_constraint(kernel_name, diff_cutoff): kern_analytical = force_en_kernel(env1[0][0], env2[0][0], d1, *args0) tol = 1e-4 + print(kern_finite_diff, kern_analytical) assert(isclose(-kern_finite_diff, kern_analytical, rtol=tol)) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('diff_cutoff', multi_cut) -def test_force_en(kernel_name, diff_cutoff): +def test_force_en(kernels, diff_cutoff): """Check that the analytical force/en kernel matches finite difference of energy kernel.""" @@ -281,19 +278,19 @@ def test_force_en(kernel_name, diff_cutoff): cell = 1e7 * np.eye(3) np.random.seed(0) - cutoffs, hyps, hm = generate_diff_hm(kernel_name, diff_cutoff) - args = from_mask_to_args(hyps, hm, cutoffs) + cutoffs, hyps, hm = generate_diff_hm(kernels, diff_cutoff) + args = from_mask_to_args(hyps, cutoffs, hm) env1 = generate_mb_envs(cutoffs, cell, delta, d1, hm) env2 = generate_mb_envs(cutoffs, cell, delta, d2, hm) - _, __, en_kernel, force_en_kernel = str_to_kernel_set(kernel_name, True) + _, __, en_kernel, force_en_kernel = str_to_kernel_set(kernels, "mc", hm) kern_analytical = force_en_kernel(env1[0][0], env2[0][0], d1, *args) kern_finite_diff = 0 - if ('mb' in kernel_name): - kernel, _, enm_kernel, efk = str_to_kernel_set('mb', True) + if ('manybody' in kernels): + kernel, _, enm_kernel, efk = str_to_kernel_set(['manybody'], "mc", hm) calc = 0 for i in range(len(env1[0])): @@ -302,34 +299,34 @@ def test_force_en(kernel_name, diff_cutoff): kern_finite_diff += (calc)/(2*delta) - if ('2' in kernel_name or '3' in kernel_name): - args23 = from_mask_to_args(hyps, hm, cutoffs[:2]) + if ('twobody' in kernels or 'threebody' in kernels): + args23 = from_mask_to_args(hyps, cutoffs, hm) - if ('2' in kernel_name): - kernel, _, en2_kernel, efk = str_to_kernel_set('2b', True) + if ('twobody' in kernels): + kernel, _, en2_kernel, efk = str_to_kernel_set(['2b'], 'mc', hm) calc1 = en2_kernel(env1[1][0], env2[0][0], *args23) calc2 = en2_kernel(env1[2][0], env2[0][0], *args23) diff2b = 4 * (calc1 - calc2) / 2.0 / delta / 2.0 kern_finite_diff += diff2b - if ('3' in kernel_name): - kernel, _, en3_kernel, efk = str_to_kernel_set('3b', True) + if ('threebody' in kernels): + kernel, _, en3_kernel, efk = str_to_kernel_set(['3b'], 'mc', hm) calc1 = en3_kernel(env1[1][0], env2[0][0], *args23) calc2 = en3_kernel(env1[2][0], env2[0][0], *args23) diff3b = 9 * (calc1 - calc2) / 3.0 / delta / 2.0 kern_finite_diff += diff3b - print("\nforce_en", kernel_name, kern_finite_diff, kern_analytical) tol = 1e-3 + print("\nforce_en", kernels, kern_finite_diff, kern_analytical) assert (isclose(-kern_finite_diff, kern_analytical, rtol=tol)) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('diff_cutoff', multi_cut) -def test_force(kernel_name, diff_cutoff): +def test_force(kernels, diff_cutoff): """Check that the analytical force kernel matches finite difference of energy kernel.""" @@ -342,12 +339,12 @@ def test_force(kernel_name, diff_cutoff): np.random.seed(10) - cutoffs, hyps, hm = generate_diff_hm(kernel_name, diff_cutoff) - kernel, kg, en_kernel, fek = str_to_kernel_set(kernel_name, True) + cutoffs, hyps, hm = generate_diff_hm(kernels, diff_cutoff) + kernel, kg, en_kernel, fek = str_to_kernel_set(kernels, 'mc', hm) nterm = 0 - for term in ['2', '3', 'mb']: - if (term in kernel_name): + for term in ['twobody', 'threebody', 'manybody']: + if (term in kernels): nterm += 1 np.random.seed(10) @@ -356,11 +353,10 @@ def test_force(kernel_name, diff_cutoff): # check force kernel kern_finite_diff = 0 - if ('mb' == kernel_name): - _, __, enm_kernel, ___ = str_to_kernel_set('mb', True) - mhyps, mhyps_mask = HyperParameterMasking.get_mb_hyps( - hyps, hm, True) - margs = from_mask_to_args(mhyps, mhyps_mask, cutoffs) + if 'manybody' in kernels and len(kernels)==1: + _, __, enm_kernel, ___ = str_to_kernel_set(['manybody'], 'mc', hm) + mhyps, mcutoffs, mhyps_mask = Parameters.get_component_mask(hm, 'manybody', hyps=hyps) + margs = from_mask_to_args(mhyps, mcutoffs, mhyps_mask) cal = 0 for i in range(3): for j in range(len(env1[0])): @@ -369,16 +365,15 @@ def test_force(kernel_name, diff_cutoff): cal -= enm_kernel(env1[1][i], env2[2][j], *margs) cal -= enm_kernel(env1[2][i], env2[1][j], *margs) kern_finite_diff += cal / (4 * delta ** 2) - else: + elif 'manybody' in kernels: # TODO: Establish why 2+3+MB fails (numerical error?) return - if ('2' in kernel_name): - nbond = 1 - _, __, en2_kernel, ___ = str_to_kernel_set('2', True) - bhyps, bhyps_mask = HyperParameterMasking.get_2b_hyps( - hyps, hm, True) - args2 = from_mask_to_args(bhyps, bhyps_mask, cutoffs[:1]) + if ('twobody' in kernels): + ntwobody = 1 + _, __, en2_kernel, ___ = str_to_kernel_set(['twobody'], 'mc', hm) + bhyps, bcutoffs, bhyps_mask = Parameters.get_component_mask(hm, 'twobody', hyps=hyps) + args2 = from_mask_to_args(bhyps, bcutoffs, bhyps_mask) calc1 = en2_kernel(env1[1][0], env2[1][0], *args2) calc2 = en2_kernel(env1[2][0], env2[2][0], *args2) @@ -386,14 +381,13 @@ def test_force(kernel_name, diff_cutoff): calc4 = en2_kernel(env1[2][0], env2[1][0], *args2) kern_finite_diff += 4 * (calc1 + calc2 - calc3 - calc4) / (4*delta**2) else: - nbond = 0 + ntwobody = 0 - if ('3' in kernel_name): - _, __, en3_kernel, ___ = str_to_kernel_set('3'+kernel_type, True) + if ('threebody' in kernels): + _, __, en3_kernel, ___ = str_to_kernel_set(['threebody'], 'mc', hm) - thyps, thyps_mask = HyperParameterMasking.get_3b_hyps( - hyps, hm, True) - args3 = from_mask_to_args(thyps, thyps_mask, cutoffs[:2]) + thyps, tcutoffs, thyps_mask = Parameters.get_component_mask(hm, 'threebody', hyps=hyps) + args3 = from_mask_to_args(thyps, tcutoffs, thyps_mask) calc1 = en3_kernel(env1[1][0], env2[1][0], *args3) calc2 = en3_kernel(env1[2][0], env2[2][0], *args3) @@ -401,25 +395,26 @@ def test_force(kernel_name, diff_cutoff): calc4 = en3_kernel(env1[2][0], env2[1][0], *args3) kern_finite_diff += 9 * (calc1 + calc2 - calc3 - calc4) / (4*delta**2) - args = from_mask_to_args(hyps, hm, cutoffs) + args = from_mask_to_args(hyps, cutoffs, hm) kern_analytical = kernel(env1[0][0], env2[0][0], d1, d2, *args) assert(isclose(kern_finite_diff, kern_analytical, rtol=tol)) -@pytest.mark.parametrize('kernel_name', list_to_test) +@pytest.mark.parametrize('kernels', list_to_test) @pytest.mark.parametrize('diff_cutoff', multi_cut) @pytest.mark.parametrize('constraint', [True, False]) -def test_hyps_grad(kernel_name, diff_cutoff, constraint): +def test_hyps_grad(kernels, diff_cutoff, constraint): delta = 1e-8 d1 = 1 d2 = 2 tol = 1e-4 - cutoffs, hyps, hm = generate_diff_hm(kernel_name, diff_cutoff, constraint=constraint) - args = from_mask_to_args(hyps, hm, cutoffs) - kernel, kernel_grad, _, __ = str_to_kernel_set(kernel_name, True) + np.random.seed(10) + cutoffs, hyps, hm = generate_diff_hm(kernels, diff_cutoff, constraint=constraint) + args = from_mask_to_args(hyps, cutoffs, hm) + kernel, kernel_grad, _, __ = str_to_kernel_set(kernels, "mc", hm) np.random.seed(0) env1 = generate_mb_envs(cutoffs, np.eye(3)*100, delta, d1) @@ -427,28 +422,26 @@ def test_hyps_grad(kernel_name, diff_cutoff, constraint): env1 = env1[0][0] env2 = env2[0][0] - # compute analytical values k, grad = kernel_grad(env1, env2, d1, d2, *args) original = kernel(env1, env2, d1, d2, *args) - nhyps = len(hyps)-1 - if ('map' in hm.keys()): - if (hm['map'][-1] != (len(hm['original'])-1)): - nhyps = len(hyps) - original_hyps = np.copy(hm['original']) + nhyps = len(hyps) + if hm['train_noise']: + nhyps -= 1 + original_hyps = Parameters.get_hyps(hm, hyps=hyps) for i in range(nhyps): newhyps = np.copy(hyps) newhyps[i] += delta if ('map' in hm.keys()): newid = hm['map'][i] - hm['original'] = np.copy(original_hyps) - hm['original'][newid] += delta - newargs = from_mask_to_args(newhyps, hm, cutoffs) + hm['original_hyps'] = np.copy(original_hyps) + hm['original_hyps'][newid] += delta + newargs = from_mask_to_args(newhyps, cutoffs, hm) hgrad = (kernel(env1, env2, d1, d2, *newargs) - original)/delta - if ('map' in hm.keys()): + if 'map' in hm: print(i, "hgrad", hgrad, grad[hm['map'][i]]) assert(isclose(grad[hm['map'][i]], hgrad, rtol=tol)) else: @@ -456,114 +449,114 @@ def test_hyps_grad(kernel_name, diff_cutoff, constraint): assert(isclose(grad[i], hgrad, rtol=tol)) -def generate_same_hm(kernel_name, multi_cutoff=False): +def generate_same_hm(kernels, multi_cutoff=False): """ generate hyperparameter and masks that are effectively the same but with single or multi group expression """ - pm1 = HyperParameterMasking(species=['H', 'He'], + pm1 = ParameterHelper(species=['H', 'He'], parameters={'noise':0.05}) - pm2 = HyperParameterMasking(species=['H', 'He'], + pm2 = ParameterHelper(species=['H', 'He'], parameters={'noise':0.05}) - if ('2' in kernel_name): + if ('twobody' in kernels): para = 2.5+0.1*random(3) - pm1.set_parameters('cutoff2b', para[-1]) - pm1.define_group('bond', 'bond0', ['*', '*'], para[:-1]) + pm1.set_parameters('cutoff_twobody', para[-1]) + pm1.define_group('twobody', 'twobody0', ['*', '*'], para[:-1]) - pm2.set_parameters('cutoff2b', para[-1]) - pm2.define_group('bond', 'bond0', ['*', '*'], para[:-1]) - pm2.define_group('bond', 'bond1', ['H', 'H'], para[:-1]) + pm2.set_parameters('cutoff_twobody', para[-1]) + pm2.define_group('twobody', 'twobody0', ['*', '*'], para[:-1]) + pm2.define_group('twobody', 'twobody1', ['H', 'H'], para[:-1]) if (multi_cutoff): - pm2.set_parameters('bond0', para) - pm2.set_parameters('bond1', para) + pm2.set_parameters('twobody0', para) + pm2.set_parameters('twobody1', para) - if ('3' in kernel_name): + if ('threebody' in kernels): para = 1.2+0.1*random(3) - pm1.set_parameters('cutoff3b', para[-1]) - pm1.define_group('triplet', 'triplet0', ['*', '*', '*'], para[:-1]) + pm1.set_parameters('cutoff_threebody', para[-1]) + pm1.define_group('threebody', 'threebody0', ['*', '*', '*'], para[:-1]) - pm2.set_parameters('cutoff3b', para[-1]) - pm2.define_group('triplet', 'triplet0', ['*', '*', '*'], para[:-1]) - pm2.define_group('triplet', 'triplet1', ['H', 'H', 'H'], para[:-1]) + pm2.set_parameters('cutoff_threebody', para[-1]) + pm2.define_group('threebody', 'threebody0', ['*', '*', '*'], para[:-1]) + pm2.define_group('threebody', 'threebody1', ['H', 'H', 'H'], para[:-1]) if (multi_cutoff): pm2.define_group('cut3b', 'c1', ['*', '*'], parameters=para) pm2.define_group('cut3b', 'c2', ['H', 'H'], parameters=para) - if ('mb' in kernel_name): + if ('manybody' in kernels): para = 1.2+0.1*random(3) - pm1.set_parameters('cutoffmb', para[-1]) - pm1.define_group('mb', 'mb0', ['*', '*'], para[:-1]) + pm1.set_parameters('cutoff_manybody', para[-1]) + pm1.define_group('manybody', 'manybody0', ['*', '*'], para[:-1]) - pm2.set_parameters('cutoffmb', para[-1]) - pm2.define_group('mb', 'mb0', ['*', '*'], para[:-1]) - pm2.define_group('mb', 'mb1', ['H', 'H'], para[:-1]) + pm2.set_parameters('cutoff_manybody', para[-1]) + pm2.define_group('manybody', 'manybody0', ['*', '*'], para[:-1]) + pm2.define_group('manybody', 'manybody1', ['H', 'H'], para[:-1]) if (multi_cutoff): - pm2.set_parameters('mb0', para) - pm2.set_parameters('mb1', para) + pm2.set_parameters('manybody0', para) + pm2.set_parameters('manybody1', para) - hm1 = pm1.generate_dict() + hm1 = pm1.as_dict() hyps1 = hm1['hyps'] cut = hm1['cutoffs'] - hm2 = pm2.generate_dict() + hm2 = pm2.as_dict() hyps2 = hm2['hyps'] cut = hm2['cutoffs'] return cut, hyps1, hyps2, hm1, hm2 -def generate_diff_hm(kernel_name, diff_cutoff=False, constraint=False): +def generate_diff_hm(kernels, diff_cutoff=False, constraint=False): - pm = HyperParameterMasking(species=['H', 'He'], + pm = ParameterHelper(species=['H', 'He'], parameters={'noise':0.05}) - if ('2' in kernel_name): + if ('twobody' in kernels): para1 = 2.5+0.1*random(3) para2 = 2.5+0.1*random(3) - pm.set_parameters('cutoff2b', para1[-1]) - pm.define_group('bond', 'bond0', ['*', '*']) - pm.set_parameters('bond0', para1[:-1], not constraint) - pm.define_group('bond', 'bond1', ['H', 'H'], para2[:-1]) + pm.set_parameters('cutoff_twobody', para1[-1]) + pm.define_group('twobody', 'twobody0', ['*', '*']) + pm.set_parameters('twobody0', para1[:-1], not constraint) + pm.define_group('twobody', 'twobody1', ['H', 'H'], para2[:-1]) if (diff_cutoff): - pm.set_parameters('bond0', para1, not constraint) - pm.set_parameters('bond1', para2) + pm.set_parameters('twobody0', para1, not constraint) + pm.set_parameters('twobody1', para2) - if ('3' in kernel_name): + if ('threebody' in kernels): para1 = 1.2+0.1*random(3) para2 = 1.2+0.1*random(3) - pm.set_parameters('cutoff3b', para1[-1]) - pm.define_group('triplet', 'triplet0', ['*', '*', '*'], para1[:-1]) - pm.set_parameters('triplet0', para1[:-1], not constraint) - pm.define_group('triplet', 'triplet1', ['H', 'H', 'H'], para2[:-1]) + pm.set_parameters('cutoff_threebody', para1[-1]) + pm.define_group('threebody', 'threebody0', ['*', '*', '*'], para1[:-1]) + pm.set_parameters('threebody0', para1[:-1], not constraint) + pm.define_group('threebody', 'threebody1', ['H', 'H', 'H'], para2[:-1]) if (diff_cutoff): pm.define_group('cut3b', 'c1', ['*', '*'], parameters=para1) pm.define_group('cut3b', 'c2', ['H', 'H'], parameters=para2) - if ('mb' in kernel_name): + if ('manybody' in kernels): para1 = 1.2+0.1*random(3) para2 = 1.2+0.1*random(3) - pm.set_parameters('cutoffmb', para1[-1]) - pm.define_group('mb', 'mb0', ['*', '*']) - pm.set_parameters('mb0', para1[:-1], not constraint) - pm.define_group('mb', 'mb1', ['H', 'H'], para2[:-1]) + pm.set_parameters('cutoff_manybody', para1[-1]) + pm.define_group('manybody', 'manybody0', ['*', '*']) + pm.set_parameters('manybody0', para1[:-1], not constraint) + pm.define_group('manybody', 'manybody1', ['H', 'H'], para2[:-1]) if (diff_cutoff): - pm.set_parameters('mb0', para1, not constraint) - pm.set_parameters('mb1', para2) + pm.set_parameters('manybody0', para1, not constraint) + pm.set_parameters('manybody1', para2) - hm = pm.generate_dict() + hm = pm.as_dict() hyps = hm['hyps'] cut = hm['cutoffs'] diff --git a/tests/test_mgp.py b/tests/test_mgp.py new file mode 100644 index 000000000..c054e0534 --- /dev/null +++ b/tests/test_mgp.py @@ -0,0 +1,403 @@ +import numpy as np +import os +import pickle +import pytest +import re +import time +import shutil + +from copy import deepcopy +from numpy import allclose, isclose + +from flare import struc, env, gp +from flare.parameters import Parameters +from flare.mgp import MappedGaussianProcess +from flare.lammps import lammps_calculator +from flare.utils.element_coder import _Z_to_mass, _Z_to_element + +from .fake_gp import get_gp, get_random_structure + +body_list = ['2', '3'] +multi_list = [False, True] +map_force_list = [False, True] +force_block_only = False + +def clean(): + for f in os.listdir("./"): + if re.search("mgp_grids", f): + shutil.rmtree(f) + if re.search("kv3", f): + os.rmdir(f) + if 'tmp' in f: + if os.path.isdir(f): + shutil.rmtree(f) + else: + os.remove(f) + + +@pytest.mark.skipif(not os.environ.get('lmp', + False), reason='lmp not found ' + 'in environment: Please install LAMMPS ' + 'and set the $lmp env. ' + 'variable to point to the executatble.') +@pytest.fixture(scope='module') +def all_gp(): + + allgp_dict = {} + np.random.seed(0) + for bodies in body_list: + for multihyps in multi_list: + gp_model = get_gp(bodies, 'mc', multihyps, cellabc=[1.5, 1, 2], + force_only=force_block_only, noa=5) #int(bodies)**2) + gp_model.parallel = True + gp_model.n_cpus = 2 + + allgp_dict[f'{bodies}{multihyps}'] = gp_model + + yield allgp_dict + del allgp_dict + +@pytest.fixture(scope='module') +def all_mgp(): + + allmgp_dict = {} + for bodies in ['2', '3', '2+3']: + for multihyps in [False, True]: + allmgp_dict[f'{bodies}{multihyps}'] = None + + yield allmgp_dict + del allmgp_dict + +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_init(bodies, multihyps, map_force, all_mgp, all_gp): + """ + test the init function + """ + + gp_model = all_gp[f'{bodies}{multihyps}'] + + # grid parameters + grid_params = {} + if ('2' in bodies): + grid_params['twobody'] = {'grid_num': [64], 'lower_bound': [0.05]} + if ('3' in bodies): + grid_params['threebody'] = {'grid_num': [24, 25, 26], 'lower_bound':[0.05]*3} + + lammps_location = f'{bodies}{multihyps}{map_force}.mgp' + data = gp_model.training_statistics + + mgp_model = MappedGaussianProcess(grid_params=grid_params, unique_species=data['species'], n_cpus=1, + map_force=map_force, lmp_file_name=lammps_location)#, mean_only=False) + all_mgp[f'{bodies}{multihyps}{map_force}'] = mgp_model + + + +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_build_map(all_gp, all_mgp, bodies, multihyps, map_force): + """ + test the mapping for mc_simple kernel + """ + gp_model = all_gp[f'{bodies}{multihyps}'] + mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] + mgp_model.build_map(gp_model) +# with open(f'grid_{bodies}_{multihyps}_{map_force}.pickle', 'wb') as f: +# pickle.dump(mgp_model, f) + + +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_write_model(all_mgp, bodies, multihyps, map_force): + """ + test the mapping for mc_simple kernel + """ + mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] + mgp_model.mean_only = True + mgp_model.write_model(f'my_mgp_{bodies}_{multihyps}_{map_force}') + + mgp_model.write_model(f'my_mgp_{bodies}_{multihyps}_{map_force}', format='pickle') + + # Ensure that user is warned when a non-mean_only + # model is serialized into a Dictionary + with pytest.warns(Warning): + mgp_model.mean_only = False + mgp_model.as_dict() + mgp_model.mean_only = True + + +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_load_model(all_mgp, bodies, multihyps, map_force): + """ + test the mapping for mc_simple kernel + """ + name = f'my_mgp_{bodies}_{multihyps}_{map_force}.json' + all_mgp[f'{bodies}{multihyps}'] = MappedGaussianProcess.from_file(name) + os.remove(name) + + name = f'my_mgp_{bodies}_{multihyps}_{map_force}.pickle' + all_mgp[f'{bodies}{multihyps}'] = MappedGaussianProcess.from_file(name) + os.remove(name) + +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_cubic_spline(all_gp, all_mgp, bodies, multihyps, map_force): + """ + test the predict for mc_simple kernel + """ + + mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] + delta = 1e-4 + + if '3' in bodies: + body_name = 'threebody' + elif '2' in bodies: + body_name = 'twobody' + + nmap = len(mgp_model.maps[body_name].maps) + print('nmap', nmap) + for i in range(nmap): + maxvalue = np.max(np.abs(mgp_model.maps[body_name].maps[i].mean.__coeffs__)) + if maxvalue >0: + comp_code = mgp_model.maps[body_name].maps[i].species_code + + if '3' in bodies: + + c_pt = np.array([[0.3, 0.4, 0.5]]) + c, cderv = mgp_model.maps[body_name].maps[i].mean(c_pt, with_derivatives=True) + cderv = cderv.reshape([-1]) + + for j in range(3): + a_pt = deepcopy(c_pt) + b_pt = deepcopy(c_pt) + a_pt[0][j]+=delta + b_pt[0][j]-=delta + a = mgp_model.maps[body_name].maps[i].mean(a_pt)[0] + b = mgp_model.maps[body_name].maps[i].mean(b_pt)[0] + num_derv = (a-b)/(2*delta) + print("spline", comp_code, num_derv, cderv[j]) + assert np.isclose(num_derv, cderv[j], rtol=1e-2) + + elif '2' in bodies: + center = np.sum(mgp_model.maps[body_name].maps[i].bounds)/2. + a_pt = np.array([[center+delta]]) + b_pt = np.array([[center-delta]]) + c_pt = np.array([[center]]) + a = mgp_model.maps[body_name].maps[i].mean(a_pt)[0] + b = mgp_model.maps[body_name].maps[i].mean(b_pt)[0] + c, cderv = mgp_model.maps[body_name].maps[i].mean(c_pt, with_derivatives=True) + cderv = cderv.reshape([-1])[0] + num_derv = (a-b)/(2*delta) + print("spline", num_derv, cderv) + assert np.isclose(num_derv, cderv, rtol=1e-2) + +def compare_triplet(mgp_model, gp_model, atom_env): + spcs, comp_r, comp_xyz = mgp_model.get_arrays(atom_env) + for i, spc in enumerate(spcs): + lengths = np.array(comp_r[i]) + xyzs = np.array(comp_xyz[i]) + + print('compare triplet spc, lengths, xyz', spc) + print(np.hstack([lengths, xyzs])) + + gp_f = [] + gp_e = [] + grid_env = get_grid_env(gp_model, spc, 3) + for l in range(lengths.shape[0]): + r1, r2, r12 = lengths[l, :] + grid_env = get_triplet_env(r1, r2, r12, grid_env) + gp_pred = np.array([gp_model.predict(grid_env, d+1) for d in range(3)]).T + gp_en, _ = gp_model.predict_local_energy_and_var(grid_env) + gp_f.append(gp_pred[0]) + gp_e.append(gp_en) + gp_force = np.sum(gp_f, axis=0) + gp_energy = np.sum(gp_e, axis=0) + print('gp_e', gp_e) + print('gp_f') + print(gp_f) + + map_ind = mgp_model.find_map_index(spc) + xyzs = np.zeros_like(xyzs) + xyzs[:, 0] = np.ones_like(xyzs[:, 0]) + f, _, _, e = mgp_model.maps[map_ind].predict(lengths, xyzs, + mgp_model.map_force, mean_only=True) + + assert np.allclose(gp_force, f, rtol=1e-2) + if not mgp_model.map_force: + assert np.allclose(gp_energy, e, rtol=1e-2) + + +def get_triplet_env(r1, r2, r12, grid_env): + grid_env.bond_array_3 = np.array([[r1, 1, 0, 0], [r2, 0, 0, 0]]) + grid_env.cross_bond_dists = np.array([[0, r12], [r12, 0]]) + print(grid_env.ctype, grid_env.etypes) + + return grid_env + + +def get_grid_env(GP, species, bodies): + if isinstance(GP.cutoffs, dict): + max_cut = np.max(list(GP.cutoffs.values())) + else: + max_cut = np.max(GP.cutoffs) + big_cell = np.eye(3) * 100 + positions = [[(i+1)/(bodies+1)*0.1, 0, 0] + for i in range(bodies)] + grid_struc = struc.Structure(big_cell, species, positions) + grid_env = env.AtomicEnvironment(grid_struc, 0, GP.cutoffs, + cutoffs_mask=GP.hyps_mask) + + return grid_env + + +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_predict(all_gp, all_mgp, bodies, multihyps, map_force): + """ + test the predict for mc_simple kernel + """ + + gp_model = all_gp[f'{bodies}{multihyps}'] + mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] + + # # debug + # filename = f'grid_{bodies}_{multihyps}_{map_force}.pickle' + # with open(filename, 'rb') as f: + # mgp_model = pickle.load(f) + + nenv = 3 + cell = 1.0 * np.eye(3) + cutoffs = gp_model.cutoffs + unique_species = gp_model.training_statistics['species'] + struc_test, f = get_random_structure(cell, unique_species, nenv) + test_envi = env.AtomicEnvironment(struc_test, 0, cutoffs, cutoffs_mask=gp_model.hyps_mask) + + if '2' in bodies: + kernel_name = 'twobody' + elif '3' in bodies: + kernel_name = 'threebody' + compare_triplet(mgp_model.maps['threebody'], gp_model, test_envi) + + assert Parameters.compare_dict(gp_model.hyps_mask, mgp_model.maps[kernel_name].hyps_mask) + + gp_pred_en, gp_pred_envar = gp_model.predict_local_energy_and_var(test_envi) + gp_pred = np.array([gp_model.predict(test_envi, d+1) for d in range(3)]).T + mgp_pred = mgp_model.predict(test_envi, mean_only=True) + + + # check mgp is within 2 meV/A of the gp + if map_force: + map_str = 'force' + gp_pred_var = gp_pred[1] + else: + map_str = 'energy' + gp_pred_var = gp_pred_envar + # TODO: energy block accuracy +# assert(np.abs(mgp_pred[3] - gp_pred_en) < 2e-3), \ +# f"{bodies} body {map_str} mapping is wrong" + +# if multihyps and ('3' in bodies): +# pytest.skip() + + print('mgp_pred', mgp_pred[0]) + print('gp_pred', gp_pred[0]) + + print("isclose?", mgp_pred[0]-gp_pred[0], gp_pred[0]) + assert(np.allclose(mgp_pred[0], gp_pred[0], atol=5e-3)), \ + f"{bodies} body {map_str} mapping is wrong" + + # TODO: energy block accuracy +# assert(np.abs(mgp_pred[1] - gp_pred_var) < 2e-3), \ +# f"{bodies} body {map_str} mapping var is wrong" + + clean() + +@pytest.mark.skipif(not os.environ.get('lmp', + False), reason='lmp not found ' + 'in environment: Please install LAMMPS ' + 'and set the $lmp env. ' + 'variable to point to the executatble.') +@pytest.mark.parametrize('bodies', body_list) +@pytest.mark.parametrize('multihyps', multi_list) +@pytest.mark.parametrize('map_force', map_force_list) +def test_lmp_predict(all_gp, all_mgp, bodies, multihyps, map_force): + """ + test the lammps implementation + """ + clean() + prefix = f'tmp{bodies}{multihyps}{map_force}' + + if ('3' in bodies) and map_force: + pytest.skip() + + mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] + gp_model = all_gp[f'{bodies}{multihyps}'] + lammps_location = mgp_model.lmp_file_name + + # lmp file is automatically written now every time MGP is constructed + mgp_model.write_lmp_file(lammps_location) + + # create test structure + cell = 5*np.eye(3) + nenv = 10 + unique_species = gp_model.training_data[0].species + cutoffs = gp_model.cutoffs + struc_test, f = get_random_structure(cell, unique_species, nenv) + atom_num = 1 + test_envi = env.AtomicEnvironment(struc_test, atom_num, cutoffs, cutoffs_mask=gp_model.hyps_mask) + + all_species=list(set(struc_test.coded_species)) + atom_types = list(np.arange(len(all_species))+1) + atom_masses=[_Z_to_mass[spec] for spec in all_species] + atom_species = [ all_species.index(spec)+1 for spec in struc_test.coded_species] + specie_symbol_list = " ".join([_Z_to_element[spec] for spec in all_species]) + + # create data file + data_file_name = f'{prefix}.data' + data_text = lammps_calculator.lammps_dat(struc_test, atom_types, + atom_masses, atom_species) + lammps_calculator.write_text(data_file_name, data_text) + + # create lammps input + by = 'no' + ty = 'no' + if '2' in bodies: + by = 'yes' + if '3' in bodies: + ty = 'yes' + + if map_force: + style_string = 'mgpf' + else: + style_string = 'mgp' + + coeff_string = f'* * {lammps_location} {specie_symbol_list} {by} {ty}' + lammps_executable = os.environ.get('lmp') + dump_file_name = f'{prefix}.dump' + input_file_name = f'{prefix}.in' + output_file_name = f'{prefix}.out' + input_text = \ + lammps_calculator.generic_lammps_input(data_file_name, style_string, + coeff_string, dump_file_name, + newton=True) + lammps_calculator.write_text(input_file_name, input_text) + + lammps_calculator.run_lammps(lammps_executable, input_file_name, + output_file_name) + + lammps_forces = lammps_calculator.lammps_parser(dump_file_name) + mgp_forces = mgp_model.predict(test_envi, mean_only=True) + + # check that lammps agrees with gp to within 1 meV/A + for i in range(3): + print("isclose? diff:", lammps_forces[atom_num, i]-mgp_forces[0][i], "mgp value", mgp_forces[0][i]) + assert np.isclose(lammps_forces[atom_num, i], mgp_forces[0][i], rtol=1e-2) + + clean() diff --git a/tests/test_mgp_unit.py b/tests/test_mgp_unit.py deleted file mode 100644 index 6479350a8..000000000 --- a/tests/test_mgp_unit.py +++ /dev/null @@ -1,277 +0,0 @@ -import numpy as np -from numpy import allclose, isclose -import time -import pytest -import os, pickle, re - -from flare import struc, env, gp -from flare import otf_parser -from flare.mgp.mgp import MappedGaussianProcess -from flare.lammps import lammps_calculator - -from .fake_gp import get_gp, get_random_structure - -body_list = ['2', '3'] -multi_list = [False, True] -map_force_list = [False, True] - -def clean(): - for f in os.listdir("./"): - if re.search(r"grid.*npy", f): - os.remove(f) - if re.search("kv3", f): - os.rmdir(f) - - -# ASSUMPTION: You have a Lammps executable with the mgp pair style with $lmp -# as the corresponding environment variable. -@pytest.mark.skipif(not os.environ.get('lmp', - False), reason='lmp not found ' - 'in environment: Please install LAMMPS ' - 'and set the $lmp env. ' - 'variable to point to the executatble.') - -@pytest.fixture(scope='module') -def all_gp(): - - allgp_dict = {} - np.random.seed(0) - for bodies in ['2', '3', '2+3']: - for multihyps in [False, True]: - gp_model = get_gp(bodies, 'mc', multihyps, cellabc=[100, 100, 100]) - gp_model.parallel = True - gp_model.n_cpus = 2 - allgp_dict[f'{bodies}{multihyps}'] = gp_model - - yield allgp_dict - del allgp_dict - -@pytest.fixture(scope='module') -def all_mgp(): - - allmgp_dict = {} - for bodies in ['2', '3', '2+3']: - for multihyps in [False, True]: - allmgp_dict[f'{bodies}{multihyps}'] = None - - yield allmgp_dict - del allmgp_dict - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -@pytest.mark.parametrize('map_force', map_force_list) -def test_init(bodies, multihyps, map_force, all_mgp, all_gp): - """ - test the init function - """ - - gp_model = all_gp[f'{bodies}{multihyps}'] - - grid_num_2 = 64 - grid_num_3 = 25 - lower_cut = 0.01 - two_cut = gp_model.cutoffs[0] - three_cut = gp_model.cutoffs[1] - if map_force: - lower_cut_3 = -1 - three_cut_3 = 1 - else: - lower_cut_3 = lower_cut - three_cut_3 = three_cut - lammps_location = f'{bodies}{multihyps}{map_force}.mgp' - - # set struc params. cell and masses arbitrary? - mapped_cell = np.eye(3) * 2 - struc_params = {'species': [1, 2], - 'cube_lat': mapped_cell, - 'mass_dict': {'0': 27, '1': 16}} - - # grid parameters - blist = [] - if ('2' in bodies): - blist += [2] - if ('3' in bodies): - blist += [3] - train_size = len(gp_model.training_data) - grid_params = {'bodies': blist, - 'cutoffs':gp_model.cutoffs, - 'bounds_2': [[lower_cut], [two_cut]], - 'bounds_3': [[lower_cut, lower_cut, lower_cut_3], - [three_cut, three_cut, three_cut_3]], - 'grid_num_2': grid_num_2, - 'grid_num_3': [grid_num_3, grid_num_3, grid_num_3], - 'svd_rank_2': 14, - 'svd_rank_3': 14, - 'load_grid': None, - 'update': False} - - struc_params = {'species': [1, 2], - 'cube_lat': np.eye(3)*2, - 'mass_dict': {'0': 27, '1': 16}} - mgp_model = MappedGaussianProcess(grid_params, struc_params, n_cpus=4, - map_force=map_force, lmp_file_name=lammps_location) - all_mgp[f'{bodies}{multihyps}{map_force}'] = mgp_model - - - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -@pytest.mark.parametrize('map_force', map_force_list) -def test_build_map(all_gp, all_mgp, bodies, multihyps, map_force): - """ - test the mapping for mc_simple kernel - """ - gp_model = all_gp[f'{bodies}{multihyps}'] - mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] - mgp_model.build_map(gp_model) - - - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -@pytest.mark.parametrize('map_force', map_force_list) -def test_write_model(all_mgp, bodies, multihyps, map_force): - """ - test the mapping for mc_simple kernel - """ - mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] - mgp_model.mean_only = True - mgp_model.write_model(f'my_mgp_{bodies}_{multihyps}_{map_force}') - - mgp_model.write_model(f'my_mgp_{bodies}_{multihyps}_{map_force}', format='pickle') - - # Ensure that user is warned when a non-mean_only - # model is serialized into a Dictionary - with pytest.warns(Warning): - mgp_model.mean_only = False - mgp_model.as_dict() - mgp_model.mean_only = True - - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -@pytest.mark.parametrize('map_force', map_force_list) -def test_load_model(all_mgp, bodies, multihyps, map_force): - """ - test the mapping for mc_simple kernel - """ - name = f'my_mgp_{bodies}_{multihyps}_{map_force}.json' - all_mgp[f'{bodies}{multihyps}'] = MappedGaussianProcess.from_file(name) - os.remove(name) - - name = f'my_mgp_{bodies}_{multihyps}_{map_force}.pickle' - all_mgp[f'{bodies}{multihyps}'] = MappedGaussianProcess.from_file(name) - os.remove(name) - - -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -@pytest.mark.parametrize('map_force', map_force_list) -def test_predict(all_gp, all_mgp, bodies, multihyps, map_force): - """ - test the predict for mc_simple kernel - """ - gp_model = all_gp[f'{bodies}{multihyps}'] - mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] - - nenv=10 - cell = np.eye(3) - cutoffs = gp_model.cutoffs - unique_species = gp_model.training_data[0].species - struc_test, f = get_random_structure(cell, unique_species, nenv) - test_envi = env.AtomicEnvironment(struc_test, 1, cutoffs) - - gp_pred_en = gp_model.predict_local_energy(test_envi) - gp_pred_f = [gp_model.predict(test_envi, d+1)[0] for d in range(3)] - mgp_pred = mgp_model.predict(test_envi, mean_only=True) - - # check mgp is within 2 meV/A of the gp - if not map_force: - assert(np.abs(mgp_pred[3] - gp_pred_en) < 2e-3), \ - f"{bodies} body energy mapping is wrong" - assert(np.abs(mgp_pred[0][0] - gp_pred_f[0]) < 2e-3), \ - f"{bodies} body mapping is wrong" - - clean() - -@pytest.mark.skipif(not os.environ.get('lmp', - False), reason='lmp not found ' - 'in environment: Please install LAMMPS ' - 'and set the $lmp env. ' - 'variable to point to the executatble.') -@pytest.mark.parametrize('bodies', body_list) -@pytest.mark.parametrize('multihyps', multi_list) -@pytest.mark.parametrize('map_force', map_force_list) -def test_lmp_predict(all_gp, all_mgp, bodies, multihyps, map_force): - """ - test the lammps implementation - """ - prefix = f'tmp{bodies}{multihyps}{map_force}' - for f in os.listdir("./"): - if prefix in f: - os.remove(f) - clean() - - mgp_model = all_mgp[f'{bodies}{multihyps}{map_force}'] - gp_model = all_gp[f'{bodies}{multihyps}'] - lammps_location = mgp_model.lmp_file_name - - # lmp file is automatically written now every time MGP is constructed - mgp_model.write_lmp_file(lammps_location) - - # create test structure - cell = np.eye(3) - nenv = 10 - unique_species = gp_model.training_data[0].species - cutoffs = gp_model.cutoffs - struc_test, f = get_random_structure(cell, unique_species, nenv) - atom_num = 1 - test_envi = env.AtomicEnvironment(struc_test, atom_num, cutoffs) - atom_types = [1, 2] - atom_masses = [108, 127] - atom_species = struc_test.coded_species - - # create data file - data_file_name = f'{prefix}.data' - data_text = lammps_calculator.lammps_dat(struc_test, atom_types, - atom_masses, atom_species) - lammps_calculator.write_text(data_file_name, data_text) - - # create lammps input - by = 'no' - ty = 'no' - if '2' in bodies: - by = 'yes' - if '3' in bodies: - ty = 'yes' - - if map_force: - style_string = 'mgpf' - else: - style_string = 'mgp' - - coeff_string = f'* * {lammps_location} H He {by} {ty}' - lammps_executable = os.environ.get('lmp') - dump_file_name = f'{prefix}.dump' - input_file_name = f'{prefix}.in' - output_file_name = f'{prefix}.out' - input_text = \ - lammps_calculator.generic_lammps_input(data_file_name, style_string, - coeff_string, dump_file_name, - newton=True) - lammps_calculator.write_text(input_file_name, input_text) - - lammps_calculator.run_lammps(lammps_executable, input_file_name, - output_file_name) - - lammps_forces = lammps_calculator.lammps_parser(dump_file_name) - mgp_forces = mgp_model.predict(test_envi, mean_only=True) - - # check that lammps agrees with gp to within 1 meV/A - for i in range(3): - assert (np.abs(lammps_forces[atom_num, i] - mgp_forces[0][i]) < 1e-3) - - for f in os.listdir("./"): - if prefix in f: - os.remove(f) - clean() diff --git a/tests/test_otf.py b/tests/test_otf.py index 2c6e6efb6..c0f8ff3d8 100644 --- a/tests/test_otf.py +++ b/tests/test_otf.py @@ -1,4 +1,5 @@ import pytest +import os import glob, os, re, shutil import numpy as np @@ -13,12 +14,16 @@ example_list = [1, 2] name_list = {1:'h2', 2:'al'} +print('running test_otf.py') +print('current working directory:') +print(os.getcwd()) + def get_gp(par=False, per_atom_par=False, n_cpus=1): hyps = np.array([1, 1, 1, 1, 1]) hyp_labels = ['Signal Std 2b', 'Length Scale 2b', 'Signal Std 3b', 'Length Scale 3b', 'Noise Std'] - cutoffs = np.array([4, 4]) + cutoffs = {'twobody':4, 'threebody':4} return GaussianProcess(\ kernel_name='23mc', hyps=hyps, cutoffs=cutoffs, hyp_labels=hyp_labels, maxiter=50, par=par, @@ -55,6 +60,10 @@ def test_otf(software, example): :return: """ + print('running test_otf.py') + print('current working directory:') + print(os.getcwd()) + outdir = f'test_outputs_{software}' if os.path.isdir(outdir): shutil.rmtree(outdir) @@ -78,15 +87,17 @@ def test_otf(software, example): gp = get_gp() - otf = OTF(dft_input, dt, number_of_steps, gp, dft_loc, - std_tolerance_factor, init_atoms=[0], + otf = OTF(dt=dt, number_of_steps=number_of_steps, + gp=gp, write_model=3, + std_tolerance_factor=std_tolerance_factor, + init_atoms=[0], calculate_energy=True, max_atoms_added=1, freeze_hyps=1, skip=1, force_source=software, + dft_input=dft_input, dft_loc=dft_loc, dft_output=dft_output, output_name=f'{casename}_otf_{software}', - store_dft_output=([dft_output, dft_input], '.'), - write_model=3) + store_dft_output=([dft_output, dft_input], '.')) otf.run() @@ -105,11 +116,16 @@ def test_otf_par(software, per_atom_par, n_cpus): Test that an otf run can survive going for more steps :return: """ + example = 1 outdir = f'test_outputs_{software}' if os.path.isdir(outdir): shutil.rmtree(outdir) + print('running test_otf.py') + print('current working directory:') + print(os.getcwd()) + if (not os.environ.get(cmd[software], False)): pytest.skip(f'{cmd[software]} not found in environment:' f' Please install the code ' @@ -133,12 +149,13 @@ def test_otf_par(software, per_atom_par, n_cpus): gp = get_gp(par=True, n_cpus=n_cpus, per_atom_par=per_atom_par) - otf = OTF(dft_input, dt, number_of_steps, gp, dft_loc, - std_tolerance_factor, init_atoms=[0], + otf = OTF(dt=dt, number_of_steps=number_of_steps, gp=gp, + std_tolerance_factor=std_tolerance_factor, init_atoms=[0], calculate_energy=True, max_atoms_added=1, - par=True, n_cpus=n_cpus, + n_cpus=n_cpus, freeze_hyps=1, skip=1, mpi="mpi", force_source=software, + dft_input=dft_input, dft_loc=dft_loc, dft_output=dft_output, output_name=f'{casename}_otf_{software}', store_dft_output=([dft_output, dft_input], '.')) diff --git a/tests/test_parameters.py b/tests/test_parameters.py new file mode 100644 index 000000000..518a30363 --- /dev/null +++ b/tests/test_parameters.py @@ -0,0 +1,169 @@ +import pytest +import numpy as np +from flare.struc import Structure +from flare.utils.parameter_helper import ParameterHelper +from flare.parameters import Parameters + + +def test_initialization(): + ''' + simplest senario + ''' + pm = ParameterHelper(kernels=['twobody', 'threebody'], + parameters={'twobody': [1, 0.5], + 'threebody': [1, 0.5], + 'cutoff_twobody': 2, + 'cutoff_threebody': 1, + 'noise': 0.05}, + verbose="DEBUG") + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +@pytest.mark.parametrize('ones', [True, False]) +def test_initialization2(ones): + ''' + simplest senario + ''' + pm = ParameterHelper(kernels=['twobody', 'threebody'], + parameters={'cutoff_twobody': 2, + 'cutoff_threebody': 1, + 'noise': 0.05}, + ones=ones, + random=not ones, + verbose="DEBUG") + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +def test_initialization3(): + pm = ParameterHelper(species=['O', 'C', 'H'], + kernels={'twobody': [['*', '*'], ['O', 'O']], + 'threebody': [['*', '*', '*'], ['O', 'O', 'O']]}, + parameters={'twobody0': [1, 0.5], 'twobody1': [2, 0.2], + 'threebody0': [1, 0.5], 'threebody1': [2, 0.2], + 'cutoff_twobody': 2, 'cutoff_threebody': 1}, + verbose="DEBUG") + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +def test_generate_by_line(): + + pm = ParameterHelper(verbose="DEBUG") + pm.define_group('specie', 'O', ['O']) + pm.define_group('specie', 'C', ['C']) + pm.define_group('specie', 'H', ['H']) + pm.define_group('twobody', '**', ['C', 'H']) + pm.define_group('twobody', 'OO', ['O', 'O']) + pm.define_group('threebody', '***', ['O', 'O', 'C']) + pm.define_group('threebody', 'OOO', ['O', 'O', 'O']) + pm.define_group('manybody', '1.5', ['C', 'H']) + pm.define_group('manybody', '1.5', ['C', 'O']) + pm.define_group('manybody', '1.5', ['O', 'H']) + pm.define_group('manybody', '2', ['O', 'O']) + pm.define_group('manybody', '2', ['H', 'O']) + pm.define_group('manybody', '2.8', ['O', 'O']) + pm.set_parameters('**', [1, 0.5]) + pm.set_parameters('OO', [1, 0.5]) + pm.set_parameters('***', [1, 0.5]) + pm.set_parameters('OOO', [1, 0.5]) + pm.set_parameters('1.5', [1, 0.5, 1.5]) + pm.set_parameters('2', [1, 0.5, 2]) + pm.set_parameters('2.8', [1, 0.5, 2.8]) + pm.set_parameters('cutoff_twobody', 5) + pm.set_parameters('cutoff_threebody', 4) + pm.set_parameters('cutoff_manybody', 3) + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +def test_generate_by_line2(): + + pm = ParameterHelper(verbose="DEBUG") + pm.define_group('specie', 'O', ['O']) + pm.define_group('specie', 'rest', ['C', 'H']) + pm.define_group('twobody', '**', ['*', '*']) + pm.define_group('twobody', 'OO', ['O', 'O']) + pm.define_group('threebody', '***', ['*', '*', '*']) + pm.define_group('threebody', 'Oall', ['O', 'O', 'O']) + pm.set_parameters('**', [1, 0.5]) + pm.set_parameters('OO', [1, 0.5]) + pm.set_parameters('Oall', [1, 0.5]) + pm.set_parameters('***', [1, 0.5]) + pm.set_parameters('cutoff_twobody', 5) + pm.set_parameters('cutoff_threebody', 4) + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +def test_generate_by_list(): + + pm = ParameterHelper(verbose="DEBUG") + pm.list_groups('specie', ['O', 'C', 'H']) + pm.list_groups('twobody', [['*', '*'], ['O', 'O']]) + pm.list_groups('threebody', [['*', '*', '*'], ['O', 'O', 'O']]) + pm.list_parameters({'twobody0': [1, 0.5], 'twobody1': [2, 0.2], + 'threebody0': [1, 0.5], 'threebody1': [2, 0.2], + 'cutoff_twobody': 2, 'cutoff_threebody': 1}) + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +def test_opt(): + pm = ParameterHelper(species=['O', 'C', 'H'], + kernels={'twobody': [['*', '*'], ['O', 'O']], + 'threebody': [['*', '*', '*'], ['O', 'O', 'O']]}, + parameters={'twobody0': [1, 0.5, 1], 'twobody1': [2, 0.2, 2], + 'threebody0': [1, 0.5], 'threebody1': [2, 0.2], + 'cutoff_twobody': 2, 'cutoff_threebody': 1}, + constraints={'twobody0': [False, True]}, + verbose="DEBUG") + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + +def test_randomization(): + pm = ParameterHelper(species=['O', 'C', 'H'], + kernels=['twobody', 'threebody'], + allseparate=True, + random=True, + parameters={'cutoff_twobody': 7, + 'cutoff_threebody': 4.5, + 'cutoff_manybody': 3}, + verbose="debug") + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + name = pm.find_group('specie', 'O') + name = pm.find_group('twobody', ['O', 'C']) + + +def test_from_dict(): + pm = ParameterHelper(species=['O', 'C', 'H'], + kernels=['twobody', 'threebody'], + allseparate=True, + random=True, + parameters={'cutoff_twobody': 7, + 'cutoff_threebody': 4.5, + 'cutoff_manybody': 3}, + verbose="debug") + hm = pm.as_dict() + Parameters.check_instantiation( + hm['hyps'], hm['cutoffs'], hm['kernels'], hm) + + pm1 = ParameterHelper.from_dict( + hm, verbose="debug", init_spec=['O', 'C', 'H']) + print("from_dict") + hm1 = pm1.as_dict() + print(hm['hyps']) + print(hm1['hyps'][:33], hm1['hyps'][33:]) + + Parameters.compare_dict(hm, hm1) diff --git a/tests/test_parse_ase_otf.py b/tests/test_parse_ase_otf.py deleted file mode 100644 index e4bff0bcd..000000000 --- a/tests/test_parse_ase_otf.py +++ /dev/null @@ -1,103 +0,0 @@ -import numpy as np -from flare.otf_parser import OtfAnalysis -from flare.env import AtomicEnvironment -from flare.predict import predict_on_structure - - -def test_parse_header(): - - header_dict = OtfAnalysis('test_files/VelocityVerlet.log').header - - assert header_dict['frames'] == 0 - assert header_dict['atoms'] == 4 - assert header_dict['species_set'] == {'Ag', 'I'} - assert header_dict['dt'] == .001 - assert header_dict['kernel_name'] == 'two_plus_three_body_mc' - assert header_dict['n_hyps'] == 5 - assert header_dict['algo'] == 'BFGS' - assert np.equal(header_dict['cell'], - np.array([[7.71, 0. , 0. ], - [0. , 3.855, 0. ], - [0. , 0. , 3.855]])).all() - - -def test_gp_parser(): - """ - Test the capability of otf parser to read GP/DFT info - :return: - """ - - parsed = OtfAnalysis('test_files/VelocityVerlet.log') - assert (parsed.gp_species_list == [['Ag', 'I']*2]) - - gp_positions = parsed.gp_position_list - assert len(gp_positions) == 1 - - pos1 = 1.819218 - pos2 = -0.141231 - assert(pos1 == gp_positions[0][-1][1]) - assert(pos2 == gp_positions[-1][0][2]) - - force1 = -0.424080 - force2 = 0.498037 - assert(force1 == parsed.gp_force_list[0][-1][1]) - assert(force2 == parsed.gp_force_list[-1][0][2]) - - -def test_md_parser(): - """ - Test the capability of otf parser to read MD info - :return: - """ - - parsed = OtfAnalysis('test_files/VelocityVerlet.log') - - pos1 = -0.172516 - assert(pos1 == parsed.position_list[0][0][2]) - assert(len(parsed.position_list[0]) == 4) - -def test_output_md_structures(): - - parsed = OtfAnalysis('test_files/VelocityVerlet.log') - - positions = parsed.position_list - forces = parsed.force_list - - structures = parsed.output_md_structures() - - assert np.isclose(structures[-1].positions, positions[-1]).all() - assert np.isclose(structures[-1].forces, forces[-1]).all() - - -def test_replicate_gp(): - """ - Based on gp_test_al.out, ensures that given hyperparameters and DFT calls - a GP model can be reproduced and correctly re-predict forces and - uncertainties - :return: - """ - - parsed = OtfAnalysis('test_files/VelocityVerlet.log') - - positions = parsed.position_list - forces = parsed.force_list - - gp_model = parsed.make_gp(kernel_name="two_plus_three_body_mc") - - structures = parsed.output_md_structures() - - assert np.isclose(structures[-1].positions, positions[-1]).all() - assert np.isclose(structures[-1].forces, forces[-1]).all() - - final_structure = structures[-1] - - pred_for, pred_stds = predict_on_structure(final_structure, gp_model) - - assert np.isclose(final_structure.forces, pred_for, rtol=1e-3).all() - assert np.isclose(final_structure.stds, pred_stds, rtol=1e-3).all() - - set_of_structures = structures[-3:-1] - for structure in set_of_structures: - pred_for, pred_stds = predict_on_structure(structure, gp_model) - assert np.isclose(structure.forces, pred_for, rtol=1e-3, atol=1e-6).all() - assert np.isclose(structure.stds, pred_stds, rtol=1e-3, atol=1e-6).all() diff --git a/tests/test_parse_otf.py b/tests/test_parse_otf.py index b2529c723..cd495b707 100644 --- a/tests/test_parse_otf.py +++ b/tests/test_parse_otf.py @@ -95,7 +95,7 @@ def test_replicate_gp(): positions = parsed.position_list forces = parsed.force_list - gp_model = parsed.make_gp(kernel_name='2+3') + gp_model = parsed.make_gp() structures = parsed.output_md_structures() diff --git a/tests/test_predict.py b/tests/test_predict.py index 11b723a9f..2fe354534 100644 --- a/tests/test_predict.py +++ b/tests/test_predict.py @@ -6,13 +6,20 @@ import numpy as np from flare.gp import GaussianProcess from flare.struc import Structure +from copy import deepcopy +from flare.predict import predict_on_structure_par, \ + predict_on_atom, predict_on_atom_en, \ + predict_on_structure_par_en -from flare.predict import predict_on_structure, predict_on_structure_par import pytest -def fake_predict(x,d): - return np.random.uniform(-1, 1), np.random.uniform(-1, 1) +def fake_predict(_, __): + return np.random.uniform(-1, 1), np.random.uniform(-1, 1) + + +def fake_predict_local_energy(_): + return np.random.uniform(-1, 1) _fake_gp = GaussianProcess(kernel_name='2_sc', cutoffs=[5], hyps=[1, 1, 1]) @@ -21,17 +28,16 @@ def fake_predict(x,d): positions=np.random.uniform(0, 1, size=(3, 3))) _fake_gp.predict = fake_predict - #lambda _, __: ( - #np.random.uniform(-1, 1), np.random.uniform(-1, 1)) +_fake_gp.predict_local_energy = fake_predict_local_energy -print(_fake_gp.predict(1, 2)) +assert isinstance(_fake_gp.predict(1, 1), tuple) +assert isinstance(_fake_gp.predict_local_energy(1), float) -@pytest.mark.parametrize('n_cpu', [1, 2]) +@pytest.mark.parametrize('n_cpu', [None, 1, 2]) def test_predict_on_structure_par(n_cpu): - # Predict only on the first atom, and make rest NAN - selective_atoms=[0] + selective_atoms = [0] skipped_atom_value = np.nan @@ -42,17 +48,15 @@ def test_predict_on_structure_par(n_cpu): selective_atoms=selective_atoms, skipped_atom_value=skipped_atom_value) - for x in forces[0][:]: - assert isinstance(x,float) + assert isinstance(x, float) for x in forces[1:]: assert np.isnan(x).all() - # Predict only on the second and third, and make rest 0 - selective_atoms = [1,2] - skipped_atom_value =0 + selective_atoms = [1, 2] + skipped_atom_value = 0 forces, stds = predict_on_structure_par(_fake_structure, _fake_gp, @@ -68,19 +72,16 @@ def test_predict_on_structure_par(n_cpu): assert np.equal(forces[0], 0).all() - - # Make selective atoms be all and ensure results are normal selective_atoms = [0, 1, 2] forces, stds = predict_on_structure_par(_fake_structure, - _fake_gp, - write_to_structure=True, - n_cpus=n_cpu, - selective_atoms=selective_atoms, - skipped_atom_value=skipped_atom_value) - + _fake_gp, + write_to_structure=True, + n_cpus=n_cpu, + selective_atoms=selective_atoms, + skipped_atom_value=skipped_atom_value) for x in forces.flatten(): assert isinstance(x, float) @@ -90,18 +91,33 @@ def test_predict_on_structure_par(n_cpu): assert np.array_equal(_fake_structure.forces, forces) assert np.array_equal(_fake_structure.stds, stds) + # Make selective atoms be nothing and ensure results are normal + + forces, stds = predict_on_structure_par(_fake_structure, + _fake_gp, + write_to_structure=True, + n_cpus=n_cpu, + selective_atoms=None, + skipped_atom_value=skipped_atom_value) + + for x in forces.flatten(): + assert isinstance(x, float) + for x in stds.flatten(): + assert isinstance(x, float) + + assert np.array_equal(_fake_structure.forces, forces) + assert np.array_equal(_fake_structure.stds, stds) # Get new examples to also test the results not being written - selective_atoms = [0,1] + selective_atoms = [0, 1] forces, stds = predict_on_structure_par(_fake_structure, - _fake_gp, - write_to_structure=True, - n_cpus=n_cpu, - selective_atoms=selective_atoms, - skipped_atom_value=skipped_atom_value) - + _fake_gp, + write_to_structure=True, + n_cpus=n_cpu, + selective_atoms=selective_atoms, + skipped_atom_value=skipped_atom_value) for x in forces.flatten(): assert isinstance(x, float) @@ -112,14 +128,76 @@ def test_predict_on_structure_par(n_cpu): assert np.array_equal(_fake_structure.forces[:2][:], forces[:2][:]) assert not np.array_equal(_fake_structure.forces[2][:], forces[2][:]) - assert np.array_equal(_fake_structure.stds[:2][:], stds[:2][:]) assert not np.array_equal(_fake_structure.stds[2][:], stds[2][:]) +def test_predict_on_atoms(): + pred_at_result = predict_on_atom((_fake_structure, 0, _fake_gp)) + assert len(pred_at_result) == 2 + assert len(pred_at_result[0]) == len(pred_at_result[1]) == 3 + + # Test results are correctly compiled into np arrays + pred_at_en_result = predict_on_atom_en((_fake_structure, 0, _fake_gp)) + assert isinstance(pred_at_en_result[0], np.ndarray) + assert isinstance(pred_at_en_result[1], np.ndarray) + + # Test 3 things are returned; two vectors of length 3 + assert len(pred_at_en_result) == 3 + assert len(pred_at_en_result[0]) == len(pred_at_result[1]) == 3 + assert isinstance(pred_at_en_result[2], float) + + +@pytest.mark.parametrize('n_cpus', [1, 2, None]) +@pytest.mark.parametrize(['write_to_structure', 'selective_atoms'], + [(True, []), + (True, [1]), + (False, []), + (False, [1])]) +def test_predict_on_structure_en(n_cpus, write_to_structure, + selective_atoms): + old_structure = deepcopy(_fake_structure) + + old_structure.forces = np.random.uniform(-1, 1, (3, 3)) + + used_structure = deepcopy(old_structure) + + forces, stds, energies = predict_on_structure_par_en( + structure=used_structure, + gp=_fake_gp, + n_cpus=n_cpus, + write_to_structure=write_to_structure, + selective_atoms=selective_atoms, + skipped_atom_value=0) + + assert np.array_equal(forces.shape, old_structure.positions.shape) + assert np.array_equal(stds.shape, old_structure.positions.shape) + + if write_to_structure: + + if selective_atoms == [1]: + assert np.array_equal(old_structure.forces[0], + used_structure.forces[0]) + assert np.array_equal(old_structure.forces[2], + used_structure.forces[2]) + assert np.array_equal(used_structure.forces[1], forces[1]) + else: + assert not np.array_equal(old_structure.forces[0], + used_structure.forces[0]) + assert not np.array_equal(old_structure.forces[2], + used_structure.forces[2]) + assert np.array_equal(forces, used_structure.forces) + # These will be unequal no matter what + assert not np.array_equal(old_structure.forces[1], + used_structure.forces[1]) + else: + assert np.array_equal(old_structure.forces, used_structure.forces) + assert np.array_equal(forces.shape, (3, 3)) + assert np.array_equal(stds.shape, (3, 3)) + assert len(energies) == len(old_structure) diff --git a/tests/test_str_to_kernel.py b/tests/test_str_to_kernel.py index a10503ccf..ccf3530ab 100644 --- a/tests/test_str_to_kernel.py +++ b/tests/test_str_to_kernel.py @@ -8,14 +8,15 @@ from flare.kernels.utils import str_to_kernel_set as stks -@pytest.mark.parametrize('kernel_name', ['2sc', '3sc', '2+3sc', - '2', '3', '2+3', - '2+3+many', '2+3mb']) -def test_stk(kernel_name): +@pytest.mark.parametrize('kernels', [['twobody'], ['threebody'], ['twobody', 'threebody'], + ['twobody', 'threebody', 'manybody']]) +@pytest.mark.parametrize('component', ['sc', 'mc']) +@pytest.mark.parametrize('nspecie', [1, 2]) +def test_stk(kernels, component, nspecie): """Check whether the str_to_kernel_set can return kernel functions properly""" try: - k, kg, ek, efk = stks(kernel_name) + k, kg, ek, efk = stks(kernels, component, {'nspecie': nspecie}) except: - raise RuntimeError(f"fail to return kernel {kernel_name}") + raise RuntimeError(f"fail to return kernel {kernels} {component} {nspecie}") diff --git a/tests/test_vasp_util.py b/tests/test_vasp_util.py index 9d529ca30..77e4d23e0 100644 --- a/tests/test_vasp_util.py +++ b/tests/test_vasp_util.py @@ -117,14 +117,18 @@ def test_vasp_input_edit(): final_structure = dft_input_to_structure(new_file) - assert np.isclose(final_structure.vec1, structure.vec1).all() + assert np.isclose(final_structure.vec1, structure.vec1, atol=1e-4).all() assert np.isclose(final_structure.positions[0], - structure.positions[0]).all() + structure.positions[0], atol=1e-4).all() os.system('rm ./POSCAR') os.system('rm ./POSCAR.bak') - +@pytest.mark.skipif(not os.environ.get('VASP_COMMAND', + False), reason='VASP_COMMAND not found ' + 'in environment: Please install VASP ' + ' and set the VASP_COMMAND env. ' + 'variable to point to cp2k.popt') def test_run_dft_par(): os.system('cp test_files/test_POSCAR ./POSCAR') test_structure = dft_input_to_structure('./POSCAR') @@ -134,11 +138,11 @@ def test_run_dft_par(): run_dft_par('POSCAR',test_structure,dft_command=dft_command, n_cpus=2) - call_string = "echo 'testing_call' > TEST_CALL_OUT" + call_string = "echo 'testing_call'" forces = run_dft_par('POSCAR', test_structure, dft_command=call_string, - n_cpus=1, serial_prefix=' ', - dft_out='test_files/test_vasprun.xml') + n_cpus=1, serial_prefix=' ', dft_out='test_files/test_vasprun.xml', + screen_out='TEST_CALL_OUT') with open("TEST_CALL_OUT", 'r') as f: assert 'testing_call' in f.readline()