def _get_options_dictionary(self, available_resources): """Returns a dictionary of options which will be serialized to a yaml file and passed to YANK. Parameters ---------- available_resources: ComputeResources The resources available to execute on. Returns ------- dict of str and Any A yaml compatible dictionary of YANK options. """ from openforcefield.utils import quantity_to_string platform_name = "CPU" if available_resources.number_of_gpus > 0: # A platform which runs on GPUs has been requested. from openff.evaluator.backends import ComputeResources toolkit_enum = ComputeResources.GPUToolkit( available_resources.preferred_gpu_toolkit ) # A platform which runs on GPUs has been requested. platform_name = ( "CUDA" if toolkit_enum == ComputeResources.GPUToolkit.CUDA else ComputeResources.GPUToolkit.OpenCL ) return { "verbose": self.verbose, "output_dir": ".", "temperature": quantity_to_string( pint_quantity_to_openmm(self.thermodynamic_state.temperature) ), "pressure": quantity_to_string( pint_quantity_to_openmm(self.thermodynamic_state.pressure) ), "minimize": True, "number_of_equilibration_iterations": self.number_of_equilibration_iterations, "default_number_of_iterations": self.number_of_iterations, "default_nsteps_per_iteration": self.steps_per_iteration, "checkpoint_interval": self.checkpoint_interval, "default_timestep": quantity_to_string( pint_quantity_to_openmm(self.timestep) ), "annihilate_electrostatics": True, "annihilate_sterics": False, "platform": platform_name, }
def setup_platform_with_resources(compute_resources, high_precision=False): """Creates an OpenMM `Platform` object which requests a set amount of compute resources (e.g with a certain number of cpus). Parameters ---------- compute_resources: ComputeResources The compute resources which describe which platform is most appropriate. high_precision: bool If true, a platform with the highest possible precision (double for CUDA and OpenCL, Reference for CPU only) will be returned. Returns ------- Platform The created platform """ from simtk.openmm import Platform # Setup the requested platform: if compute_resources.number_of_gpus > 0: # TODO: Make sure use mixing precision - CUDA, OpenCL. # TODO: Deterministic forces = True from openff.evaluator.backends import ComputeResources toolkit_enum = ComputeResources.GPUToolkit( compute_resources.preferred_gpu_toolkit) # A platform which runs on GPUs has been requested. platform_name = ("CUDA" if toolkit_enum == ComputeResources.GPUToolkit.CUDA else ComputeResources.GPUToolkit.OpenCL) # noinspection PyCallByClass,PyTypeChecker platform = Platform.getPlatformByName(platform_name) if compute_resources.gpu_device_indices is not None: property_platform_name = platform_name if toolkit_enum == ComputeResources.GPUToolkit.CUDA: property_platform_name = platform_name.lower().capitalize() platform.setPropertyDefaultValue( property_platform_name + "DeviceIndex", compute_resources.gpu_device_indices, ) if high_precision: platform.setPropertyDefaultValue("Precision", "double") logger.info("Setting up an openmm platform on GPU {}".format( compute_resources.gpu_device_indices or 0)) else: if not high_precision: # noinspection PyCallByClass,PyTypeChecker platform = Platform.getPlatformByName("CPU") platform.setPropertyDefaultValue( "Threads", str(compute_resources.number_of_threads)) else: # noinspection PyCallByClass,PyTypeChecker platform = Platform.getPlatformByName("Reference") logger.info("Setting up a simulation with {} threads".format( compute_resources.number_of_threads)) return platform