├── .gitignore ├── CMakeLists.txt ├── FindOpenCL.cmake ├── README.md ├── openmmapi ├── include │ ├── NeuralNetworkForce.h │ ├── NeuralNetworkKernels.h │ └── internal │ │ ├── NeuralNetworkForceImpl.h │ │ └── windowsExportNN.h └── src │ ├── NeuralNetworkForce.cpp │ └── NeuralNetworkForceImpl.cpp ├── platforms ├── cuda │ ├── CMakeLists.txt │ ├── EncodeCUDAFiles.cmake │ ├── include │ │ └── CudaNeuralNetworkKernelFactory.h │ ├── src │ │ ├── CudaNeuralNetworkKernelFactory.cpp │ │ ├── CudaNeuralNetworkKernelSources.cpp.in │ │ ├── CudaNeuralNetworkKernelSources.h.in │ │ ├── CudaNeuralNetworkKernels.cpp │ │ ├── CudaNeuralNetworkKernels.h │ │ └── kernels │ │ │ └── neuralNetworkForce.cu │ └── tests │ │ ├── CMakeLists.txt │ │ └── TestCudaNeuralNetworkForce.cpp ├── opencl │ ├── CMakeLists.txt │ ├── EncodeCLFiles.cmake │ ├── include │ │ └── OpenCLNeuralNetworkKernelFactory.h │ ├── src │ │ ├── OpenCLNeuralNetworkKernelFactory.cpp │ │ ├── OpenCLNeuralNetworkKernelSources.cpp.in │ │ ├── OpenCLNeuralNetworkKernelSources.h.in │ │ ├── OpenCLNeuralNetworkKernels.cpp │ │ ├── OpenCLNeuralNetworkKernels.h │ │ ├── cl.hpp │ │ └── kernels │ │ │ └── neuralNetworkForce.cl │ └── tests │ │ ├── CMakeLists.txt │ │ └── TestOpenCLNeuralNetworkForce.cpp └── reference │ ├── CMakeLists.txt │ ├── include │ └── ReferenceNeuralNetworkKernelFactory.h │ ├── src │ ├── ReferenceNeuralNetworkKernelFactory.cpp │ ├── ReferenceNeuralNetworkKernels.cpp │ └── ReferenceNeuralNetworkKernels.h │ └── tests │ ├── CMakeLists.txt │ └── TestReferenceNeuralNetworkForce.cpp ├── python ├── CMakeLists.txt ├── openmmnn.i ├── setup.py └── tests │ └── TestNeuralNetworkForce.py ├── serialization ├── include │ └── NeuralNetworkForceProxy.h ├── src │ ├── NeuralNetworkForceProxy.cpp │ └── NeuralNetworkSerializationProxyRegistration.cpp └── tests │ ├── CMakeLists.txt │ └── TestSerializeNeuralNetworkForce.cpp └── tests ├── central.pb └── periodic.pb /.gitignore: -------------------------------------------------------------------------------- 1 | # Prerequisites 2 | *.d 3 | 4 | # Compiled Object files 5 | *.slo 6 | *.lo 7 | *.o 8 | *.obj 9 | 10 | # Precompiled Headers 11 | *.gch 12 | *.pch 13 | 14 | # Compiled Dynamic libraries 15 | *.so 16 | *.dylib 17 | *.dll 18 | 19 | # Fortran module files 20 | *.mod 21 | *.smod 22 | 23 | # Compiled Static libraries 24 | *.lai 25 | *.la 26 | *.a 27 | *.lib 28 | 29 | # Executables 30 | *.exe 31 | *.out 32 | *.app 33 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #--------------------------------------------------- 2 | # OpenMM NeuralNetwork Plugin 3 | #---------------------------------------------------- 4 | 5 | CMAKE_MINIMUM_REQUIRED(VERSION 3.5) 6 | 7 | # We need to know where OpenMM is installed so we can access the headers and libraries. 8 | SET(OPENMM_DIR "/usr/local/openmm" CACHE PATH "Where OpenMM is installed") 9 | INCLUDE_DIRECTORIES("${OPENMM_DIR}/include") 10 | LINK_DIRECTORIES("${OPENMM_DIR}/lib" "${OPENMM_DIR}/lib/plugins") 11 | 12 | # We need to know where TensorFlow is installed so we can access the headers and libraries. 13 | SET(TENSORFLOW_DIR "/usr/local" CACHE PATH "Where TensorFlow is installed") 14 | INCLUDE_DIRECTORIES("${TENSORFLOW_DIR}/include/tensorflow/c") 15 | LINK_DIRECTORIES("${TENSORFLOW_DIR}/lib") 16 | 17 | # Specify the C++ version we are building for. 18 | SET (CMAKE_CXX_STANDARD 11) 19 | 20 | # Set flags for linking on mac 21 | IF(APPLE) 22 | SET (CMAKE_INSTALL_NAME_DIR "@rpath") 23 | SET(EXTRA_COMPILE_FLAGS "-msse2 -stdlib=libc++") 24 | ENDIF(APPLE) 25 | 26 | # Select where to install 27 | IF(${CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT}) 28 | IF(WIN32) 29 | SET(CMAKE_INSTALL_PREFIX "$ENV{ProgramFiles}/OpenMM" CACHE PATH "Where to install the plugin" FORCE) 30 | ELSE(WIN32) 31 | SET(CMAKE_INSTALL_PREFIX "/usr/local/openmm" CACHE PATH "Where to install the plugin" FORCE) 32 | ENDIF(WIN32) 33 | ENDIF(${CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT}) 34 | 35 | # Put all the tests and libraries in a single output directory. 36 | IF(NOT EXECUTABLE_OUTPUT_PATH) 37 | SET(EXECUTABLE_OUTPUT_PATH ${PROJECT_BINARY_DIR} 38 | CACHE INTERNAL "Single output directory for building all executables.") 39 | ENDIF() 40 | IF(NOT LIBRARY_OUTPUT_PATH) 41 | SET(LIBRARY_OUTPUT_PATH ${PROJECT_BINARY_DIR} 42 | CACHE INTERNAL "Single output directory for building all libraries.") 43 | ENDIF() 44 | SET(${PROJECT_NAME}_EXECUTABLE_DIR ${EXECUTABLE_OUTPUT_PATH}/${CMAKE_CFG_INTDIR}) 45 | SET(${PROJECT_NAME}_LIBRARY_DIR ${LIBRARY_OUTPUT_PATH}/${CMAKE_CFG_INTDIR}) 46 | 47 | # The source is organized into subdirectories, but we handle them all from 48 | # this CMakeLists file rather than letting CMake visit them as SUBDIRS. 49 | SET(NN_PLUGIN_SOURCE_SUBDIRS openmmapi serialization) 50 | 51 | # Set the library name 52 | SET(NN_LIBRARY_NAME OpenMMNN) 53 | SET(SHARED_NN_TARGET ${NN_LIBRARY_NAME}) 54 | 55 | # These are all the places to search for header files which are to be part of the API. 56 | SET(API_INCLUDE_DIRS "openmmapi/include" "openmmapi/include/internal") 57 | 58 | # Locate header files. 59 | SET(API_INCLUDE_FILES) 60 | FOREACH(dir ${API_INCLUDE_DIRS}) 61 | FILE(GLOB fullpaths ${dir}/*.h) 62 | SET(API_INCLUDE_FILES ${API_INCLUDE_FILES} ${fullpaths}) 63 | ENDFOREACH(dir) 64 | 65 | # Collect up source files 66 | SET(SOURCE_FILES) # empty 67 | SET(SOURCE_INCLUDE_FILES) 68 | FOREACH(subdir ${NN_PLUGIN_SOURCE_SUBDIRS}) 69 | FILE(GLOB src_files ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.cpp) 70 | FILE(GLOB incl_files ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.h) 71 | SET(SOURCE_FILES ${SOURCE_FILES} ${src_files}) #append 72 | SET(SOURCE_INCLUDE_FILES ${SOURCE_INCLUDE_FILES} ${incl_files}) 73 | 74 | ## Make sure we find these locally before looking in OpenMM/include if 75 | ## OpenMM was previously installed there. 76 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/include) 77 | ENDFOREACH(subdir) 78 | 79 | # Create the library. 80 | 81 | ADD_LIBRARY(${SHARED_NN_TARGET} SHARED ${SOURCE_FILES} ${SOURCE_INCLUDE_FILES} ${API_INCLUDE_FILES}) 82 | SET_TARGET_PROPERTIES(${SHARED_NN_TARGET} 83 | PROPERTIES COMPILE_FLAGS "-DNN_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" 84 | LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") 85 | TARGET_LINK_LIBRARIES(${SHARED_NN_TARGET} OpenMM) 86 | TARGET_LINK_LIBRARIES(${SHARED_NN_TARGET} tensorflow) 87 | INSTALL_TARGETS(/lib RUNTIME_DIRECTORY /lib ${SHARED_NN_TARGET}) 88 | 89 | # install headers 90 | FILE(GLOB API_ONLY_INCLUDE_FILES "openmmapi/include/*.h") 91 | INSTALL (FILES ${API_ONLY_INCLUDE_FILES} DESTINATION include) 92 | FILE(GLOB API_ONLY_INCLUDE_FILES_INTERNAL "openmmapi/include/internal/*.h") 93 | INSTALL (FILES ${API_ONLY_INCLUDE_FILES_INTERNAL} DESTINATION include/internal) 94 | 95 | # Enable testing 96 | 97 | ENABLE_TESTING() 98 | ADD_SUBDIRECTORY(serialization/tests) 99 | 100 | # Copy test files to the build directory. 101 | 102 | file(GLOB_RECURSE TEST_FILES RELATIVE "${CMAKE_SOURCE_DIR}" 103 | "${CMAKE_SOURCE_DIR}/tests/*.pb" 104 | ) 105 | file(MAKE_DIRECTORY ${CMAKE_BINARY_DIR}/tests) 106 | set(COPIED_TEST_FILES) 107 | foreach(TEST_FILE ${TEST_FILES}) 108 | set(infile "${CMAKE_SOURCE_DIR}/${TEST_FILE}") 109 | set(outfile "${CMAKE_BINARY_DIR}/${TEST_FILE}") 110 | add_custom_command( 111 | OUTPUT "${outfile}" 112 | COMMAND "${CMAKE_COMMAND}" -E copy_if_different "${infile}" "${outfile}" 113 | DEPENDS "${infile}" 114 | COMMENT "CMake-copying file ${infile} to ${outfile}") 115 | set(COPIED_TEST_FILES ${COPIED_TEST_FILES} "${outfile}") 116 | endforeach() 117 | add_custom_target(CopyTestFiles ALL DEPENDS ${COPIED_TEST_FILES}) 118 | 119 | # Build the implementations for different platforms 120 | 121 | ADD_SUBDIRECTORY(platforms/reference) 122 | 123 | SET(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_SOURCE_DIR}") 124 | FIND_PACKAGE(OpenCL QUIET) 125 | IF(OPENCL_FOUND) 126 | SET(NN_BUILD_OPENCL_LIB ON CACHE BOOL "Build implementation for OpenCL") 127 | ELSE(OPENCL_FOUND) 128 | SET(NN_BUILD_OPENCL_LIB OFF CACHE BOOL "Build implementation for OpenCL") 129 | ENDIF(OPENCL_FOUND) 130 | IF(NN_BUILD_OPENCL_LIB) 131 | ADD_SUBDIRECTORY(platforms/opencl) 132 | ENDIF(NN_BUILD_OPENCL_LIB) 133 | 134 | FIND_PACKAGE(CUDA QUIET) 135 | IF(CUDA_FOUND) 136 | SET(NN_BUILD_CUDA_LIB ON CACHE BOOL "Build implementation for CUDA") 137 | ELSE(CUDA_FOUND) 138 | SET(NN_BUILD_CUDA_LIB OFF CACHE BOOL "Build implementation for CUDA") 139 | ENDIF(CUDA_FOUND) 140 | IF(NN_BUILD_CUDA_LIB) 141 | ADD_SUBDIRECTORY(platforms/cuda) 142 | ENDIF(NN_BUILD_CUDA_LIB) 143 | 144 | # Build the Python API 145 | 146 | FIND_PROGRAM(PYTHON_EXECUTABLE python) 147 | FIND_PROGRAM(SWIG_EXECUTABLE swig) 148 | IF(PYTHON_EXECUTABLE AND SWIG_EXECUTABLE) 149 | SET(NN_BUILD_PYTHON_WRAPPERS ON CACHE BOOL "Build wrappers for Python") 150 | ELSE(PYTHON_EXECUTABLE AND SWIG_EXECUTABLE) 151 | SET(NN_BUILD_PYTHON_WRAPPERS OFF CACHE BOOL "Build wrappers for Python") 152 | ENDIF(PYTHON_EXECUTABLE AND SWIG_EXECUTABLE) 153 | IF(NN_BUILD_PYTHON_WRAPPERS) 154 | ADD_SUBDIRECTORY(python) 155 | ENDIF(NN_BUILD_PYTHON_WRAPPERS) 156 | 157 | -------------------------------------------------------------------------------- /FindOpenCL.cmake: -------------------------------------------------------------------------------- 1 | 2 | ### OPENCL_INCLUDE_DIR ### 3 | # Try OPENCL_DIR variable before looking elsewhere 4 | find_path(OPENCL_INCLUDE_DIR 5 | NAMES OpenCL/opencl.h CL/opencl.h 6 | PATHS $ENV{OPENCL_DIR} 7 | PATH_SUFFIXES "include" 8 | NO_DEFAULT_PATH 9 | ) 10 | # Next look in environment variables set by OpenCL SDK installations 11 | find_path(OPENCL_INCLUDE_DIR 12 | NAMES OpenCL/opencl.h CL/opencl.h 13 | PATHS 14 | $ENV{CUDA_PATH} 15 | $ENV{AMDAPPSDKROOT} 16 | PATH_SUFFIXES "include" 17 | NO_DEFAULT_PATH 18 | ) 19 | # On Macs, look inside the platform SDK 20 | if(DEFINED CMAKE_OSX_SYSROOT) 21 | find_path(OPENCL_INCLUDE_DIR 22 | NAMES opencl.h opencl.h 23 | PATHS 24 | "${CMAKE_OSX_SYSROOT}/System/Library/Frameworks/OpenCL.framework/Headers" 25 | NO_DEFAULT_PATH 26 | ) 27 | endif(DEFINED CMAKE_OSX_SYSROOT) 28 | # As a last resort, look in default system areas followed by other possible locations 29 | find_path(OPENCL_INCLUDE_DIR 30 | NAMES OpenCL/opencl.h CL/opencl.h 31 | PATHS 32 | "C:/CUDA" 33 | "/usr/local/cuda" 34 | "/usr/local/streamsdk" 35 | "/usr" 36 | PATH_SUFFIXES "include" 37 | ) 38 | 39 | ### OPENCL_LIBRARY ### 40 | if("${CMAKE_SYSTEM_NAME}" MATCHES "Linux") 41 | if("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64") 42 | set(path_suffixes "lib/x86_64") 43 | else("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64") 44 | set(path_suffixes "lib/x86") 45 | endif("${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "x86_64") 46 | elseif(MSVC) 47 | if(CMAKE_CL_64) 48 | set(path_suffixes "lib/x64" "lib/x86_64") 49 | else(CMAKE_CL_64) 50 | set(path_suffixes "lib/Win32" "lib/x86") 51 | endif(CMAKE_CL_64) 52 | else(MSVC) 53 | set(path_suffixes "lib") 54 | endif("${CMAKE_SYSTEM_NAME}" MATCHES "Linux") 55 | # Try OPENCL_DIR variable before looking elsewhere 56 | find_library(OPENCL_LIBRARY 57 | NAMES OpenCL 58 | PATHS 59 | $ENV{OPENCL_DIR} 60 | ${OPENCL_LIB_SEARCH_PATH} 61 | PATH_SUFFIXES ${path_suffixes} 62 | NO_DEFAULT_PATH 63 | ) 64 | # Next look in environment variables set by OpenCL SDK installations 65 | find_library(OPENCL_LIBRARY 66 | NAMES OpenCL 67 | PATHS 68 | $ENV{CUDA_PATH} 69 | $ENV{AMDAPPSDKROOT} 70 | PATH_SUFFIXES ${path_suffixes} 71 | NO_DEFAULT_PATH 72 | ) 73 | # As a last resort, look in default system areas followed by other possible locations 74 | find_library(OPENCL_LIBRARY 75 | NAMES OpenCL 76 | PATHS 77 | "C:/CUDA" 78 | "/usr/local/cuda" 79 | "/usr/local/streamsdk" 80 | "/usr" 81 | PATH_SUFFIXES ${path_suffixes} "lib" 82 | ) 83 | 84 | include(FindPackageHandleStandardArgs) 85 | find_package_handle_standard_args(OPENCL DEFAULT_MSG OPENCL_LIBRARY OPENCL_INCLUDE_DIR) 86 | 87 | if(OPENCL_FOUND) 88 | set(OPENCL_LIBRARIES ${OPENCL_LIBRARY}) 89 | mark_as_advanced(CLEAR OPENCL_INCLUDE_DIR) 90 | mark_as_advanced(CLEAR OPENCL_LIBRARY) 91 | else(OPENCL_FOUND) 92 | set(OPENCL_LIBRARIES) 93 | mark_as_advanced(OPENCL_INCLUDE_DIR) 94 | mark_as_advanced(OPENCL_LIBRARY) 95 | endif(OPENCL_FOUND) 96 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | OpenMM Neural Network Plugin 2 | ============================ 3 | 4 | This is a plugin for [OpenMM](http://openmm.org) that allows neural networks 5 | to be used for defining forces. It is implemented with [TensorFlow](https://www.tensorflow.org/). 6 | To use it, you create a TensorFlow graph that takes particle positions as input 7 | and produces forces and energy as output. This plugin uses the graph to apply 8 | forces to particles during a simulation. 9 | 10 | Installation 11 | ============ 12 | 13 | At present this plugin must be compiled from source. It uses CMake as its build 14 | system. Before compiling you must install the TensorFlow C API by following the 15 | instructions at https://www.tensorflow.org/install/install_c. You can then 16 | follow these steps. 17 | 18 | 1. Create a directory in which to build the plugin. 19 | 20 | 2. Run the CMake GUI or ccmake, specifying your new directory as the build directory and the top 21 | level directory of this project as the source directory. 22 | 23 | 3. Press "Configure". 24 | 25 | 4. Set OPENMM_DIR to point to the directory where OpenMM is installed. This is needed to locate 26 | the OpenMM header files and libraries. 27 | 28 | 5. Set TENSORFLOW_DIR to point to the directory where you installed the TensorFlow C API. 29 | 30 | 6. Set CMAKE_INSTALL_PREFIX to the directory where the plugin should be installed. Usually, 31 | this will be the same as OPENMM_DIR, so the plugin will be added to your OpenMM installation. 32 | 33 | 7. If you plan to build the OpenCL platform, make sure that OPENCL_INCLUDE_DIR and 34 | OPENCL_LIBRARY are set correctly, and that NN_BUILD_OPENCL_LIB is selected. 35 | 36 | 8. If you plan to build the CUDA platform, make sure that CUDA_TOOLKIT_ROOT_DIR is set correctly 37 | and that NN_BUILD_CUDA_LIB is selected. 38 | 39 | 9. Press "Configure" again if necessary, then press "Generate". 40 | 41 | 10. Use the build system you selected to build and install the plugin. For example, if you 42 | selected Unix Makefiles, type `make install` to install the plugin, and `make PythonInstall` to 43 | install the Python wrapper. 44 | 45 | Usage 46 | ===== 47 | 48 | The first step is to create a TensorFlow graph defining the calculation to 49 | perform. It should have an input called `positions` which will be set to the 50 | particle positions. It should produce two outputs: one called `forces` 51 | containing the forces to apply to the particles, and one called `energy` with 52 | the potential energy. This graph must then be saved to a binary protocol 53 | buffer file. Here is an example of Python code that does this for a very 54 | simple calculation (a harmonic force attracting every particle to the origin). 55 | 56 | ```python 57 | import tensorflow as tf 58 | graph = tf.Graph() 59 | with graph.as_default(): 60 | positions = tf.placeholder(tf.float32, [None, 3], 'positions') 61 | energy = tf.reduce_sum(positions**2, name='energy') 62 | forces = tf.identity(tf.gradients(-energy, positions), name='forces') 63 | tf.train.write_graph(graph, '.', 'graph.pb', as_text=False) 64 | ``` 65 | 66 | The data types of the graph's input and output tensors may be either `float32` 67 | or `float64`. 68 | 69 | To use the graph in a simulation, create a `NeuralNetworkForce` object and add 70 | it to your `System`. The constructor takes the path to the saved graph as an 71 | argument. For example, 72 | 73 | ```python 74 | from openmmnn import * 75 | f = NeuralNetworkForce('graph.pb') 76 | system.addForce(f) 77 | ``` 78 | 79 | Alternatively, in Python (but not C++) you can directly pass a `tf.Graph` to the 80 | constructor without writing it to a file first: 81 | 82 | ```python 83 | f = NeuralNetworkForce(graph) 84 | ``` 85 | 86 | If the graph includes any variables, pass a `tf.Session` as the second argument. 87 | It will use TensorFlow's `freeze_graph` utility to create a frozen version of the 88 | graph in which variables have been replaced with values taken from the session: 89 | 90 | ```python 91 | f = NeuralNetworkForce(graph, session) 92 | ``` 93 | 94 | When defining the graph to perform a calculation, you may want to apply 95 | periodic boundary conditions. To do this, call `setUsesPeriodicBoundaryConditions(True)` 96 | on the `NeuralNetworkForce`. The graph is then expected to contain an input 97 | called `boxvectors` which will contain the current periodic box vectors. You 98 | can make use of them in whatever way you want for computing the force. For 99 | example, the following code applies periodic boundary conditions to each 100 | particle position to translate all of them into a single periodic cell. 101 | 102 | ```python 103 | positions = tf.placeholder(tf.float32, [None, 3], 'positions') 104 | boxvectors = tf.placeholder(tf.float32, [3, 3], 'boxvectors') 105 | boxsize = tf.diag_part(boxvectors) 106 | periodicPositions = positions - tf.floor(positions/boxsize)*boxsize 107 | ``` 108 | 109 | Note that this code assumes a rectangular box. Applying periodic boundary 110 | conditions with a triclinic box requires a slightly more complicated 111 | calculation. 112 | 113 | License 114 | ======= 115 | 116 | This is part of the OpenMM molecular simulation toolkit originating from 117 | Simbios, the NIH National Center for Physics-Based Simulation of 118 | Biological Structures at Stanford, funded under the NIH Roadmap for 119 | Medical Research, grant U54 GM072970. See https://simtk.org. 120 | 121 | Portions copyright (c) 2018 Stanford University and the Authors. 122 | 123 | Authors: Peter Eastman 124 | 125 | Contributors: 126 | 127 | Permission is hereby granted, free of charge, to any person obtaining a 128 | copy of this software and associated documentation files (the "Software"), 129 | to deal in the Software without restriction, including without limitation 130 | the rights to use, copy, modify, merge, publish, distribute, sublicense, 131 | and/or sell copies of the Software, and to permit persons to whom the 132 | Software is furnished to do so, subject to the following conditions: 133 | 134 | The above copyright notice and this permission notice shall be included in 135 | all copies or substantial portions of the Software. 136 | 137 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 138 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 139 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 140 | THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, 141 | DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR 142 | OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE 143 | USE OR OTHER DEALINGS IN THE SOFTWARE. 144 | 145 | -------------------------------------------------------------------------------- /openmmapi/include/NeuralNetworkForce.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_NEURAL_NETWORKFORCE_H_ 2 | #define OPENMM_NEURAL_NETWORKFORCE_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "openmm/Context.h" 36 | #include "openmm/Force.h" 37 | #include 38 | #include "internal/windowsExportNN.h" 39 | 40 | namespace NNPlugin { 41 | 42 | /** 43 | * This class implements forces that are defined by user-supplied neural networks. 44 | * It uses the TensorFlow library to perform the computations. */ 45 | 46 | class OPENMM_EXPORT_NN NeuralNetworkForce : public OpenMM::Force { 47 | public: 48 | /** 49 | * Create a NeuralNetworkForce. The network is defined by a TensorFlow graph saved 50 | * to a binary protocol buffer file. 51 | * 52 | * @param file the path to the file containing the network 53 | */ 54 | NeuralNetworkForce(const std::string& file); 55 | /** 56 | * Get the path to the file containing the graph. 57 | */ 58 | const std::string& getFile() const; 59 | /** 60 | * Get the content of the protocol buffer defining the graph. 61 | */ 62 | const std::string& getGraphProto() const; 63 | /** 64 | * Set whether this force makes use of periodic boundary conditions. If this is set 65 | * to true, the TensorFlow graph must include a 3x3 tensor called "boxvectors", which 66 | * is set to the current periodic box vectors. 67 | */ 68 | void setUsesPeriodicBoundaryConditions(bool periodic); 69 | /** 70 | * Get whether this force makes use of periodic boundary conditions. 71 | */ 72 | bool usesPeriodicBoundaryConditions() const; 73 | protected: 74 | OpenMM::ForceImpl* createImpl() const; 75 | private: 76 | std::string file, graphProto; 77 | bool usePeriodic; 78 | }; 79 | 80 | } // namespace NNPlugin 81 | 82 | #endif /*OPENMM_NEURAL_NETWORKFORCE_H_*/ 83 | -------------------------------------------------------------------------------- /openmmapi/include/NeuralNetworkKernels.h: -------------------------------------------------------------------------------- 1 | #ifndef NEURAL_NETWORK_KERNELS_H_ 2 | #define NEURAL_NETWORK_KERNELS_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "NeuralNetworkForce.h" 36 | #include "openmm/KernelImpl.h" 37 | #include "openmm/Platform.h" 38 | #include "openmm/System.h" 39 | #include 40 | #include 41 | 42 | namespace NNPlugin { 43 | 44 | /** 45 | * This kernel is invoked by NeuralNetworkForce to calculate the forces acting on the system and the energy of the system. 46 | */ 47 | class CalcNeuralNetworkForceKernel : public OpenMM::KernelImpl { 48 | public: 49 | static std::string Name() { 50 | return "CalcNeuralNetworkForce"; 51 | } 52 | CalcNeuralNetworkForceKernel(std::string name, const OpenMM::Platform& platform) : OpenMM::KernelImpl(name, platform) { 53 | } 54 | /** 55 | * Initialize the kernel. 56 | * 57 | * @param system the System this kernel will be applied to 58 | * @param force the NeuralNetworkForce this kernel will be used for 59 | * @param session the TensorFlow session in which to do calculations 60 | * @param graph the TensorFlow graph to use for computing forces and energy 61 | * @param positionsType the data type of the "positions" tensor 62 | * @param boxType the data type of the "boxvectors" tensor 63 | * @param energyType the data type of the "energy" tensor 64 | * @param forcesType the data type of the "forces" tensor 65 | */ 66 | virtual void initialize(const OpenMM::System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 67 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType) = 0; 68 | /** 69 | * Execute the kernel to calculate the forces and/or energy. 70 | * 71 | * @param context the context in which to execute this kernel 72 | * @param includeForces true if forces should be calculated 73 | * @param includeEnergy true if the energy should be calculated 74 | * @return the potential energy due to the force 75 | */ 76 | virtual double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy) = 0; 77 | }; 78 | 79 | } // namespace NNPlugin 80 | 81 | #endif /*NEURAL_NETWORK_KERNELS_H_*/ 82 | -------------------------------------------------------------------------------- /openmmapi/include/internal/NeuralNetworkForceImpl.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_NEURAL_NETWORK_FORCE_IMPL_H_ 2 | #define OPENMM_NEURAL_NETWORK_FORCE_IMPL_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "NeuralNetworkForce.h" 36 | #include "openmm/internal/ForceImpl.h" 37 | #include "openmm/Kernel.h" 38 | #include 39 | #include 40 | #include 41 | #include 42 | 43 | namespace NNPlugin { 44 | 45 | class System; 46 | 47 | /** 48 | * This is the internal implementation of NeuralNetworkForce. 49 | */ 50 | 51 | class OPENMM_EXPORT_NN NeuralNetworkForceImpl : public OpenMM::ForceImpl { 52 | public: 53 | NeuralNetworkForceImpl(const NeuralNetworkForce& owner); 54 | ~NeuralNetworkForceImpl(); 55 | void initialize(OpenMM::ContextImpl& context); 56 | const NeuralNetworkForce& getOwner() const { 57 | return owner; 58 | } 59 | void updateContextState(OpenMM::ContextImpl& context, bool& forcesInvalid) { 60 | // This force field doesn't update the state directly. 61 | } 62 | double calcForcesAndEnergy(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy, int groups); 63 | std::map getDefaultParameters() { 64 | return std::map(); // This force field doesn't define any parameters. 65 | } 66 | std::vector getKernelNames(); 67 | private: 68 | const NeuralNetworkForce& owner; 69 | OpenMM::Kernel kernel; 70 | TF_Graph* graph; 71 | TF_Session* session; 72 | TF_Status* status; 73 | }; 74 | 75 | } // namespace NNPlugin 76 | 77 | #endif /*OPENMM_NEURAL_NETWORK_FORCE_IMPL_H_*/ 78 | -------------------------------------------------------------------------------- /openmmapi/include/internal/windowsExportNN.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_WINDOWSEXPORTNN_H_ 2 | #define OPENMM_WINDOWSEXPORTNN_H_ 3 | 4 | /* 5 | * Shared libraries are messy in Visual Studio. We have to distinguish three 6 | * cases: 7 | * (1) this header is being used to build the OpenMM shared library 8 | * (dllexport) 9 | * (2) this header is being used by a *client* of the OpenMM shared 10 | * library (dllimport) 11 | * (3) we are building the OpenMM static library, or the client is 12 | * being compiled with the expectation of linking with the 13 | * OpenMM static library (nothing special needed) 14 | * In the CMake script for building this library, we define one of the symbols 15 | * NN_BUILDING_{SHARED|STATIC}_LIBRARY 16 | * Client code normally has no special symbol defined, in which case we'll 17 | * assume it wants to use the shared library. However, if the client defines 18 | * the symbol OPENMM_USE_STATIC_LIBRARIES we'll suppress the dllimport so 19 | * that the client code can be linked with static libraries. Note that 20 | * the client symbol is not library dependent, while the library symbols 21 | * affect only the OpenMM library, meaning that other libraries can 22 | * be clients of this one. However, we are assuming all-static or all-shared. 23 | */ 24 | 25 | #ifdef _MSC_VER 26 | // We don't want to hear about how sprintf is "unsafe". 27 | #pragma warning(disable:4996) 28 | // Keep MS VC++ quiet about lack of dll export of private members. 29 | #pragma warning(disable:4251) 30 | #if defined(NN_BUILDING_SHARED_LIBRARY) 31 | #define OPENMM_EXPORT_NN __declspec(dllexport) 32 | #elif defined(NN_BUILDING_STATIC_LIBRARY) || defined(NN_USE_STATIC_LIBRARIES) 33 | #define OPENMM_EXPORT_NN 34 | #else 35 | #define OPENMM_EXPORT_NN __declspec(dllimport) // i.e., a client of a shared library 36 | #endif 37 | #else 38 | #define OPENMM_EXPORT_NN // Linux, Mac 39 | #endif 40 | 41 | #endif // OPENMM_WINDOWSEXPORTNN_H_ 42 | -------------------------------------------------------------------------------- /openmmapi/src/NeuralNetworkForce.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "NeuralNetworkForce.h" 33 | #include "internal/NeuralNetworkForceImpl.h" 34 | #include "openmm/OpenMMException.h" 35 | #include "openmm/internal/AssertionUtilities.h" 36 | #include 37 | 38 | using namespace NNPlugin; 39 | using namespace OpenMM; 40 | using namespace std; 41 | 42 | NeuralNetworkForce::NeuralNetworkForce(const std::string& file) : file(file), usePeriodic(false) { 43 | ifstream graphFile(file); 44 | graphProto = string((istreambuf_iterator(graphFile)), istreambuf_iterator()); 45 | } 46 | 47 | const string& NeuralNetworkForce::getFile() const { 48 | return file; 49 | } 50 | 51 | const string& NeuralNetworkForce::getGraphProto() const { 52 | return graphProto; 53 | } 54 | 55 | ForceImpl* NeuralNetworkForce::createImpl() const { 56 | return new NeuralNetworkForceImpl(*this); 57 | } 58 | 59 | void NeuralNetworkForce::setUsesPeriodicBoundaryConditions(bool periodic) { 60 | usePeriodic = periodic; 61 | } 62 | 63 | bool NeuralNetworkForce::usesPeriodicBoundaryConditions() const { 64 | return usePeriodic; 65 | } 66 | -------------------------------------------------------------------------------- /openmmapi/src/NeuralNetworkForceImpl.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "internal/NeuralNetworkForceImpl.h" 33 | #include "NeuralNetworkKernels.h" 34 | #include "openmm/OpenMMException.h" 35 | #include "openmm/internal/ContextImpl.h" 36 | 37 | using namespace NNPlugin; 38 | using namespace OpenMM; 39 | using namespace std; 40 | 41 | NeuralNetworkForceImpl::NeuralNetworkForceImpl(const NeuralNetworkForce& owner) : owner(owner), graph(NULL), session(NULL), status(TF_NewStatus()) { 42 | } 43 | 44 | NeuralNetworkForceImpl::~NeuralNetworkForceImpl() { 45 | if (session != NULL) { 46 | TF_CloseSession(session, status); 47 | TF_DeleteSession(session, status); 48 | } 49 | if (graph != NULL) 50 | TF_DeleteGraph(graph); 51 | TF_DeleteStatus(status); 52 | } 53 | 54 | void NeuralNetworkForceImpl::initialize(ContextImpl& context) { 55 | // Load the graph from the file. 56 | 57 | string graphProto = owner.getGraphProto(); 58 | TF_Buffer* buffer = TF_NewBufferFromString(graphProto.c_str(), graphProto.size()); 59 | graph = TF_NewGraph(); 60 | TF_ImportGraphDefOptions* importOptions = TF_NewImportGraphDefOptions(); 61 | TF_GraphImportGraphDef(graph, buffer, importOptions, status); 62 | if (TF_GetCode(status) != TF_OK) 63 | throw OpenMMException(string("Error loading TensorFlow graph: ")+TF_Message(status)); 64 | TF_DeleteImportGraphDefOptions(importOptions); 65 | TF_DeleteBuffer(buffer); 66 | 67 | // Check that the graph contains all the expected elements and that their types 68 | // are supported. 69 | 70 | TF_Output positions = {TF_GraphOperationByName(graph, "positions"), 0}; 71 | if (positions.oper == NULL) 72 | throw OpenMMException("NeuralNetworkForce: the graph does not have a 'positions' input"); 73 | TF_DataType positionsType = TF_OperationOutputType(positions); 74 | if (positionsType != TF_FLOAT && positionsType != TF_DOUBLE) 75 | throw OpenMMException("NeuralNetworkForce: 'positions' must have type float32 or float64"); 76 | TF_DataType boxType = TF_FLOAT; 77 | if (owner.usesPeriodicBoundaryConditions()) { 78 | TF_Output boxvectors = {TF_GraphOperationByName(graph, "boxvectors"), 0}; 79 | if (boxvectors.oper == NULL) 80 | throw OpenMMException("NeuralNetworkForce: the graph does not have a 'boxvectors' input"); 81 | boxType = TF_OperationOutputType(boxvectors); 82 | if (boxType != TF_FLOAT && boxType != TF_DOUBLE) 83 | throw OpenMMException("NeuralNetworkForce: 'boxvectors' must have type float32 or float64"); 84 | } 85 | TF_Output energy = {TF_GraphOperationByName(graph, "energy"), 0}; 86 | if (energy.oper == NULL) 87 | throw OpenMMException("NeuralNetworkForce: the graph does not have an 'energy' output"); 88 | TF_DataType energyType = TF_OperationOutputType(energy); 89 | if (energyType != TF_FLOAT && energyType != TF_DOUBLE) 90 | throw OpenMMException("NeuralNetworkForce: 'energy' must have type float32 or float64"); 91 | TF_Output forces = {TF_GraphOperationByName(graph, "forces"), 0}; 92 | if (forces.oper == NULL) 93 | throw OpenMMException("NeuralNetworkForce: the graph does not have a 'forces' output"); 94 | TF_DataType forcesType = TF_OperationOutputType(forces); 95 | if (forcesType != TF_FLOAT && forcesType != TF_DOUBLE) 96 | throw OpenMMException("NeuralNetworkForce: 'forces' must have type float32 or float64"); 97 | 98 | // Create the TensorFlow Session. 99 | 100 | TF_SessionOptions* sessionOptions = TF_NewSessionOptions(); 101 | session = TF_NewSession(graph, sessionOptions, status); 102 | if (TF_GetCode(status) != TF_OK) 103 | throw OpenMMException(string("Error creating TensorFlow session: ")+TF_Message(status)); 104 | TF_DeleteSessionOptions(sessionOptions); 105 | 106 | // Create the kernel. 107 | 108 | kernel = context.getPlatform().createKernel(CalcNeuralNetworkForceKernel::Name(), context); 109 | kernel.getAs().initialize(context.getSystem(), owner, session, graph, 110 | positionsType, boxType, energyType, forcesType); 111 | } 112 | 113 | double NeuralNetworkForceImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) { 114 | if ((groups&(1<().execute(context, includeForces, includeEnergy); 116 | return 0.0; 117 | } 118 | 119 | std::vector NeuralNetworkForceImpl::getKernelNames() { 120 | std::vector names; 121 | names.push_back(CalcNeuralNetworkForceKernel::Name()); 122 | return names; 123 | } 124 | -------------------------------------------------------------------------------- /platforms/cuda/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #--------------------------------------------------- 2 | # OpenMM Neural Network Plugin CUDA Platform 3 | #---------------------------------------------------- 4 | 5 | # Collect up information about the version of the OpenMM library we're building 6 | # and make it available to the code so it can be built into the binaries. 7 | 8 | SET(NN_CUDA_LIBRARY_NAME OpenMMNNCUDA) 9 | 10 | SET(SHARED_TARGET ${NN_CUDA_LIBRARY_NAME}) 11 | 12 | 13 | # These are all the places to search for header files which are 14 | # to be part of the API. 15 | SET(API_INCLUDE_DIRS "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/include/internal") 16 | 17 | # Locate header files. 18 | SET(API_INCLUDE_FILES) 19 | FOREACH(dir ${API_INCLUDE_DIRS}) 20 | FILE(GLOB fullpaths ${dir}/*.h) 21 | SET(API_INCLUDE_FILES ${API_INCLUDE_FILES} ${fullpaths}) 22 | ENDFOREACH(dir) 23 | 24 | # collect up source files 25 | SET(SOURCE_FILES) # empty 26 | SET(SOURCE_INCLUDE_FILES) 27 | 28 | FILE(GLOB_RECURSE src_files ${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.c) 29 | FILE(GLOB incl_files ${CMAKE_CURRENT_SOURCE_DIR}/src/*.h) 30 | SET(SOURCE_FILES ${SOURCE_FILES} ${src_files}) #append 31 | SET(SOURCE_INCLUDE_FILES ${SOURCE_INCLUDE_FILES} ${incl_files}) 32 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/include) 33 | 34 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/src) 35 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/cuda/include) 36 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/cuda/src) 37 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_BINARY_DIR}/platforms/cuda/src) 38 | 39 | # Set variables needed for encoding kernel sources into a C++ class 40 | 41 | SET(CUDA_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src) 42 | SET(CUDA_SOURCE_CLASS CudaNeuralNetworkKernelSources) 43 | SET(CUDA_KERNELS_CPP ${CMAKE_CURRENT_BINARY_DIR}/src/${CUDA_SOURCE_CLASS}.cpp) 44 | SET(CUDA_KERNELS_H ${CMAKE_CURRENT_BINARY_DIR}/src/${CUDA_SOURCE_CLASS}.h) 45 | SET(SOURCE_FILES ${SOURCE_FILES} ${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H}) 46 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_BINARY_DIR}/src) 47 | 48 | # Create the library 49 | 50 | INCLUDE_DIRECTORIES(${CUDA_TOOLKIT_INCLUDE}) 51 | 52 | FILE(GLOB CUDA_KERNELS ${CUDA_SOURCE_DIR}/kernels/*.cu) 53 | ADD_CUSTOM_COMMAND(OUTPUT ${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H} 54 | COMMAND ${CMAKE_COMMAND} 55 | ARGS -D CUDA_SOURCE_DIR=${CUDA_SOURCE_DIR} -D CUDA_KERNELS_CPP=${CUDA_KERNELS_CPP} -D CUDA_KERNELS_H=${CUDA_KERNELS_H} -D CUDA_SOURCE_CLASS=${CUDA_SOURCE_CLASS} -P ${CMAKE_SOURCE_DIR}/platforms/cuda/EncodeCUDAFiles.cmake 56 | DEPENDS ${CUDA_KERNELS} 57 | ) 58 | SET_SOURCE_FILES_PROPERTIES(${CUDA_KERNELS_CPP} ${CUDA_KERNELS_H} PROPERTIES GENERATED TRUE) 59 | ADD_LIBRARY(${SHARED_TARGET} SHARED ${SOURCE_FILES} ${SOURCE_INCLUDE_FILES} ${API_INCLUDE_FILES}) 60 | 61 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} ${CUDA_LIBRARIES}) 62 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} OpenMM) 63 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} OpenMMCUDA) 64 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} ${NN_LIBRARY_NAME}) 65 | SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES 66 | COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" 67 | LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") 68 | IF (APPLE) 69 | SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES LINK_FLAGS "-F/Library/Frameworks -framework CUDA ${EXTRA_COMPILE_FLAGS}") 70 | ENDIF (APPLE) 71 | 72 | INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) 73 | # Ensure that links to the main library will be resolved. 74 | IF (APPLE) 75 | INSTALL(CODE "EXECUTE_PROCESS(COMMAND install_name_tool -add_rpath @loader_path/.. ${CMAKE_INSTALL_PREFIX}/lib/plugins/lib${SHARED_TARGET}.dylib)") 76 | ENDIF (APPLE) 77 | 78 | SUBDIRS (tests) 79 | -------------------------------------------------------------------------------- /platforms/cuda/EncodeCUDAFiles.cmake: -------------------------------------------------------------------------------- 1 | FILE(GLOB CUDA_KERNELS ${CUDA_SOURCE_DIR}/kernels/*.cu) 2 | SET(CUDA_FILE_DECLARATIONS) 3 | SET(CUDA_FILE_DEFINITIONS) 4 | CONFIGURE_FILE(${CUDA_SOURCE_DIR}/${CUDA_SOURCE_CLASS}.cpp.in ${CUDA_KERNELS_CPP}) 5 | FOREACH(file ${CUDA_KERNELS}) 6 | # Load the file contents and process it. 7 | FILE(STRINGS ${file} file_content NEWLINE_CONSUME) 8 | # Replace all backslashes by double backslashes as they are being put in a C string. 9 | # Be careful not to replace the backslash before a semicolon as that is the CMAKE 10 | # internal escaping of a semicolon to prevent it from acting as a list seperator. 11 | STRING(REGEX REPLACE "\\\\([^;])" "\\\\\\\\\\1" file_content "${file_content}") 12 | # Escape double quotes as being put in a C string. 13 | STRING(REPLACE "\"" "\\\"" file_content "${file_content}") 14 | # Split in separate C strings for each line. 15 | STRING(REPLACE "\n" "\\n\"\n\"" file_content "${file_content}") 16 | 17 | # Determine a name for the variable that will contain this file's contents 18 | FILE(RELATIVE_PATH filename ${CUDA_SOURCE_DIR}/kernels ${file}) 19 | STRING(LENGTH ${filename} filename_length) 20 | MATH(EXPR filename_length ${filename_length}-3) 21 | STRING(SUBSTRING ${filename} 0 ${filename_length} variable_name) 22 | 23 | # Record the variable declaration and definition. 24 | SET(CUDA_FILE_DECLARATIONS ${CUDA_FILE_DECLARATIONS}static\ const\ std::string\ ${variable_name};\n) 25 | FILE(APPEND ${CUDA_KERNELS_CPP} const\ string\ ${CUDA_SOURCE_CLASS}::${variable_name}\ =\ \"${file_content}\"\;\n) 26 | ENDFOREACH(file) 27 | CONFIGURE_FILE(${CUDA_SOURCE_DIR}/${CUDA_SOURCE_CLASS}.h.in ${CUDA_KERNELS_H}) 28 | -------------------------------------------------------------------------------- /platforms/cuda/include/CudaNeuralNetworkKernelFactory.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_CUDA_NEURAL_NETWORK_KERNEL_FACTORY_H_ 2 | #define OPENMM_CUDA_NEURAL_NETWORK_KERNEL_FACTORY_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "openmm/KernelFactory.h" 36 | 37 | namespace OpenMM { 38 | 39 | /** 40 | * This KernelFactory creates kernels for the CUDA implementation of the neural network plugin. 41 | */ 42 | 43 | class CudaNeuralNetworkKernelFactory : public KernelFactory { 44 | public: 45 | KernelImpl* createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const; 46 | }; 47 | 48 | } // namespace OpenMM 49 | 50 | #endif /*OPENMM_CUDA_NEURAL_NETWORK_KERNEL_FACTORY_H_*/ 51 | -------------------------------------------------------------------------------- /platforms/cuda/src/CudaNeuralNetworkKernelFactory.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM-NN * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include 33 | 34 | #include "CudaNeuralNetworkKernelFactory.h" 35 | #include "CudaNeuralNetworkKernels.h" 36 | #include "openmm/internal/windowsExport.h" 37 | #include "openmm/internal/ContextImpl.h" 38 | #include "openmm/OpenMMException.h" 39 | #include 40 | 41 | using namespace NNPlugin; 42 | using namespace OpenMM; 43 | using namespace std; 44 | 45 | extern "C" OPENMM_EXPORT void registerPlatforms() { 46 | } 47 | 48 | extern "C" OPENMM_EXPORT void registerKernelFactories() { 49 | try { 50 | int argc = 0; 51 | vector argv = {NULL}; 52 | Platform& platform = Platform::getPlatformByName("CUDA"); 53 | CudaNeuralNetworkKernelFactory* factory = new CudaNeuralNetworkKernelFactory(); 54 | platform.registerKernelFactory(CalcNeuralNetworkForceKernel::Name(), factory); 55 | } 56 | catch (std::exception ex) { 57 | // Ignore 58 | } 59 | } 60 | 61 | extern "C" OPENMM_EXPORT void registerNeuralNetworkCudaKernelFactories() { 62 | try { 63 | Platform::getPlatformByName("CUDA"); 64 | } 65 | catch (...) { 66 | Platform::registerPlatform(new CudaPlatform()); 67 | } 68 | registerKernelFactories(); 69 | } 70 | 71 | KernelImpl* CudaNeuralNetworkKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const { 72 | CudaContext& cu = *static_cast(context.getPlatformData())->contexts[0]; 73 | if (name == CalcNeuralNetworkForceKernel::Name()) 74 | return new CudaCalcNeuralNetworkForceKernel(name, platform, cu); 75 | throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); 76 | } 77 | -------------------------------------------------------------------------------- /platforms/cuda/src/CudaNeuralNetworkKernelSources.cpp.in: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "CudaNeuralNetworkKernelSources.h" 33 | 34 | using namespace NNPlugin; 35 | using namespace std; 36 | 37 | -------------------------------------------------------------------------------- /platforms/cuda/src/CudaNeuralNetworkKernelSources.h.in: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_CUDA_NEURAL_NETWORK_KERNEL_SOURCES_H_ 2 | #define OPENMM_CUDA_NEURAL_NETWORK_KERNEL_SOURCES_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include 36 | 37 | namespace NNPlugin { 38 | 39 | /** 40 | * This class is a central holding place for the source code of CUDA kernels. 41 | * The CMake build script inserts declarations into it based on the .cu files in the 42 | * kernels subfolder. 43 | */ 44 | 45 | class CudaNeuralNetworkKernelSources { 46 | public: 47 | @CUDA_FILE_DECLARATIONS@ 48 | }; 49 | 50 | } // namespace NNPlugin 51 | 52 | #endif /*OPENMM_CUDA_NEURAL_NETWORK_KERNEL_SOURCES_H_*/ 53 | -------------------------------------------------------------------------------- /platforms/cuda/src/CudaNeuralNetworkKernels.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "CudaNeuralNetworkKernels.h" 33 | #include "CudaNeuralNetworkKernelSources.h" 34 | #include "openmm/internal/ContextImpl.h" 35 | #include 36 | 37 | using namespace NNPlugin; 38 | using namespace OpenMM; 39 | using namespace std; 40 | 41 | CudaCalcNeuralNetworkForceKernel::~CudaCalcNeuralNetworkForceKernel() { 42 | if (positionsTensor != NULL) 43 | TF_DeleteTensor(positionsTensor); 44 | if (boxVectorsTensor != NULL) 45 | TF_DeleteTensor(boxVectorsTensor); 46 | } 47 | 48 | void CudaCalcNeuralNetworkForceKernel::initialize(const System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 49 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType) { 50 | cu.setAsCurrent(); 51 | this->session = session; 52 | this->graph = graph; 53 | this->positionsType = positionsType; 54 | this->boxType = boxType; 55 | this->energyType = energyType; 56 | this->forcesType = forcesType; 57 | usePeriodic = force.usesPeriodicBoundaryConditions(); 58 | int numParticles = system.getNumParticles(); 59 | 60 | // Construct input tensors. 61 | 62 | int64_t positionsDims[] = {numParticles, 3}; 63 | positionsTensor = TF_AllocateTensor(positionsType, positionsDims, 2, numParticles*3*TF_DataTypeSize(positionsType)); 64 | if (usePeriodic) { 65 | int64_t boxVectorsDims[] = {3, 3}; 66 | boxVectorsTensor = TF_AllocateTensor(boxType, boxVectorsDims, 2, 9*TF_DataTypeSize(boxType)); 67 | } 68 | 69 | // Inititalize CUDA objects. 70 | 71 | networkForces.initialize(cu, 3*numParticles, TF_DataTypeSize(forcesType), "networkForces"); 72 | map defines; 73 | if (forcesType == TF_FLOAT) 74 | defines["FORCES_TYPE"] = "float"; 75 | else 76 | defines["FORCES_TYPE"] = "double"; 77 | CUmodule module = cu.createModule(CudaNeuralNetworkKernelSources::neuralNetworkForce, defines); 78 | addForcesKernel = cu.getKernel(module, "addForces"); 79 | } 80 | 81 | double CudaCalcNeuralNetworkForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) { 82 | vector pos; 83 | context.getPositions(pos); 84 | int numParticles = cu.getNumAtoms(); 85 | if (positionsType == TF_FLOAT) { 86 | float* positions = reinterpret_cast(TF_TensorData(positionsTensor)); 87 | for (int i = 0; i < numParticles; i++) { 88 | positions[3*i] = pos[i][0]; 89 | positions[3*i+1] = pos[i][1]; 90 | positions[3*i+2] = pos[i][2]; 91 | } 92 | } 93 | else { 94 | double* positions = reinterpret_cast(TF_TensorData(positionsTensor)); 95 | for (int i = 0; i < numParticles; i++) { 96 | positions[3*i] = pos[i][0]; 97 | positions[3*i+1] = pos[i][1]; 98 | positions[3*i+2] = pos[i][2]; 99 | } 100 | } 101 | if (usePeriodic) { 102 | Vec3 box[3]; 103 | cu.getPeriodicBoxVectors(box[0], box[1], box[2]); 104 | if (boxType == TF_FLOAT) { 105 | float* boxVectors = reinterpret_cast(TF_TensorData(boxVectorsTensor)); 106 | for (int i = 0; i < 3; i++) 107 | for (int j = 0; j < 3; j++) 108 | boxVectors[3*i+j] = box[i][j]; 109 | } 110 | else { 111 | double* boxVectors = reinterpret_cast(TF_TensorData(boxVectorsTensor)); 112 | for (int i = 0; i < 3; i++) 113 | for (int j = 0; j < 3; j++) 114 | boxVectors[3*i+j] = box[i][j]; 115 | } 116 | } 117 | vector inputs, outputs; 118 | int forceOutputIndex = 0; 119 | if (includeEnergy) 120 | outputs.push_back({TF_GraphOperationByName(graph, "energy"), 0}); 121 | if (includeForces) { 122 | forceOutputIndex = outputs.size(); 123 | outputs.push_back({TF_GraphOperationByName(graph, "forces"), 0}); 124 | } 125 | vector inputTensors, outputTensors(outputs.size()); 126 | inputs.push_back({TF_GraphOperationByName(graph, "positions"), 0}); 127 | inputTensors.push_back(positionsTensor); 128 | if (usePeriodic) { 129 | inputs.push_back({TF_GraphOperationByName(graph, "boxvectors"), 0}); 130 | inputTensors.push_back(boxVectorsTensor); 131 | } 132 | TF_Status* status = TF_NewStatus(); 133 | TF_SessionRun(session, NULL, &inputs[0], &inputTensors[0], inputs.size(), 134 | &outputs[0], &outputTensors[0], outputs.size(), 135 | NULL, 0, NULL, status); 136 | if (TF_GetCode(status) != TF_OK) 137 | throw OpenMMException(string("Error running TensorFlow session: ")+TF_Message(status)); 138 | TF_DeleteStatus(status); 139 | double energy = 0.0; 140 | if (includeEnergy) { 141 | if (energyType == TF_FLOAT) 142 | energy = reinterpret_cast(TF_TensorData(outputTensors[0]))[0]; 143 | else 144 | energy = reinterpret_cast(TF_TensorData(outputTensors[0]))[0]; 145 | } 146 | if (includeForces) { 147 | const void* data = TF_TensorData(outputTensors[forceOutputIndex]); 148 | networkForces.upload(data); 149 | int paddedNumAtoms = cu.getPaddedNumAtoms(); 150 | void* args[] = {&networkForces.getDevicePointer(), &cu.getForce().getDevicePointer(), &cu.getAtomIndexArray().getDevicePointer(), &numParticles, &paddedNumAtoms}; 151 | cu.executeKernel(addForcesKernel, args, numParticles); 152 | } 153 | return energy; 154 | } 155 | -------------------------------------------------------------------------------- /platforms/cuda/src/CudaNeuralNetworkKernels.h: -------------------------------------------------------------------------------- 1 | #ifndef CUDA_NEURAL_NETWORK_KERNELS_H_ 2 | #define CUDA_NEURAL_NETWORK_KERNELS_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "NeuralNetworkKernels.h" 36 | #include "openmm/cuda/CudaContext.h" 37 | #include "openmm/cuda/CudaArray.h" 38 | 39 | namespace NNPlugin { 40 | 41 | /** 42 | * This kernel is invoked by NeuralNetworkForce to calculate the forces acting on the system and the energy of the system. 43 | */ 44 | class CudaCalcNeuralNetworkForceKernel : public CalcNeuralNetworkForceKernel { 45 | public: 46 | CudaCalcNeuralNetworkForceKernel(std::string name, const OpenMM::Platform& platform, OpenMM::CudaContext& cu) : 47 | CalcNeuralNetworkForceKernel(name, platform), hasInitializedKernel(false), cu(cu), 48 | positionsTensor(NULL), boxVectorsTensor(NULL) { 49 | } 50 | ~CudaCalcNeuralNetworkForceKernel(); 51 | /** 52 | * Initialize the kernel. 53 | * 54 | * @param system the System this kernel will be applied to 55 | * @param force the NeuralNetworkForce this kernel will be used for 56 | * @param session the TensorFlow session in which to do calculations 57 | * @param graph the TensorFlow graph to use for computing forces and energy 58 | * @param positionsType the data type of the "positions" tensor 59 | * @param boxType the data type of the "boxvectors" tensor 60 | * @param energyType the data type of the "energy" tensor 61 | * @param forcesType the data type of the "forces" tensor 62 | */ 63 | void initialize(const OpenMM::System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 64 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType); 65 | /** 66 | * Execute the kernel to calculate the forces and/or energy. 67 | * 68 | * @param context the context in which to execute this kernel 69 | * @param includeForces true if forces should be calculated 70 | * @param includeEnergy true if the energy should be calculated 71 | * @return the potential energy due to the force 72 | */ 73 | double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy); 74 | private: 75 | bool hasInitializedKernel; 76 | OpenMM::CudaContext& cu; 77 | TF_Session* session; 78 | TF_Graph* graph; 79 | TF_Tensor* positionsTensor; 80 | TF_Tensor* boxVectorsTensor; 81 | TF_DataType positionsType, boxType, energyType, forcesType; 82 | bool usePeriodic; 83 | OpenMM::CudaArray networkForces; 84 | CUfunction addForcesKernel; 85 | }; 86 | 87 | } // namespace NNPlugin 88 | 89 | #endif /*CUDA_NEURAL_NETWORK_KERNELS_H_*/ 90 | -------------------------------------------------------------------------------- /platforms/cuda/src/kernels/neuralNetworkForce.cu: -------------------------------------------------------------------------------- 1 | extern "C" __global__ 2 | void addForces(const FORCES_TYPE* __restrict__ forces, long long* __restrict__ forceBuffers, int* __restrict__ atomIndex, int numAtoms, int paddedNumAtoms) { 3 | for (int atom = blockIdx.x*blockDim.x+threadIdx.x; atom < numAtoms; atom += blockDim.x*gridDim.x) { 4 | int index = atomIndex[atom]; 5 | forceBuffers[atom] += (long long) (forces[3*index]*0x100000000); 6 | forceBuffers[atom+paddedNumAtoms] += (long long) (forces[3*index+1]*0x100000000); 7 | forceBuffers[atom+2*paddedNumAtoms] += (long long) (forces[3*index+2]*0x100000000); 8 | } 9 | } 10 | 11 | -------------------------------------------------------------------------------- /platforms/cuda/tests/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # Testing 3 | # 4 | 5 | INCLUDE_DIRECTORIES(${CUDA_INCLUDE_DIR}) 6 | 7 | # Automatically create tests using files named "Test*.cpp" 8 | FILE(GLOB TEST_PROGS "*Test*.cpp") 9 | FOREACH(TEST_PROG ${TEST_PROGS}) 10 | GET_FILENAME_COMPONENT(TEST_ROOT ${TEST_PROG} NAME_WE) 11 | 12 | # Link with shared library 13 | ADD_EXECUTABLE(${TEST_ROOT} ${TEST_PROG}) 14 | TARGET_LINK_LIBRARIES(${TEST_ROOT} ${SHARED_EXAMPLE_TARGET} ${SHARED_TARGET}) 15 | IF (APPLE) 16 | SET_TARGET_PROPERTIES(${TEST_ROOT} PROPERTIES LINK_FLAGS "${EXTRA_COMPILE_FLAGS} -F/Library/Frameworks -framework CUDA" COMPILE_FLAGS "${EXTRA_COMPILE_FLAGS}") 17 | ELSE (APPLE) 18 | SET_TARGET_PROPERTIES(${TEST_ROOT} PROPERTIES LINK_FLAGS "${EXTRA_COMPILE_FLAGS}" COMPILE_FLAGS "${EXTRA_COMPILE_FLAGS}") 19 | ENDIF (APPLE) 20 | ADD_TEST(NAME ${TEST_ROOT}Single COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} single WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 21 | ADD_TEST(NAME ${TEST_ROOT}Mixed COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} mixed WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 22 | ADD_TEST(NAME ${TEST_ROOT}Double COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} double WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 23 | 24 | ENDFOREACH(TEST_PROG ${TEST_PROGS}) 25 | -------------------------------------------------------------------------------- /platforms/cuda/tests/TestCudaNeuralNetworkForce.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | /** 33 | * This tests the CUDA implementation of NeuralNetworkForce. 34 | */ 35 | 36 | #include "NeuralNetworkForce.h" 37 | #include "openmm/internal/AssertionUtilities.h" 38 | #include "openmm/Context.h" 39 | #include "openmm/Platform.h" 40 | #include "openmm/System.h" 41 | #include "openmm/VerletIntegrator.h" 42 | #include "sfmt/SFMT.h" 43 | #include 44 | #include 45 | #include 46 | 47 | using namespace NNPlugin; 48 | using namespace OpenMM; 49 | using namespace std; 50 | 51 | extern "C" OPENMM_EXPORT void registerNeuralNetworkCudaKernelFactories(); 52 | 53 | void testForce() { 54 | // Create a random cloud of particles. 55 | 56 | const int numParticles = 10; 57 | System system; 58 | vector positions(numParticles); 59 | OpenMM_SFMT::SFMT sfmt; 60 | init_gen_rand(0, sfmt); 61 | for (int i = 0; i < numParticles; i++) { 62 | system.addParticle(1.0); 63 | positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10; 64 | } 65 | NeuralNetworkForce* force = new NeuralNetworkForce("tests/central.pb"); 66 | system.addForce(force); 67 | 68 | // Compute the forces and energy. 69 | 70 | VerletIntegrator integ(1.0); 71 | Platform& platform = Platform::getPlatformByName("CUDA"); 72 | Context context(system, integ, platform); 73 | context.setPositions(positions); 74 | State state = context.getState(State::Energy | State::Forces); 75 | 76 | // See if the energy is correct. The network defines a potential of the form E(r) = |r|^2 77 | 78 | double expectedEnergy = 0; 79 | for (int i = 0; i < numParticles; i++) { 80 | Vec3 pos = positions[i]; 81 | double r = sqrt(pos.dot(pos)); 82 | expectedEnergy += r*r; 83 | ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5); 84 | } 85 | ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5); 86 | } 87 | 88 | void testPeriodicForce() { 89 | // Create a random cloud of particles. 90 | 91 | const int numParticles = 10; 92 | System system; 93 | system.setDefaultPeriodicBoxVectors(Vec3(2, 0, 0), Vec3(0, 3, 0), Vec3(0, 0, 4)); 94 | vector positions(numParticles); 95 | OpenMM_SFMT::SFMT sfmt; 96 | init_gen_rand(0, sfmt); 97 | for (int i = 0; i < numParticles; i++) { 98 | system.addParticle(1.0); 99 | positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10; 100 | } 101 | NeuralNetworkForce* force = new NeuralNetworkForce("tests/periodic.pb"); 102 | force->setUsesPeriodicBoundaryConditions(true); 103 | system.addForce(force); 104 | 105 | // Compute the forces and energy. 106 | 107 | VerletIntegrator integ(1.0); 108 | Platform& platform = Platform::getPlatformByName("CUDA"); 109 | Context context(system, integ, platform); 110 | context.setPositions(positions); 111 | State state = context.getState(State::Energy | State::Forces); 112 | 113 | // See if the energy is correct. The network defines a potential of the form E(r) = |r|^2 114 | 115 | double expectedEnergy = 0; 116 | for (int i = 0; i < numParticles; i++) { 117 | Vec3 pos = positions[i]; 118 | pos[0] -= floor(pos[0]/2.0)*2.0; 119 | pos[1] -= floor(pos[1]/3.0)*3.0; 120 | pos[2] -= floor(pos[2]/4.0)*4.0; 121 | double r = sqrt(pos.dot(pos)); 122 | expectedEnergy += r*r; 123 | ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5); 124 | } 125 | ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5); 126 | } 127 | 128 | int main(int argc, char* argv[]) { 129 | try { 130 | registerNeuralNetworkCudaKernelFactories(); 131 | if (argc > 1) 132 | Platform::getPlatformByName("CUDA").setPropertyDefaultValue("Precision", string(argv[1])); 133 | testForce(); 134 | testPeriodicForce(); 135 | } 136 | catch(const std::exception& e) { 137 | std::cout << "exception: " << e.what() << std::endl; 138 | return 1; 139 | } 140 | std::cout << "Done" << std::endl; 141 | return 0; 142 | } 143 | -------------------------------------------------------------------------------- /platforms/opencl/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #--------------------------------------------------- 2 | # OpenMM NeuralNetwork Plugin OpenCL Platform 3 | #---------------------------------------------------- 4 | 5 | # Collect up information about the version of the OpenMM library we're building 6 | # and make it available to the code so it can be built into the binaries. 7 | 8 | SET(NN_OPENCL_LIBRARY_NAME OpenMMNNOpenCL) 9 | 10 | SET(SHARED_TARGET ${NN_OPENCL_LIBRARY_NAME}) 11 | 12 | 13 | # These are all the places to search for header files which are 14 | # to be part of the API. 15 | SET(API_INCLUDE_DIRS "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/include/internal") 16 | 17 | # Locate header files. 18 | SET(API_INCLUDE_FILES) 19 | FOREACH(dir ${API_INCLUDE_DIRS}) 20 | FILE(GLOB fullpaths ${dir}/*.h) 21 | SET(API_INCLUDE_FILES ${API_INCLUDE_FILES} ${fullpaths}) 22 | ENDFOREACH(dir) 23 | 24 | # collect up source files 25 | SET(SOURCE_FILES) # empty 26 | SET(SOURCE_INCLUDE_FILES) 27 | 28 | FILE(GLOB_RECURSE src_files ${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.c) 29 | FILE(GLOB incl_files ${CMAKE_CURRENT_SOURCE_DIR}/src/*.h) 30 | SET(SOURCE_FILES ${SOURCE_FILES} ${src_files}) #append 31 | SET(SOURCE_INCLUDE_FILES ${SOURCE_INCLUDE_FILES} ${incl_files}) 32 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/include) 33 | 34 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/src) 35 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/opencl/include) 36 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/opencl/src) 37 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_BINARY_DIR}/platforms/opencl/src) 38 | 39 | # Set variables needed for encoding kernel sources into a C++ class 40 | 41 | SET(CL_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/src) 42 | SET(CL_SOURCE_CLASS OpenCLNeuralNetworkKernelSources) 43 | SET(CL_KERNELS_CPP ${CMAKE_CURRENT_BINARY_DIR}/src/${CL_SOURCE_CLASS}.cpp) 44 | SET(CL_KERNELS_H ${CMAKE_CURRENT_BINARY_DIR}/src/${CL_SOURCE_CLASS}.h) 45 | SET(SOURCE_FILES ${SOURCE_FILES} ${CL_KERNELS_CPP} ${CL_KERNELS_H}) 46 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_BINARY_DIR}/src) 47 | 48 | # Create the library 49 | 50 | INCLUDE_DIRECTORIES(${OPENCL_INCLUDE_DIR}) 51 | 52 | FILE(GLOB OPENCL_KERNELS ${CL_SOURCE_DIR}/kernels/*.cl) 53 | ADD_CUSTOM_COMMAND(OUTPUT ${CL_KERNELS_CPP} ${CL_KERNELS_H} 54 | COMMAND ${CMAKE_COMMAND} 55 | ARGS -D CL_SOURCE_DIR=${CL_SOURCE_DIR} -D CL_KERNELS_CPP=${CL_KERNELS_CPP} -D CL_KERNELS_H=${CL_KERNELS_H} -D CL_SOURCE_CLASS=${CL_SOURCE_CLASS} -P ${CMAKE_SOURCE_DIR}/platforms/opencl/EncodeCLFiles.cmake 56 | DEPENDS ${OPENCL_KERNELS} 57 | ) 58 | SET_SOURCE_FILES_PROPERTIES(${CL_KERNELS_CPP} ${CL_KERNELS_H} PROPERTIES GENERATED TRUE) 59 | ADD_LIBRARY(${SHARED_TARGET} SHARED ${SOURCE_FILES} ${SOURCE_INCLUDE_FILES} ${API_INCLUDE_FILES}) 60 | 61 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} ${OPENCL_LIBRARIES}) 62 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} OpenMM) 63 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} OpenMMOpenCL) 64 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} ${SHARED_NN_TARGET}) 65 | SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES 66 | COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" 67 | LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") 68 | 69 | INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) 70 | # Ensure that links to the main library will be resolved. 71 | IF (APPLE) 72 | INSTALL(CODE "EXECUTE_PROCESS(COMMAND install_name_tool -add_rpath @loader_path/.. ${CMAKE_INSTALL_PREFIX}/lib/plugins/lib${SHARED_TARGET}.dylib)") 73 | ENDIF (APPLE) 74 | 75 | SUBDIRS (tests) 76 | 77 | -------------------------------------------------------------------------------- /platforms/opencl/EncodeCLFiles.cmake: -------------------------------------------------------------------------------- 1 | FILE(GLOB OPENCL_KERNELS ${CL_SOURCE_DIR}/kernels/*.cl) 2 | SET(CL_FILE_DECLARATIONS) 3 | SET(CL_FILE_DEFINITIONS) 4 | CONFIGURE_FILE(${CL_SOURCE_DIR}/${CL_SOURCE_CLASS}.cpp.in ${CL_KERNELS_CPP}) 5 | FOREACH(file ${OPENCL_KERNELS}) 6 | # Load the file contents and process it. 7 | FILE(STRINGS ${file} file_content NEWLINE_CONSUME) 8 | # Replace all backslashes by double backslashes as they are being put in a C string. 9 | # Be careful not to replace the backslash before a semicolon as that is the CMAKE 10 | # internal escaping of a semicolon to prevent it from acting as a list seperator. 11 | STRING(REGEX REPLACE "\\\\([^;])" "\\\\\\\\\\1" file_content "${file_content}") 12 | # Escape double quotes as being put in a C string. 13 | STRING(REPLACE "\"" "\\\"" file_content "${file_content}") 14 | # Split in separate C strings for each line. 15 | STRING(REPLACE "\n" "\\n\"\n\"" file_content "${file_content}") 16 | 17 | # Determine a name for the variable that will contain this file's contents 18 | FILE(RELATIVE_PATH filename ${CL_SOURCE_DIR}/kernels ${file}) 19 | STRING(LENGTH ${filename} filename_length) 20 | MATH(EXPR filename_length ${filename_length}-3) 21 | STRING(SUBSTRING ${filename} 0 ${filename_length} variable_name) 22 | 23 | # Record the variable declaration and definition. 24 | SET(CL_FILE_DECLARATIONS ${CL_FILE_DECLARATIONS}static\ const\ std::string\ ${variable_name};\n) 25 | FILE(APPEND ${CL_KERNELS_CPP} const\ string\ ${CL_SOURCE_CLASS}::${variable_name}\ =\ \"${file_content}\"\;\n) 26 | ENDFOREACH(file) 27 | CONFIGURE_FILE(${CL_SOURCE_DIR}/${CL_SOURCE_CLASS}.h.in ${CL_KERNELS_H}) 28 | -------------------------------------------------------------------------------- /platforms/opencl/include/OpenCLNeuralNetworkKernelFactory.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_OPENCL_NEURAL_NETWORK_KERNEL_FACTORY_H_ 2 | #define OPENMM_OPENCL_NEURAL_NETWORK_KERNEL_FACTORY_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "openmm/KernelFactory.h" 36 | 37 | namespace NNPlugin { 38 | 39 | /** 40 | * This KernelFactory creates kernels for the OpenCL implementation of the neural network plugin. 41 | */ 42 | 43 | class OpenCLNeuralNetworkKernelFactory : public OpenMM::KernelFactory { 44 | public: 45 | OpenMM::KernelImpl* createKernelImpl(std::string name, const OpenMM::Platform& platform, OpenMM::ContextImpl& context) const; 46 | }; 47 | 48 | } // namespace NNPlugin 49 | 50 | #endif /*OPENMM_OPENCL_NEURAL_NETWORK_KERNEL_FACTORY_H_*/ 51 | -------------------------------------------------------------------------------- /platforms/opencl/src/OpenCLNeuralNetworkKernelFactory.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM-NN * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include 33 | 34 | #include "OpenCLNeuralNetworkKernelFactory.h" 35 | #include "OpenCLNeuralNetworkKernels.h" 36 | #include "openmm/internal/windowsExport.h" 37 | #include "openmm/internal/ContextImpl.h" 38 | #include "openmm/OpenMMException.h" 39 | #include 40 | 41 | using namespace NNPlugin; 42 | using namespace OpenMM; 43 | using namespace std; 44 | 45 | extern "C" OPENMM_EXPORT void registerPlatforms() { 46 | } 47 | 48 | extern "C" OPENMM_EXPORT void registerKernelFactories() { 49 | try { 50 | int argc = 0; 51 | vector argv = {NULL}; 52 | Platform& platform = Platform::getPlatformByName("OpenCL"); 53 | OpenCLNeuralNetworkKernelFactory* factory = new OpenCLNeuralNetworkKernelFactory(); 54 | platform.registerKernelFactory(CalcNeuralNetworkForceKernel::Name(), factory); 55 | } 56 | catch (std::exception ex) { 57 | // Ignore 58 | } 59 | } 60 | 61 | extern "C" OPENMM_EXPORT void registerNeuralNetworkOpenCLKernelFactories() { 62 | try { 63 | Platform::getPlatformByName("OpenCL"); 64 | } 65 | catch (...) { 66 | Platform::registerPlatform(new OpenCLPlatform()); 67 | } 68 | registerKernelFactories(); 69 | } 70 | 71 | KernelImpl* OpenCLNeuralNetworkKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const { 72 | OpenCLContext& cl = *static_cast(context.getPlatformData())->contexts[0]; 73 | if (name == CalcNeuralNetworkForceKernel::Name()) 74 | return new OpenCLCalcNeuralNetworkForceKernel(name, platform, cl); 75 | throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); 76 | } 77 | -------------------------------------------------------------------------------- /platforms/opencl/src/OpenCLNeuralNetworkKernelSources.cpp.in: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "OpenCLNeuralNetworkKernelSources.h" 33 | 34 | using namespace NNPlugin; 35 | using namespace std; 36 | 37 | -------------------------------------------------------------------------------- /platforms/opencl/src/OpenCLNeuralNetworkKernelSources.h.in: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_OPENCL_NEURAL_NETWORK_KERNEL_SOURCES_H_ 2 | #define OPENMM_OPENCL_NEURAL_NETWORK_KERNEL_SOURCES_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include 36 | 37 | namespace NNPlugin { 38 | 39 | /** 40 | * This class is a central holding place for the source code of OpenCL kernels. 41 | * The CMake build script inserts declarations into it based on the .cl files in the 42 | * kernels subfolder. 43 | */ 44 | 45 | class OpenCLNeuralNetworkKernelSources { 46 | public: 47 | @CL_FILE_DECLARATIONS@ 48 | }; 49 | 50 | } // namespace NNePlugin 51 | 52 | #endif /*OPENMM_OPENCL_NEURAL_NETWORK_KERNEL_SOURCES_H_*/ 53 | -------------------------------------------------------------------------------- /platforms/opencl/src/OpenCLNeuralNetworkKernels.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "OpenCLNeuralNetworkKernels.h" 33 | #include "OpenCLNeuralNetworkKernelSources.h" 34 | #include "openmm/internal/ContextImpl.h" 35 | #include 36 | 37 | using namespace NNPlugin; 38 | using namespace OpenMM; 39 | using namespace std; 40 | 41 | OpenCLCalcNeuralNetworkForceKernel::~OpenCLCalcNeuralNetworkForceKernel() { 42 | if (positionsTensor != NULL) 43 | TF_DeleteTensor(positionsTensor); 44 | if (boxVectorsTensor != NULL) 45 | TF_DeleteTensor(boxVectorsTensor); 46 | } 47 | 48 | void OpenCLCalcNeuralNetworkForceKernel::initialize(const System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 49 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType) { 50 | this->session = session; 51 | this->graph = graph; 52 | this->positionsType = positionsType; 53 | this->boxType = boxType; 54 | this->energyType = energyType; 55 | this->forcesType = forcesType; 56 | usePeriodic = force.usesPeriodicBoundaryConditions(); 57 | int numParticles = system.getNumParticles(); 58 | 59 | // Construct input tensors. 60 | 61 | int64_t positionsDims[] = {numParticles, 3}; 62 | positionsTensor = TF_AllocateTensor(positionsType, positionsDims, 2, numParticles*3*TF_DataTypeSize(positionsType)); 63 | if (usePeriodic) { 64 | int64_t boxVectorsDims[] = {3, 3}; 65 | boxVectorsTensor = TF_AllocateTensor(boxType, boxVectorsDims, 2, 9*TF_DataTypeSize(boxType)); 66 | } 67 | 68 | // Inititalize OpenCL objects. 69 | 70 | networkForces.initialize(cl, 3*numParticles, TF_DataTypeSize(forcesType), "networkForces"); 71 | map defines; 72 | if (forcesType == TF_FLOAT) 73 | defines["FORCES_TYPE"] = "float"; 74 | else 75 | defines["FORCES_TYPE"] = "double"; 76 | cl::Program program = cl.createProgram(OpenCLNeuralNetworkKernelSources::neuralNetworkForce, defines); 77 | addForcesKernel = cl::Kernel(program, "addForces"); 78 | } 79 | 80 | double OpenCLCalcNeuralNetworkForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) { 81 | vector pos; 82 | context.getPositions(pos); 83 | int numParticles = cl.getNumAtoms(); 84 | if (positionsType == TF_FLOAT) { 85 | float* positions = reinterpret_cast(TF_TensorData(positionsTensor)); 86 | for (int i = 0; i < numParticles; i++) { 87 | positions[3*i] = pos[i][0]; 88 | positions[3*i+1] = pos[i][1]; 89 | positions[3*i+2] = pos[i][2]; 90 | } 91 | } 92 | else { 93 | double* positions = reinterpret_cast(TF_TensorData(positionsTensor)); 94 | for (int i = 0; i < numParticles; i++) { 95 | positions[3*i] = pos[i][0]; 96 | positions[3*i+1] = pos[i][1]; 97 | positions[3*i+2] = pos[i][2]; 98 | } 99 | } 100 | if (usePeriodic) { 101 | Vec3 box[3]; 102 | cl.getPeriodicBoxVectors(box[0], box[1], box[2]); 103 | if (boxType == TF_FLOAT) { 104 | float* boxVectors = reinterpret_cast(TF_TensorData(boxVectorsTensor)); 105 | for (int i = 0; i < 3; i++) 106 | for (int j = 0; j < 3; j++) 107 | boxVectors[3*i+j] = box[i][j]; 108 | } 109 | else { 110 | double* boxVectors = reinterpret_cast(TF_TensorData(boxVectorsTensor)); 111 | for (int i = 0; i < 3; i++) 112 | for (int j = 0; j < 3; j++) 113 | boxVectors[3*i+j] = box[i][j]; 114 | } 115 | } 116 | vector inputs, outputs; 117 | int forceOutputIndex = 0; 118 | if (includeEnergy) 119 | outputs.push_back({TF_GraphOperationByName(graph, "energy"), 0}); 120 | if (includeForces) { 121 | forceOutputIndex = outputs.size(); 122 | outputs.push_back({TF_GraphOperationByName(graph, "forces"), 0}); 123 | } 124 | vector inputTensors, outputTensors(outputs.size()); 125 | inputs.push_back({TF_GraphOperationByName(graph, "positions"), 0}); 126 | inputTensors.push_back(positionsTensor); 127 | if (usePeriodic) { 128 | inputs.push_back({TF_GraphOperationByName(graph, "boxvectors"), 0}); 129 | inputTensors.push_back(boxVectorsTensor); 130 | } 131 | TF_Status* status = TF_NewStatus(); 132 | TF_SessionRun(session, NULL, &inputs[0], &inputTensors[0], inputs.size(), 133 | &outputs[0], &outputTensors[0], outputs.size(), 134 | NULL, 0, NULL, status); 135 | if (TF_GetCode(status) != TF_OK) 136 | throw OpenMMException(string("Error running TensorFlow session: ")+TF_Message(status)); 137 | TF_DeleteStatus(status); 138 | double energy = 0.0; 139 | if (includeEnergy) { 140 | if (energyType == TF_FLOAT) 141 | energy = reinterpret_cast(TF_TensorData(outputTensors[0]))[0]; 142 | else 143 | energy = reinterpret_cast(TF_TensorData(outputTensors[0]))[0]; 144 | } 145 | if (includeForces) { 146 | const void* data = TF_TensorData(outputTensors[forceOutputIndex]); 147 | networkForces.upload(data); 148 | addForcesKernel.setArg(0, networkForces.getDeviceBuffer()); 149 | addForcesKernel.setArg(1, cl.getForceBuffers().getDeviceBuffer()); 150 | addForcesKernel.setArg(2, cl.getAtomIndexArray().getDeviceBuffer()); 151 | addForcesKernel.setArg(3, numParticles); 152 | cl.executeKernel(addForcesKernel, numParticles); 153 | } 154 | return energy; 155 | } 156 | 157 | -------------------------------------------------------------------------------- /platforms/opencl/src/OpenCLNeuralNetworkKernels.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENCL_NEURAL_NETWORK_KERNELS_H_ 2 | #define OPENCL_NEURAL_NETWORK_KERNELS_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "NeuralNetworkKernels.h" 36 | #include "openmm/opencl/OpenCLContext.h" 37 | #include "openmm/opencl/OpenCLArray.h" 38 | 39 | namespace NNPlugin { 40 | 41 | /** 42 | * This kernel is invoked by NeuralNetworkForce to calculate the forces acting on the system and the energy of the system. 43 | */ 44 | class OpenCLCalcNeuralNetworkForceKernel : public CalcNeuralNetworkForceKernel { 45 | public: 46 | OpenCLCalcNeuralNetworkForceKernel(std::string name, const OpenMM::Platform& platform, OpenMM::OpenCLContext& cl) : 47 | CalcNeuralNetworkForceKernel(name, platform), hasInitializedKernel(false), cl(cl), 48 | positionsTensor(NULL), boxVectorsTensor(NULL) { 49 | } 50 | ~OpenCLCalcNeuralNetworkForceKernel(); 51 | /** 52 | * Initialize the kernel. 53 | * 54 | * @param system the System this kernel will be applied to 55 | * @param force the NeuralNetworkForce this kernel will be used for 56 | * @param session the TensorFlow session in which to do calculations 57 | * @param graph the TensorFlow graph to use for computing forces and energy 58 | * @param positionsType the data type of the "positions" tensor 59 | * @param boxType the data type of the "boxvectors" tensor 60 | * @param energyType the data type of the "energy" tensor 61 | * @param forcesType the data type of the "forces" tensor 62 | */ 63 | void initialize(const OpenMM::System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 64 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType); 65 | /** 66 | * Execute the kernel to calculate the forces and/or energy. 67 | * 68 | * @param context the context in which to execute this kernel 69 | * @param includeForces true if forces should be calculated 70 | * @param includeEnergy true if the energy should be calculated 71 | * @return the potential energy due to the force 72 | */ 73 | double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy); 74 | private: 75 | bool hasInitializedKernel; 76 | OpenMM::OpenCLContext& cl; 77 | TF_Session* session; 78 | TF_Graph* graph; 79 | TF_Tensor* positionsTensor; 80 | TF_Tensor* boxVectorsTensor; 81 | TF_DataType positionsType, boxType, energyType, forcesType; 82 | bool usePeriodic; 83 | OpenMM::OpenCLArray networkForces; 84 | cl::Kernel addForcesKernel; 85 | }; 86 | 87 | } // namespace NNPlugin 88 | 89 | #endif /*OPENCL_NEURAL_NETWORK_KERNELS_H_*/ 90 | -------------------------------------------------------------------------------- /platforms/opencl/src/kernels/neuralNetworkForce.cl: -------------------------------------------------------------------------------- 1 | __kernel void addForces(__global const FORCES_TYPE* restrict forces, __global real4* restrict forceBuffers, __global int* restrict atomIndex, int numAtoms) { 2 | for (int atom = get_global_id(0); atom < numAtoms; atom += get_global_size(0)) { 3 | int index = atomIndex[atom]; 4 | real4 f = forceBuffers[atom]; 5 | f.xyz += (real3) (forces[3*index], forces[3*index+1], forces[3*index+2]); 6 | forceBuffers[atom] = f; 7 | } 8 | } 9 | 10 | -------------------------------------------------------------------------------- /platforms/opencl/tests/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # Testing 3 | # 4 | 5 | INCLUDE_DIRECTORIES(${OPENCL_INCLUDE_DIR}) 6 | 7 | # Automatically create tests using files named "Test*.cpp" 8 | FILE(GLOB TEST_PROGS "*Test*.cpp") 9 | FOREACH(TEST_PROG ${TEST_PROGS}) 10 | GET_FILENAME_COMPONENT(TEST_ROOT ${TEST_PROG} NAME_WE) 11 | 12 | # Link with shared library 13 | ADD_EXECUTABLE(${TEST_ROOT} ${TEST_PROG}) 14 | TARGET_LINK_LIBRARIES(${TEST_ROOT} ${SHARED_NN_TARGET} ${SHARED_TARGET}) 15 | SET_TARGET_PROPERTIES(${TEST_ROOT} PROPERTIES LINK_FLAGS "${EXTRA_COMPILE_FLAGS}" COMPILE_FLAGS "${EXTRA_COMPILE_FLAGS}") 16 | ADD_TEST(NAME ${TEST_ROOT}Single COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} single WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 17 | ADD_TEST(NAME ${TEST_ROOT}Mixed COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} mixed WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 18 | ADD_TEST(NAME ${TEST_ROOT}Double COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} double WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 19 | 20 | ENDFOREACH(TEST_PROG ${TEST_PROGS}) 21 | -------------------------------------------------------------------------------- /platforms/opencl/tests/TestOpenCLNeuralNetworkForce.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | /** 33 | * This tests the OpenCL implementation of NeuralNetworkForce. 34 | */ 35 | 36 | #include "NeuralNetworkForce.h" 37 | #include "openmm/internal/AssertionUtilities.h" 38 | #include "openmm/Context.h" 39 | #include "openmm/Platform.h" 40 | #include "openmm/System.h" 41 | #include "openmm/VerletIntegrator.h" 42 | #include "sfmt/SFMT.h" 43 | #include 44 | #include 45 | #include 46 | 47 | using namespace NNPlugin; 48 | using namespace OpenMM; 49 | using namespace std; 50 | 51 | extern "C" OPENMM_EXPORT void registerNeuralNetworkOpenCLKernelFactories(); 52 | 53 | void testForce() { 54 | // Create a random cloud of particles. 55 | 56 | const int numParticles = 10; 57 | System system; 58 | vector positions(numParticles); 59 | OpenMM_SFMT::SFMT sfmt; 60 | init_gen_rand(0, sfmt); 61 | for (int i = 0; i < numParticles; i++) { 62 | system.addParticle(1.0); 63 | positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10; 64 | } 65 | NeuralNetworkForce* force = new NeuralNetworkForce("tests/central.pb"); 66 | system.addForce(force); 67 | 68 | // Compute the forces and energy. 69 | 70 | VerletIntegrator integ(1.0); 71 | Platform& platform = Platform::getPlatformByName("OpenCL"); 72 | Context context(system, integ, platform); 73 | context.setPositions(positions); 74 | State state = context.getState(State::Energy | State::Forces); 75 | 76 | // See if the energy is correct. The network defines a potential of the form E(r) = |r|^2 77 | 78 | double expectedEnergy = 0; 79 | for (int i = 0; i < numParticles; i++) { 80 | Vec3 pos = positions[i]; 81 | double r = sqrt(pos.dot(pos)); 82 | expectedEnergy += r*r; 83 | ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5); 84 | } 85 | ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5); 86 | } 87 | 88 | void testPeriodicForce() { 89 | // Create a random cloud of particles. 90 | 91 | const int numParticles = 10; 92 | System system; 93 | system.setDefaultPeriodicBoxVectors(Vec3(2, 0, 0), Vec3(0, 3, 0), Vec3(0, 0, 4)); 94 | vector positions(numParticles); 95 | OpenMM_SFMT::SFMT sfmt; 96 | init_gen_rand(0, sfmt); 97 | for (int i = 0; i < numParticles; i++) { 98 | system.addParticle(1.0); 99 | positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10; 100 | } 101 | NeuralNetworkForce* force = new NeuralNetworkForce("tests/periodic.pb"); 102 | force->setUsesPeriodicBoundaryConditions(true); 103 | system.addForce(force); 104 | 105 | // Compute the forces and energy. 106 | 107 | VerletIntegrator integ(1.0); 108 | Platform& platform = Platform::getPlatformByName("OpenCL"); 109 | Context context(system, integ, platform); 110 | context.setPositions(positions); 111 | State state = context.getState(State::Energy | State::Forces); 112 | 113 | // See if the energy is correct. The network defines a potential of the form E(r) = |r|^2 114 | 115 | double expectedEnergy = 0; 116 | for (int i = 0; i < numParticles; i++) { 117 | Vec3 pos = positions[i]; 118 | pos[0] -= floor(pos[0]/2.0)*2.0; 119 | pos[1] -= floor(pos[1]/3.0)*3.0; 120 | pos[2] -= floor(pos[2]/4.0)*4.0; 121 | double r = sqrt(pos.dot(pos)); 122 | expectedEnergy += r*r; 123 | ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5); 124 | } 125 | ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5); 126 | } 127 | 128 | int main(int argc, char* argv[]) { 129 | try { 130 | registerNeuralNetworkOpenCLKernelFactories(); 131 | if (argc > 1) 132 | Platform::getPlatformByName("OpenCL").setPropertyDefaultValue("Precision", string(argv[1])); 133 | testForce(); 134 | testPeriodicForce(); 135 | } 136 | catch(const std::exception& e) { 137 | std::cout << "exception: " << e.what() << std::endl; 138 | return 1; 139 | } 140 | std::cout << "Done" << std::endl; 141 | return 0; 142 | } 143 | -------------------------------------------------------------------------------- /platforms/reference/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | #--------------------------------------------------- 2 | # OpenMM Neural Network Plugin Reference Platform 3 | #---------------------------------------------------- 4 | 5 | # Collect up information about the version of the OpenMM library we're building 6 | # and make it available to the code so it can be built into the binaries. 7 | 8 | SET(NN_REFERENCE_LIBRARY_NAME OpenMMNNReference) 9 | 10 | SET(SHARED_TARGET ${NN_REFERENCE_LIBRARY_NAME}) 11 | 12 | 13 | # These are all the places to search for header files which are 14 | # to be part of the API. 15 | SET(API_INCLUDE_DIRS "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/include/internal") 16 | 17 | # Locate header files. 18 | SET(API_INCLUDE_FILES) 19 | FOREACH(dir ${API_INCLUDE_DIRS}) 20 | FILE(GLOB fullpaths ${dir}/*.h) 21 | SET(API_INCLUDE_FILES ${API_INCLUDE_FILES} ${fullpaths}) 22 | ENDFOREACH(dir) 23 | 24 | # collect up source files 25 | SET(SOURCE_FILES) # empty 26 | SET(SOURCE_INCLUDE_FILES) 27 | 28 | FILE(GLOB_RECURSE src_files ${CMAKE_CURRENT_SOURCE_DIR}/src/*.cpp ${CMAKE_CURRENT_SOURCE_DIR}/${subdir}/src/*.c) 29 | FILE(GLOB incl_files ${CMAKE_CURRENT_SOURCE_DIR}/src/*.h) 30 | SET(SOURCE_FILES ${SOURCE_FILES} ${src_files}) #append 31 | SET(SOURCE_INCLUDE_FILES ${SOURCE_INCLUDE_FILES} ${incl_files}) 32 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/include) 33 | 34 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/src) 35 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/reference/include) 36 | INCLUDE_DIRECTORIES(BEFORE ${CMAKE_SOURCE_DIR}/platforms/reference/src) 37 | 38 | # Create the library 39 | 40 | INCLUDE_DIRECTORIES(${REFERENCE_INCLUDE_DIR}) 41 | 42 | ADD_LIBRARY(${SHARED_TARGET} SHARED ${SOURCE_FILES} ${SOURCE_INCLUDE_FILES} ${API_INCLUDE_FILES}) 43 | 44 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} OpenMM) 45 | TARGET_LINK_LIBRARIES(${SHARED_TARGET} debug ${SHARED_NN_TARGET} optimized ${SHARED_NN_TARGET}) 46 | SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES 47 | COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" 48 | LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") 49 | 50 | INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) 51 | # Ensure that links to the main library will be resolved. 52 | IF (APPLE) 53 | INSTALL(CODE "EXECUTE_PROCESS(COMMAND install_name_tool -add_rpath @loader_path/.. ${CMAKE_INSTALL_PREFIX}/lib/plugins/lib${SHARED_TARGET}.dylib)") 54 | ENDIF (APPLE) 55 | 56 | SUBDIRS (tests) 57 | -------------------------------------------------------------------------------- /platforms/reference/include/ReferenceNeuralNetworkKernelFactory.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_REFERENCE_NEURAL_NETWORK_KERNEL_FACTORY_H_ 2 | #define OPENMM_REFERENCE_NEURAL_NETWORK_KERNEL_FACTORY_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "openmm/KernelFactory.h" 36 | 37 | namespace OpenMM { 38 | 39 | /** 40 | * This KernelFactory creates kernels for the reference implementation of the neural network plugin. 41 | */ 42 | 43 | class ReferenceNeuralNetworkKernelFactory : public KernelFactory { 44 | public: 45 | KernelImpl* createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const; 46 | }; 47 | 48 | } // namespace OpenMM 49 | 50 | #endif /*OPENMM_REFERENCE_NEURAL_NETWORK_KERNEL_FACTORY_H_*/ 51 | -------------------------------------------------------------------------------- /platforms/reference/src/ReferenceNeuralNetworkKernelFactory.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM-NN * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "ReferenceNeuralNetworkKernelFactory.h" 33 | #include "ReferenceNeuralNetworkKernels.h" 34 | #include "openmm/reference/ReferencePlatform.h" 35 | #include "openmm/internal/ContextImpl.h" 36 | #include "openmm/OpenMMException.h" 37 | #include 38 | 39 | using namespace NNPlugin; 40 | using namespace OpenMM; 41 | using namespace std; 42 | 43 | extern "C" OPENMM_EXPORT void registerPlatforms() { 44 | } 45 | 46 | extern "C" OPENMM_EXPORT void registerKernelFactories() { 47 | int argc = 0; 48 | vector argv = {NULL}; 49 | for (int i = 0; i < Platform::getNumPlatforms(); i++) { 50 | Platform& platform = Platform::getPlatform(i); 51 | if (dynamic_cast(&platform) != NULL) { 52 | ReferenceNeuralNetworkKernelFactory* factory = new ReferenceNeuralNetworkKernelFactory(); 53 | platform.registerKernelFactory(CalcNeuralNetworkForceKernel::Name(), factory); 54 | } 55 | } 56 | } 57 | 58 | extern "C" OPENMM_EXPORT void registerNeuralNetworkReferenceKernelFactories() { 59 | registerKernelFactories(); 60 | } 61 | 62 | KernelImpl* ReferenceNeuralNetworkKernelFactory::createKernelImpl(std::string name, const Platform& platform, ContextImpl& context) const { 63 | ReferencePlatform::PlatformData& data = *static_cast(context.getPlatformData()); 64 | if (name == CalcNeuralNetworkForceKernel::Name()) 65 | return new ReferenceCalcNeuralNetworkForceKernel(name, platform); 66 | throw OpenMMException((std::string("Tried to create kernel with illegal kernel name '")+name+"'").c_str()); 67 | } 68 | -------------------------------------------------------------------------------- /platforms/reference/src/ReferenceNeuralNetworkKernels.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "ReferenceNeuralNetworkKernels.h" 33 | #include "NeuralNetworkForce.h" 34 | #include "openmm/OpenMMException.h" 35 | #include "openmm/internal/ContextImpl.h" 36 | #include "openmm/reference/ReferencePlatform.h" 37 | 38 | using namespace NNPlugin; 39 | using namespace OpenMM; 40 | using namespace std; 41 | 42 | static vector& extractPositions(ContextImpl& context) { 43 | ReferencePlatform::PlatformData* data = reinterpret_cast(context.getPlatformData()); 44 | return *((vector*) data->positions); 45 | } 46 | 47 | static vector& extractForces(ContextImpl& context) { 48 | ReferencePlatform::PlatformData* data = reinterpret_cast(context.getPlatformData()); 49 | return *((vector*) data->forces); 50 | } 51 | 52 | static Vec3* extractBoxVectors(ContextImpl& context) { 53 | ReferencePlatform::PlatformData* data = reinterpret_cast(context.getPlatformData()); 54 | return (Vec3*) data->periodicBoxVectors; 55 | } 56 | 57 | ReferenceCalcNeuralNetworkForceKernel::~ReferenceCalcNeuralNetworkForceKernel() { 58 | if (positionsTensor != NULL) 59 | TF_DeleteTensor(positionsTensor); 60 | if (boxVectorsTensor != NULL) 61 | TF_DeleteTensor(boxVectorsTensor); 62 | } 63 | 64 | void ReferenceCalcNeuralNetworkForceKernel::initialize(const System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 65 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType) { 66 | this->session = session; 67 | this->graph = graph; 68 | this->positionsType = positionsType; 69 | this->boxType = boxType; 70 | this->energyType = energyType; 71 | this->forcesType = forcesType; 72 | usePeriodic = force.usesPeriodicBoundaryConditions(); 73 | int numParticles = system.getNumParticles(); 74 | 75 | // Construct input tensors. 76 | 77 | int64_t positionsDims[] = {numParticles, 3}; 78 | positionsTensor = TF_AllocateTensor(positionsType, positionsDims, 2, numParticles*3*TF_DataTypeSize(positionsType)); 79 | if (usePeriodic) { 80 | int64_t boxVectorsDims[] = {3, 3}; 81 | boxVectorsTensor = TF_AllocateTensor(boxType, boxVectorsDims, 2, 9*TF_DataTypeSize(boxType)); 82 | } 83 | } 84 | 85 | double ReferenceCalcNeuralNetworkForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) { 86 | vector& pos = extractPositions(context); 87 | vector& force = extractForces(context); 88 | int numParticles = pos.size(); 89 | if (positionsType == TF_FLOAT) { 90 | float* positions = reinterpret_cast(TF_TensorData(positionsTensor)); 91 | for (int i = 0; i < numParticles; i++) { 92 | positions[3*i] = pos[i][0]; 93 | positions[3*i+1] = pos[i][1]; 94 | positions[3*i+2] = pos[i][2]; 95 | } 96 | } 97 | else { 98 | double* positions = reinterpret_cast(TF_TensorData(positionsTensor)); 99 | for (int i = 0; i < numParticles; i++) { 100 | positions[3*i] = pos[i][0]; 101 | positions[3*i+1] = pos[i][1]; 102 | positions[3*i+2] = pos[i][2]; 103 | } 104 | } 105 | if (usePeriodic) { 106 | Vec3* box = extractBoxVectors(context); 107 | if (boxType == TF_FLOAT) { 108 | float* boxVectors = reinterpret_cast(TF_TensorData(boxVectorsTensor)); 109 | for (int i = 0; i < 3; i++) 110 | for (int j = 0; j < 3; j++) 111 | boxVectors[3*i+j] = box[i][j]; 112 | } 113 | else { 114 | double* boxVectors = reinterpret_cast(TF_TensorData(boxVectorsTensor)); 115 | for (int i = 0; i < 3; i++) 116 | for (int j = 0; j < 3; j++) 117 | boxVectors[3*i+j] = box[i][j]; 118 | } 119 | } 120 | vector inputs, outputs; 121 | int forceOutputIndex = 0; 122 | if (includeEnergy) 123 | outputs.push_back({TF_GraphOperationByName(graph, "energy"), 0}); 124 | if (includeForces) { 125 | forceOutputIndex = outputs.size(); 126 | outputs.push_back({TF_GraphOperationByName(graph, "forces"), 0}); 127 | } 128 | vector inputTensors, outputTensors(outputs.size()); 129 | inputs.push_back({TF_GraphOperationByName(graph, "positions"), 0}); 130 | inputTensors.push_back(positionsTensor); 131 | if (usePeriodic) { 132 | inputs.push_back({TF_GraphOperationByName(graph, "boxvectors"), 0}); 133 | inputTensors.push_back(boxVectorsTensor); 134 | } 135 | TF_Status* status = TF_NewStatus(); 136 | TF_SessionRun(session, NULL, &inputs[0], &inputTensors[0], inputs.size(), 137 | &outputs[0], &outputTensors[0], outputs.size(), 138 | NULL, 0, NULL, status); 139 | if (TF_GetCode(status) != TF_OK) 140 | throw OpenMMException(string("Error running TensorFlow session: ")+TF_Message(status)); 141 | TF_DeleteStatus(status); 142 | double energy = 0.0; 143 | if (includeEnergy) { 144 | if (energyType == TF_FLOAT) 145 | energy = reinterpret_cast(TF_TensorData(outputTensors[0]))[0]; 146 | else 147 | energy = reinterpret_cast(TF_TensorData(outputTensors[0]))[0]; 148 | } 149 | if (includeForces) { 150 | if (forcesType == TF_FLOAT) { 151 | const float* data = reinterpret_cast(TF_TensorData(outputTensors[forceOutputIndex])); 152 | for (int i = 0; i < numParticles; i++) { 153 | force[i][0] += data[3*i]; 154 | force[i][1] += data[3*i+1]; 155 | force[i][2] += data[3*i+2]; 156 | } 157 | } 158 | else { 159 | const double* data = reinterpret_cast(TF_TensorData(outputTensors[forceOutputIndex])); 160 | for (int i = 0; i < numParticles; i++) { 161 | force[i][0] += data[3*i]; 162 | force[i][1] += data[3*i+1]; 163 | force[i][2] += data[3*i+2]; 164 | } 165 | } 166 | } 167 | return energy; 168 | } 169 | -------------------------------------------------------------------------------- /platforms/reference/src/ReferenceNeuralNetworkKernels.h: -------------------------------------------------------------------------------- 1 | #ifndef REFERENCE_NEURAL_NETWORK_KERNELS_H_ 2 | #define REFERENCE_NEURAL_NETWORK_KERNELS_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "NeuralNetworkKernels.h" 36 | #include "openmm/Platform.h" 37 | #include 38 | 39 | namespace NNPlugin { 40 | 41 | /** 42 | * This kernel is invoked by NeuralNetworkForce to calculate the forces acting on the system and the energy of the system. 43 | */ 44 | class ReferenceCalcNeuralNetworkForceKernel : public CalcNeuralNetworkForceKernel { 45 | public: 46 | ReferenceCalcNeuralNetworkForceKernel(std::string name, const OpenMM::Platform& platform) : CalcNeuralNetworkForceKernel(name, platform), 47 | positionsTensor(NULL), boxVectorsTensor(NULL) { 48 | } 49 | ~ReferenceCalcNeuralNetworkForceKernel(); 50 | /** 51 | * Initialize the kernel. 52 | * 53 | * @param system the System this kernel will be applied to 54 | * @param force the NeuralNetworkForce this kernel will be used for 55 | * @param session the TensorFlow session in which to do calculations 56 | * @param graph the TensorFlow graph to use for computing forces and energy 57 | * @param positionsType the data type of the "positions" tensor 58 | * @param boxType the data type of the "boxvectors" tensor 59 | * @param energyType the data type of the "energy" tensor 60 | * @param forcesType the data type of the "forces" tensor 61 | */ 62 | void initialize(const OpenMM::System& system, const NeuralNetworkForce& force, TF_Session* session, TF_Graph* graph, 63 | TF_DataType positionsType, TF_DataType boxType, TF_DataType energyType, TF_DataType forcesType); 64 | /** 65 | * Execute the kernel to calculate the forces and/or energy. 66 | * 67 | * @param context the context in which to execute this kernel 68 | * @param includeForces true if forces should be calculated 69 | * @param includeEnergy true if the energy should be calculated 70 | * @return the potential energy due to the force 71 | */ 72 | double execute(OpenMM::ContextImpl& context, bool includeForces, bool includeEnergy); 73 | private: 74 | TF_Session* session; 75 | TF_Graph* graph; 76 | TF_Tensor* positionsTensor; 77 | TF_Tensor* boxVectorsTensor; 78 | TF_DataType positionsType, boxType, energyType, forcesType; 79 | std::vector positions, boxVectors; 80 | bool usePeriodic; 81 | }; 82 | 83 | } // namespace NNPlugin 84 | 85 | #endif /*REFERENCE_NEURAL_NETWORK_KERNELS_H_*/ 86 | -------------------------------------------------------------------------------- /platforms/reference/tests/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # Testing 3 | # 4 | 5 | # Automatically create tests using files named "Test*.cpp" 6 | FILE(GLOB TEST_PROGS "*Test*.cpp") 7 | FOREACH(TEST_PROG ${TEST_PROGS}) 8 | GET_FILENAME_COMPONENT(TEST_ROOT ${TEST_PROG} NAME_WE) 9 | 10 | # Link with shared library 11 | 12 | ADD_EXECUTABLE(${TEST_ROOT} ${TEST_PROG}) 13 | TARGET_LINK_LIBRARIES(${TEST_ROOT} ${SHARED_TARGET}) 14 | SET_TARGET_PROPERTIES(${TEST_ROOT} PROPERTIES LINK_FLAGS "${EXTRA_COMPILE_FLAGS}" COMPILE_FLAGS "${EXTRA_COMPILE_FLAGS}") 15 | ADD_TEST(NAME ${TEST_ROOT} COMMAND ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT} single WORKING_DIRECTORY ${CMAKE_BINARY_DIR}) 16 | 17 | ENDFOREACH(TEST_PROG ${TEST_PROGS}) 18 | -------------------------------------------------------------------------------- /platforms/reference/tests/TestReferenceNeuralNetworkForce.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | /** 33 | * This tests the Reference implementation of NeuralNetworkForce. 34 | */ 35 | 36 | #include "NeuralNetworkForce.h" 37 | #include "openmm/internal/AssertionUtilities.h" 38 | #include "openmm/Context.h" 39 | #include "openmm/Platform.h" 40 | #include "openmm/System.h" 41 | #include "openmm/VerletIntegrator.h" 42 | #include "sfmt/SFMT.h" 43 | #include 44 | #include 45 | #include 46 | 47 | using namespace NNPlugin; 48 | using namespace OpenMM; 49 | using namespace std; 50 | 51 | extern "C" OPENMM_EXPORT void registerNeuralNetworkReferenceKernelFactories(); 52 | 53 | void testForce() { 54 | // Create a random cloud of particles. 55 | 56 | const int numParticles = 10; 57 | System system; 58 | vector positions(numParticles); 59 | OpenMM_SFMT::SFMT sfmt; 60 | init_gen_rand(0, sfmt); 61 | for (int i = 0; i < numParticles; i++) { 62 | system.addParticle(1.0); 63 | positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10; 64 | } 65 | NeuralNetworkForce* force = new NeuralNetworkForce("tests/central.pb"); 66 | system.addForce(force); 67 | 68 | // Compute the forces and energy. 69 | 70 | VerletIntegrator integ(1.0); 71 | Platform& platform = Platform::getPlatformByName("Reference"); 72 | Context context(system, integ, platform); 73 | context.setPositions(positions); 74 | State state = context.getState(State::Energy | State::Forces); 75 | 76 | // See if the energy is correct. The network defines a potential of the form E(r) = |r|^2 77 | 78 | double expectedEnergy = 0; 79 | for (int i = 0; i < numParticles; i++) { 80 | Vec3 pos = positions[i]; 81 | double r = sqrt(pos.dot(pos)); 82 | expectedEnergy += r*r; 83 | ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5); 84 | } 85 | ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5); 86 | } 87 | 88 | void testPeriodicForce() { 89 | // Create a random cloud of particles. 90 | 91 | const int numParticles = 10; 92 | System system; 93 | system.setDefaultPeriodicBoxVectors(Vec3(2, 0, 0), Vec3(0, 3, 0), Vec3(0, 0, 4)); 94 | vector positions(numParticles); 95 | OpenMM_SFMT::SFMT sfmt; 96 | init_gen_rand(0, sfmt); 97 | for (int i = 0; i < numParticles; i++) { 98 | system.addParticle(1.0); 99 | positions[i] = Vec3(genrand_real2(sfmt), genrand_real2(sfmt), genrand_real2(sfmt))*10; 100 | } 101 | NeuralNetworkForce* force = new NeuralNetworkForce("tests/periodic.pb"); 102 | force->setUsesPeriodicBoundaryConditions(true); 103 | system.addForce(force); 104 | 105 | // Compute the forces and energy. 106 | 107 | VerletIntegrator integ(1.0); 108 | Platform& platform = Platform::getPlatformByName("Reference"); 109 | Context context(system, integ, platform); 110 | context.setPositions(positions); 111 | State state = context.getState(State::Energy | State::Forces); 112 | 113 | // See if the energy is correct. The network defines a potential of the form E(r) = |r|^2 114 | 115 | double expectedEnergy = 0; 116 | for (int i = 0; i < numParticles; i++) { 117 | Vec3 pos = positions[i]; 118 | pos[0] -= floor(pos[0]/2.0)*2.0; 119 | pos[1] -= floor(pos[1]/3.0)*3.0; 120 | pos[2] -= floor(pos[2]/4.0)*4.0; 121 | double r = sqrt(pos.dot(pos)); 122 | expectedEnergy += r*r; 123 | ASSERT_EQUAL_VEC(pos*(-2.0), state.getForces()[i], 1e-5); 124 | } 125 | ASSERT_EQUAL_TOL(expectedEnergy, state.getPotentialEnergy(), 1e-5); 126 | } 127 | 128 | int main() { 129 | try { 130 | registerNeuralNetworkReferenceKernelFactories(); 131 | testForce(); 132 | testPeriodicForce(); 133 | } 134 | catch(const std::exception& e) { 135 | std::cout << "exception: " << e.what() << std::endl; 136 | return 1; 137 | } 138 | std::cout << "Done" << std::endl; 139 | return 0; 140 | } 141 | -------------------------------------------------------------------------------- /python/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | set(WRAP_FILE NNPluginWrapper.cpp) 2 | set(MODULE_NAME openmmnn) 3 | 4 | # Execute SWIG to generate source code for the Python module. 5 | 6 | add_custom_command( 7 | OUTPUT "${WRAP_FILE}" 8 | COMMAND "${SWIG_EXECUTABLE}" 9 | -python -c++ 10 | -o "${WRAP_FILE}" 11 | "-I${OPENMM_DIR}/include" 12 | "${CMAKE_CURRENT_SOURCE_DIR}/openmmnn.i" 13 | DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/openmmnn.i" 14 | WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" 15 | ) 16 | 17 | # Compile the Python module. 18 | 19 | add_custom_target(PythonInstall DEPENDS "${WRAP_FILE}") 20 | set(NN_PLUGIN_HEADER_DIR "${CMAKE_SOURCE_DIR}/openmmapi/include") 21 | set(NN_PLUGIN_LIBRARY_DIR "${CMAKE_BINARY_DIR}") 22 | configure_file(${CMAKE_CURRENT_SOURCE_DIR}/setup.py ${CMAKE_CURRENT_BINARY_DIR}/setup.py) 23 | add_custom_command(TARGET PythonInstall 24 | COMMAND "${PYTHON_EXECUTABLE}" setup.py build 25 | COMMAND "${PYTHON_EXECUTABLE}" setup.py install 26 | WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" 27 | ) 28 | 29 | -------------------------------------------------------------------------------- /python/openmmnn.i: -------------------------------------------------------------------------------- 1 | %module openmmnn 2 | 3 | %import(module="simtk.openmm") "swig/OpenMMSwigHeaders.i" 4 | %include "swig/typemaps.i" 5 | %include 6 | 7 | %{ 8 | #include "NeuralNetworkForce.h" 9 | #include "OpenMM.h" 10 | #include "OpenMMAmoeba.h" 11 | #include "OpenMMDrude.h" 12 | #include "openmm/RPMDIntegrator.h" 13 | #include "openmm/RPMDMonteCarloBarostat.h" 14 | %} 15 | 16 | /* 17 | * Convert C++ exceptions to Python exceptions. 18 | */ 19 | %exception { 20 | try { 21 | $action 22 | } catch (std::exception &e) { 23 | PyErr_SetString(PyExc_Exception, const_cast(e.what())); 24 | return NULL; 25 | } 26 | } 27 | 28 | %feature("shadow") NNPlugin::NeuralNetworkForce::NeuralNetworkForce %{ 29 | def __init__(self, *args): 30 | if len(args) == 1 and isinstance(args[0], str): 31 | this = _openmmnn.new_NeuralNetworkForce(args[0]) 32 | else: 33 | import tensorflow as tf 34 | import os 35 | import tempfile 36 | graph = args[0] 37 | if len(args) == 2: 38 | session = args[1] 39 | graph = tf.compat.v1.graph_util.convert_variables_to_constants(session, graph.as_graph_def(), ['energy', 'forces']) 40 | with tempfile.TemporaryDirectory() as dir: 41 | tf.io.write_graph(graph, dir, 'graph.pb', as_text=False) 42 | file = os.path.join(dir, 'graph.pb') 43 | this = _openmmnn.new_NeuralNetworkForce(file) 44 | try: 45 | self.this.append(this) 46 | except Exception: 47 | self.this = this 48 | %} 49 | 50 | namespace NNPlugin { 51 | 52 | class NeuralNetworkForce : public OpenMM::Force { 53 | public: 54 | NeuralNetworkForce(const std::string& file); 55 | const std::string& getFile() const; 56 | void setUsesPeriodicBoundaryConditions(bool periodic); 57 | bool usesPeriodicBoundaryConditions() const; 58 | 59 | /* 60 | * Add methods for casting a Force to a NeuralNetworkForce. 61 | */ 62 | %extend { 63 | static NNPlugin::NeuralNetworkForce& cast(OpenMM::Force& force) { 64 | return dynamic_cast(force); 65 | } 66 | 67 | static bool isinstance(OpenMM::Force& force) { 68 | return (dynamic_cast(&force) != NULL); 69 | } 70 | } 71 | }; 72 | 73 | } 74 | -------------------------------------------------------------------------------- /python/setup.py: -------------------------------------------------------------------------------- 1 | from distutils.core import setup 2 | from distutils.extension import Extension 3 | import os 4 | import sys 5 | import platform 6 | 7 | openmm_dir = '@OPENMM_DIR@' 8 | nn_plugin_header_dir = '@NN_PLUGIN_HEADER_DIR@' 9 | nn_plugin_library_dir = '@NN_PLUGIN_LIBRARY_DIR@' 10 | 11 | # setup extra compile and link arguments on Mac 12 | extra_compile_args = ['-std=c++11'] 13 | extra_link_args = [] 14 | 15 | if platform.system() == 'Darwin': 16 | extra_compile_args += ['-stdlib=libc++', '-mmacosx-version-min=10.7'] 17 | extra_link_args += ['-stdlib=libc++', '-mmacosx-version-min=10.7', '-Wl', '-rpath', openmm_dir+'/lib'] 18 | 19 | extension = Extension(name='_openmmnn', 20 | sources=['NNPluginWrapper.cpp'], 21 | libraries=['OpenMM', 'OpenMMNN'], 22 | include_dirs=[os.path.join(openmm_dir, 'include'), nn_plugin_header_dir], 23 | library_dirs=[os.path.join(openmm_dir, 'lib'), nn_plugin_library_dir], 24 | runtime_library_dirs=[os.path.join(openmm_dir, 'lib')], 25 | extra_compile_args=extra_compile_args, 26 | extra_link_args=extra_link_args 27 | ) 28 | 29 | setup(name='openmmnn', 30 | version='1.0', 31 | py_modules=['openmmnn'], 32 | ext_modules=[extension], 33 | ) 34 | -------------------------------------------------------------------------------- /python/tests/TestNeuralNetworkForce.py: -------------------------------------------------------------------------------- 1 | import simtk.openmm as mm 2 | import simtk.unit as unit 3 | import openmmnn as nn 4 | import tensorflow as tf 5 | import unittest 6 | 7 | class TestNeuralNetworkForce(unittest.TestCase): 8 | 9 | def testFreezeGraph(self): 10 | graph = tf.Graph() 11 | with graph.as_default(): 12 | positions = tf.placeholder(tf.float32, [None, 3], 'positions') 13 | scale = tf.Variable(5.0) 14 | energy = tf.multiply(scale, tf.reduce_sum(positions**2), name='energy') 15 | forces = tf.identity(tf.gradients(-energy, positions), name='forces') 16 | session = tf.Session() 17 | session.run(tf.global_variables_initializer()) 18 | force = nn.NeuralNetworkForce(graph, session) 19 | system = mm.System() 20 | for i in range(3): 21 | system.addParticle(1.0) 22 | system.addForce(force) 23 | integrator = mm.VerletIntegrator(0.001) 24 | context = mm.Context(system, integrator) 25 | positions = [mm.Vec3(3, 0, 0), mm.Vec3(0, 4, 0), mm.Vec3(3, 4, 0)] 26 | context.setPositions(positions) 27 | assert context.getState(getEnergy=True).getPotentialEnergy() == 250.0*unit.kilojoules_per_mole 28 | 29 | 30 | if __name__ == '__main__': 31 | unittest.main() 32 | -------------------------------------------------------------------------------- /serialization/include/NeuralNetworkForceProxy.h: -------------------------------------------------------------------------------- 1 | #ifndef OPENMM_NEURAL_NETWORK_FORCE_PROXY_H_ 2 | #define OPENMM_NEURAL_NETWORK_FORCE_PROXY_H_ 3 | 4 | /* -------------------------------------------------------------------------- * 5 | * OpenMM-NN * 6 | * -------------------------------------------------------------------------- * 7 | * This is part of the OpenMM molecular simulation toolkit originating from * 8 | * Simbios, the NIH National Center for Physics-Based Simulation of * 9 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 10 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 11 | * * 12 | * Portions copyright (c) 2018 Stanford University and the Authors. * 13 | * Authors: Peter Eastman * 14 | * Contributors: * 15 | * * 16 | * Permission is hereby granted, free of charge, to any person obtaining a * 17 | * copy of this software and associated documentation files (the "Software"), * 18 | * to deal in the Software without restriction, including without limitation * 19 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 20 | * and/or sell copies of the Software, and to permit persons to whom the * 21 | * Software is furnished to do so, subject to the following conditions: * 22 | * * 23 | * The above copyright notice and this permission notice shall be included in * 24 | * all copies or substantial portions of the Software. * 25 | * * 26 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 27 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 28 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 29 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 30 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 31 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 32 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 33 | * -------------------------------------------------------------------------- */ 34 | 35 | #include "internal/windowsExportNN.h" 36 | #include "openmm/serialization/SerializationProxy.h" 37 | 38 | namespace OpenMM { 39 | 40 | /** 41 | * This is a proxy for serializing NeuralNetworkForce objects. 42 | */ 43 | 44 | class OPENMM_EXPORT_NN NeuralNetworkForceProxy : public SerializationProxy { 45 | public: 46 | NeuralNetworkForceProxy(); 47 | void serialize(const void* object, SerializationNode& node) const; 48 | void* deserialize(const SerializationNode& node) const; 49 | }; 50 | 51 | } // namespace OpenMM 52 | 53 | #endif /*OPENMM_NEURAL_NETWORK_FORCE_PROXY_H_*/ 54 | -------------------------------------------------------------------------------- /serialization/src/NeuralNetworkForceProxy.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM-NN * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "NeuralNetworkForceProxy.h" 33 | #include "NeuralNetworkForce.h" 34 | #include "openmm/serialization/SerializationNode.h" 35 | #include 36 | 37 | using namespace NNPlugin; 38 | using namespace OpenMM; 39 | using namespace std; 40 | 41 | NeuralNetworkForceProxy::NeuralNetworkForceProxy() : SerializationProxy("NeuralNetworkForce") { 42 | } 43 | 44 | void NeuralNetworkForceProxy::serialize(const void* object, SerializationNode& node) const { 45 | node.setIntProperty("version", 1); 46 | const NeuralNetworkForce& force = *reinterpret_cast(object); 47 | node.setStringProperty("file", force.getFile()); 48 | } 49 | 50 | void* NeuralNetworkForceProxy::deserialize(const SerializationNode& node) const { 51 | if (node.getIntProperty("version") != 1) 52 | throw OpenMMException("Unsupported version number"); 53 | NeuralNetworkForce* force = new NeuralNetworkForce(node.getStringProperty("file")); 54 | return force; 55 | } 56 | -------------------------------------------------------------------------------- /serialization/src/NeuralNetworkSerializationProxyRegistration.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM-NN * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #ifdef WIN32 33 | #include 34 | #include 35 | #else 36 | #include 37 | #include 38 | #include 39 | #endif 40 | 41 | #include "NeuralNetworkForce.h" 42 | #include "NeuralNetworkForceProxy.h" 43 | #include "openmm/serialization/SerializationProxy.h" 44 | 45 | #if defined(WIN32) 46 | #include 47 | extern "C" OPENMM_EXPORT_NN void registerNeuralNetworkSerializationProxies(); 48 | BOOL WINAPI DllMain(HANDLE hModule, DWORD ul_reason_for_call, LPVOID lpReserved) { 49 | if (ul_reason_for_call == DLL_PROCESS_ATTACH) 50 | registerNeuralNetworkSerializationProxies(); 51 | return TRUE; 52 | } 53 | #else 54 | extern "C" void __attribute__((constructor)) registerNeuralNetworkSerializationProxies(); 55 | #endif 56 | 57 | using namespace NNPlugin; 58 | using namespace OpenMM; 59 | 60 | extern "C" OPENMM_EXPORT_NN void registerNeuralNetworkSerializationProxies() { 61 | SerializationProxy::registerProxy(typeid(NeuralNetworkForce), new NeuralNetworkForceProxy()); 62 | } 63 | -------------------------------------------------------------------------------- /serialization/tests/CMakeLists.txt: -------------------------------------------------------------------------------- 1 | # 2 | # Testing 3 | # 4 | 5 | # Automatically create tests using files named "Test*.cpp" 6 | FILE(GLOB TEST_PROGS "*Test*.cpp") 7 | FOREACH(TEST_PROG ${TEST_PROGS}) 8 | GET_FILENAME_COMPONENT(TEST_ROOT ${TEST_PROG} NAME_WE) 9 | 10 | # Link with shared library 11 | 12 | ADD_EXECUTABLE(${TEST_ROOT} ${TEST_PROG}) 13 | TARGET_LINK_LIBRARIES(${TEST_ROOT} ${SHARED_NN_TARGET}) 14 | SET_TARGET_PROPERTIES(${TEST_ROOT} PROPERTIES LINK_FLAGS "${EXTRA_COMPILE_FLAGS}" COMPILE_FLAGS "${EXTRA_COMPILE_FLAGS}") 15 | ADD_TEST(${TEST_ROOT} ${EXECUTABLE_OUTPUT_PATH}/${TEST_ROOT}) 16 | 17 | ENDFOREACH(TEST_PROG ${TEST_PROGS}) 18 | -------------------------------------------------------------------------------- /serialization/tests/TestSerializeNeuralNetworkForce.cpp: -------------------------------------------------------------------------------- 1 | /* -------------------------------------------------------------------------- * 2 | * OpenMM-NN * 3 | * -------------------------------------------------------------------------- * 4 | * This is part of the OpenMM molecular simulation toolkit originating from * 5 | * Simbios, the NIH National Center for Physics-Based Simulation of * 6 | * Biological Structures at Stanford, funded under the NIH Roadmap for * 7 | * Medical Research, grant U54 GM072970. See https://simtk.org. * 8 | * * 9 | * Portions copyright (c) 2018 Stanford University and the Authors. * 10 | * Authors: Peter Eastman * 11 | * Contributors: * 12 | * * 13 | * Permission is hereby granted, free of charge, to any person obtaining a * 14 | * copy of this software and associated documentation files (the "Software"), * 15 | * to deal in the Software without restriction, including without limitation * 16 | * the rights to use, copy, modify, merge, publish, distribute, sublicense, * 17 | * and/or sell copies of the Software, and to permit persons to whom the * 18 | * Software is furnished to do so, subject to the following conditions: * 19 | * * 20 | * The above copyright notice and this permission notice shall be included in * 21 | * all copies or substantial portions of the Software. * 22 | * * 23 | * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * 24 | * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * 25 | * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL * 26 | * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, * 27 | * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR * 28 | * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE * 29 | * USE OR OTHER DEALINGS IN THE SOFTWARE. * 30 | * -------------------------------------------------------------------------- */ 31 | 32 | #include "NeuralNetworkForce.h" 33 | #include "openmm/Platform.h" 34 | #include "openmm/internal/AssertionUtilities.h" 35 | #include "openmm/serialization/XmlSerializer.h" 36 | #include 37 | #include 38 | 39 | using namespace NNPlugin; 40 | using namespace OpenMM; 41 | using namespace std; 42 | 43 | extern "C" void registerNeuralNetworkSerializationProxies(); 44 | 45 | void testSerialization() { 46 | // Create a Force. 47 | 48 | NeuralNetworkForce force("graph.pb"); 49 | 50 | // Serialize and then deserialize it. 51 | 52 | stringstream buffer; 53 | XmlSerializer::serialize(&force, "Force", buffer); 54 | NeuralNetworkForce* copy = XmlSerializer::deserialize(buffer); 55 | 56 | // Compare the two forces to see if they are identical. 57 | 58 | NeuralNetworkForce& force2 = *copy; 59 | ASSERT_EQUAL(force.getFile(), force2.getFile()); 60 | } 61 | 62 | int main() { 63 | try { 64 | registerNeuralNetworkSerializationProxies(); 65 | testSerialization(); 66 | } 67 | catch(const exception& e) { 68 | cout << "exception: " << e.what() << endl; 69 | return 1; 70 | } 71 | cout << "Done" << endl; 72 | return 0; 73 | } 74 | -------------------------------------------------------------------------------- /tests/central.pb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openmm/openmm-tensorflow/c9a0e09cdb2c9449fca5cf02fe9e000c4757a721/tests/central.pb -------------------------------------------------------------------------------- /tests/periodic.pb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openmm/openmm-tensorflow/c9a0e09cdb2c9449fca5cf02fe9e000c4757a721/tests/periodic.pb --------------------------------------------------------------------------------