import multiprocessing as mp
import os
import sys
import time
import importlib
import numpy as np
__all__ = ['Potential', 'Potential_NoMP', 'NN_Potential', 'Potential_Direct']
[docs]
class Potential:
"""
A potential handler that is able to call python functions that
call .so files, either generated by f2py or loaded in by ctypes.
:param potential_function: The name of a python function (user specified) that will take in a n x m x 3 stack of geometries and return a 1D numpy array filled with potential values in hartrees.
:type potential_function: str
:param potential_dir: The *absolute path* to the directory that contains the .so file and .py file. If it"s a python function, then just the absolute path to your .py file.
:type: str
:param num_cores: Will create a pool of <num_cores> processes using Python"s multiprocessing module. This should never be larger than the number of processors on the machine this code is run.
:type: int
"""
def __init__(self,
potential_function,
potential_directory,
python_file,
num_cores=1,
pass_timestep=False,
pot_kwargs=None,
):
self.potential_function = potential_function
self.python_file = python_file
self.potential_directory = potential_directory
self.num_cores = num_cores
self.pass_timestep = pass_timestep
self.pot_kwargs = pot_kwargs
if self.pass_timestep:
self.ct = 0
if self.pot_kwargs is None:
self.pot_kwargs = {'timestep': 0}
else:
self.pot_kwargs['timestep'] = 0
self._init_pool()
def _init_pot(self):
"""
Sets _pot
"""
# Go to potential directory that houses python function and assign a self._pot variable to it
self._curdir = os.getcwd()
os.chdir(self.potential_directory)
sys.path.insert(0, os.getcwd())
module = self.python_file.split(".")[0]
x = importlib.import_module(module)
self._pot = getattr(x, self.potential_function)
# leave pool workers there
def _init_pool(self):
if self.num_cores <= 0:
print('Weird number of cores specified. Defaulting to 1...')
self.num_cores = 1
self._potPool = mp.Pool(self.num_cores, initializer=self._init_pot())
os.chdir(self._curdir)
@property
def pool(self):
"""Returns the potential manager's pool so that it can be used internally with Imp Samp or with the user elsewhere"""
return self._potPool
[docs]
def getpot(self, cds, timeit=False):
"""
Uses the potential function we got to call potential
:param cds: A stack of geometries (nxmx3, n=num_geoms;m=num_atoms;3=x,y,z) whose energies we need
:type cds: np.ndarray
:param timeit: The logger telling the potential manager whether or not to time the potential call
:type timeit: bool
"""
if timeit:
start = time.time()
if self._potPool is not None:
from itertools import repeat
cds = np.array_split(cds, self.num_cores)
if self.pot_kwargs is not None:
res = self._potPool.starmap(self._pot, zip(cds, repeat(self.pot_kwargs, len(cds))))
else:
res = self._potPool.map(self._pot, cds)
v = np.concatenate(res)
else:
v = self._pot(cds)
if self.pass_timestep:
self.pot_kwargs['timestep'] += 1
if timeit:
elapsed = time.time() - start
return v, elapsed
else:
return v
[docs]
def mp_close(self):
if self._potPool is not None:
self._potPool.close()
self._potPool.join()
[docs]
class Potential_NoMP:
"""
Version of Potential where no multiprocessing is included. As such, it does not leave a worker in the potential
directory of interest. If you need to cd into the directory to call a potential (for example, the potential
loads in a file using a relative path or something), then use ch_dir=True. Not the default since it is
computationally inefficient to cd during each potential call when not necessary.
"""
def __init__(self,
potential_function,
potential_directory,
python_file,
pass_timestep=False,
ch_dir=False,
pot_kwargs=None):
self.potential_function = potential_function
self.python_file = python_file
self.pass_timestep = pass_timestep
self.potential_directory = potential_directory
self.pot_kwargs = pot_kwargs
if self.pass_timestep:
if self.pot_kwargs is None:
self.pot_kwargs = {'timestep': 0}
else:
self.pot_kwargs['timestep'] = 0
self.ch_dir = ch_dir
self._init_pot()
def _init_pot(self):
self._curdir = os.getcwd()
os.chdir(self.potential_directory)
sys.path.insert(0, os.getcwd())
module = self.python_file.split(".")[0]
x = importlib.import_module(module)
self._pot = getattr(x, self.potential_function)
os.chdir(self._curdir)
[docs]
def getpot(self, cds, timeit=False):
"""
Uses the potential function we got to call potential
:param cds: A stack of geometries (nxmx3, n=num_geoms;m=num_atoms;3=x,y,z) whose energies we need
:type cds: np.ndarray
:param timeit: The logger telling the potential manager whether or not to time the potential call
:type timeit: bool
"""
if timeit:
start = time.time()
if not self.ch_dir:
if self.pot_kwargs is not None:
v = self._pot(cds, self.pot_kwargs)
else:
v = self._pot(cds)
else:
# Change to the potential directory and then change back after call
os.chdir(self.potential_directory)
if self.pot_kwargs is not None:
v = self._pot(cds, self.pot_kwargs)
else:
v = self._pot(cds)
os.chdir(self._curdir)
if self.pass_timestep:
self.pot_kwargs['timestep'] += 1
if timeit:
elapsed = time.time() - start
return v, elapsed
else:
return v
[docs]
class NN_Potential(Potential_NoMP):
"""
Subclass of Potential_NoMP, where a model argument is provided for convenience
"""
def __init__(self,
potential_function,
potential_directory,
python_file,
model,
ch_dir=False,
pot_kwargs=None,
pass_timestep=False):
super().__init__(potential_function, potential_directory, python_file, pass_timestep, ch_dir, pot_kwargs)
self.model = model
# Kwargs: if passed to NN_Potential, override Potential_NoMP default values vv
self.ch_dir = ch_dir
self.pot_kwargs = pot_kwargs
[docs]
def getpot(self, cds, timeit=False):
"""
Subclass of Potential_NoMP, but with an explicit argument passed for the NN model for convenience.
"""
if timeit:
start = time.time()
if self.pot_kwargs is not None:
v = self._pot(cds, self.model, self.pot_kwargs)
else:
v = self._pot(cds, self.model)
if self.pass_timestep:
self.pot_kwargs['timestep'] += 1
if timeit:
elapsed = time.time() - start
return v, elapsed
else:
return v
[docs]
class Potential_Direct:
"""
Version of Potential where the actual Python function is passed rather than
being imported from an external file. At the request of Mark.
"""
def __init__(self,
potential_function,
pot_kwargs=None,
pass_timestep=False):
self.potential_function = potential_function
self.pot_kwargs = pot_kwargs
self.pass_timestep = pass_timestep
if self.pass_timestep:
self.ct = 0
if self.pot_kwargs is None:
self.pot_kwargs = {'timestep': 0}
else:
self.pot_kwargs['timestep'] = 0
[docs]
def getpot(self, cds, timeit=False):
if timeit:
start = time.time()
if self.pot_kwargs is not None:
v = self.potential_function(cds, self.pot_kwargs)
else:
v = self.potential_function(cds)
if self.pass_timestep:
self.pot_kwargs['timestep'] += 1
if timeit:
elapsed = time.time() - start
return v, elapsed
else:
return v