Commit 7c1b5c5a authored by André Melo's avatar André Melo

Test loop without precomputing

parent e56b6067
Pipeline #44922 passed with stage
in 2 minutes and 26 seconds
from miniff.ml import NNPotential
from miniff.kernel import NeighborWrapper
from torch import load
from dfttools.simple import parse
from pathlib import Path
script_location = Path(__file__).parent
data_location = script_location / ".." / "datasets"
potentials = tuple(map(NNPotential.from_state_dict, load(data_location / "bise.pt")))
cutoff = max(i.cutoff for i in potentials)
with open(data_location / "bise.json", 'r') as f:
c = parse(f, "unit-cell")
w = NeighborWrapper(c, x=(2, 2, 2))
w.compute_distances(cutoff)
for i in range(10):
print(w.total(potentials), w.grad(potentials))
This diff is collapsed.
......@@ -3,6 +3,8 @@ import numpy as np
cimport numpy as np
cimport cython
from libc.math cimport exp, cos, sin, pi, pow
from cython.parallel import prange
# Template potential-2.pyx
# A two-point potential: depends on the distance between pairs of atoms
......@@ -155,7 +157,7 @@ def kernel_general_3(
# (r1, r2)
for dim in range(3):
r12_cos += (cartesian_col[col1, dim] - cartesian_row[row, dim]) * (cartesian_col[col2, dim]- cartesian_row[row, dim])
r12_cos /= r1 * r2
r12_cos = r12_cos / (r1 * r2)
# --- before ---
# (no 'before' statements)
# --- kernel ---
......@@ -223,7 +225,7 @@ def kernel_g_general_3(
# (r1, r2)
for dim in range(3):
r12_cos += (cartesian_col[col1, dim] - cartesian_row[row, dim]) * (cartesian_col[col2, dim]- cartesian_row[row, dim])
r12_cos /= r1 * r2
r12_cos = r12_cos / (r1 * r2)
# --- before ---
# (no 'before' statements)
......@@ -363,7 +365,7 @@ def kernel_lj(
double[::1] r_data,
double[:, ::1] cartesian_row,
double[:, ::1] cartesian_col,
double a,
double a,
int[::1] species_row,
int[::1] species_mask,
double[::1] out,
......@@ -409,7 +411,7 @@ def kernel_g_lj(
double[::1] r_data,
double[:, ::1] cartesian_row,
double[:, ::1] cartesian_col,
double a,
double a,
int[::1] species_row,
int[::1] species_mask,
double[:, :, ::1] out,
......@@ -604,7 +606,7 @@ def kernel_sw_phi3(
# (r1, r2)
for dim in range(3):
r12_cos += (cartesian_col[col1, dim] - cartesian_row[row, dim]) * (cartesian_col[col2, dim]- cartesian_row[row, dim])
r12_cos /= r1 * r2
r12_cos = r12_cos / (r1 * r2)
# --- before ---
# (no 'before' statements)
# --- kernel ---
......@@ -672,7 +674,7 @@ def kernel_g_sw_phi3(
# (r1, r2)
for dim in range(3):
r12_cos += (cartesian_col[col1, dim] - cartesian_row[row, dim]) * (cartesian_col[col2, dim]- cartesian_row[row, dim])
r12_cos /= r1 * r2
r12_cos = r12_cos / (r1 * r2)
# --- before ---
# (no 'before' statements)
......@@ -708,6 +710,10 @@ def kernel_g_sw_phi3(
# Template potential-2.pyx
# A two-point potential: depends on the distance between pairs of atoms
@cython.cdivision(True)
@cython.boundscheck(False)
@cython.nonecheck(False)
@cython.wraparound(False)
def kernel_mlsf_g2(
int[::1] r_indptr,
int[::1] r_indices,
......@@ -733,7 +739,7 @@ def kernel_mlsf_g2(
# (no 'preamble' statements)
# ----------------
for row in range(nrows):
for row in prange(nrows, nogil=True, schedule='static'):
if species_row[row] == row_mask:
ptr_fr = r_indptr[row]
ptr_to = r_indptr[row + 1]
......@@ -754,6 +760,10 @@ def kernel_mlsf_g2(
# Template potential-2-g.pyx
# A two-point potential: depends on the distance between pairs of atoms
@cython.cdivision(True)
@cython.boundscheck(False)
@cython.nonecheck(False)
@cython.wraparound(False)
def kernel_g_mlsf_g2(
int[::1] r_indptr,
int[::1] r_indices,
......@@ -778,7 +788,7 @@ def kernel_g_mlsf_g2(
# (no 'preamble' statements)
# ----------------
for row in range(nrows):
for row in prange(nrows, nogil=True, schedule='static'):
if species_row[row] == row_mask:
ptr_fr = r_indptr[row]
ptr_to = r_indptr[row + 1]
......@@ -807,6 +817,10 @@ def kernel_g_mlsf_g2(
# Template potential-3.pyx
# A three-point potential: depends on the distance between two pairs of atoms
# sharing the same atom at origin and the cosine of the angle formed
@cython.cdivision(True)
@cython.boundscheck(False)
@cython.nonecheck(False)
@cython.wraparound(False)
def kernel_mlsf_g5(
int[::1] r_indptr,
int[::1] r_indices,
......@@ -838,7 +852,7 @@ def kernel_mlsf_g5(
r_cutoff_fn[ptr] = .5 + cos(pi * r_data[ptr] / a) / 2
# ----------------
for row in range(nrows):
for row in prange(nrows, nogil=True, schedule='static'):
if species_row[row] == row_mask:
ptr_fr = r_indptr[row]
ptr_to = r_indptr[row + 1]
......@@ -862,7 +876,7 @@ def kernel_mlsf_g5(
# (r1, r2)
for dim in range(3):
r12_cos += (cartesian_col[col1, dim] - cartesian_row[row, dim]) * (cartesian_col[col2, dim]- cartesian_row[row, dim])
r12_cos /= r1 * r2
r12_cos = r12_cos / (r1 * r2)
# --- before ---
_fn_cutoff2 = r_cutoff_fn[ptr2]
_fn_exponent = _prefactor * exp(- eta * (r1 * r1 + r2 * r2))
......@@ -877,6 +891,10 @@ def kernel_mlsf_g5(
# Template potential-3-g.pyx
# A three-point potential: depends on the distance between two pairs of atoms
# sharing the same atom at origin and the cosine of the angle formed
@cython.cdivision(True)
@cython.boundscheck(False)
@cython.nonecheck(False)
@cython.wraparound(False)
def kernel_g_mlsf_g5(
int[::1] r_indptr,
int[::1] r_indices,
......@@ -914,7 +932,7 @@ def kernel_g_mlsf_g5(
r_cutoff_fp[ptr] = - pi * sin(pi * r_data[ptr] / a) / 2 / a
# ----------------
for row in range(nrows):
for row in prange(nrows, nogil=True, schedule='static'):
if species_row[row] == row_mask:
ptr_fr = r_indptr[row]
ptr_to = r_indptr[row + 1]
......@@ -941,7 +959,7 @@ def kernel_g_mlsf_g5(
# (r1, r2)
for dim in range(3):
r12_cos += (cartesian_col[col1, dim] - cartesian_row[row, dim]) * (cartesian_col[col2, dim]- cartesian_row[row, dim])
r12_cos /= r1 * r2
r12_cos = r12_cos / (r1 * r2)
# --- before ---
_fn_cutoff2 = r_cutoff_fn[ptr2]
......
......@@ -99,7 +99,8 @@ class ext_modules_lazy(list):
def eval(self):
self.ext_modules = cythonize([
Extension("miniff._potentials", ["miniff/_potentials.pyx"], include_dirs=[numpy.get_include()]),
Extension("miniff._potentials", ["miniff/_potentials.pyx"], include_dirs=[numpy.get_include()],
extra_compile_args=['-fopenmp'], extra_link_args=['-fopenmp']),
Extension("miniff._util", ["miniff/_util.pyx"], include_dirs=[numpy.get_include()]),
], compiler_directives={"profile": True}, annotate=True)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment