fix pyx file location
This commit is contained in:
0
src/MLIAP/mliap_model_python_couple.pyx
Executable file → Normal file
0
src/MLIAP/mliap_model_python_couple.pyx
Executable file → Normal file
@ -1,120 +0,0 @@
|
||||
# cython: language_level=3
|
||||
# distutils: language = c++
|
||||
|
||||
cimport cython
|
||||
|
||||
import pickle
|
||||
|
||||
# For converting C arrays to numpy arrays
|
||||
import numpy as np
|
||||
|
||||
# For converting void * to integer for tracking object identity
|
||||
from libc.stdint cimport uintptr_t
|
||||
|
||||
from libcpp.string cimport string
|
||||
|
||||
|
||||
cdef extern from "mliap_data.h" namespace "LAMMPS_NS":
|
||||
cdef cppclass MLIAPData:
|
||||
# Array shapes
|
||||
int natoms
|
||||
int ndescriptors
|
||||
|
||||
# Input data
|
||||
int * ielems # types for all atoms in list
|
||||
double ** descriptors # descriptors for all atoms in list
|
||||
|
||||
# Output data to write to
|
||||
double ** betas # betas for all atoms in list
|
||||
double * eatoms # energy for all atoms in list
|
||||
double energy
|
||||
|
||||
cdef extern from "mliap_model_python.h" namespace "LAMMPS_NS":
|
||||
cdef cppclass MLIAPModelPython:
|
||||
void connect_param_counts()
|
||||
|
||||
|
||||
class MLIAPPYModelNotLinked(Exception): pass
|
||||
|
||||
|
||||
LOADED_MODELS = {}
|
||||
|
||||
cdef object c_id(MLIAPModelPython * c_model):
|
||||
"""
|
||||
Use python-style id of object to keep track of identity.
|
||||
Note, this is probably not a perfect general strategy but it should work fine with LAMMPS pair styles.
|
||||
"""
|
||||
return int(<uintptr_t> c_model)
|
||||
|
||||
cdef object retrieve(MLIAPModelPython * c_model) with gil:
|
||||
try:
|
||||
model = LOADED_MODELS[c_id(c_model)]
|
||||
except KeyError as ke:
|
||||
raise KeyError("Model has not been loaded.") from ke
|
||||
if model is None:
|
||||
raise MLIAPPYModelNotLinked("Model not linked, connect the model from the python side.")
|
||||
return model
|
||||
|
||||
cdef public int MLIAPPY_load_model(MLIAPModelPython * c_model, char* fname) with gil:
|
||||
str_fname = fname.decode('utf-8') # Python 3 only; not Python 2 not supported.
|
||||
if str_fname == "LATER":
|
||||
model = None
|
||||
returnval = 0
|
||||
else:
|
||||
if str_fname.endswith(".pt") or str_fname.endswith('.pth'):
|
||||
import torch
|
||||
model = torch.load(str_fname)
|
||||
else:
|
||||
with open(str_fname,'rb') as pfile:
|
||||
model = pickle.load(pfile)
|
||||
returnval = 1
|
||||
LOADED_MODELS[c_id(c_model)] = model
|
||||
return returnval
|
||||
|
||||
def load_from_python(model):
|
||||
unloaded_models = [k for k, v in LOADED_MODELS.items() if v is None]
|
||||
num_models = len(unloaded_models)
|
||||
cdef MLIAPModelPython * lmp_model
|
||||
|
||||
if num_models == 0:
|
||||
raise ValueError("No model in the waiting area.")
|
||||
elif num_models > 1:
|
||||
raise ValueError("Model is amibguous, more than one model in waiting area.")
|
||||
else:
|
||||
c_id = unloaded_models[0]
|
||||
LOADED_MODELS[c_id]=model
|
||||
lmp_model = <MLIAPModelPython *> <uintptr_t> c_id
|
||||
lmp_model.connect_param_counts()
|
||||
|
||||
|
||||
cdef public void MLIAPPY_unload_model(MLIAPModelPython * c_model) with gil:
|
||||
del LOADED_MODELS[c_id(c_model)]
|
||||
|
||||
cdef public int MLIAPPY_nparams(MLIAPModelPython * c_model) with gil:
|
||||
return int(retrieve(c_model).n_params)
|
||||
|
||||
cdef public int MLIAPPY_nelements(MLIAPModelPython * c_model) with gil:
|
||||
return int(retrieve(c_model).n_elements)
|
||||
|
||||
cdef public int MLIAPPY_ndescriptors(MLIAPModelPython * c_model) with gil:
|
||||
return int(retrieve(c_model).n_descriptors)
|
||||
|
||||
cdef public void MLIAPPY_compute_gradients(MLIAPModelPython * c_model, MLIAPData * data) with gil:
|
||||
model = retrieve(c_model)
|
||||
|
||||
n_d = data.ndescriptors
|
||||
n_a = data.natoms
|
||||
|
||||
# Make numpy arrays from pointers
|
||||
beta_np = np.asarray(<double[:n_a,:n_d] > &data.betas[0][0])
|
||||
desc_np = np.asarray(<double[:n_a,:n_d]> &data.descriptors[0][0])
|
||||
elem_np = np.asarray(<int[:n_a]> &data.ielems[0])
|
||||
en_np = np.asarray(<double[:n_a]> &data.eatoms[0])
|
||||
|
||||
# Invoke python model on numpy arrays.
|
||||
model(elem_np,desc_np,beta_np,en_np)
|
||||
|
||||
# Get the total energy from the atom energy.
|
||||
energy = np.sum(en_np)
|
||||
data.energy = <double> energy
|
||||
return
|
||||
Reference in New Issue
Block a user