From 4e47715afe917ce829a33ee6840628b6ae09ed88 Mon Sep 17 00:00:00 2001 From: alvarom Date: Mon, 23 Feb 2026 02:06:46 +0000 Subject: [PATCH 01/29] adding sycl headers --- libsymmetrix/source/mace_kokkos.cpp | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/libsymmetrix/source/mace_kokkos.cpp b/libsymmetrix/source/mace_kokkos.cpp index 80ba9bd..ba9728c 100644 --- a/libsymmetrix/source/mace_kokkos.cpp +++ b/libsymmetrix/source/mace_kokkos.cpp @@ -9,6 +9,10 @@ #include "sphericart.hpp" #include "sphericart_cuda.hpp" +#ifdef SYMMETRIX_SPHERICART_SYCL +#include "sphericart_sycl.hpp" +#endif + #include "tools_kokkos.hpp" #include "mace_kokkos.hpp" @@ -191,7 +195,7 @@ void MACEKokkos::compute_R1( template void MACEKokkos::compute_Y(Kokkos::View xyz) { -#ifndef SYMMETRIX_SPHERICART_CUDA +#if !defined (SYMMETRIX_SPHERICART_CUDA) && !defined (SYMMETRIX_SPHERICART_SYCL) const int num = xyz.extent(0) / 3; if (Y.extent(0) < num*num_lm) { @@ -248,7 +252,7 @@ void MACEKokkos::compute_Y(Kokkos::View xyz) { }); Kokkos::fence(); -#else // SYMMETRIX_SPHERICART_CUDA +#else // SYMMETRIX_SPHERICART_CUDA or SYMMETRIX_SPHERICART_SYCL const int num = xyz.extent(0) / 3; const int num_lm = (l_max+1)*(l_max+1); @@ -271,7 +275,14 @@ void MACEKokkos::compute_Y(Kokkos::View xyz) { Kokkos::fence(); // call sphericart +#if defined (SYMMETRIX_SPHERICART_CUDA) +#error "here 2" sphericart::cuda::SphericalHarmonics sphericart(l_max); +#elif defined (SYMMETRIX_SPHERICART_SYCL) + sphericart::intel::SphericalHarmonics sphericart(l_max); +#else +#error "NO GPU defined" +#endif sphericart.compute_with_gradients(xyz_shuffled.data(), num, Y.data(), Y_grad.data()); // unshuffle gradient From e024c39f514581fdf35b37884f97a35d12896caf Mon Sep 17 00:00:00 2001 From: alvarom Date: Mon, 23 Feb 2026 02:07:04 +0000 Subject: [PATCH 02/29] adding sycl macro --- libsymmetrix/CMakeLists.txt | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/libsymmetrix/CMakeLists.txt b/libsymmetrix/CMakeLists.txt index de21bea..f061d42 100644 --- a/libsymmetrix/CMakeLists.txt +++ b/libsymmetrix/CMakeLists.txt @@ -74,4 +74,9 @@ if (SYMMETRIX_SPHERICART_CUDA) message(STATUS "Symmetrix: Will use CUDA version of sphericart.") target_compile_definitions(symmetrix PRIVATE SYMMETRIX_SPHERICART_CUDA) endif() +option(SYMMETRIX_SPHERICART_SYCL OFF) +if (SYMMETRIX_SPHERICART_SYCL) + message(STATUS "Symmetrix: Will use SYCL version of sphericart.") + target_compile_definitions(symmetrix PRIVATE SYMMETRIX_SPHERICART_SYCL) +endif() From ea7761470dbaeb84b283468ee0d5904239a30de5 Mon Sep 17 00:00:00 2001 From: alvarom Date: Mon, 23 Feb 2026 08:12:21 +0000 Subject: [PATCH 03/29] update namespace sycl --- libsymmetrix/source/mace_kokkos.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libsymmetrix/source/mace_kokkos.cpp b/libsymmetrix/source/mace_kokkos.cpp index ba9728c..03fc13f 100644 --- a/libsymmetrix/source/mace_kokkos.cpp +++ b/libsymmetrix/source/mace_kokkos.cpp @@ -279,7 +279,7 @@ void MACEKokkos::compute_Y(Kokkos::View xyz) { #error "here 2" sphericart::cuda::SphericalHarmonics sphericart(l_max); #elif defined (SYMMETRIX_SPHERICART_SYCL) - sphericart::intel::SphericalHarmonics sphericart(l_max); + sphericart::sycl::SphericalHarmonics sphericart(l_max); #else #error "NO GPU defined" #endif From cf1bcfe41b7b63ad83e17bbe566545a8b45ff96b Mon Sep 17 00:00:00 2001 From: G-Seaford Date: Fri, 10 Apr 2026 11:17:17 +0000 Subject: [PATCH 04/29] Fixed issue with mixed devices being used when extracting MACE models --- .../source/symmetrix/extract_mace_data.py | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/symmetrix/source/symmetrix/extract_mace_data.py b/symmetrix/source/symmetrix/extract_mace_data.py index 143c878..c86a4ce 100755 --- a/symmetrix/source/symmetrix/extract_mace_data.py +++ b/symmetrix/source/symmetrix/extract_mace_data.py @@ -35,11 +35,13 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): ------- output_data: dict with symmetrix model data """ + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = torch.load( model, - map_location=torch.device('cpu'), + map_location=device, weights_only=False - ).to(torch.float64) + ).to(device=device, dtype=torch.float64) + model.eval() if species is None: species = [] @@ -63,7 +65,8 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): if hasattr(model, 'heads') and len(model.heads) != 1: torch.set_default_dtype(next(model.parameters()).dtype) - model = remove_pt_head(model, head) + model = remove_pt_head(model, head).to(device=device, dtype=torch.float64) + model.eval() ### ----- CHECK FOR COMPATIBILITY ----- @@ -96,7 +99,7 @@ def linear_simplify(linear): Irreps(linear.irreps_out).simplify()) simplified.weight = linear.weight simplified.bias = linear.bias - return simplified + return simplified.to(device=device, dtype=torch.float64) ### ----- BASIC MODEL INFO ----- @@ -154,10 +157,10 @@ def linear_simplify(linear): model_i = model.atomic_numbers.tolist().index(a_i) model_j = model.atomic_numbers.tolist().index(a_j) bessels = model.radial_embedding( - torch.tensor(r, dtype=torch.get_default_dtype()).unsqueeze(-1), - torch.eye(len(model.atomic_numbers)), - torch.tensor([[model_i],[model_j]], dtype=torch.int64), - model.atomic_numbers) + torch.tensor(r, dtype=torch.get_default_dtype(), device=device).unsqueeze(-1), + torch.eye(len(model.atomic_numbers), device=device), + torch.tensor([[model_i],[model_j]], dtype=torch.int64, device=device), + model.atomic_numbers.to(device)) if isinstance(bessels, tuple): bessels = bessels[0] # newer versions return (bessels, cutoffs) # radial basis for interaction 0 @@ -211,10 +214,10 @@ def linear_simplify(linear): model_i = model.atomic_numbers.tolist().index(a_i) model_j = model.atomic_numbers.tolist().index(a_j) bessels = model.radial_embedding( - torch.tensor(r, dtype=torch.get_default_dtype()).unsqueeze(-1), - torch.eye(len(model.atomic_numbers)), - torch.tensor([[model_i],[model_j]], dtype=torch.int64), - model.atomic_numbers) + torch.tensor(r, dtype=torch.get_default_dtype(), device=device).unsqueeze(-1), + torch.eye(len(model.atomic_numbers), device=device), + torch.tensor([[model_i],[model_j]], dtype=torch.int64, device=device), + model.atomic_numbers.to(device)) if isinstance(bessels, tuple): bessels = bessels[0] # newer versions return (bessels, cutoffs) R = torch.tanh(model.interactions[0].density_fn(bessels)**2).numpy(force=True) @@ -347,16 +350,16 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): lelm1lm2 += (2*l1+1)*(2*l2+1) l, l1, l2 = (Phi1_l[le], Phi1_l1[le], Phi1_l2[le]) return lelm1lm2 + (l1+m1)*(2*l2+1) + l2+m2 - tp = model.interactions[1].conv_tp + tp = model.interactions[1].conv_tp.to(device) for l1 in range(l_max+1): for m1 in range(-l1,l1+1): lm1 = l1*l1+l1+m1 for l2 in range(L_max+1): for m2 in range(-l2,l2+1): - R = torch.ones([1,len(tp.instructions)*num_channels],dtype=torch.double) - Y = torch.zeros([1,num_lm1], dtype=torch.double) + R = torch.ones([1,len(tp.instructions)*num_channels],dtype=torch.double, device=device) + Y = torch.zeros([1,num_lm1], dtype=torch.double, device=device) Y[0,lm1] = 1.0 - H = torch.zeros([1,num_lm2*num_channels],dtype=torch.double) + H = torch.zeros([1,num_lm2*num_channels],dtype=torch.double, device=device) H[0,sum([2*p+1 for p in range(l2)])*num_channels+l2+m2] = 1.0 Phi = tp(H, Y, R) # extract Phi values for k=0 @@ -396,10 +399,10 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): model_i = model.atomic_numbers.tolist().index(a_i) model_j = model.atomic_numbers.tolist().index(a_j) bessels = model.radial_embedding( - torch.tensor(r, dtype=torch.get_default_dtype()).unsqueeze(-1), - torch.eye(len(model.atomic_numbers)), - torch.tensor([[model_i],[model_j]], dtype=torch.int64), - model.atomic_numbers) + torch.tensor(r, dtype=torch.get_default_dtype(), device=device).unsqueeze(-1), + torch.eye(len(model.atomic_numbers), device=device), + torch.tensor([[model_i],[model_j]], dtype=torch.int64, device=device), + model.atomic_numbers.to(device)) if isinstance(bessels, tuple): bessels = bessels[0] # newer versions return (bessels, cutoffs) R = torch.tanh(model.interactions[1].density_fn(bessels)**2).numpy(force=True) From d5367e1702affdb19a4705ce9c02b68b26fd3b58 Mon Sep 17 00:00:00 2001 From: Gianluca Seaford <147088212+G-Seaford@users.noreply.github.com> Date: Fri, 10 Apr 2026 16:47:44 +0100 Subject: [PATCH 05/29] Change device to CPU for model loading --- symmetrix/source/symmetrix/extract_mace_data.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/symmetrix/source/symmetrix/extract_mace_data.py b/symmetrix/source/symmetrix/extract_mace_data.py index c86a4ce..40bc2d9 100755 --- a/symmetrix/source/symmetrix/extract_mace_data.py +++ b/symmetrix/source/symmetrix/extract_mace_data.py @@ -35,7 +35,7 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): ------- output_data: dict with symmetrix model data """ - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = torch.device("cpu") model = torch.load( model, map_location=device, From b5ac8ec3f599c480a4aacc9200e490f50d5bbe06 Mon Sep 17 00:00:00 2001 From: Noam Bernstein Date: Fri, 8 May 2026 14:00:40 -0400 Subject: [PATCH 06/29] better error checking, including float issues, for spline bounds --- libsymmetrix/source/cubic_spline.cpp | 24 ++++++++++++------------ libsymmetrix/source/cubic_spline.hpp | 2 ++ 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index f53eb70..304fb26 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -1,5 +1,7 @@ #include #include +#include +#include #include "cubic_spline.hpp" @@ -12,12 +14,16 @@ CubicSpline::CubicSpline( { } +int CubicSpline::get_i(double r) +{ + if (r<0 or r>h*c.size()/4) + throw std::invalid_argument("Out of bounds in CubicSpline::evaluate. r=" + std::to_string(r)); + return std::clamp(static_cast(r / h), 0, c.size()/4 - 1); +} + double CubicSpline::evaluate(double r) { - const int i = static_cast(r / h); - // TODO: something better with this bounds checking - if (i<0 or i>=c.size()/4) - throw std::invalid_argument("Out of bounds in CubicSpline::evaluate."); + const int i = get_i(r) const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; @@ -28,10 +34,7 @@ double CubicSpline::evaluate(double r) std::tuple CubicSpline::evaluate_deriv(double r) { - const int i = static_cast(r / h); - // TODO: something better with this bounds checking - if (i<0 or i>=c.size()/4) - throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv."); + const int i = get_i(r) const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; @@ -42,10 +45,7 @@ std::tuple CubicSpline::evaluate_deriv(double r) std::tuple CubicSpline::evaluate_deriv_divided(double r) { - const int i = static_cast(r / h); - // TODO: something better with this bounds checking - if (i<0 or i>=c.size()/4) - throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv."); + const int i = get_i(r) const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; diff --git a/libsymmetrix/source/cubic_spline.hpp b/libsymmetrix/source/cubic_spline.hpp index 32cc28c..ee648b9 100644 --- a/libsymmetrix/source/cubic_spline.hpp +++ b/libsymmetrix/source/cubic_spline.hpp @@ -20,6 +20,8 @@ auto evaluate_deriv_divided(double r) -> std::tuple; double h; std::vector c; +auto get_i(double r) -> int; + auto generate_coefficients( double h, std::vector nodal_values, From a825dc2f7ea021031c923b936a426b365ab5f015 Mon Sep 17 00:00:00 2001 From: Noam Bernstein Date: Fri, 8 May 2026 14:19:18 -0400 Subject: [PATCH 07/29] fix bug in clamp arg type, and spline upper bound check --- libsymmetrix/source/cubic_spline.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index 304fb26..e07d061 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -16,9 +16,9 @@ CubicSpline::CubicSpline( int CubicSpline::get_i(double r) { - if (r<0 or r>h*c.size()/4) + if (r<0 or r>=h*c.size()/4) throw std::invalid_argument("Out of bounds in CubicSpline::evaluate. r=" + std::to_string(r)); - return std::clamp(static_cast(r / h), 0, c.size()/4 - 1); + return std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); } double CubicSpline::evaluate(double r) From e012645e4f58f20649d667f0c539673560200088 Mon Sep 17 00:00:00 2001 From: Noam Bernstein Date: Fri, 8 May 2026 14:33:27 -0400 Subject: [PATCH 08/29] missing ; --- libsymmetrix/source/cubic_spline.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index e07d061..64ed963 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -23,7 +23,7 @@ int CubicSpline::get_i(double r) double CubicSpline::evaluate(double r) { - const int i = get_i(r) + const int i = get_i(r); const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; @@ -34,7 +34,7 @@ double CubicSpline::evaluate(double r) std::tuple CubicSpline::evaluate_deriv(double r) { - const int i = get_i(r) + const int i = get_i(r); const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; @@ -45,7 +45,7 @@ std::tuple CubicSpline::evaluate_deriv(double r) std::tuple CubicSpline::evaluate_deriv_divided(double r) { - const int i = get_i(r) + const int i = get_i(r); const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; From 40942d10c7569c225dcd38aacc5306268439bd25 Mon Sep 17 00:00:00 2001 From: Noam Bernstein Date: Fri, 8 May 2026 16:12:14 -0400 Subject: [PATCH 09/29] update tests for new cubic spline out of bounds exception message --- symmetrix/test/test_cubic_spline.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/symmetrix/test/test_cubic_spline.py b/symmetrix/test/test_cubic_spline.py index e25f17a..53670f9 100644 --- a/symmetrix/test/test_cubic_spline.py +++ b/symmetrix/test/test_cubic_spline.py @@ -28,7 +28,7 @@ def test_evaluate(): for r in [-1.0, 5.0, 9.0]: with raises(ValueError) as exception: spl.evaluate(r) - assert str(exception.value) == "Out of bounds in CubicSpline::evaluate." + assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") def test_evaluate_deriv(): @@ -50,6 +50,11 @@ def test_evaluate_deriv(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2)) + for r in [-1.0, 5.0, 9.0]: + with raises(ValueError) as exception: + _ = spl.evaluate_deriv(r) + assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") + def test_evaluate_deriv_divided(): # generate data @@ -68,3 +73,8 @@ def test_evaluate_deriv_divided(): f2[i], d2[i] = spl.evaluate_deriv_divided(ri) assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) + + for r in [-1.0, 5.0, 9.0]: + with raises(ValueError) as exception: + _ = spl.evaluate_deriv(r) + assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") From 571debc8d10a95fa90d39f9add92fa33cec8c944 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 16:25:29 -0400 Subject: [PATCH 10/29] tidy. --- libsymmetrix/source/mace_kokkos.cpp | 3 --- 1 file changed, 3 deletions(-) diff --git a/libsymmetrix/source/mace_kokkos.cpp b/libsymmetrix/source/mace_kokkos.cpp index 03fc13f..1a4f8cd 100644 --- a/libsymmetrix/source/mace_kokkos.cpp +++ b/libsymmetrix/source/mace_kokkos.cpp @@ -276,12 +276,9 @@ void MACEKokkos::compute_Y(Kokkos::View xyz) { // call sphericart #if defined (SYMMETRIX_SPHERICART_CUDA) -#error "here 2" sphericart::cuda::SphericalHarmonics sphericart(l_max); #elif defined (SYMMETRIX_SPHERICART_SYCL) sphericart::sycl::SphericalHarmonics sphericart(l_max); -#else -#error "NO GPU defined" #endif sphericart.compute_with_gradients(xyz_shuffled.data(), num, Y.data(), Y_grad.data()); From f4a3ba3fc864be902411d10f022aff9d1b5a2748 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 16:26:29 -0400 Subject: [PATCH 11/29] spaces for consistency. --- libsymmetrix/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libsymmetrix/CMakeLists.txt b/libsymmetrix/CMakeLists.txt index f061d42..e7e7de7 100644 --- a/libsymmetrix/CMakeLists.txt +++ b/libsymmetrix/CMakeLists.txt @@ -76,7 +76,7 @@ if (SYMMETRIX_SPHERICART_CUDA) endif() option(SYMMETRIX_SPHERICART_SYCL OFF) if (SYMMETRIX_SPHERICART_SYCL) - message(STATUS "Symmetrix: Will use SYCL version of sphericart.") - target_compile_definitions(symmetrix PRIVATE SYMMETRIX_SPHERICART_SYCL) + message(STATUS "Symmetrix: Will use SYCL version of sphericart.") + target_compile_definitions(symmetrix PRIVATE SYMMETRIX_SPHERICART_SYCL) endif() From 9b1cdd62f71e65d4a81f78a98707053c8a8b05b3 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 16:32:48 -0400 Subject: [PATCH 12/29] bump sphericart to v2.0.3. --- libsymmetrix/external/sphericart | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libsymmetrix/external/sphericart b/libsymmetrix/external/sphericart index 1866da4..67643cf 160000 --- a/libsymmetrix/external/sphericart +++ b/libsymmetrix/external/sphericart @@ -1 +1 @@ -Subproject commit 1866da4162f14491c44a299dd9991a0816f15e62 +Subproject commit 67643cfd14d2d0cb78a9dfd1c264ca9bc111b724 From dc144898adaec9a391c90ddfd9517c204dfe7445 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 17:10:26 -0400 Subject: [PATCH 13/29] tweak for correct name in error message. --- libsymmetrix/source/cubic_spline.cpp | 19 +++++++++---------- libsymmetrix/source/cubic_spline.hpp | 2 -- symmetrix/test/test_cubic_spline.py | 19 +++++++++---------- 3 files changed, 18 insertions(+), 22 deletions(-) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index 64ed963..12a270e 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -14,16 +14,11 @@ CubicSpline::CubicSpline( { } -int CubicSpline::get_i(double r) -{ - if (r<0 or r>=h*c.size()/4) - throw std::invalid_argument("Out of bounds in CubicSpline::evaluate. r=" + std::to_string(r)); - return std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); -} - double CubicSpline::evaluate(double r) { - const int i = get_i(r); + if (r<0 or r>h*c.size()/4) + throw std::invalid_argument("Out of bounds in CubicSpline::evaluate. r=" + std::to_string(r)); + const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; @@ -34,7 +29,9 @@ double CubicSpline::evaluate(double r) std::tuple CubicSpline::evaluate_deriv(double r) { - const int i = get_i(r); + if (r<0 or r>h*c.size()/4) + throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv. r=" + std::to_string(r)); + const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; @@ -45,7 +42,9 @@ std::tuple CubicSpline::evaluate_deriv(double r) std::tuple CubicSpline::evaluate_deriv_divided(double r) { - const int i = get_i(r); + if (r<0 or r>h*c.size()/4) + throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv_divided. r=" + std::to_string(r)); + const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; diff --git a/libsymmetrix/source/cubic_spline.hpp b/libsymmetrix/source/cubic_spline.hpp index ee648b9..32cc28c 100644 --- a/libsymmetrix/source/cubic_spline.hpp +++ b/libsymmetrix/source/cubic_spline.hpp @@ -20,8 +20,6 @@ auto evaluate_deriv_divided(double r) -> std::tuple; double h; std::vector c; -auto get_i(double r) -> int; - auto generate_coefficients( double h, std::vector nodal_values, diff --git a/symmetrix/test/test_cubic_spline.py b/symmetrix/test/test_cubic_spline.py index 53670f9..d24744a 100644 --- a/symmetrix/test/test_cubic_spline.py +++ b/symmetrix/test/test_cubic_spline.py @@ -18,14 +18,13 @@ def test_evaluate(): d = scipy_spl.derivative()(r) spl = symmetrix.CubicSpline(h, f, d) # test equivalence - # NOTE: endpoint=False avoids out of bounds (see exception tests below) - r2 = np.linspace(0, r_cut, 1000, endpoint=False) + r2 = np.linspace(0, r_cut, 1000) f2 = np.zeros(len(r2)) for i, ri in enumerate(r2): f2[i] = spl.evaluate(ri) assert f2 == approx(scipy_spl(r2)) # test out of bounds errors - for r in [-1.0, 5.0, 9.0]: + for r in [-1.0, np.nextafter(r_cut, np.inf), 9.0]: with raises(ValueError) as exception: spl.evaluate(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") @@ -42,7 +41,7 @@ def test_evaluate_deriv(): d = scipy_spl.derivative()(r) spl = symmetrix.CubicSpline(h, f, d) # test equivalence - r2 = np.linspace(0, r_cut, 1000, endpoint=False) + r2 = np.linspace(0, r_cut, 1000) f2 = np.zeros(len(r2)) d2 = np.zeros(len(r2)) for i, ri in enumerate(r2): @@ -50,10 +49,10 @@ def test_evaluate_deriv(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2)) - for r in [-1.0, 5.0, 9.0]: + for r in [-1.0, np.nextafter(r_cut, np.inf), 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv(r) - assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") + assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv.") def test_evaluate_deriv_divided(): @@ -66,7 +65,7 @@ def test_evaluate_deriv_divided(): d = scipy_spl.derivative()(r) spl = symmetrix.CubicSpline(h, f, d) # test equivalence - r2 = np.linspace(1e-6, r_cut, 1000, endpoint=False) + r2 = np.linspace(1e-6, r_cut, 1000) f2 = np.zeros(len(r2)) d2 = np.zeros(len(r2)) for i, ri in enumerate(r2): @@ -74,7 +73,7 @@ def test_evaluate_deriv_divided(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) - for r in [-1.0, 5.0, 9.0]: + for r in [-1.0, np.nextafter(r_cut, np.inf), 9.0]: with raises(ValueError) as exception: - _ = spl.evaluate_deriv(r) - assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") + _ = spl.evaluate_deriv_divided(r) + assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv_divided.") From ebfb8b101e6e610cc5738306e8e39782b0f3136c Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 17:29:57 -0400 Subject: [PATCH 14/29] fix outdated syntax. --- symmetrix/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/symmetrix/pyproject.toml b/symmetrix/pyproject.toml index 5f48686..4caf83a 100644 --- a/symmetrix/pyproject.toml +++ b/symmetrix/pyproject.toml @@ -3,7 +3,7 @@ requires = ["scikit-build-core", "pybind11", "cmake>=3.27"] build-backend = "scikit_build_core.build" [tool.scikit-build] -cmake.minimum-version = "3.27" +cmake.version = ">=3.27" build-dir = "build" wheel.exclude = ["lib/", "include/"] wheel.packages = ["source/symmetrix"] From ac934aae2d5e2e6322c3709c86a9bec633efc824 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 17:43:58 -0400 Subject: [PATCH 15/29] check for nan. --- libsymmetrix/source/cubic_spline.cpp | 7 ++++--- symmetrix/test/test_cubic_spline.py | 6 +++--- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index 12a270e..d3fabbd 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include "cubic_spline.hpp" @@ -16,7 +17,7 @@ CubicSpline::CubicSpline( double CubicSpline::evaluate(double r) { - if (r<0 or r>h*c.size()/4) + if (r<0 or r>h*c.size()/4 or std::isnan(r)) throw std::invalid_argument("Out of bounds in CubicSpline::evaluate. r=" + std::to_string(r)); const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; @@ -29,7 +30,7 @@ double CubicSpline::evaluate(double r) std::tuple CubicSpline::evaluate_deriv(double r) { - if (r<0 or r>h*c.size()/4) + if (r<0 or r>h*c.size()/4 or std::isnan(r)) throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv. r=" + std::to_string(r)); const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; @@ -42,7 +43,7 @@ std::tuple CubicSpline::evaluate_deriv(double r) std::tuple CubicSpline::evaluate_deriv_divided(double r) { - if (r<0 or r>h*c.size()/4) + if (r<0 or r>h*c.size()/4 or std::isnan(r)) throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv_divided. r=" + std::to_string(r)); const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; diff --git a/symmetrix/test/test_cubic_spline.py b/symmetrix/test/test_cubic_spline.py index d24744a..d4196b2 100644 --- a/symmetrix/test/test_cubic_spline.py +++ b/symmetrix/test/test_cubic_spline.py @@ -24,7 +24,7 @@ def test_evaluate(): f2[i] = spl.evaluate(ri) assert f2 == approx(scipy_spl(r2)) # test out of bounds errors - for r in [-1.0, np.nextafter(r_cut, np.inf), 9.0]: + for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: with raises(ValueError) as exception: spl.evaluate(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") @@ -49,7 +49,7 @@ def test_evaluate_deriv(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2)) - for r in [-1.0, np.nextafter(r_cut, np.inf), 9.0]: + for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: with raises(ValueError) as exception: _ = spl.evaluate_deriv(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv.") @@ -73,7 +73,7 @@ def test_evaluate_deriv_divided(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) - for r in [-1.0, np.nextafter(r_cut, np.inf), 9.0]: + for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: with raises(ValueError) as exception: _ = spl.evaluate_deriv_divided(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv_divided.") From b2476124348e5bd8503a1d3f6baccdd6b6ba5695 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 17:48:38 -0400 Subject: [PATCH 16/29] sanity check constructor. --- libsymmetrix/source/cubic_spline.cpp | 5 +++++ symmetrix/test/test_cubic_spline.py | 10 ++++++++++ 2 files changed, 15 insertions(+) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index d3fabbd..d58713e 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -60,6 +60,11 @@ auto CubicSpline::generate_coefficients( std::vector nodal_derivs) -> std::vector { + if (h<=0 or not std::isfinite(h)) + throw std::invalid_argument("CubicSpline requires positive finite spacing."); + if (nodal_values.size()<2 or nodal_values.size()!=nodal_derivs.size()) + throw std::invalid_argument("CubicSpline requires at least two values and matching derivatives."); + auto c = std::vector(4*(nodal_values.size()-1), 0.0); for (int i=0; i Date: Thu, 23 Jul 2026 17:50:06 -0400 Subject: [PATCH 17/29] special check for deriv_divided. --- libsymmetrix/source/cubic_spline.cpp | 2 +- symmetrix/test/test_cubic_spline.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libsymmetrix/source/cubic_spline.cpp b/libsymmetrix/source/cubic_spline.cpp index d58713e..6ccf2fe 100644 --- a/libsymmetrix/source/cubic_spline.cpp +++ b/libsymmetrix/source/cubic_spline.cpp @@ -43,7 +43,7 @@ std::tuple CubicSpline::evaluate_deriv(double r) std::tuple CubicSpline::evaluate_deriv_divided(double r) { - if (r<0 or r>h*c.size()/4 or std::isnan(r)) + if (r<=0 or r>h*c.size()/4 or std::isnan(r)) throw std::invalid_argument("Out of bounds in CubicSpline::evaluate_deriv_divided. r=" + std::to_string(r)); const int i = std::clamp(static_cast(r / h), 0, static_cast(c.size()/4 - 1)); const double x = r - h*i; diff --git a/symmetrix/test/test_cubic_spline.py b/symmetrix/test/test_cubic_spline.py index aec698e..315e438 100644 --- a/symmetrix/test/test_cubic_spline.py +++ b/symmetrix/test/test_cubic_spline.py @@ -83,7 +83,7 @@ def test_evaluate_deriv_divided(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) - for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-np.inf, -1.0, 0.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: with raises(ValueError) as exception: _ = spl.evaluate_deriv_divided(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv_divided.") From c6283c9ed9592de31b076f469a5632923932fc12 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 17:56:47 -0400 Subject: [PATCH 18/29] harden splines for kokkos. --- libsymmetrix/source/cubic_spline_kokkos.cpp | 23 ++++++++++----------- symmetrix/test/test_cubic_spline_kokkos.py | 21 +++++++++++++------ 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/libsymmetrix/source/cubic_spline_kokkos.cpp b/libsymmetrix/source/cubic_spline_kokkos.cpp index 882788b..e212309 100644 --- a/libsymmetrix/source/cubic_spline_kokkos.cpp +++ b/libsymmetrix/source/cubic_spline_kokkos.cpp @@ -1,6 +1,8 @@ #include #include #include +#include +#include #include "cubic_spline_kokkos.hpp" CubicSplineKokkos::CubicSplineKokkos( @@ -25,10 +27,9 @@ CubicSplineKokkos::CubicSplineKokkos( double CubicSplineKokkos::evaluate(double r) { - const int i = static_cast(r / h); - // TODO: something better with this bounds checking - if (i < 0 || i >= num_coeffs / 4) - throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate."); + if (r<0 or r>h*num_coeffs/4 or std::isnan(r)) + throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate. r=" + std::to_string(r)); + const int i = std::clamp(static_cast(r / h), 0, static_cast(num_coeffs/4 - 1)); const double x = r - h * i; const double xx = x * x; @@ -49,10 +50,9 @@ double CubicSplineKokkos::evaluate(double r) std::tuple CubicSplineKokkos::evaluate_deriv(double r) { - const int i = static_cast(r / h); - // TODO: something better with this bounds checking - if (i < 0 || i > num_coeffs / 4) - throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate_deriv."); + if (r<0 or r>h*num_coeffs/4 or std::isnan(r)) + throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate_deriv. r=" + std::to_string(r)); + const int i = std::clamp(static_cast(r / h), 0, static_cast(num_coeffs/4 - 1)); const double x = r - h * i; const double xx = x * x; @@ -74,10 +74,9 @@ std::tuple CubicSplineKokkos::evaluate_deriv(double r) std::tuple CubicSplineKokkos::evaluate_deriv_divided(double r) { - const int i = static_cast(r / h); - // TODO: something better with this bounds checking - if (i<0 or i> num_coeffs) - throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate_deriv."); + if (r<=0 or r>h*num_coeffs/4 or std::isnan(r)) + throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate_deriv_divided. r=" + std::to_string(r)); + const int i = std::clamp(static_cast(r / h), 0, static_cast(num_coeffs/4 - 1)); const double x = r - h*i; const double xx = x*x; diff --git a/symmetrix/test/test_cubic_spline_kokkos.py b/symmetrix/test/test_cubic_spline_kokkos.py index 858b1c2..e45d412 100644 --- a/symmetrix/test/test_cubic_spline_kokkos.py +++ b/symmetrix/test/test_cubic_spline_kokkos.py @@ -21,17 +21,16 @@ def test_evaluate(): d = scipy_spl.derivative()(r) spl = symmetrix.CubicSplineKokkos(h, f, d) # test equivalence - # NOTE: endpoint=False avoids out of bounds (see exception tests below) - r2 = np.linspace(0, r_cut, 1000, endpoint=False) + r2 = np.linspace(0, r_cut, 1000) f2 = np.zeros(len(r2)) for i, ri in enumerate(r2): f2[i] = spl.evaluate(ri) assert f2 == approx(scipy_spl(r2)) # test out of bounds errors - for r in [-1.0, 5.0, 9.0]: + for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: with raises(ValueError) as exception: spl.evaluate(r) - assert str(exception.value) == "Out of bounds in CubicSplineKokkos::evaluate." + assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate.") def test_evaluate_deriv(): @@ -44,7 +43,7 @@ def test_evaluate_deriv(): d = scipy_spl.derivative()(r) spl = symmetrix.CubicSplineKokkos(h, f, d) # test equivalence - r2 = np.linspace(0, r_cut, 1000, endpoint=False) + r2 = np.linspace(0, r_cut, 1000) f2 = np.zeros(len(r2)) d2 = np.zeros(len(r2)) for i, ri in enumerate(r2): @@ -52,6 +51,11 @@ def test_evaluate_deriv(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2)) + for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + with raises(ValueError) as exception: + _ = spl.evaluate_deriv(r) + assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate_deriv.") + def test_evaluate_deriv_divided(): # generate data @@ -63,10 +67,15 @@ def test_evaluate_deriv_divided(): d = scipy_spl.derivative()(r) spl = symmetrix.CubicSplineKokkos(h, f, d) # test equivalence - r2 = np.linspace(1e-6, r_cut, 1000, endpoint=False) + r2 = np.linspace(1e-6, r_cut, 1000) f2 = np.zeros(len(r2)) d2 = np.zeros(len(r2)) for i, ri in enumerate(r2): f2[i], d2[i] = spl.evaluate_deriv_divided(ri) assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) + + for r in [-np.inf, -1.0, 0.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + with raises(ValueError) as exception: + _ = spl.evaluate_deriv_divided(r) + assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate_deriv_divided.") From 047936fba896aadf54110fb0b7fda57374285630 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Thu, 23 Jul 2026 18:03:45 -0400 Subject: [PATCH 19/29] improve constructor in cubic_spline_kokkos. --- libsymmetrix/source/cubic_spline_kokkos.cpp | 14 ++++++++++++-- symmetrix/test/test_cubic_spline_kokkos.py | 11 +++++++++++ 2 files changed, 23 insertions(+), 2 deletions(-) diff --git a/libsymmetrix/source/cubic_spline_kokkos.cpp b/libsymmetrix/source/cubic_spline_kokkos.cpp index e212309..9226d57 100644 --- a/libsymmetrix/source/cubic_spline_kokkos.cpp +++ b/libsymmetrix/source/cubic_spline_kokkos.cpp @@ -9,8 +9,13 @@ CubicSplineKokkos::CubicSplineKokkos( double h, std::vector nodal_values, std::vector nodal_derivs) - : h(h), num_coeffs(4*(nodal_values.size() - 1)) + : h(h) { + if (h<=0 or not std::isfinite(h)) + throw std::invalid_argument("CubicSplineKokkos requires positive finite spacing."); + if (nodal_values.size()<2 or nodal_values.size()!=nodal_derivs.size()) + throw std::invalid_argument("CubicSplineKokkos requires at least two values and matching derivatives."); + num_coeffs = 4*(nodal_values.size() - 1); c = Kokkos::View("coeffs",num_coeffs); generate_coefficients(h, nodal_values, nodal_derivs); } @@ -19,8 +24,13 @@ CubicSplineKokkos::CubicSplineKokkos( double h, Kokkos::View nodal_values, Kokkos::View nodal_derivs) - : h(h), num_coeffs(4*(nodal_values.size() - 1)) + : h(h) { + if (h<=0 or not std::isfinite(h)) + throw std::invalid_argument("CubicSplineKokkos requires positive finite spacing."); + if (nodal_values.size()<2 or nodal_values.size()!=nodal_derivs.size()) + throw std::invalid_argument("CubicSplineKokkos requires at least two values and matching derivatives."); + num_coeffs = 4*(nodal_values.size() - 1); c = Kokkos::View("coeffs",num_coeffs); generate_coefficients(h, nodal_values, nodal_derivs); } diff --git a/symmetrix/test/test_cubic_spline_kokkos.py b/symmetrix/test/test_cubic_spline_kokkos.py index e45d412..071ed86 100644 --- a/symmetrix/test/test_cubic_spline_kokkos.py +++ b/symmetrix/test/test_cubic_spline_kokkos.py @@ -10,6 +10,17 @@ if not symmetrix._kokkos_is_initialized(): symmetrix._init_kokkos() + +def test_invalid_input(): + for h in [0.0, -1.0, np.inf, np.nan]: + with raises(ValueError): + symmetrix.CubicSplineKokkos(h, [0.0, 1.0], [0.0, 1.0]) + + for values, derivs in [([], []), ([0.0], [0.0]), ([0.0, 1.0], [0.0])]: + with raises(ValueError): + symmetrix.CubicSplineKokkos(1.0, values, derivs) + + def test_evaluate(): # generate data From 4fe7eb3fc9a2b8094fbf97bf3dc7f2bf9cae3d73 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Fri, 24 Jul 2026 11:13:01 -0400 Subject: [PATCH 20/29] drop tests that trigger with fast-math. --- symmetrix/test/test_cubic_spline.py | 8 ++++---- symmetrix/test/test_cubic_spline_kokkos.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/symmetrix/test/test_cubic_spline.py b/symmetrix/test/test_cubic_spline.py index 315e438..2304111 100644 --- a/symmetrix/test/test_cubic_spline.py +++ b/symmetrix/test/test_cubic_spline.py @@ -8,7 +8,7 @@ def test_invalid_input(): - for h in [0.0, -1.0, np.inf, np.nan]: + for h in [0.0, -1.0]: with raises(ValueError): symmetrix.CubicSpline(h, [0.0, 1.0], [0.0, 1.0]) @@ -34,7 +34,7 @@ def test_evaluate(): f2[i] = spl.evaluate(ri) assert f2 == approx(scipy_spl(r2)) # test out of bounds errors - for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: spl.evaluate(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") @@ -59,7 +59,7 @@ def test_evaluate_deriv(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2)) - for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv.") @@ -83,7 +83,7 @@ def test_evaluate_deriv_divided(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) - for r in [-np.inf, -1.0, 0.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-1.0, 0.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv_divided(r) assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv_divided.") diff --git a/symmetrix/test/test_cubic_spline_kokkos.py b/symmetrix/test/test_cubic_spline_kokkos.py index 071ed86..e46bc04 100644 --- a/symmetrix/test/test_cubic_spline_kokkos.py +++ b/symmetrix/test/test_cubic_spline_kokkos.py @@ -12,7 +12,7 @@ def test_invalid_input(): - for h in [0.0, -1.0, np.inf, np.nan]: + for h in [0.0, -1.0]: with raises(ValueError): symmetrix.CubicSplineKokkos(h, [0.0, 1.0], [0.0, 1.0]) @@ -38,7 +38,7 @@ def test_evaluate(): f2[i] = spl.evaluate(ri) assert f2 == approx(scipy_spl(r2)) # test out of bounds errors - for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: spl.evaluate(r) assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate.") @@ -62,7 +62,7 @@ def test_evaluate_deriv(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2)) - for r in [-np.inf, -1.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv(r) assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate_deriv.") @@ -86,7 +86,7 @@ def test_evaluate_deriv_divided(): assert f2 == approx(scipy_spl(r2)) assert d2 == approx(scipy_spl.derivative()(r2) / r2) - for r in [-np.inf, -1.0, 0.0, r_cut + 1e-12, 9.0, np.inf, np.nan]: + for r in [-1.0, 0.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv_divided(r) assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate_deriv_divided.") From d53eace2386ab5c1f823b4ebfdd3ae700e06f7b3 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Mon, 27 Jul 2026 10:07:30 -0400 Subject: [PATCH 21/29] bump kokkos to 5.1.1. --- libsymmetrix/external/kokkos | 2 +- libsymmetrix/external/kokkos-kernels | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/libsymmetrix/external/kokkos b/libsymmetrix/external/kokkos index dfceb77..267ebc2 160000 --- a/libsymmetrix/external/kokkos +++ b/libsymmetrix/external/kokkos @@ -1 +1 @@ -Subproject commit dfceb77277e23361cccfbcf72c0468fd8475153e +Subproject commit 267ebc25fc5c8b96bb321f34d78f97b1c30f8830 diff --git a/libsymmetrix/external/kokkos-kernels b/libsymmetrix/external/kokkos-kernels index 7226191..c41b623 160000 --- a/libsymmetrix/external/kokkos-kernels +++ b/libsymmetrix/external/kokkos-kernels @@ -1 +1 @@ -Subproject commit 722619104a9c99da0d497e65d8dfdff07e25145a +Subproject commit c41b62335a4fe3363ff5f2f6559ab4db08ff0b5b From 6b9b831f692e0f32383a0de131596dae82c28f35 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Mon, 27 Jul 2026 10:43:38 -0400 Subject: [PATCH 22/29] finalize kokkos correctly when testing. --- symmetrix/test/conftest.py | 10 ++++++++++ symmetrix/test/test_zbl_kokkos.py | 21 +++++++++++++++------ 2 files changed, 25 insertions(+), 6 deletions(-) diff --git a/symmetrix/test/conftest.py b/symmetrix/test/conftest.py index 5c3f9df..3ccb059 100644 --- a/symmetrix/test/conftest.py +++ b/symmetrix/test/conftest.py @@ -10,6 +10,16 @@ } +@pytest.fixture(scope="session", autouse=True) +def finalize_kokkos_after_tests(): + yield + + import symmetrix + + if symmetrix._kokkos_is_initialized(): + symmetrix._finalize_kokkos() + + @pytest.fixture(scope="session") def model_cache(): cache_dir = Path(__file__).parent / "model-cache" diff --git a/symmetrix/test/test_zbl_kokkos.py b/symmetrix/test/test_zbl_kokkos.py index fd204be..1ad69e6 100644 --- a/symmetrix/test/test_zbl_kokkos.py +++ b/symmetrix/test/test_zbl_kokkos.py @@ -14,20 +14,26 @@ covalent_radii = [0.2, 0.31, 0.28, 1.28, 0.96, 0.84, 0.76, 0.71, 0.66, 0.57, 0.58] r_max = covalent_radii[Z_u] + covalent_radii[Z_v] -zbl = symmetrix.ZBLKokkos( - 0.3, - 0.4543, - [0.1818, 0.5099, 0.2802, 0.02817], - covalent_radii, - 6) + +def make_zbl(): + return symmetrix.ZBLKokkos( + 0.3, + 0.4543, + [0.1818, 0.5099, 0.2802, 0.02817], + covalent_radii, + 6) + def test_compute_envelope(): + zbl = make_zbl() assert zbl.compute_envelope(r_max, r_max, 6) == 0.0 assert zbl.compute_envelope(1.0, r_max, 6) == pytest.approx(0.4374788794430078) + def test_compute_envelope_grad(): + zbl = make_zbl() assert zbl.compute_envelope_gradient(r_max, r_max, 6) == 0.0 @@ -39,12 +45,15 @@ def test_compute_envelope_grad(): assert zbl.compute_envelope_gradient(x, r_max, 6) == pytest.approx((v_p - v_m) / (2.0 * dx)) def test_compute_ZBL(): + zbl = make_zbl() assert zbl.compute(5, 10, r_max) == 0.0 assert zbl.compute(5, 10, 1.0) == pytest.approx(0.3166652764835175) + def test_compute_ZBL_grad(): + zbl = make_zbl() assert zbl.compute_gradient(Z_u, Z_v, r_max) == 0.0 From 9bf6e33b9cfa9c218f7e4160afb46a91f08f1217 Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Mon, 27 Jul 2026 17:02:01 -0400 Subject: [PATCH 23/29] modernize python packaging and ci. --- .github/workflows/ci.yaml | 83 +++++++++---------- .pre-commit-config.yaml | 8 ++ README.md | 10 +++ symmetrix/README.md | 11 ++- symmetrix/pyproject.toml | 14 ++++ symmetrix/source/symmetrix/__init__.py | 2 +- .../{symmetrix_calc.py => calculator.py} | 2 +- .../source/symmetrix/extract_mace_data.py | 3 +- 8 files changed, 84 insertions(+), 49 deletions(-) create mode 100644 .pre-commit-config.yaml rename symmetrix/source/symmetrix/{symmetrix_calc.py => calculator.py} (99%) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 3ea27fb..987b5dc 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -3,37 +3,40 @@ name: CI on: [pull_request, workflow_dispatch] jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - name: Clone repo + uses: actions/checkout@v6 + - name: Set up uv + uses: astral-sh/setup-uv@v8.3.2 + - name: Run pre-commit + run: uvx pre-commit run --all-files symmetrix: runs-on: ubuntu-latest steps: - name: Clone repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: submodules: recursive - name: Install dependencies run: | - sudo apt-get install -y build-essential cmake git libblas-dev liblapack-dev - - name: Set up python - uses: actions/setup-python@v5 + sudo apt-get update + sudo apt-get install -y build-essential cmake git libblas-dev liblapack-dev ninja-build + - name: Set up uv + uses: astral-sh/setup-uv@v8.3.2 with: python-version: '3.12' - - name: Create python venv + - name: Create python venv and install run: | - python -m venv venv - source venv/bin/activate - pip install ase cmake-build-extension[all] numpy pytest setuptools scipy wheel cmake-build-extension[all] - deactivate - - name: Build and package symmetrix - run: | - source venv/bin/activate - cd symmetrix - pip install . - cd .. + uv venv --clear + source .venv/bin/activate + uv pip install ./symmetrix[test] deactivate - name: Run tests run: | - source venv/bin/activate + source .venv/bin/activate python -c "import symmetrix; print(symmetrix.__version__)" cd symmetrix/test pytest @@ -44,25 +47,26 @@ jobs: runs-on: ubuntu-latest steps: - name: Clone repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: submodules: recursive - name: Install dependencies run: | + sudo apt-get update sudo apt-get install -y build-essential cmake git libfftw3-dev libopenmpi-dev mpi-default-bin mpi-default-dev libblas-dev liblapack-dev - - name: Set up python - uses: actions/setup-python@v5 + - name: Set up uv + uses: astral-sh/setup-uv@v8.3.2 with: python-version: '3.12' - - name: Create python venv + - name: Create python venv and install run: | - python -m venv venv - source venv/bin/activate - pip install numpy pytest + uv venv --clear + source .venv/bin/activate + uv pip install pip numpy pytest deactivate - name: Clone and build LAMMPS run: | - source venv/bin/activate + source .venv/bin/activate git clone -b release --depth 1 https://github.com/lammps/lammps.git cd pair_symmetrix chmod +x install.sh @@ -82,13 +86,13 @@ jobs: -D Kokkos_ENABLE_AGGRESSIVE_VECTORIZATION=ON \ -D SYMMETRIX_KOKKOS=ON \ cmake - cmake --build build -j 2 + cmake --build build -j 4 cd build make install-python cd ../../.. - name: Run tests run: | - source venv/bin/activate + source .venv/bin/activate cd pair_symmetrix/test/ python -m pytest @@ -100,33 +104,26 @@ jobs: mace-torch: ["mace-torch==0.3.10", "mace-torch"] steps: - name: Clone repo - uses: actions/checkout@v4 + uses: actions/checkout@v6 with: submodules: recursive - name: Install dependencies run: | - sudo apt-get install -y build-essential cmake git libblas-dev liblapack-dev - - name: Set up python - uses: actions/setup-python@v5 + sudo apt-get update + sudo apt-get install -y build-essential cmake git libblas-dev liblapack-dev ninja-build + - name: Set up uv + uses: astral-sh/setup-uv@v8.3.2 with: python-version: '3.12' - - name: Create python venv - run: | - python -m venv venv - source venv/bin/activate - pip install ${{ matrix.mace-torch }} - deactivate - - name: Build and package symmetrix + - name: Create python venv and install run: | - source venv/bin/activate - cd symmetrix - pip install . - cd .. + uv venv --clear + source .venv/bin/activate + uv pip install ${{ matrix.mace-torch }} ./symmetrix deactivate - name: Test model extraction run: | - source venv/bin/activate - pip install mace-torch + source .venv/bin/activate # check for valid symmetrix json from atomic numbers and default filename wget https://github.com/ACEsuit/mace-off/raw/refs/heads/main/mace_off23/MACE-OFF23_small.model symmetrix_extract_mace --model MACE-OFF23_small.model --atomic-numbers 1 8 diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..af907b0 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,8 @@ +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.5.0 + hooks: + # TODO: Restore trailing-whitespace and end-of-file-fixer with the deferred + # repository-wide formatting changes. + - id: check-yaml + - id: check-added-large-files diff --git a/README.md b/README.md index 5132761..0be5422 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,16 @@ See the `symmetrix` [README](symmetrix/README.md) to build and use the Python pa See the `pair_symmetrix` [README](pair_symmetrix/README.md) for use from LAMMPS. +### Development Setup + +Use `uv` to create a virtual environment and install the package with its test dependencies: + +```bash +uv venv +source .venv/bin/activate +uv pip install -e "./symmetrix[test]" +``` + ### Citing Symmetrix The earliest `symmetrix` results are reported in: diff --git a/symmetrix/README.md b/symmetrix/README.md index 1b586a6..7a6cbd1 100644 --- a/symmetrix/README.md +++ b/symmetrix/README.md @@ -33,7 +33,14 @@ pip install --verbose . \ ### Generating Symmetrix `.json` model files -Once the Python package is installed, use +Install the Python package with the optional MACE dependencies: + +``` +pip install ".[mace]" +``` + +Then use: + ``` symmetrix_extract_mace my-mace.model --atomic-numbers 1 8 ``` @@ -47,5 +54,5 @@ One can import the ASE calculator with ``` from symmetrix import Symmetrix ``` -See [the source code](source/symmetrix/symmetrix_calc.py) and [this test](test/test_symmetrix_calc.py) +See [the source code](source/symmetrix/calculator.py) and [this test](test/test_symmetrix_calc.py) for additional details. diff --git a/symmetrix/pyproject.toml b/symmetrix/pyproject.toml index 4caf83a..1186049 100644 --- a/symmetrix/pyproject.toml +++ b/symmetrix/pyproject.toml @@ -12,6 +12,20 @@ wheel.packages = ["source/symmetrix"] name = "symmetrix" version = "0.0.1" description = "Symmetrix — a package for functions equivariant under: translation, rotation, inversion, and exchange of particles." +dependencies = [ + "numpy", + "scipy", + "ase", +] + +[project.optional-dependencies] +mace = [ + "torch", + "mace-torch", +] +test = [ + "pytest", +] [project.scripts] symmetrix_extract_mace = "symmetrix.cli.extract_mace:main" diff --git a/symmetrix/source/symmetrix/__init__.py b/symmetrix/source/symmetrix/__init__.py index 20b559e..21eef3f 100644 --- a/symmetrix/source/symmetrix/__init__.py +++ b/symmetrix/source/symmetrix/__init__.py @@ -6,4 +6,4 @@ _sym.__all__ = [n for n in vars(_sym) if not (n.startswith('__') and n.endswith('__'))] from .symmetrix import * -from .symmetrix_calc import Symmetrix +from .calculator import Symmetrix diff --git a/symmetrix/source/symmetrix/symmetrix_calc.py b/symmetrix/source/symmetrix/calculator.py similarity index 99% rename from symmetrix/source/symmetrix/symmetrix_calc.py rename to symmetrix/source/symmetrix/calculator.py index 067af9a..829d47e 100755 --- a/symmetrix/source/symmetrix/symmetrix_calc.py +++ b/symmetrix/source/symmetrix/calculator.py @@ -11,7 +11,7 @@ try: from matscipy.neighbours import neighbour_list as neighbor_list -except: +except ImportError: logging.warning("Symmetrix using slow ase.neighborlist.neighbor_list") from ase.neighborlist import neighbor_list diff --git a/symmetrix/source/symmetrix/extract_mace_data.py b/symmetrix/source/symmetrix/extract_mace_data.py index 40bc2d9..842e9a3 100755 --- a/symmetrix/source/symmetrix/extract_mace_data.py +++ b/symmetrix/source/symmetrix/extract_mace_data.py @@ -1,5 +1,4 @@ import torch -torch.serialization.add_safe_globals([slice]) import os import logging @@ -55,7 +54,7 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): try: Z = chemical_symbols.index(sp) except ValueError as exc: - raise ValueError("Failed to parse {sp} as atomic number or chemical species") from exc + raise ValueError(f"Failed to parse {sp} as atomic number or chemical species") from exc atomic_numbers.append(Z) # ensure that splines goes smoothly to 0 at outer cutoff From 34d72cec17f43c59b3ce46b0599060de47f992aa Mon Sep 17 00:00:00 2001 From: Yuan Chiang Date: Mon, 27 Jul 2026 17:14:56 -0400 Subject: [PATCH 24/29] restore ruff workflow and python formatting --- .pre-commit-config.yaml | 7 + README.md | 2 + pair_symmetrix/test/conftest.py | 3 +- .../test/test_pair_symmetrix_mace.py | 123 ++--- symmetrix/pyproject.toml | 7 + symmetrix/source/symmetrix/__init__.py | 7 +- symmetrix/source/symmetrix/calculator.py | 77 ++- .../source/symmetrix/cli/extract_mace.py | 36 +- .../source/symmetrix/extract_mace_data.py | 456 +++++++++++------- symmetrix/test/test_cubic_spline.py | 24 +- symmetrix/test/test_cubic_spline_kokkos.py | 26 +- symmetrix/test/test_cubic_spline_set.py | 33 +- .../test/test_cubic_spline_set_kokkos.py | 36 +- symmetrix/test/test_lammpslib.py | 121 +++-- symmetrix/test/test_mace.py | 258 ++++++---- symmetrix/test/test_multilayer_perceptron.py | 113 +++-- .../test/test_multilayer_perceptron_kokkos.py | 90 ++-- .../test/test_multivariate_polynomial.py | 20 +- symmetrix/test/test_symmetrix_calc.py | 65 +-- symmetrix/test/test_tools.py | 50 +- symmetrix/test/test_zbl.py | 25 +- symmetrix/test/test_zbl_kokkos.py | 18 +- 22 files changed, 949 insertions(+), 648 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index af907b0..517e2c3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,3 +6,10 @@ repos: # repository-wide formatting changes. - id: check-yaml - id: check-added-large-files + +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.3.0 + hooks: + - id: ruff + args: [ --fix ] + - id: ruff-format diff --git a/README.md b/README.md index 0be5422..07a8a2d 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,8 @@ source .venv/bin/activate uv pip install -e "./symmetrix[test]" ``` +Run Python formatting and lint checks with `uvx pre-commit run --all-files`. + ### Citing Symmetrix The earliest `symmetrix` results are reported in: diff --git a/pair_symmetrix/test/conftest.py b/pair_symmetrix/test/conftest.py index eca5189..8a801c3 100644 --- a/pair_symmetrix/test/conftest.py +++ b/pair_symmetrix/test/conftest.py @@ -1,4 +1,5 @@ -import os, sys +import os + # Hacky, prevents bad exit due that doesn't seem directly related to the tests def pytest_sessionfinish(session, exitstatus): diff --git a/pair_symmetrix/test/test_pair_symmetrix_mace.py b/pair_symmetrix/test/test_pair_symmetrix_mace.py index 33b1e5c..2cf90c6 100644 --- a/pair_symmetrix/test/test_pair_symmetrix_mace.py +++ b/pair_symmetrix/test/test_pair_symmetrix_mace.py @@ -6,16 +6,22 @@ if not os.path.exists("MACE-OFF23_small-1-8.json"): - urlretrieve("https://www.dropbox.com/scl/fi/zbg122s1zeeb1j6ogheok/MACE-OFF23_small-1-8.json?rlkey=mqb7cje9y3l0smwf75cfoahr7&st=iabk9093&dl=1", - "MACE-OFF23_small-1-8.json") + urlretrieve( + "https://www.dropbox.com/scl/fi/zbg122s1zeeb1j6ogheok/MACE-OFF23_small-1-8.json?rlkey=mqb7cje9y3l0smwf75cfoahr7&st=iabk9093&dl=1", + "MACE-OFF23_small-1-8.json", + ) if not os.path.exists("mace-mp-0b3-medium-1-8.json"): - urlretrieve("https://www.dropbox.com/scl/fi/ymzotmy9nw2lp7pvv2awc/mace-mp-0b3-medium-1-8.json?rlkey=3y2y42ieo79ekjwpt8zbfjgoe&st=91o13eux&dl=1", - "mace-mp-0b3-medium-1-8.json") + urlretrieve( + "https://www.dropbox.com/scl/fi/ymzotmy9nw2lp7pvv2awc/mace-mp-0b3-medium-1-8.json?rlkey=3y2y42ieo79ekjwpt8zbfjgoe&st=91o13eux&dl=1", + "mace-mp-0b3-medium-1-8.json", + ) if not os.path.exists("mace-mp-0b3-medium-hea.json"): - urlretrieve("https://www.dropbox.com/scl/fi/gexhyg8sqy39m5j0mnsnv/mace-mp-0b3-medium-hea.json?rlkey=9cz9g3oxrbsek9a599ul2kdvc&st=fqsyv5yb&dl=1", - "mace-mp-0b3-medium-hea.json") + urlretrieve( + "https://www.dropbox.com/scl/fi/gexhyg8sqy39m5j0mnsnv/mace-mp-0b3-medium-hea.json?rlkey=9cz9g3oxrbsek9a599ul2kdvc&st=fqsyv5yb&dl=1", + "mace-mp-0b3-medium-hea.json", + ) @pytest.mark.parametrize( @@ -23,7 +29,7 @@ [ ["-screen", "none"], ["-screen", "none", "-k", "on", "-sf", "kk"], # kokkos - ] + ], ) @pytest.mark.parametrize( "pair_style", @@ -35,11 +41,10 @@ "symmetrix/mace/float32", "symmetrix/mace/float32 no_domain_decomposition", "symmetrix/mace/float32 mpi_message_passing", - "symmetrix/mace/float32 no_mpi_message_passing" - ] + "symmetrix/mace/float32 no_mpi_message_passing", + ], ) def test_h20(cmdargs, pair_style): - if "float32" in pair_style: pytest.skip("Skipping float32 lammps tests.") if "float32" in pair_style and "kk" not in cmdargs: @@ -47,7 +52,8 @@ def test_h20(cmdargs, pair_style): # ----- setup ----- lmp = lammps(cmdargs=cmdargs) - lmp.commands_string(""" + lmp.commands_string( + """ clear units metal atom_style atomic @@ -66,7 +72,8 @@ def test_h20(cmdargs, pair_style): pair_coeff * * MACE-OFF23_small-1-8.json H O run 0 - """.format(pair_style)) + """.format(pair_style) + ) # ----- energy ----- e = lmp.get_thermo("pe") @@ -82,18 +89,18 @@ def test_h20(cmdargs, pair_style): h = 1e-4 x = lmp.numpy.extract_atom("x", nelem=3, dim=3) f = lmp.numpy.extract_atom("f", nelem=3, dim=3) - f_num = np.zeros([3,3]) - for i in range(0,3): - for j in range(0,3): - x[i,j] += h + f_num = np.zeros([3, 3]) + for i in range(0, 3): + for j in range(0, 3): + x[i, j] += h lmp.command("run 0") ep = lmp.get_thermo("pe") - x[i,j] -= 2*h + x[i, j] -= 2 * h lmp.command("run 0") em = lmp.get_thermo("pe") - x[i,j] += h + x[i, j] += h lmp.command("run 0") - f_num[i,j] = -(ep-em)/(2*h) + f_num[i, j] = -(ep - em) / (2 * h) assert np.allclose(f, f_num, atol=1e-5) # ----- teardown ----- @@ -105,7 +112,7 @@ def test_h20(cmdargs, pair_style): [ ["-screen", "none"], ["-screen", "none", "-k", "on", "-sf", "kk"], # kokkos - ] + ], ) @pytest.mark.parametrize( "pair_style", @@ -117,11 +124,10 @@ def test_h20(cmdargs, pair_style): "symmetrix/mace/float32", "symmetrix/mace/float32 no_domain_decomposition", "symmetrix/mace/float32 mpi_message_passing", - "symmetrix/mace/float32 no_mpi_message_passing" - ] + "symmetrix/mace/float32 no_mpi_message_passing", + ], ) def test_h20_zbl(cmdargs, pair_style): - if "float32" in pair_style: pytest.skip("Skipping float32 lammps tests.") if "float32" in pair_style and "kk" not in cmdargs: @@ -129,7 +135,8 @@ def test_h20_zbl(cmdargs, pair_style): # ----- setup ----- lmp = lammps(cmdargs=cmdargs) - lmp.commands_string(""" + lmp.commands_string( + """ clear units metal atom_style atomic @@ -148,7 +155,8 @@ def test_h20_zbl(cmdargs, pair_style): pair_coeff * * mace-mp-0b3-medium-1-8.json H O run 0 - """.format(pair_style)) + """.format(pair_style) + ) # ----- energy ----- e = lmp.get_thermo("pe") @@ -158,29 +166,30 @@ def test_h20_zbl(cmdargs, pair_style): h = 1e-4 x = lmp.numpy.extract_atom("x", nelem=3, dim=3) f = lmp.numpy.extract_atom("f", nelem=3, dim=3) - f_num = np.zeros([3,3]) - for i in range(0,3): - for j in range(0,3): - x[i,j] += h + f_num = np.zeros([3, 3]) + for i in range(0, 3): + for j in range(0, 3): + x[i, j] += h lmp.command("run 0") ep = lmp.get_thermo("pe") - x[i,j] -= 2*h + x[i, j] -= 2 * h lmp.command("run 0") em = lmp.get_thermo("pe") - x[i,j] += h + x[i, j] += h lmp.command("run 0") - f_num[i,j] = -(ep-em)/(2*h) + f_num[i, j] = -(ep - em) / (2 * h) assert np.allclose(f, f_num, rtol=1e-4, atol=1e-6) # ----- teardown ----- lmp.close() + @pytest.mark.parametrize( "cmdargs", [ ["-screen", "none"], ["-screen", "none", "-k", "on", "-sf", "kk"], # kokkos - ] + ], ) @pytest.mark.parametrize( "pair_style", @@ -192,11 +201,10 @@ def test_h20_zbl(cmdargs, pair_style): "symmetrix/mace/float32", "symmetrix/mace/float32 no_domain_decomposition", "symmetrix/mace/float32 mpi_message_passing", - "symmetrix/mace/float32 no_mpi_message_passing" - ] + "symmetrix/mace/float32 no_mpi_message_passing", + ], ) def test_water(cmdargs, pair_style): - if "float32" in pair_style: pytest.skip("Skipping float32 lammps tests.") if "float32" in pair_style and "kk" not in cmdargs: @@ -204,7 +212,8 @@ def test_water(cmdargs, pair_style): # ----- setup ----- lmp = lammps(cmdargs=cmdargs) - lmp.commands_string(""" + lmp.commands_string( + """ clear units metal boundary p p p @@ -252,23 +261,26 @@ def test_water(cmdargs, pair_style): compute peratom all pe/atom fix f1 all nve run 0 - """.format(pair_style)) + """.format(pair_style) + ) # ----- test energy and stress ----- assert lmp.get_thermo("pe") == pytest.approx(-16649.784441, abs=1e-6) assert lmp.get_thermo("pxx") == pytest.approx(-69407.514290, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pyy") == pytest.approx(-69407.514290, abs=1e-8, rel=1e-4) - assert lmp.get_thermo("pzz") == pytest.approx( 18042.601669, abs=1e-8, rel=1e-4) + assert lmp.get_thermo("pzz") == pytest.approx(18042.601669, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pxy") == pytest.approx(-55297.126324, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pxz") == pytest.approx(0.0, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pyz") == pytest.approx(0.0, abs=1e-8, rel=1e-4) # ----- run 10 steps, test again ----- lmp.command("run 10") - assert lmp.get_thermo("pe") == pytest.approx(-16649.988675, abs=1e-4) # note lower tolerance + assert lmp.get_thermo("pe") == pytest.approx( + -16649.988675, abs=1e-4 + ) # note lower tolerance assert lmp.get_thermo("pxx") == pytest.approx(-56913.479676, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pyy") == pytest.approx(-56913.479676, abs=1e-8, rel=1e-4) - assert lmp.get_thermo("pzz") == pytest.approx( 17756.761767, abs=1e-8, rel=1e-4) + assert lmp.get_thermo("pzz") == pytest.approx(17756.761767, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pxy") == pytest.approx(-50938.320172, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pxz") == pytest.approx(0.0, abs=1e-8, rel=1e-4) assert lmp.get_thermo("pyz") == pytest.approx(0.0, abs=1e-8, rel=1e-4) @@ -282,7 +294,7 @@ def test_water(cmdargs, pair_style): [ ["-screen", "none"], ["-screen", "none", "-k", "on", "-sf", "kk"], # kokkos - ] + ], ) @pytest.mark.parametrize( "pair_style", @@ -294,11 +306,10 @@ def test_water(cmdargs, pair_style): "symmetrix/mace/float32", "symmetrix/mace/float32 no_domain_decomposition", "symmetrix/mace/float32 mpi_message_passing", - "symmetrix/mace/float32 no_mpi_message_passing" - ] + "symmetrix/mace/float32 no_mpi_message_passing", + ], ) def test_hea(cmdargs, pair_style): - if "float32" in pair_style: pytest.skip("Skipping float32 lammps tests.") if "float32" in pair_style and "kk" not in cmdargs: @@ -306,7 +317,8 @@ def test_hea(cmdargs, pair_style): # ----- setup ----- lmp = lammps(cmdargs=cmdargs) - lmp.commands_string(""" + lmp.commands_string( + """ clear units metal boundary p p p @@ -368,7 +380,8 @@ def test_hea(cmdargs, pair_style): compute peratom all pe/atom fix f1 all nve run 0 - """.format(pair_style)) + """.format(pair_style) + ) # ----- test energy and stress ----- assert lmp.get_thermo("pe") == pytest.approx(-105.640759, abs=1e-3) @@ -381,13 +394,15 @@ def test_hea(cmdargs, pair_style): # ----- run 10 steps, test again ----- lmp.command("run 10") - assert lmp.get_thermo("pe") == pytest.approx(-105.642334, abs=1e-3) # note lower tolerance - assert lmp.get_thermo("pxx") == pytest.approx(-85880.067428, abs=1e-8, rel=1e-2) - assert lmp.get_thermo("pyy") == pytest.approx(-75813.607646, abs=1e-8, rel=1e-2) - assert lmp.get_thermo("pzz") == pytest.approx(-93780.229278, abs=1e-8, rel=1e-2) - assert lmp.get_thermo("pxy") == pytest.approx(0.0, abs=1e-8, rel=1e-2) - assert lmp.get_thermo("pxz") == pytest.approx(0.0, abs=1e-8, rel=1e-2) - assert lmp.get_thermo("pyz") == pytest.approx(0.0, abs=1e-8, rel=1e-2) + assert lmp.get_thermo("pe") == pytest.approx( + -105.642334, abs=1e-3 + ) # note lower tolerance + assert lmp.get_thermo("pxx") == pytest.approx(-85880.067428, abs=1e-8, rel=1e-2) + assert lmp.get_thermo("pyy") == pytest.approx(-75813.607646, abs=1e-8, rel=1e-2) + assert lmp.get_thermo("pzz") == pytest.approx(-93780.229278, abs=1e-8, rel=1e-2) + assert lmp.get_thermo("pxy") == pytest.approx(0.0, abs=1e-8, rel=1e-2) + assert lmp.get_thermo("pxz") == pytest.approx(0.0, abs=1e-8, rel=1e-2) + assert lmp.get_thermo("pyz") == pytest.approx(0.0, abs=1e-8, rel=1e-2) # ----- teardown ----- lmp.close() diff --git a/symmetrix/pyproject.toml b/symmetrix/pyproject.toml index 1186049..c2b502a 100644 --- a/symmetrix/pyproject.toml +++ b/symmetrix/pyproject.toml @@ -29,3 +29,10 @@ test = [ [project.scripts] symmetrix_extract_mace = "symmetrix.cli.extract_mace:main" + +[tool.ruff] +lint.ignore = ["E741"] + +[tool.ruff.lint.per-file-ignores] +"__init__.py" = ["F403", "E402"] +"test_*.py" = ["F811"] diff --git a/symmetrix/source/symmetrix/__init__.py b/symmetrix/source/symmetrix/__init__.py index 21eef3f..60116d5 100644 --- a/symmetrix/source/symmetrix/__init__.py +++ b/symmetrix/source/symmetrix/__init__.py @@ -2,8 +2,7 @@ # TODO: very hacky, to import underscored names import importlib -_sym = importlib.import_module('.symmetrix', __name__) -_sym.__all__ = [n for n in vars(_sym) if not (n.startswith('__') and n.endswith('__'))] -from .symmetrix import * -from .calculator import Symmetrix +_sym = importlib.import_module(".symmetrix", __name__) +_sym.__all__ = [n for n in vars(_sym) if not (n.startswith("__") and n.endswith("__"))] +from .symmetrix import * diff --git a/symmetrix/source/symmetrix/calculator.py b/symmetrix/source/symmetrix/calculator.py index 829d47e..104dd80 100755 --- a/symmetrix/source/symmetrix/calculator.py +++ b/symmetrix/source/symmetrix/calculator.py @@ -4,6 +4,7 @@ This file was written and publicly released by Dr. Noam Bernstein as part of his work for the U. S. Government, and is not subject to copyright. """ + import json import logging from tempfile import NamedTemporaryFile @@ -15,13 +16,14 @@ logging.warning("Symmetrix using slow ase.neighborlist.neighbor_list") from ase.neighborlist import neighbor_list -from ase.calculators.calculator import Calculator, PropertyNotImplementedError, all_changes +from ase.calculators.calculator import Calculator, all_changes from ase.stress import full_3x3_to_voigt_6_stress from . import symmetrix + class Symmetrix(Calculator): - """ASE Calculator using symmetrix library to evaluate equivariant graph neural network + """ASE Calculator using symmetrix library to evaluate equivariant graph neural network potential energy functions Parameters @@ -33,34 +35,44 @@ class Symmetrix(Calculator): ----- Wraps symmetrix library from https://github.com/wcwitt/symmetrix via python interface at https://pypi.org/project/symmetrix/ """ - implemented_properties = ['energy', 'free_energy', 'energies', 'forces', 'stress'] + implemented_properties = ["energy", "free_energy", "energies", "forces", "stress"] def __init__(self, model_file, dtype="float64", use_kokkos=True, **kwargs): Calculator.__init__(self, **kwargs) if dtype not in ["float32", "float64"]: - raise ValueError(f"Unsupported dtype '{dtype}'. Supported dtypes are 'float64' and 'float32'.") + raise ValueError( + f"Unsupported dtype '{dtype}'. Supported dtypes are 'float64' and 'float32'." + ) if use_kokkos and not hasattr(symmetrix, "MACEKokkos"): raise RuntimeError("Symmetrix was built without Kokkos support.") self.use_kokkos = use_kokkos if self.use_kokkos: if not symmetrix._kokkos_is_initialized(): symmetrix._init_kokkos() - MACE = symmetrix.MACEKokkos if dtype == "float64" else symmetrix.MACEKokkosFloat + MACE = ( + symmetrix.MACEKokkos + if dtype == "float64" + else symmetrix.MACEKokkosFloat + ) else: if dtype == "float32": raise ValueError(f"dtype '{dtype}' requires `use_kokkos = True`") MACE = symmetrix.MACE try: self.evaluator = MACE(str(model_file)) - except RuntimeError: # expecting json.exception.parse_error.101 + except RuntimeError: # expecting json.exception.parse_error.101 # import this here so that torch/mace support isn't needed if file is already symmetrix json from .extract_mace_data import extract_mace_data - kwargs_extract = {k: v for k, v in kwargs.items() - if k in ['species', - 'head', - 'num_spline_points']} - logging.warning(f"Converting model from pytorch model to symmetrix dict with {kwargs_extract}") + + kwargs_extract = { + k: v + for k, v in kwargs.items() + if k in ["species", "head", "num_spline_points"] + } + logging.warning( + f"Converting model from pytorch model to symmetrix dict with {kwargs_extract}" + ) data = extract_mace_data(model_file, **kwargs_extract) with NamedTemporaryFile("w") as fout: logging.warning(f"Converting via NamedTemporaryFile {fout.name}") @@ -69,34 +81,49 @@ def __init__(self, model_file, dtype="float64", use_kokkos=True, **kwargs): self.cutoff = self.evaluator.r_cut - def calculate(self, atoms=None, properties=['energy'], system_changes=all_changes): + def calculate(self, atoms=None, properties=["energy"], system_changes=all_changes): Calculator.calculate(self, atoms, properties, system_changes) ase_atomic_numbers = self.atoms.get_atomic_numbers().tolist() mace_atomic_numbers = self.evaluator.atomic_numbers - i_list, j_list, r, xyz = neighbor_list('ijdD', self.atoms, self.cutoff) + i_list, j_list, r, xyz = neighbor_list("ijdD", self.atoms, self.cutoff) num_nodes = np.max(i_list) + 1 - node_types = [mace_atomic_numbers.index(ase_atomic_numbers[i]) for i in range(num_nodes)] + node_types = [ + mace_atomic_numbers.index(ase_atomic_numbers[i]) for i in range(num_nodes) + ] num_neigh = np.bincount(j_list, minlength=num_nodes) neigh_types = [mace_atomic_numbers.index(ase_atomic_numbers[j]) for j in j_list] self.evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r + ) - self.results['energy'] = self.results['free_energy'] = np.sum(self.evaluator.node_energies) - self.results['energies'] = np.asarray(self.evaluator.node_energies) + self.results["energy"] = self.results["free_energy"] = np.sum( + self.evaluator.node_energies + ) + self.results["energies"] = np.asarray(self.evaluator.node_energies) pair_forces = np.asarray(self.evaluator.node_forces).reshape((-1, 3)) - pair_forces = pair_forces[:len(i_list), :] # currently, `evaluator.node_forces` is a container - # which can grow larger than the actual number of pairs + pair_forces = pair_forces[ + : len(i_list), : + ] # currently, `evaluator.node_forces` is a container + # which can grow larger than the actual number of pairs # atom forces from pair_forces N_atoms = len(self.atoms) atom_forces = np.zeros((N_atoms, 3)) - atom_forces[:, 0] = np.bincount(j_list, weights=pair_forces[:, 0], minlength=N_atoms) - np.bincount(i_list, weights=pair_forces[:, 0], minlength=N_atoms) - atom_forces[:, 1] = np.bincount(j_list, weights=pair_forces[:, 1], minlength=N_atoms) - np.bincount(i_list, weights=pair_forces[:, 1], minlength=N_atoms) - atom_forces[:, 2] = np.bincount(j_list, weights=pair_forces[:, 2], minlength=N_atoms) - np.bincount(i_list, weights=pair_forces[:, 2], minlength=N_atoms) - - self.results['forces'] = atom_forces + atom_forces[:, 0] = np.bincount( + j_list, weights=pair_forces[:, 0], minlength=N_atoms + ) - np.bincount(i_list, weights=pair_forces[:, 0], minlength=N_atoms) + atom_forces[:, 1] = np.bincount( + j_list, weights=pair_forces[:, 1], minlength=N_atoms + ) - np.bincount(i_list, weights=pair_forces[:, 1], minlength=N_atoms) + atom_forces[:, 2] = np.bincount( + j_list, weights=pair_forces[:, 2], minlength=N_atoms + ) - np.bincount(i_list, weights=pair_forces[:, 2], minlength=N_atoms) + + self.results["forces"] = atom_forces # stress from pair_forces - self.results['stress'] = full_3x3_to_voigt_6_stress((-pair_forces.T @ xyz) / self.atoms.get_volume()) + self.results["stress"] = full_3x3_to_voigt_6_stress( + (-pair_forces.T @ xyz) / self.atoms.get_volume() + ) diff --git a/symmetrix/source/symmetrix/cli/extract_mace.py b/symmetrix/source/symmetrix/cli/extract_mace.py index a9d35aa..1ad75d8 100755 --- a/symmetrix/source/symmetrix/cli/extract_mace.py +++ b/symmetrix/source/symmetrix/cli/extract_mace.py @@ -2,16 +2,34 @@ from argparse import ArgumentParser -from ..extract_mace_data import extract_mace_data +from ..extract_mace_data import extract_mace_data + def main(): parser = ArgumentParser() parser.add_argument("--model", "-m", required=True, help="Torch model file.") group = parser.add_mutually_exclusive_group() - group.add_argument("--atomic-numbers", "-Z", "-z", nargs="+", help="Atomic numbers to extract.", default=[]) - group.add_argument("--chemical-symbols", "-s", nargs="+", help="Chemical symbols to extract.", default=[]) - parser.add_argument("--head", "-H", help="Head to keep, ignored unless model is multihead. " - "Defaults to first non-PT head, same as mace.tools.script_utils.remove_pt_head") + group.add_argument( + "--atomic-numbers", + "-Z", + "-z", + nargs="+", + help="Atomic numbers to extract.", + default=[], + ) + group.add_argument( + "--chemical-symbols", + "-s", + nargs="+", + help="Chemical symbols to extract.", + default=[], + ) + parser.add_argument( + "--head", + "-H", + help="Head to keep, ignored unless model is multihead. " + "Defaults to first non-PT head, same as mace.tools.script_utils.remove_pt_head", + ) parser.add_argument("--output", "-o", help="Output filename.") args = parser.parse_args() @@ -23,13 +41,17 @@ def main(): else: model_name = Path(args.model).stem - species = args.atomic_numbers if args.chemical_symbols == [] else args.chemical_symbols + species = ( + args.atomic_numbers if args.chemical_symbols == [] else args.chemical_symbols + ) output = extract_mace_data(args.model, species, args.head) ### ----- WRITE JSON ----- if args.output is None: - args.output = model_name + '-' + '-'.join(str(a) for a in sorted(species)) + '.json' + args.output = ( + model_name + "-" + "-".join(str(a) for a in sorted(species)) + ".json" + ) print("WRITING JSON TO", args.output) with open(args.output, "w") as f: json.dump(output, f, indent=4) diff --git a/symmetrix/source/symmetrix/extract_mace_data.py b/symmetrix/source/symmetrix/extract_mace_data.py index 842e9a3..dfba74a 100755 --- a/symmetrix/source/symmetrix/extract_mace_data.py +++ b/symmetrix/source/symmetrix/extract_mace_data.py @@ -1,20 +1,19 @@ import torch -import os import logging import itertools import numpy as np -import matplotlib.pyplot as plt from scipy.interpolate import CubicSpline -from e3nn.o3 import Irrep, Irreps, Linear, wigner_3j +from e3nn.o3 import Irreps, Linear from mace.modules.radial import ZBLBasis from mace.tools.cg import U_matrix_real from mace.tools.scripts_utils import remove_pt_head from ase.data import chemical_symbols + def extract_mace_data(model, species, head=None, num_spline_points=256): """Extract data from pytorch model file into structure that can be written as symmetrix JSON data file @@ -35,11 +34,9 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): output_data: dict with symmetrix model data """ device = torch.device("cpu") - model = torch.load( - model, - map_location=device, - weights_only=False - ).to(device=device, dtype=torch.float64) + model = torch.load(model, map_location=device, weights_only=False).to( + device=device, dtype=torch.float64 + ) model.eval() if species is None: @@ -54,7 +51,9 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): try: Z = chemical_symbols.index(sp) except ValueError as exc: - raise ValueError(f"Failed to parse {sp} as atomic number or chemical species") from exc + raise ValueError( + f"Failed to parse {sp} as atomic number or chemical species" + ) from exc atomic_numbers.append(Z) # ensure that splines goes smoothly to 0 at outer cutoff @@ -62,7 +61,7 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): ### ----- EXTRACT SINGLE HEAD ----- - if hasattr(model, 'heads') and len(model.heads) != 1: + if hasattr(model, "heads") and len(model.heads) != 1: torch.set_default_dtype(next(model.parameters()).dtype) model = remove_pt_head(model, head).to(device=device, dtype=torch.float64) model.eval() @@ -72,30 +71,45 @@ def extract_mace_data(model, species, head=None, num_spline_points=256): if len(model.interactions) != 2: raise RuntimeError("Currently, symmetrix only supports two-layer MACE models.") - from mace.modules.blocks import RealAgnosticInteractionBlock, RealAgnosticDensityInteractionBlock - if (not isinstance(model.interactions[0], RealAgnosticInteractionBlock) - and - not isinstance(model.interactions[0], RealAgnosticDensityInteractionBlock)): + from mace.modules.blocks import ( + RealAgnosticInteractionBlock, + RealAgnosticDensityInteractionBlock, + ) + + if not isinstance( + model.interactions[0], RealAgnosticInteractionBlock + ) and not isinstance(model.interactions[0], RealAgnosticDensityInteractionBlock): raise RuntimeError( "Currently, symmetrix only supports MACE models whose first interaction is " - "RealAgnosticInteractionBlock or RealAgnosticDensityInteractionBlock.") + "RealAgnosticInteractionBlock or RealAgnosticDensityInteractionBlock." + ) - from mace.modules.blocks import RealAgnosticResidualInteractionBlock, RealAgnosticDensityResidualInteractionBlock - if (not isinstance(model.interactions[1], RealAgnosticResidualInteractionBlock) - and - not isinstance(model.interactions[1], RealAgnosticDensityResidualInteractionBlock)): + from mace.modules.blocks import ( + RealAgnosticResidualInteractionBlock, + RealAgnosticDensityResidualInteractionBlock, + ) + + if not isinstance( + model.interactions[1], RealAgnosticResidualInteractionBlock + ) and not isinstance( + model.interactions[1], RealAgnosticDensityResidualInteractionBlock + ): raise RuntimeError( "Currently, symmetrix only supports MACE models whose second interaction is " - "RealAgnosticResidualInteractionBlock or RealAgnosticDensityResidualInteractionBlock.") + "RealAgnosticResidualInteractionBlock or RealAgnosticDensityResidualInteractionBlock." + ) - if (model.spherical_harmonics._lmax != 3): - raise RuntimeError("Currently, symmetrix only supports MACE models with l_max=3.") + if model.spherical_harmonics._lmax != 3: + raise RuntimeError( + "Currently, symmetrix only supports MACE models with l_max=3." + ) ### ----- HELPER FUNCTION ----- def linear_simplify(linear): - simplified = Linear(Irreps(linear.irreps_in).simplify(), - Irreps(linear.irreps_out).simplify()) + simplified = Linear( + Irreps(linear.irreps_in).simplify(), Irreps(linear.irreps_out).simplify() + ) simplified.weight = linear.weight simplified.bias = linear.bias return simplified.to(device=device, dtype=torch.float64) @@ -105,12 +119,12 @@ def linear_simplify(linear): num_channels = model.node_embedding.linear.irreps_out.count("0e") r_cut = model.r_max.item() l_max = model.spherical_harmonics._lmax - L_max = model.products[0].linear.irreps_out.lmax + L_max = model.products[0].linear.irreps_out.lmax output = {} - output['num_channels'] = num_channels - output['r_cut'] = r_cut - output['l_max'] = l_max - output['L_max'] = L_max + output["num_channels"] = num_channels + output["r_cut"] = r_cut + output["l_max"] = l_max + output["L_max"] = L_max ### ----- ATOMIC NUMBERS AND ENERGIES ----- @@ -119,32 +133,37 @@ def linear_simplify(linear): atomic_numbers = sorted(model.atomic_numbers.tolist()) logging.warning(f"No atomic_numbers, including all: {atomic_numbers}") atomic_energies = [ - torch.atleast_1d(model.atomic_energies_fn.atomic_energies.squeeze())[model.atomic_numbers.tolist().index(a)].item() - + model.scale_shift.shift.item() - for a in atomic_numbers] - output['atomic_numbers'] = atomic_numbers - output['num_elements'] = len(atomic_numbers) - output['atomic_energies'] = atomic_energies + torch.atleast_1d(model.atomic_energies_fn.atomic_energies.squeeze())[ + model.atomic_numbers.tolist().index(a) + ].item() + + model.scale_shift.shift.item() + for a in atomic_numbers + ] + output["atomic_numbers"] = atomic_numbers + output["num_elements"] = len(atomic_numbers) + output["atomic_energies"] = atomic_energies ### --- ZBL --- if hasattr(model, "pair_repulsion") and model.pair_repulsion: if not isinstance(model.pair_repulsion_fn, ZBLBasis): raise Exception("Only ZBL pair_repulsion is supported.") - output['has_zbl'] = True + output["has_zbl"] = True zbl = model.pair_repulsion_fn - output['zbl_a_exp'] = zbl.a_exp.item() - output['zbl_a_prefactor'] = zbl.a_prefactor.item() - output['zbl_c'] = (model.scale_shift.scale.item() * zbl.c.numpy(force=True)).tolist() - output['zbl_covalent_radii'] = zbl.covalent_radii.numpy(force=True).tolist() - output['zbl_p'] = zbl.p.item() + output["zbl_a_exp"] = zbl.a_exp.item() + output["zbl_a_prefactor"] = zbl.a_prefactor.item() + output["zbl_c"] = ( + model.scale_shift.scale.item() * zbl.c.numpy(force=True) + ).tolist() + output["zbl_covalent_radii"] = zbl.covalent_radii.numpy(force=True).tolist() + output["zbl_p"] = zbl.p.item() else: - output['has_zbl'] = False + output["has_zbl"] = False ### ----- RADIAL SPLINES ----- logging.info("R0+R1") - r,h = np.linspace(1e-12, r_cut, num_spline_points, retstep=True) + r, h = np.linspace(1e-12, r_cut, num_spline_points, retstep=True) spline_values_0 = [] spline_derivatives_0 = [] spline_values_1 = [] @@ -156,42 +175,56 @@ def linear_simplify(linear): model_i = model.atomic_numbers.tolist().index(a_i) model_j = model.atomic_numbers.tolist().index(a_j) bessels = model.radial_embedding( - torch.tensor(r, dtype=torch.get_default_dtype(), device=device).unsqueeze(-1), + torch.tensor( + r, dtype=torch.get_default_dtype(), device=device + ).unsqueeze(-1), torch.eye(len(model.atomic_numbers), device=device), - torch.tensor([[model_i],[model_j]], dtype=torch.int64, device=device), - model.atomic_numbers.to(device)) + torch.tensor([[model_i], [model_j]], dtype=torch.int64, device=device), + model.atomic_numbers.to(device), + ) if isinstance(bessels, tuple): bessels = bessels[0] # newer versions return (bessels, cutoffs) # radial basis for interaction 0 R = model.interactions[0].conv_tp_weights(bessels).numpy(force=True) - spl_0 = [CubicSpline(r, R[:,k], bc_type=spline_bc_type) for k in range(R.shape[1])] + spl_0 = [ + CubicSpline(r, R[:, k], bc_type=spline_bc_type) + for k in range(R.shape[1]) + ] spline_values_0.append([spl(r).tolist() for spl in spl_0]) spline_derivatives_0.append([spl.derivative()(r).tolist() for spl in spl_0]) # radial basis for interaction 1 R = model.interactions[1].conv_tp_weights(bessels).numpy(force=True) - spl_1 = [CubicSpline(r, R[:,k], bc_type=spline_bc_type) for k in range(R.shape[1])] + spl_1 = [ + CubicSpline(r, R[:, k], bc_type=spline_bc_type) + for k in range(R.shape[1]) + ] spline_values_1.append([spl(r).tolist() for spl in spl_1]) spline_derivatives_1.append([spl.derivative()(r).tolist() for spl in spl_1]) - output['radial_spline_h'] = float(h) - output['radial_spline_values_0'] = spline_values_0 - output['radial_spline_derivs_0'] = spline_derivatives_0 - output['radial_spline_values_1'] = spline_values_1 - output['radial_spline_derivs_1'] = spline_derivatives_1 + output["radial_spline_h"] = float(h) + output["radial_spline_values_0"] = spline_values_0 + output["radial_spline_derivs_0"] = spline_derivatives_0 + output["radial_spline_values_1"] = spline_values_1 + output["radial_spline_derivs_1"] = spline_derivatives_1 ### ----- H0 ----- logging.info("H0") H0_weights = ( - np.reshape(model.node_embedding.linear.weight.numpy(force=True), - [len(model.atomic_numbers),num_channels]) / np.sqrt(len(model.atomic_numbers)) - @ - np.reshape(model.interactions[0].linear_up.weight.numpy(force=True), - [num_channels,num_channels]) / np.sqrt(num_channels) + np.reshape( + model.node_embedding.linear.weight.numpy(force=True), + [len(model.atomic_numbers), num_channels], ) + / np.sqrt(len(model.atomic_numbers)) + @ np.reshape( + model.interactions[0].linear_up.weight.numpy(force=True), + [num_channels, num_channels], + ) + / np.sqrt(num_channels) + ) indices = [model.atomic_numbers.tolist().index(a) for a in atomic_numbers] - H0_weights = H0_weights[indices,:] - output['H0_weights'] = H0_weights.flatten().tolist() + H0_weights = H0_weights[indices, :] + output["H0_weights"] = H0_weights.flatten().tolist() ### ----- Phi0 ----- @@ -201,9 +234,9 @@ def linear_simplify(linear): logging.info("A0") A0_scaled = True if ("Density" in type(model.interactions[0]).__name__) else False - output['A0_scaled'] = A0_scaled + output["A0_scaled"] = A0_scaled if A0_scaled: - r,h = np.linspace(1e-12, r_cut, num_spline_points, retstep=True) + r, h = np.linspace(1e-12, r_cut, num_spline_points, retstep=True) A0_spline_values = [] A0_spline_derivs = [] for a_i in atomic_numbers: @@ -213,28 +246,43 @@ def linear_simplify(linear): model_i = model.atomic_numbers.tolist().index(a_i) model_j = model.atomic_numbers.tolist().index(a_j) bessels = model.radial_embedding( - torch.tensor(r, dtype=torch.get_default_dtype(), device=device).unsqueeze(-1), + torch.tensor( + r, dtype=torch.get_default_dtype(), device=device + ).unsqueeze(-1), torch.eye(len(model.atomic_numbers), device=device), - torch.tensor([[model_i],[model_j]], dtype=torch.int64, device=device), - model.atomic_numbers.to(device)) + torch.tensor( + [[model_i], [model_j]], dtype=torch.int64, device=device + ), + model.atomic_numbers.to(device), + ) if isinstance(bessels, tuple): bessels = bessels[0] # newer versions return (bessels, cutoffs) - R = torch.tanh(model.interactions[0].density_fn(bessels)**2).numpy(force=True) - spl = CubicSpline(r, R[:,0], bc_type=spline_bc_type) + R = torch.tanh(model.interactions[0].density_fn(bessels) ** 2).numpy( + force=True + ) + spl = CubicSpline(r, R[:, 0], bc_type=spline_bc_type) A0_spline_values.append(spl(r).tolist()) A0_spline_derivs.append(spl.derivative()(r).tolist()) - output['A0_spline_h'] = float(h) - output['A0_spline_values'] = A0_spline_values - output['A0_spline_derivs'] = A0_spline_derivs + output["A0_spline_h"] = float(h) + output["A0_spline_values"] = A0_spline_values + output["A0_spline_derivs"] = A0_spline_derivs A0_weights = [] for i, a in enumerate(atomic_numbers): model_i = model.atomic_numbers.tolist().index(a) A0_weights.append([]) - for l,_,w in model.interactions[0].skip_tp.weight_views(yield_instruction=True): - w_linear = model.interactions[0].linear.weight_view_for_instruction(l).numpy(force=True) / np.sqrt(num_channels) + for l, _, w in model.interactions[0].skip_tp.weight_views( + yield_instruction=True + ): + w_linear = model.interactions[0].linear.weight_view_for_instruction( + l + ).numpy(force=True) / np.sqrt(num_channels) if not A0_scaled: w_linear /= model.interactions[0].avg_num_neighbors - fused = w_linear @ w[:,model_i,:].numpy(force=True) / np.sqrt(len(model.atomic_numbers)*num_channels) + fused = ( + w_linear + @ w[:, model_i, :].numpy(force=True) + / np.sqrt(len(model.atomic_numbers) * num_channels) + ) A0_weights[i].append(fused.flatten().tolist()) output["A0_weights"] = A0_weights @@ -242,16 +290,19 @@ def linear_simplify(linear): logging.info("M0") correlation = model.products[0].symmetric_contractions.contractions[0].correlation + ### Computes U_{lm\eta, l1m1 l2m2 ...} # * `irrep_out` is essentially l_out # * `irreps_in` essentially provides l_in_max # * `corr_in_max` is the max correlation order def U_sparse(irrep_out, irreps_in, corr_in_max): U = [[]] # list of lists because U[0] should be empty - for corr in range(1,corr_in_max+1): + for corr in range(1, corr_in_max + 1): # get U matrix for this correlation order try: - U_matrix = U_matrix_real(irreps_in, [irrep_out], corr, use_cueq_cg=False)[1] + U_matrix = U_matrix_real( + irreps_in, [irrep_out], corr, use_cueq_cg=False + )[1] except TypeError: U_matrix = U_matrix_real(irreps_in, [irrep_out], corr)[1] if irrep_out.l == 0: # makes U_matrix.shape consistent with l>0 cases @@ -259,10 +310,12 @@ def U_sparse(irrep_out, irreps_in, corr_in_max): num_eta = U_matrix.shape[-1] U_matrix = U_matrix.flatten() # extract sparse U for this correlation order - U_sparse_corr = [[{} for _ in range(num_eta)] for _ in range(2*irrep_out.l+1)] + U_sparse_corr = [ + [{} for _ in range(num_eta)] for _ in range(2 * irrep_out.l + 1) + ] j = 0 - for m in range(2*irrep_out.l+1): - for lm_list in itertools.product(range((l_max+1)**2), repeat=corr): + for m in range(2 * irrep_out.l + 1): + for lm_list in itertools.product(range((l_max + 1) ** 2), repeat=corr): for eta in range(num_eta): if abs(U_matrix[j]) > 1e-12: lm_tuple_sorted = tuple(sorted(lm_list)) @@ -272,6 +325,7 @@ def U_sparse(irrep_out, irreps_in, corr_in_max): j += 1 U.append(U_sparse_corr) return U + irreps_in = [ir[1] for ir in model.products[0].symmetric_contractions.irreps_in] irreps_out = [ir[1] for ir in model.products[0].symmetric_contractions.irreps_out] C = {} @@ -285,38 +339,54 @@ def U_sparse(irrep_out, irreps_in, corr_in_max): # extract weights from model # warning: slightly odd order of the contractions weights due to reverse countdown W = [[]] # list of lists because W[0] should be empty - W.append(model.products[0].symmetric_contractions.contractions[l].weights[1].numpy(force=True)) - W.append(model.products[0].symmetric_contractions.contractions[l].weights[0].numpy(force=True)) - W.append(model.products[0].symmetric_contractions.contractions[l].weights_max.numpy(force=True)) + W.append( + model.products[0] + .symmetric_contractions.contractions[l] + .weights[1] + .numpy(force=True) + ) + W.append( + model.products[0] + .symmetric_contractions.contractions[l] + .weights[0] + .numpy(force=True) + ) + W.append( + model.products[0] + .symmetric_contractions.contractions[l] + .weights_max.numpy(force=True) + ) # combine U and W into polynomial-like terms for recursive evaluator - for m in range(-l,l+1): - lm = l*(l+1)+m + for m in range(-l, l + 1): + lm = l * (l + 1) + m C[i][lm] = {} for k in range(num_channels): P_lmk = {} - for corr in range(1,correlation+1): - for eta in range(len(U[corr][l+m])): - for key,value in U[corr][l+m][eta].items(): + for corr in range(1, correlation + 1): + for eta in range(len(U[corr][l + m])): + for key, value in U[corr][l + m][eta].items(): if key not in P_lmk.keys(): P_lmk[key] = 0.0 - P_lmk[key] += float(W[corr][model_i,eta,k]) * value + P_lmk[key] += float(W[corr][model_i, eta, k]) * value C[i][lm][k] = list(P_lmk.values()) M[lm] = [list(key) for key in P_lmk.keys()] - output['M0_weights'] = C - output['M0_monomials'] = M + output["M0_weights"] = C + output["M0_monomials"] = M ### ----- H1 ----- logging.info("H1") - H1_weights = np.zeros([L_max+1, num_channels, num_channels]) + H1_weights = np.zeros([L_max + 1, num_channels, num_channels]) weights_0 = np.reshape( model.products[0].linear.weight.numpy(force=True), - [L_max+1,num_channels,num_channels]) / np.sqrt(num_channels) + [L_max + 1, num_channels, num_channels], + ) / np.sqrt(num_channels) weights_1 = np.reshape( model.interactions[1].linear_up.weight.numpy(force=True), - [L_max+1,num_channels,num_channels]) / np.sqrt(num_channels) - for l in range(L_max+1): - H1_weights[l,:,:] = weights_0[l,:,:] @ weights_1[l,:,:] + [L_max + 1, num_channels, num_channels], + ) / np.sqrt(num_channels) + for l in range(L_max + 1): + H1_weights[l, :, :] = weights_0[l, :, :] @ weights_1[l, :, :] output["H1_weights"] = H1_weights.flatten().tolist() ### ----- Phi1 ----- @@ -328,53 +398,69 @@ def U_sparse(irrep_out, irreps_in, corr_in_max): Phi1_clebsch_gordan = [] Phi1_lme = [] Phi1_lelm1lm2 = [] - num_lm1 = (l_max+1)**2 - num_lm2 = (L_max+1)**2 + num_lm1 = (l_max + 1) ** 2 + num_lm2 = (L_max + 1) ** 2 + def compute_lem(le, l, m): lem = 0 for j in range(le): - lem += 2*Phi1_l[j]+1 - return lem+l+m + lem += 2 * Phi1_l[j] + 1 + return lem + l + m + def compute_lme(le, l, m): - e = le - int(sum(np.array(Phi1_l) 1e-12: - Phi1_lme.append(compute_lme(le,l,m)) + Phi1_lme.append(compute_lme(le, l, m)) Phi1_clebsch_gordan.append(Phi_0[lem]) - Phi1_lelm1lm2.append(compute_lelm1lm2(le,l1,m1,l2,m2)) + Phi1_lelm1lm2.append( + compute_lelm1lm2(le, l1, m1, l2, m2) + ) output["Phi1_l"] = Phi1_l output["Phi1_l1"] = Phi1_l1 output["Phi1_l2"] = Phi1_l2 @@ -386,9 +472,9 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): logging.info("A1") A1_scaled = True if ("Density" in type(model.interactions[1]).__name__) else False - output['A1_scaled'] = A1_scaled + output["A1_scaled"] = A1_scaled if A1_scaled: - r,h = np.linspace(1e-12, r_cut, num_spline_points, retstep=True) + r, h = np.linspace(1e-12, r_cut, num_spline_points, retstep=True) A1_spline_values = [] A1_spline_derivs = [] for a_i in atomic_numbers: @@ -398,24 +484,33 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): model_i = model.atomic_numbers.tolist().index(a_i) model_j = model.atomic_numbers.tolist().index(a_j) bessels = model.radial_embedding( - torch.tensor(r, dtype=torch.get_default_dtype(), device=device).unsqueeze(-1), + torch.tensor( + r, dtype=torch.get_default_dtype(), device=device + ).unsqueeze(-1), torch.eye(len(model.atomic_numbers), device=device), - torch.tensor([[model_i],[model_j]], dtype=torch.int64, device=device), - model.atomic_numbers.to(device)) + torch.tensor( + [[model_i], [model_j]], dtype=torch.int64, device=device + ), + model.atomic_numbers.to(device), + ) if isinstance(bessels, tuple): bessels = bessels[0] # newer versions return (bessels, cutoffs) - R = torch.tanh(model.interactions[1].density_fn(bessels)**2).numpy(force=True) - spl = CubicSpline(r, R[:,0], bc_type=spline_bc_type) + R = torch.tanh(model.interactions[1].density_fn(bessels) ** 2).numpy( + force=True + ) + spl = CubicSpline(r, R[:, 0], bc_type=spline_bc_type) A1_spline_values.append(spl(r).tolist()) A1_spline_derivs.append(spl.derivative()(r).tolist()) - output['A1_spline_h'] = float(h) - output['A1_spline_values'] = A1_spline_values - output['A1_spline_derivs'] = A1_spline_derivs + output["A1_spline_h"] = float(h) + output["A1_spline_values"] = A1_spline_values + output["A1_spline_derivs"] = A1_spline_derivs A1_weights = [] - num_eta = [sum([l==ll for ll in Phi1_l]) for l in range(l_max+1)] + num_eta = [sum([l == ll for ll in Phi1_l]) for l in range(l_max + 1)] A1_linear = linear_simplify(model.interactions[1].linear) - for l in range(l_max+1): - w_linear = A1_linear.weight_view_for_instruction(l).numpy(force=True) / np.sqrt(num_eta[l]*num_channels) + for l in range(l_max + 1): + w_linear = A1_linear.weight_view_for_instruction(l).numpy(force=True) / np.sqrt( + num_eta[l] * num_channels + ) if not A1_scaled: w_linear /= model.interactions[1].avg_num_neighbors w_linear = np.reshape(w_linear, (num_eta[l], num_channels, num_channels)) @@ -431,7 +526,7 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): irreps_out = Irreps("0e") # extract U in sparse format U = [[]] - for corr in range(1,4): + for corr in range(1, 4): # get U matrix for this correlation order try: U_matrix = U_matrix_real(irreps_in, irreps_out, corr, use_cueq_cg=False)[1] @@ -442,7 +537,7 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): # extract sparse U for this correlation order U_sparse = [{} for _ in range(num_nu)] j = 0 - for lm_list in itertools.product(range((l_max+1)**2), repeat=corr): + for lm_list in itertools.product(range((l_max + 1) ** 2), repeat=corr): for nu in range(num_nu): if abs(U_matrix[j]) > 1e-12: lm_tuple_sorted = tuple(sorted(lm_list)) @@ -454,9 +549,23 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): # extract weights from model # warning: slightly odd order of the contractions weights due to reverse countdown W = [[]] - W.append(model.products[1].symmetric_contractions.contractions[0].weights[1].numpy(force=True)) - W.append(model.products[1].symmetric_contractions.contractions[0].weights[0].numpy(force=True)) - W.append(model.products[1].symmetric_contractions.contractions[0].weights_max.numpy(force=True)) + W.append( + model.products[1] + .symmetric_contractions.contractions[0] + .weights[1] + .numpy(force=True) + ) + W.append( + model.products[1] + .symmetric_contractions.contractions[0] + .weights[0] + .numpy(force=True) + ) + W.append( + model.products[1] + .symmetric_contractions.contractions[0] + .weights_max.numpy(force=True) + ) # combine U and W into polynomial-like terms for recursive evaluator C = {} for i, a in enumerate(atomic_numbers): @@ -464,58 +573,75 @@ def compute_lelm1lm2(le,l1,m1,l2,m2): C[i] = {} for k in range(num_channels): P_ik = {} - for corr in range(1,4): + for corr in range(1, 4): for nu in range(len(U[corr])): - for key,value in U[corr][nu].items(): + for key, value in U[corr][nu].items(): if key not in P_ik.keys(): P_ik[key] = 0.0 - P_ik[key] += float(W[corr][model_i,nu,k]) * value + P_ik[key] += float(W[corr][model_i, nu, k]) * value C[i][k] = list(P_ik.values()) M = [list(key) for key in P_ik.keys()] - output['M1_weights'] = C - output['M1_monomials'] = M + output["M1_weights"] = C + output["M1_monomials"] = M ### ----- H2 ----- logging.info("H2") - weights_to_fuse = model.interactions[1].linear_up.weight_view_for_instruction(0).numpy(force=True) / np.sqrt(num_channels) + weights_to_fuse = model.interactions[1].linear_up.weight_view_for_instruction( + 0 + ).numpy(force=True) / np.sqrt(num_channels) weights_to_fuse_rank = np.linalg.matrix_rank(weights_to_fuse) if weights_to_fuse_rank < num_channels: - raise RuntimeError('ERROR: fusing weights have too low rank {weights_to_fuse_rank} < {num_channels}') + raise RuntimeError( + "ERROR: fusing weights have too low rank {weights_to_fuse_rank} < {num_channels}" + ) # H2 weights for H1 H2_weights_for_H1 = [] for i, a in enumerate(atomic_numbers): model_i = model.atomic_numbers.tolist().index(a) - w = model.interactions[1].skip_tp.weight_view_for_instruction(0)[:,model_i,:].numpy(force=True) - H2_weights_for_H1.append(w / np.sqrt(len(model.atomic_numbers)*num_channels)) + w = ( + model.interactions[1] + .skip_tp.weight_view_for_instruction(0)[:, model_i, :] + .numpy(force=True) + ) + H2_weights_for_H1.append(w / np.sqrt(len(model.atomic_numbers) * num_channels)) H2_weights_for_H1[i] = np.linalg.inv(weights_to_fuse) @ H2_weights_for_H1[i] H2_weights_for_H1[i] = H2_weights_for_H1[i].flatten().tolist() output["H2_weights_for_H1"] = H2_weights_for_H1 # H2 weights for M1 output["H2_weights_for_M1"] = ( - model.products[1].linear.weight.numpy(force=True) / np.sqrt(num_channels)).tolist() + model.products[1].linear.weight.numpy(force=True) / np.sqrt(num_channels) + ).tolist() ### ----- READOUTS ----- # linear readout weights_to_fuse = np.reshape( - model.interactions[1].linear_up.weight.numpy(force=True) / np.sqrt(num_channels), - [L_max+1,num_channels,num_channels]) - readout_1_weights = model.readouts[0].linear.weight.numpy(force=True) / np.sqrt(num_channels) - readout_1_weights = np.linalg.inv(weights_to_fuse[0,:,:]) @ readout_1_weights - output['readout_1_weights'] = (readout_1_weights * model.scale_shift.scale.item()).tolist() + model.interactions[1].linear_up.weight.numpy(force=True) + / np.sqrt(num_channels), + [L_max + 1, num_channels, num_channels], + ) + readout_1_weights = model.readouts[0].linear.weight.numpy(force=True) / np.sqrt( + num_channels + ) + readout_1_weights = np.linalg.inv(weights_to_fuse[0, :, :]) @ readout_1_weights + output["readout_1_weights"] = ( + readout_1_weights * model.scale_shift.scale.item() + ).tolist() # nonlinear readout - #output["mlp_hidden_layers"] = 16 // TODO + # output["mlp_hidden_layers"] = 16 // TODO output["readout_2_weights_1"] = ( - torch.reshape( - model.readouts[1].linear_1.weight, - (num_channels,16) - ).T.numpy(force=True).flatten() / np.sqrt(num_channels) - ).tolist() + torch.reshape(model.readouts[1].linear_1.weight, (num_channels, 16)) + .T.numpy(force=True) + .flatten() + / np.sqrt(num_channels) + ).tolist() output["readout_2_weights_2"] = ( - model.readouts[1].linear_2.weight.numpy(force=True).flatten() / np.sqrt(16) - * model.scale_shift.scale.item()).tolist() + model.readouts[1].linear_2.weight.numpy(force=True).flatten() + / np.sqrt(16) + * model.scale_shift.scale.item() + ).tolist() output["readout_2_scale_factor"] = model.readouts[1].non_linearity.acts[0].cst return output diff --git a/symmetrix/test/test_cubic_spline.py b/symmetrix/test/test_cubic_spline.py index 2304111..0cc24e1 100644 --- a/symmetrix/test/test_cubic_spline.py +++ b/symmetrix/test/test_cubic_spline.py @@ -1,8 +1,6 @@ import numpy as np -import os from pytest import approx, raises from scipy.interpolate import CubicSpline -import sys import symmetrix @@ -18,11 +16,10 @@ def test_invalid_input(): def test_evaluate(): - # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f = np.sin(r) * r**2 * (r_cut-r)**2 + f = np.sin(r) * r**2 * (r_cut - r) ** 2 # create splines scipy_spl = CubicSpline(r, f) d = scipy_spl.derivative()(r) @@ -37,15 +34,16 @@ def test_evaluate(): for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: spl.evaluate(r) - assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate.") + assert str(exception.value).startswith( + "Out of bounds in CubicSpline::evaluate." + ) def test_evaluate_deriv(): - # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f = np.sin(r) * r**2 * (r_cut-r)**2 + f = np.sin(r) * r**2 * (r_cut - r) ** 2 # create splines scipy_spl = CubicSpline(r, f) d = scipy_spl.derivative()(r) @@ -62,14 +60,16 @@ def test_evaluate_deriv(): for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv(r) - assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv.") + assert str(exception.value).startswith( + "Out of bounds in CubicSpline::evaluate_deriv." + ) -def test_evaluate_deriv_divided(): +def test_evaluate_deriv_divided(): # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f = np.sin(r) * r**2 * (r_cut-r)**2 + f = np.sin(r) * r**2 * (r_cut - r) ** 2 # create splines scipy_spl = CubicSpline(r, f) d = scipy_spl.derivative()(r) @@ -86,4 +86,6 @@ def test_evaluate_deriv_divided(): for r in [-1.0, 0.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv_divided(r) - assert str(exception.value).startswith("Out of bounds in CubicSpline::evaluate_deriv_divided.") + assert str(exception.value).startswith( + "Out of bounds in CubicSpline::evaluate_deriv_divided." + ) diff --git a/symmetrix/test/test_cubic_spline_kokkos.py b/symmetrix/test/test_cubic_spline_kokkos.py index e46bc04..68c3522 100644 --- a/symmetrix/test/test_cubic_spline_kokkos.py +++ b/symmetrix/test/test_cubic_spline_kokkos.py @@ -1,9 +1,6 @@ import numpy as np -import os from pytest import approx, raises from scipy.interpolate import CubicSpline -import sys -import pytest import symmetrix @@ -22,11 +19,10 @@ def test_invalid_input(): def test_evaluate(): - # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f = np.sin(r) * r**2 * (r_cut-r)**2 + f = np.sin(r) * r**2 * (r_cut - r) ** 2 # create splines scipy_spl = CubicSpline(r, f) d = scipy_spl.derivative()(r) @@ -41,14 +37,16 @@ def test_evaluate(): for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: spl.evaluate(r) - assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate.") + assert str(exception.value).startswith( + "Out of bounds in CubicSplineKokkos::evaluate." + ) -def test_evaluate_deriv(): +def test_evaluate_deriv(): # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f = np.sin(r) * r**2 * (r_cut-r)**2 + f = np.sin(r) * r**2 * (r_cut - r) ** 2 # create splines scipy_spl = CubicSpline(r, f) d = scipy_spl.derivative()(r) @@ -65,14 +63,16 @@ def test_evaluate_deriv(): for r in [-1.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv(r) - assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate_deriv.") + assert str(exception.value).startswith( + "Out of bounds in CubicSplineKokkos::evaluate_deriv." + ) -def test_evaluate_deriv_divided(): +def test_evaluate_deriv_divided(): # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f = np.sin(r) * r**2 * (r_cut-r)**2 + f = np.sin(r) * r**2 * (r_cut - r) ** 2 # create splines scipy_spl = CubicSpline(r, f) d = scipy_spl.derivative()(r) @@ -89,4 +89,6 @@ def test_evaluate_deriv_divided(): for r in [-1.0, 0.0, r_cut + 1e-12, 9.0]: with raises(ValueError) as exception: _ = spl.evaluate_deriv_divided(r) - assert str(exception.value).startswith("Out of bounds in CubicSplineKokkos::evaluate_deriv_divided.") + assert str(exception.value).startswith( + "Out of bounds in CubicSplineKokkos::evaluate_deriv_divided." + ) diff --git a/symmetrix/test/test_cubic_spline_set.py b/symmetrix/test/test_cubic_spline_set.py index a484e04..0eafa55 100644 --- a/symmetrix/test/test_cubic_spline_set.py +++ b/symmetrix/test/test_cubic_spline_set.py @@ -1,20 +1,17 @@ import numpy as np -import os from pytest import approx from scipy.interpolate import CubicSpline -import sys import symmetrix def test_evaluate(): - # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f1 = np.sin(r) * r**2 * (r_cut-r)**2 - f2 = np.cos(r) * r**2 * (r_cut-r)**2 - f3 = np.sin(r)*np.cos(r) * r**2 * (r_cut-r)**2 + f1 = np.sin(r) * r**2 * (r_cut - r) ** 2 + f2 = np.cos(r) * r**2 * (r_cut - r) ** 2 + f3 = np.sin(r) * np.cos(r) * r**2 * (r_cut - r) ** 2 # create splines spl1 = CubicSpline(r, f1) spl2 = CubicSpline(r, f2) @@ -22,26 +19,26 @@ def test_evaluate(): d1 = spl1.derivative()(r) d2 = spl2.derivative()(r) d3 = spl3.derivative()(r) - spl_set = symmetrix.CubicSplineSet(h, [f1,f2,f3], [d1,d2,d3]) + spl_set = symmetrix.CubicSplineSet(h, [f1, f2, f3], [d1, d2, d3]) # test equivalence r = np.linspace(0, r_cut, 1000, endpoint=False) - f1,f2,f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) + f1, f2, f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) values = np.zeros(3) for i, ri in enumerate(r): spl_set.evaluate(ri, values) - f1[i],f2[i],f3[i] = values + f1[i], f2[i], f3[i] = values assert f1 == approx(spl1(r)) assert f2 == approx(spl2(r)) assert f3 == approx(spl3(r)) -def test_evaluate_derivs(): +def test_evaluate_derivs(): # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f1 = np.sin(r) * r**2 * (r_cut-r)**2 - f2 = np.cos(r) * r**2 * (r_cut-r)**2 - f3 = np.sin(r)*np.cos(r) * r**2 * (r_cut-r)**2 + f1 = np.sin(r) * r**2 * (r_cut - r) ** 2 + f2 = np.cos(r) * r**2 * (r_cut - r) ** 2 + f3 = np.sin(r) * np.cos(r) * r**2 * (r_cut - r) ** 2 # create splines spl1 = CubicSpline(r, f1) spl2 = CubicSpline(r, f2) @@ -49,17 +46,17 @@ def test_evaluate_derivs(): d1 = spl1.derivative()(r) d2 = spl2.derivative()(r) d3 = spl3.derivative()(r) - spl_set = symmetrix.CubicSplineSet(h, [f1,f2,f3], [d1,d2,d3]) + spl_set = symmetrix.CubicSplineSet(h, [f1, f2, f3], [d1, d2, d3]) # test equivalence r = np.linspace(0, r_cut, 1000, endpoint=False) - f1,f2,f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) - d1,d2,d3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) + f1, f2, f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) + d1, d2, d3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) values = np.zeros(3) derivs = np.zeros(3) for i, ri in enumerate(r): spl_set.evaluate_derivs(ri, values, derivs) - f1[i],f2[i],f3[i] = values - d1[i],d2[i],d3[i] = derivs + f1[i], f2[i], f3[i] = values + d1[i], d2[i], d3[i] = derivs assert f1 == approx(spl1(r)) assert d1 == approx(spl1.derivative()(r)) assert f2 == approx(spl2(r)) diff --git a/symmetrix/test/test_cubic_spline_set_kokkos.py b/symmetrix/test/test_cubic_spline_set_kokkos.py index 1c4b0cc..43a4bf4 100644 --- a/symmetrix/test/test_cubic_spline_set_kokkos.py +++ b/symmetrix/test/test_cubic_spline_set_kokkos.py @@ -1,23 +1,20 @@ import numpy as np -import os -import pytest from pytest import approx from scipy.interpolate import CubicSpline -import sys import symmetrix + if not symmetrix._kokkos_is_initialized(): symmetrix._init_kokkos() - + def test_evaluate(): - # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f1 = np.sin(r) * r**2 * (r_cut-r)**2 - f2 = np.cos(r) * r**2 * (r_cut-r)**2 - f3 = np.sin(r)*np.cos(r) * r**2 * (r_cut-r)**2 + f1 = np.sin(r) * r**2 * (r_cut - r) ** 2 + f2 = np.cos(r) * r**2 * (r_cut - r) ** 2 + f3 = np.sin(r) * np.cos(r) * r**2 * (r_cut - r) ** 2 # create splines spl1 = CubicSpline(r, f1) spl2 = CubicSpline(r, f2) @@ -25,27 +22,26 @@ def test_evaluate(): d1 = spl1.derivative()(r) d2 = spl2.derivative()(r) d3 = spl3.derivative()(r) - spl_set = symmetrix.CubicSplineSetKokkos(h, [f1,f2,f3], [d1,d2,d3]) + spl_set = symmetrix.CubicSplineSetKokkos(h, [f1, f2, f3], [d1, d2, d3]) # test equivalence r = np.linspace(0, r_cut, 1000, endpoint=False) - f1,f2,f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) + f1, f2, f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) values = np.zeros(3) for i, ri in enumerate(r): spl_set.evaluate(ri, values) - f1[i],f2[i],f3[i] = values + f1[i], f2[i], f3[i] = values assert f1 == approx(spl1(r)) assert f2 == approx(spl2(r)) assert f3 == approx(spl3(r)) def test_evaluate_derivs(): - # generate data r_cut = 5 r, h = np.linspace(0, r_cut, 20, retstep=True) - f1 = np.sin(r) * r**2 * (r_cut-r)**2 - f2 = np.cos(r) * r**2 * (r_cut-r)**2 - f3 = np.sin(r)*np.cos(r) * r**2 * (r_cut-r)**2 + f1 = np.sin(r) * r**2 * (r_cut - r) ** 2 + f2 = np.cos(r) * r**2 * (r_cut - r) ** 2 + f3 = np.sin(r) * np.cos(r) * r**2 * (r_cut - r) ** 2 # create splines spl1 = CubicSpline(r, f1) spl2 = CubicSpline(r, f2) @@ -53,17 +49,17 @@ def test_evaluate_derivs(): d1 = spl1.derivative()(r) d2 = spl2.derivative()(r) d3 = spl3.derivative()(r) - spl_set = symmetrix.CubicSplineSetKokkos(h, [f1,f2,f3], [d1,d2,d3]) + spl_set = symmetrix.CubicSplineSetKokkos(h, [f1, f2, f3], [d1, d2, d3]) # test equivalence r = np.linspace(0, r_cut, 1000, endpoint=False) - f1,f2,f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) - d1,d2,d3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) + f1, f2, f3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) + d1, d2, d3 = (np.zeros(len(r)), np.zeros(len(r)), np.zeros(len(r))) values = np.zeros(3) derivs = np.zeros(3) for i, ri in enumerate(r): spl_set.evaluate_derivs(ri, values, derivs) - f1[i],f2[i],f3[i] = values - d1[i],d2[i],d3[i] = derivs + f1[i], f2[i], f3[i] = values + d1[i], d2[i], d3[i] = derivs assert f1 == approx(spl1(r)) assert d1 == approx(spl1.derivative()(r)) assert f2 == approx(spl2(r)) diff --git a/symmetrix/test/test_lammpslib.py b/symmetrix/test/test_lammpslib.py index 955b0d5..861dde0 100755 --- a/symmetrix/test/test_lammpslib.py +++ b/symmetrix/test/test_lammpslib.py @@ -8,94 +8,113 @@ from symmetrix import Symmetrix except ModuleNotFoundError as exc: if "No module named 'symmetrix.symmetrix'" in str(exc): - raise RuntimeError("Can't import symmetrix.symmetrix, probably need to run pytest in venv " - "and install version to be tested with " - "'(cd /path/to/repo && python3 -m pip install -e .)'") from exc + raise RuntimeError( + "Can't import symmetrix.symmetrix, probably need to run pytest in venv " + "and install version to be tested with " + "'(cd /path/to/repo && python3 -m pip install -e .)'" + ) from exc else: raise try: - import lammps + import lammps # noqa: F401 except ImportError: pytest.skip("No lammps python package available", allow_module_level=True) from ase.calculators.lammpslib import LAMMPSlib def test_lammpslib_map(model_cache): - rng = np.random.default_rng(3) - calc_symmetrix = Symmetrix(model_cache["mace-mp-0b3-medium-1-8.json"]) species = ["H", "O"] for map_type in [None, "yes", "array", "hash"]: - lammpslib_kwargs = dict(lmpcmds = [ "pair_style symmetrix/mace", - f"pair_coeff * * {model_cache['mace-mp-0b3-medium-1-8.json']} " + " ".join(species)], - atom_types = {sp: sp_i+1 for sp_i, sp in enumerate(species)}, - lammps_header = ['units metal', 'atom_style atomic', 'atom_modify sort 0 0']) + lammpslib_kwargs = dict( + lmpcmds=[ + "pair_style symmetrix/mace", + f"pair_coeff * * {model_cache['mace-mp-0b3-medium-1-8.json']} " + + " ".join(species), + ], + atom_types={sp: sp_i + 1 for sp_i, sp in enumerate(species)}, + lammps_header=["units metal", "atom_style atomic", "atom_modify sort 0 0"], + ) if map_type is None: calc_lammpslib = LAMMPSlib(**lammpslib_kwargs) - atoms = Atoms('OH', cell=[4, 2, 2], positions=[[0, 0, 0], [2, 0, 0]], pbc=[True] * 3) + atoms = Atoms( + "OH", cell=[4, 2, 2], positions=[[0, 0, 0], [2, 0, 0]], pbc=[True] * 3 + ) atoms.calc = calc_lammpslib with pytest.raises(Exception): _ = atoms.get_potential_energy() continue - lammpslib_kwargs['lammps_header'].append("atom_modify map " + map_type) + lammpslib_kwargs["lammps_header"].append("atom_modify map " + map_type) compare_lammpslib_symmetrix(calc_symmetrix, lammpslib_kwargs) def test_lammpslib_default_header(model_cache): - rng = np.random.default_rng(3) - calc_symmetrix = Symmetrix(model_cache["mace-mp-0b3-medium-1-8.json"]) species = ["H", "O"] # use default lammps_header, should do "atom_modify map array sort 0 0" - lammpslib_kwargs = dict(lmpcmds = [ "pair_style symmetrix/mace", - f"pair_coeff * * {model_cache['mace-mp-0b3-medium-1-8.json']} " + " ".join(species)], - atom_types = {sp: sp_i+1 for sp_i, sp in enumerate(species)}) + lammpslib_kwargs = dict( + lmpcmds=[ + "pair_style symmetrix/mace", + f"pair_coeff * * {model_cache['mace-mp-0b3-medium-1-8.json']} " + + " ".join(species), + ], + atom_types={sp: sp_i + 1 for sp_i, sp in enumerate(species)}, + ) compare_lammpslib_symmetrix(calc_symmetrix, lammpslib_kwargs) # confirm that without "atom_modify sort 0 0" it fails - lammpslib_kwargs = dict(lmpcmds = [ "pair_style symmetrix/mace", - f"pair_coeff * * {model_cache['mace-mp-0b3-medium-1-8.json']} " + " ".join(species)], - atom_types = {sp: sp_i+1 for sp_i, sp in enumerate(species)}, - lammps_header = ['units metal', 'atom_style atomic', 'atom_modify map yes']) + lammpslib_kwargs = dict( + lmpcmds=[ + "pair_style symmetrix/mace", + f"pair_coeff * * {model_cache['mace-mp-0b3-medium-1-8.json']} " + + " ".join(species), + ], + atom_types={sp: sp_i + 1 for sp_i, sp in enumerate(species)}, + lammps_header=["units metal", "atom_style atomic", "atom_modify map yes"], + ) with pytest.raises(AssertionError): compare_lammpslib_symmetrix(calc_symmetrix, lammpslib_kwargs) def compare_lammpslib_symmetrix(calc_symmetrix, lammpslib_kwargs): - calc_lammpslib = LAMMPSlib(**lammpslib_kwargs) - - atoms_list = [] - for n1 in range(4, 2, -1): # when size gets smaller over different calls LAMMPSlib fails without 'sort 0 0' - for n2 in range(2, 4): - atoms = Atoms('OH', cell=[4, 2, 2], positions=[[0, 0, 0], [2, 0, 0]], pbc=[True] * 3) - atoms *= (1, n1 + 1, n2 + 1) - rng = np.random.default_rng(5) - atoms.rattle(rng=rng) - - rng.shuffle(atoms.numbers) - atoms_list.append(atoms) - - E_lammpslib = [] - F_lammpslib = [] - for atoms in atoms_list: - atoms.calc = calc_lammpslib - E_lammpslib.append(atoms.get_potential_energy()) - F_lammpslib.append(atoms.get_forces()) - - E_symmetrix = [] - F_symmetrix = [] - for atoms in atoms_list: - atoms.calc = calc_symmetrix - E_symmetrix.append(atoms.get_potential_energy()) - F_symmetrix.append(atoms.get_forces()) - - assert np.allclose(E_symmetrix, E_lammpslib) - for F_s, F_l in zip(F_symmetrix, F_lammpslib): - assert np.allclose(F_s, F_l, rtol=1e-2) + calc_lammpslib = LAMMPSlib(**lammpslib_kwargs) + + atoms_list = [] + for n1 in range( + 4, 2, -1 + ): # when size gets smaller over different calls LAMMPSlib fails without 'sort 0 0' + for n2 in range(2, 4): + atoms = Atoms( + "OH", cell=[4, 2, 2], positions=[[0, 0, 0], [2, 0, 0]], pbc=[True] * 3 + ) + atoms *= (1, n1 + 1, n2 + 1) + rng = np.random.default_rng(5) + atoms.rattle(rng=rng) + + rng.shuffle(atoms.numbers) + atoms_list.append(atoms) + + E_lammpslib = [] + F_lammpslib = [] + for atoms in atoms_list: + atoms.calc = calc_lammpslib + E_lammpslib.append(atoms.get_potential_energy()) + F_lammpslib.append(atoms.get_forces()) + + E_symmetrix = [] + F_symmetrix = [] + for atoms in atoms_list: + atoms.calc = calc_symmetrix + E_symmetrix.append(atoms.get_potential_energy()) + F_symmetrix.append(atoms.get_forces()) + + assert np.allclose(E_symmetrix, E_lammpslib) + for F_s, F_l in zip(F_symmetrix, F_lammpslib): + assert np.allclose(F_s, F_l, rtol=1e-2) diff --git a/symmetrix/test/test_mace.py b/symmetrix/test/test_mace.py index d03f84c..c837f99 100644 --- a/symmetrix/test/test_mace.py +++ b/symmetrix/test/test_mace.py @@ -3,28 +3,31 @@ import numpy as np import os import pytest -import sys from urllib.request import urlretrieve import symmetrix if not os.path.exists("MACE-OFF23_small-1-8.json"): - urlretrieve("https://www.dropbox.com/scl/fi/7rz3vh5mhacofp5w2u8cu/MACE-OFF23_small-1-8.json?rlkey=rubpqlut6uhjf4w9pej54alu7&st=w23fcknx&dl=1", - "MACE-OFF23_small-1-8.json") + urlretrieve( + "https://www.dropbox.com/scl/fi/7rz3vh5mhacofp5w2u8cu/MACE-OFF23_small-1-8.json?rlkey=rubpqlut6uhjf4w9pej54alu7&st=w23fcknx&dl=1", + "MACE-OFF23_small-1-8.json", + ) if not os.path.exists("mace-mp-0b3-medium-1-8.json"): - urlretrieve("https://www.dropbox.com/scl/fi/3lydfgta1lijymq98pgal/mace-mp-0b3-medium-1-8.json?rlkey=7wofp9gznqt5b3wmk5ybbj76z&st=w7cd09x6&dl=1", - "mace-mp-0b3-medium-1-8.json") + urlretrieve( + "https://www.dropbox.com/scl/fi/3lydfgta1lijymq98pgal/mace-mp-0b3-medium-1-8.json?rlkey=7wofp9gznqt5b3wmk5ybbj76z&st=w7cd09x6&dl=1", + "mace-mp-0b3-medium-1-8.json", + ) model = "mace-off-small" -#model = "mace-off-medium" -#model = "mace-off-large" -#model = "mace-mp-small" -#model = "mace-mp-medium" -#model = "mace-mp-large" -#model = "mace-mpa-medium" -#model = "mace-mp-0b3-medium" -#model = "mace-omat-0-medium" +# model = "mace-off-medium" +# model = "mace-off-large" +# model = "mace-mp-small" +# model = "mace-mp-medium" +# model = "mace-mp-large" +# model = "mace-mpa-medium" +# model = "mace-mp-0b3-medium" +# model = "mace-omat-0-medium" kokkos = False if kokkos: @@ -33,7 +36,7 @@ symmetrix._init_kokkos() else: MACE = symmetrix.MACE - + # load model if model == "mace-off-small": evaluator = MACE("MACE-OFF23_small-1-8.json") @@ -55,20 +58,21 @@ evaluator = MACE("mace-omat-0-medium-1-8.json") # prepare for tests -atoms = ase.Atoms('OHH', - positions=[[0.0, -2.0, 0.0], - [1.0, 0.0, 0.0], - [0.0, 1.0, 0.0]]) +atoms = ase.Atoms("OHH", positions=[[0.0, -2.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) ase_atomic_numbers = atoms.get_atomic_numbers().tolist() mace_atomic_numbers = evaluator.atomic_numbers -i_list, j_list, r, xyz = neighbor_list('ijdD', atoms, 5.0) -xyz = -xyz # TODO: why exactly is this necessary!? -num_nodes = np.max(i_list)+1 -node_types = [mace_atomic_numbers.index(ase_atomic_numbers[i]) for i in range(num_nodes)] +i_list, j_list, r, xyz = neighbor_list("ijdD", atoms, 5.0) +xyz = -xyz # TODO: why exactly is this necessary!? +num_nodes = np.max(i_list) + 1 +node_types = [ + mace_atomic_numbers.index(ase_atomic_numbers[i]) for i in range(num_nodes) +] num_neigh = [sum(j_list == i) for i in range(num_nodes)] neigh_types = [mace_atomic_numbers.index(ase_atomic_numbers[j]) for j in j_list] evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r +) + def numerical_gradient(f, x): h = 1e-4 @@ -76,24 +80,24 @@ def numerical_gradient(f, x): for i in range(len(x)): x[i] += h fp = f(x) - x[i] -= 2*h + x[i] -= 2 * h fm = f(x) x[i] += h - grad[i] = (fp-fm)/(2*h) + grad[i] = (fp - fm) / (2 * h) f(x) # reverts side effects of applying f(x+h) return grad -#def test_Y(): - # TODO +# def test_Y(): +# TODO -#def test_R0(): +# def test_R0(): # # ### FORWARD # evaluator.compute_R0(num_nodes, node_types, num_neigh, neigh_types, r) # if model == "mace-off-small": -# R0_sum = 21.20435772182849 +# R0_sum = 21.20435772182849 # elif model == "mace-off-medium": # R0_sum = -21.480436737535587 # elif model == "mace-off-large": @@ -105,7 +109,7 @@ def numerical_gradient(f, x): # elif model == "mace-mp-large": # R0_sum = -132.59276393205272 # elif model == "mace-mpa-medium": -# R0_sum = -39.38919552881708 +# R0_sum = -39.38919552881708 # elif model == "mace-mp-0b3-medium": # R0_sum = 72.27732642745596 # elif model == "mace-omat-0-medium": @@ -117,7 +121,6 @@ def numerical_gradient(f, x): def test_A0(): - ### FORWARD if model == "mace-off-small": A0_sum = 0.5140705628937071 @@ -134,7 +137,7 @@ def test_A0(): elif model == "mace-mpa-medium": A0_sum = -7.991043151569456 elif model == "mace-mp-0b3-medium": - A0_sum = 10.329510160691235 + A0_sum = 10.329510160691235 elif model == "mace-omat-0-medium": A0_sum = -21.88004976955362 evaluator.compute_R0(num_nodes, node_types, num_neigh, neigh_types, r) @@ -143,16 +146,21 @@ def test_A0(): assert sum(evaluator.A0) == pytest.approx(A0_sum, abs=1e-4) def f(xyz_flat): - r = np.sqrt(np.sum(np.reshape(xyz_flat*xyz_flat, [xyz_flat.size//3,3]), axis=1)) + r = np.sqrt( + np.sum(np.reshape(xyz_flat * xyz_flat, [xyz_flat.size // 3, 3]), axis=1) + ) evaluator.compute_R0(num_nodes, node_types, num_neigh, neigh_types, r) evaluator.compute_Y(xyz_flat) evaluator.compute_A0(num_nodes, node_types, num_neigh, neigh_types) return np.sum(evaluator.A0) + # compute analytical forces f(xyz.flatten()) evaluator.A0_adj = np.ones(len(evaluator.A0)) evaluator.node_forces = np.zeros(xyz.size) - evaluator.reverse_A0(num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r) + evaluator.reverse_A0( + num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r + ) node_forces = evaluator.node_forces # compare with numerical forces node_forces_num = -numerical_gradient(f, xyz.flatten()) @@ -160,7 +168,6 @@ def f(xyz_flat): def test_A0_scaled(): - # store A0_unscaled evaluator.compute_A0(num_nodes, node_types, num_neigh, neigh_types) A0_unscaled = np.array(evaluator.A0) @@ -183,42 +190,51 @@ def test_A0_scaled(): elif model == "mace-mp-0b3-medium": A0_sum = 9.440126671650006 elif model == "mace-omat-0-medium": - A0_sum = -20.379975414199357 + A0_sum = -20.379975414199357 evaluator.compute_A0_scaled(num_nodes, node_types, num_neigh, neigh_types, r) assert sum(evaluator.A0) == pytest.approx(A0_sum, abs=1e-4) ### REVERSE # define f(xyz)=sum(A0), used to test dA0/dxyz def f(xyz_flat): - r = np.sqrt(np.sum(np.reshape(xyz_flat*xyz_flat, [xyz_flat.size//3,3]), axis=1)) + r = np.sqrt( + np.sum(np.reshape(xyz_flat * xyz_flat, [xyz_flat.size // 3, 3]), axis=1) + ) evaluator.A0 = A0_unscaled evaluator.compute_A0_scaled(num_nodes, node_types, num_neigh, neigh_types, r) return np.sum(evaluator.A0) + # compute analytical forces f(xyz.flatten()) evaluator.node_forces = np.zeros(len(xyz.flatten())) evaluator.A0_adj = np.ones(len(evaluator.A0)) - evaluator.reverse_A0_scaled(num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r) + evaluator.reverse_A0_scaled( + num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r + ) node_forces = evaluator.node_forces # compare with numerical forces node_forces_num = -numerical_gradient(f, xyz.flatten()) assert np.allclose(node_forces, node_forces_num, rtol=1e-4, atol=1e-6) + # define f(A0)=sum(A0_scaled), used to test dA0_scaled/dA0 def f(A0): evaluator.A0 = A0 evaluator.compute_A0_scaled(num_nodes, node_types, num_neigh, neigh_types, r) return np.sum(evaluator.A0) + # compute analytical gradient f(A0_unscaled) evaluator.A0_adj = np.ones(len(evaluator.A0)) - evaluator.reverse_A0_scaled(num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r) + evaluator.reverse_A0_scaled( + num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r + ) g = evaluator.A0_adj # compare with numerical gradient g_num = numerical_gradient(f, A0_unscaled) assert np.allclose(g, g_num, rtol=1e-4, atol=1e-6) -def test_M0(): +def test_M0(): ### FORWARD if model == "mace-off-small": M0_sum = -0.22431147864930753 @@ -247,6 +263,7 @@ def f(A0): evaluator.A0 = A0 evaluator.compute_M0(num_nodes, node_types) return np.sum(evaluator.M0) + # compute analytical gradient f(evaluator.A0) evaluator.M0_adj = np.ones(len(evaluator.M0)) @@ -258,7 +275,6 @@ def f(A0): def test_H1(): - ### FORWARD if model == "mace-off-small": H1_sum = 3.030576444359059 @@ -288,6 +304,7 @@ def f(M0): evaluator.M0 = M0 evaluator.compute_H1(num_nodes) return np.sum(evaluator.H1) + # compute analytical gradient f(evaluator.M0) evaluator.H1_adj = np.ones(len(evaluator.H1)) @@ -299,10 +316,9 @@ def f(M0): def test_R1(): - ### FORWARD if model == "mace-off-small": - R1_sum = 35.985092839298346 + R1_sum = 35.985092839298346 elif model == "mace-off-medium": R1_sum = -118.43918364181435 elif model == "mace-off-large": @@ -327,12 +343,11 @@ def test_R1(): def test_Phi1(): - ### FORWARD if model == "mace-off-small": Phi1_sum = 1.414641743513294 elif model == "mace-off-medium": - Phi1_sum = -11.473894785398894 + Phi1_sum = -11.473894785398894 elif model == "mace-off-large": Phi1_sum = -5.03354101033154 elif model == "mace-mp-small": @@ -342,7 +357,7 @@ def test_Phi1(): elif model == "mace-mp-large": Phi1_sum = -8.217173578481036 elif model == "mace-mpa-medium": - Phi1_sum = 47.18438578544733 + Phi1_sum = 47.18438578544733 elif model == "mace-mp-0b3-medium": Phi1_sum = 12.53997938705195 elif model == "mace-omat-0-medium": @@ -355,11 +370,14 @@ def test_Phi1(): ### REVERSE # define f(xyz)=sum(Phi1), used to test dPhi1/dxyz def f(xyz_flat): - r = np.sqrt(np.sum(np.reshape(xyz_flat*xyz_flat, [xyz_flat.size//3,3]), axis=1)) + r = np.sqrt( + np.sum(np.reshape(xyz_flat * xyz_flat, [xyz_flat.size // 3, 3]), axis=1) + ) evaluator.compute_R1(num_nodes, node_types, num_neigh, neigh_types, r) evaluator.compute_Y(xyz_flat) evaluator.compute_Phi1(num_nodes, num_neigh, j_list) return np.sum(evaluator.Phi1) + # compute analytical forces f(xyz.flatten()) evaluator.Phi1_adj = np.ones(len(evaluator.Phi1)) @@ -369,11 +387,13 @@ def f(xyz_flat): # compare with numerical forces node_forces_num = -numerical_gradient(f, xyz.flatten()) assert np.allclose(node_forces, node_forces_num, rtol=1e-4, atol=1e-6) + # define f(H1)=sum(Phi1), used to test dPhi1/dH1 def f(H1): evaluator.H1 = H1 evaluator.compute_Phi1(num_nodes, num_neigh, j_list) return np.sum(evaluator.Phi1) + # compute analytical gradient f(evaluator.H1) evaluator.Phi1_adj = np.ones(len(evaluator.Phi1)) @@ -386,10 +406,9 @@ def f(H1): def test_A1(): - ### FORWARD if model == "mace-off-small": - A1_sum = 0.2854230509340355 + A1_sum = 0.2854230509340355 elif model == "mace-off-medium": A1_sum = 1.0645744085666808 elif model == "mace-off-large": @@ -399,7 +418,7 @@ def test_A1(): elif model == "mace-mp-medium": A1_sum = 4.865592770455143 elif model == "mace-mp-large": - A1_sum = -2.9562330922326314 + A1_sum = -2.9562330922326314 elif model == "mace-mpa-medium": A1_sum = -0.9342584160946092 elif model == "mace-mp-0b3-medium": @@ -415,6 +434,7 @@ def f(Phi1): evaluator.Phi1 = Phi1 evaluator.compute_A1(num_nodes) return np.sum(evaluator.A1) + # compute analytical gradient f(evaluator.Phi1) evaluator.A1_adj = np.ones(len(evaluator.A1)) @@ -424,15 +444,15 @@ def f(Phi1): g_num = numerical_gradient(f, evaluator.Phi1) assert np.allclose(g, g_num, rtol=1e-4, atol=1e-6) -def test_A1_scaled(): +def test_A1_scaled(): # store A1_unscaled evaluator.compute_A1(num_nodes) A1_unscaled = np.array(evaluator.A1) ### FORWARD if model == "mace-off-small": - A1_sum = 0.2854230509340355 + A1_sum = 0.2854230509340355 elif model == "mace-off-medium": A1_sum = 1.0645744085666808 elif model == "mace-off-large": @@ -444,7 +464,7 @@ def test_A1_scaled(): elif model == "mace-mp-large": A1_sum = -2.9306745971117203 elif model == "mace-mpa-medium": - A1_sum = -0.7384077299498335 + A1_sum = -0.7384077299498335 elif model == "mace-mp-0b3-medium": A1_sum = 5.1042645420231745 elif model == "mace-omat-0-medium": @@ -455,35 +475,44 @@ def test_A1_scaled(): ### REVERSE # define f(xyz)=sum(A1), used to test dA1/dxyz def f(xyz_flat): - r = np.sqrt(np.sum(np.reshape(xyz_flat*xyz_flat, [xyz_flat.size//3,3]), axis=1)) + r = np.sqrt( + np.sum(np.reshape(xyz_flat * xyz_flat, [xyz_flat.size // 3, 3]), axis=1) + ) evaluator.A1 = A1_unscaled evaluator.compute_A1_scaled(num_nodes, node_types, num_neigh, neigh_types, r) return np.sum(evaluator.A1) + # compute analytical forces f(xyz.flatten()) evaluator.node_forces = np.zeros(len(xyz.flatten())) evaluator.A1_adj = np.ones(len(evaluator.A1)) - evaluator.reverse_A1_scaled(num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r) + evaluator.reverse_A1_scaled( + num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r + ) node_forces = evaluator.node_forces # compare with numerical forces node_forces_num = -numerical_gradient(f, xyz.flatten()) assert np.allclose(node_forces, node_forces_num, rtol=1e-4, atol=1e-6) + # define f(A1)=sum(A1_scaled), used to test dA1_scaled/dA1 def f(A1): evaluator.A1 = A1 evaluator.compute_A1_scaled(num_nodes, node_types, num_neigh, neigh_types, r) return np.sum(evaluator.A1) + # compute analytical gradient f(A1_unscaled) evaluator.A1_adj = np.ones(len(evaluator.A1)) - evaluator.reverse_A1_scaled(num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r) + evaluator.reverse_A1_scaled( + num_nodes, node_types, num_neigh, neigh_types, xyz.flatten(), r + ) g = evaluator.A1_adj # compare with numerical gradient g_num = numerical_gradient(f, A1_unscaled) assert np.allclose(g, g_num, rtol=1e-4, atol=1e-6) -def test_M1(): +def test_M1(): ### FORWARD if model == "mace-off-small": M1_sum = -0.05168003793600238 @@ -496,9 +525,9 @@ def test_M1(): elif model == "mace-mp-medium": M1_sum = -0.41865181071865554 elif model == "mace-mp-large": - M1_sum = -0.8190295085983332 + M1_sum = -0.8190295085983332 elif model == "mace-mp-large": - M1_sum = -0.8190295085983332 + M1_sum = -0.8190295085983332 elif model == "mace-mpa-medium": M1_sum = 0.8098763169906235 elif model == "mace-mp-0b3-medium": @@ -524,6 +553,7 @@ def f(A1): evaluator.A1 = A1 evaluator.compute_M1(num_nodes, node_types) return np.sum(evaluator.M1) + # compute analytical gradient f(evaluator.A1) evaluator.M1_adj = np.ones(len(evaluator.M1)) @@ -535,7 +565,6 @@ def f(A1): def test_H2(): - ### FORWARD if model == "mace-off-small": H2_sum = 0.8844547738634937 @@ -575,6 +604,7 @@ def f(H1): evaluator.H1 = H1 evaluator.compute_H2(num_nodes, node_types) return np.sum(evaluator.H2) + # compute analytical gradient f(evaluator.H1) evaluator.H2_adj = np.ones(len(evaluator.H2)) @@ -583,11 +613,13 @@ def f(H1): # compare with numerical gradient g_num = numerical_gradient(f, evaluator.H1) assert np.allclose(g, g_num, rtol=1e-4, atol=1e-6) + # define f(M1)=sum(H2), used to test dH2/dM1 def f(M1): evaluator.M1 = M1 evaluator.compute_H2(num_nodes, node_types) return np.sum(evaluator.H2) + # compute analytical gradient f(evaluator.M1) evaluator.H2_adj = np.ones(len(evaluator.H2)) @@ -599,7 +631,6 @@ def f(M1): def test_readouts(): - ### FORWARD if model == "mace-off-small": readout = -2071.839005822318 @@ -608,9 +639,9 @@ def test_readouts(): elif model == "mace-off-large": readout = -2074.154047738083 elif model == "mace-mp-small": - readout = -5.998523387682857 + readout = -5.998523387682857 elif model == "mace-mp-medium": - readout = -5.355659696240375 + readout = -5.355659696240375 elif model == "mace-mp-large": readout = -5.766392385834781 elif model == "mace-mpa-medium": @@ -620,7 +651,8 @@ def test_readouts(): elif model == "mace-omat-0-medium": readout = -5.356825090124245 evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r + ) evaluator.node_energies = np.zeros(num_nodes) evaluator.compute_readouts(num_nodes, node_types) assert sum(evaluator.node_energies) == pytest.approx(readout, abs=1e-4) @@ -632,18 +664,21 @@ def f(H1): evaluator.node_energies = np.zeros(num_nodes) evaluator.compute_readouts(num_nodes, node_types) return sum(evaluator.node_energies) + # compute analytical gradient f(evaluator.H1) g = evaluator.H1_adj # compare with numerical gradient g_num = numerical_gradient(f, evaluator.H1) assert np.allclose(g, g_num, rtol=1e-4, atol=1e-6) + # define readout as function of H2 def f(H2): evaluator.H2 = H2 evaluator.node_energies = np.zeros(num_nodes) evaluator.compute_readouts(num_nodes, node_types) return sum(evaluator.node_energies) + # compute analytical gradient f(evaluator.H2) g = evaluator.H2_adj @@ -653,10 +688,10 @@ def f(H2): def test_compute_node_energies_forces(): - ### FORWARD evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r + ) e = sum(evaluator.node_energies) f = evaluator.node_forces if model == "mace-off-small": @@ -666,13 +701,13 @@ def test_compute_node_energies_forces(): elif model == "mace-off-large": exact_e = -2074.154047738083 elif model == "mace-mp-small": - exact_e = -5.998523387682857 + exact_e = -5.998523387682857 elif model == "mace-mp-medium": exact_e = -5.355659696240375 elif model == "mace-mp-large": exact_e = -5.766392385834781 elif model == "mace-mpa-medium": - exact_e = -5.089426502695993 + exact_e = -5.089426502695993 elif model == "mace-mp-0b3-medium": exact_e = -4.920488393882309 elif model == "mace-omat-0-medium": @@ -687,45 +722,60 @@ def test_compute_node_energies_forces(): for i in range(num_nodes): for j in range(num_neigh[i]): for w in range(3): - xyz[ij,w] += h - r[ij] = np.sqrt(xyz[ij,:].dot(xyz[ij,:])) + xyz[ij, w] += h + r[ij] = np.sqrt(xyz[ij, :].dot(xyz[ij, :])) evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, + node_types, + num_neigh, + j_list, + neigh_types, + xyz.flatten(), + r, + ) ep = sum(evaluator.node_energies) - xyz[ij,w] -= 2*h - r[ij] = np.sqrt(xyz[ij,:].dot(xyz[ij,:])) + xyz[ij, w] -= 2 * h + r[ij] = np.sqrt(xyz[ij, :].dot(xyz[ij, :])) evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, + node_types, + num_neigh, + j_list, + neigh_types, + xyz.flatten(), + r, + ) em = sum(evaluator.node_energies) - xyz[ij,w] += h - r[ij] = np.sqrt(xyz[ij,:].dot(xyz[ij,:])) - f_num[3*ij+w] = -(ep-em)/(2*h) + xyz[ij, w] += h + r[ij] = np.sqrt(xyz[ij, :].dot(xyz[ij, :])) + f_num[3 * ij + w] = -(ep - em) / (2 * h) ij += 1 assert np.allclose(f, f_num, rtol=1e-4, atol=1e-6) def test_zbl(): - evaluator = MACE("mace-mp-0b3-medium-1-8.json") - atoms = ase.Atoms('OHH', - positions=[[0.0, -0.5, 0.0], - [0.5, 0.0, 0.0], - [0.0, 0.5, 0.0]]) + atoms = ase.Atoms( + "OHH", positions=[[0.0, -0.5, 0.0], [0.5, 0.0, 0.0], [0.0, 0.5, 0.0]] + ) ase_atomic_numbers = atoms.get_atomic_numbers().tolist() mace_atomic_numbers = evaluator.atomic_numbers - i_list, j_list, r, xyz = neighbor_list('ijdD', atoms, 5.0) - xyz = -xyz # TODO: why exactly is this necessary!? - num_nodes = np.max(i_list)+1 - node_types = [mace_atomic_numbers.index(ase_atomic_numbers[i]) for i in range(num_nodes)] + i_list, j_list, r, xyz = neighbor_list("ijdD", atoms, 5.0) + xyz = -xyz # TODO: why exactly is this necessary!? + num_nodes = np.max(i_list) + 1 + node_types = [ + mace_atomic_numbers.index(ase_atomic_numbers[i]) for i in range(num_nodes) + ] num_neigh = [sum(j_list == i) for i in range(num_nodes)] neigh_types = [mace_atomic_numbers.index(ase_atomic_numbers[j]) for j in j_list] ### FORWARD evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r + ) e = sum(evaluator.node_energies) f = evaluator.node_forces - exact_e = -5.003106904473648 + exact_e = -5.003106904473648 assert e == pytest.approx(exact_e, abs=1e-3) ### REVERSE @@ -736,18 +786,32 @@ def test_zbl(): for i in range(num_nodes): for j in range(num_neigh[i]): for w in range(3): - xyz[ij,w] += h - r[ij] = np.sqrt(xyz[ij,:].dot(xyz[ij,:])) + xyz[ij, w] += h + r[ij] = np.sqrt(xyz[ij, :].dot(xyz[ij, :])) evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, + node_types, + num_neigh, + j_list, + neigh_types, + xyz.flatten(), + r, + ) ep = sum(evaluator.node_energies) - xyz[ij,w] -= 2*h - r[ij] = np.sqrt(xyz[ij,:].dot(xyz[ij,:])) + xyz[ij, w] -= 2 * h + r[ij] = np.sqrt(xyz[ij, :].dot(xyz[ij, :])) evaluator.compute_node_energies_forces( - num_nodes, node_types, num_neigh, j_list, neigh_types, xyz.flatten(), r) + num_nodes, + node_types, + num_neigh, + j_list, + neigh_types, + xyz.flatten(), + r, + ) em = sum(evaluator.node_energies) - xyz[ij,w] += h - r[ij] = np.sqrt(xyz[ij,:].dot(xyz[ij,:])) - f_num[3*ij+w] = -(ep-em)/(2*h) + xyz[ij, w] += h + r[ij] = np.sqrt(xyz[ij, :].dot(xyz[ij, :])) + f_num[3 * ij + w] = -(ep - em) / (2 * h) ij += 1 assert np.allclose(f, f_num, rtol=1e-4, atol=1e-6) diff --git a/symmetrix/test/test_multilayer_perceptron.py b/symmetrix/test/test_multilayer_perceptron.py index b20ab06..d383ceb 100644 --- a/symmetrix/test/test_multilayer_perceptron.py +++ b/symmetrix/test/test_multilayer_perceptron.py @@ -1,22 +1,23 @@ -import os import numpy as np import pytest -import sys from symmetrix import MultilayerPerceptron -def test_evaluate(): +def test_evaluate(): ### One hidden layer, one output shape = [3, 6, 1] - w0 = np.random.random([3,6]).T - w1 = np.random.random([6,1]).T + w0 = np.random.random([3, 6]).T + w1 = np.random.random([6, 1]).T weights = [w0.flatten(), w1.flatten()] scale = 1.1 + def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w1.dot(act(w0.dot(x))) + # create MultilayerPerceptron mlp = MultilayerPerceptron(shape, weights, scale) # test @@ -25,16 +26,19 @@ def MLP(x): ### Three hidden layers, one output shape = [3, 8, 8, 4, 1] - w0 = np.random.random([3,8]).T - w1 = np.random.random([8,8]).T - w2 = np.random.random([8,4]).T - w3 = np.random.random([4,1]).T + w0 = np.random.random([3, 8]).T + w1 = np.random.random([8, 8]).T + w2 = np.random.random([8, 4]).T + w3 = np.random.random([4, 1]).T weights = [w0.flatten(), w1.flatten(), w2.flatten(), w3.flatten()] scale = 1.3 + def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w3.dot(act(w2.dot(act(w1.dot(act(w0.dot(x))))))) + # create MultilayerPerceptron mlp = MultilayerPerceptron(shape, weights, scale) # test @@ -43,28 +47,31 @@ def MLP(x): ### Two hidden layers, three outputs shape = [5, 8, 4, 3] - w0 = np.random.random([5,8]).T - w1 = np.random.random([8,4]).T - w2 = np.random.random([4,3]).T + w0 = np.random.random([5, 8]).T + w1 = np.random.random([8, 4]).T + w2 = np.random.random([4, 3]).T weights = [w0.flatten(), w1.flatten(), w2.flatten()] scale = 0.9 + def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w2.dot(act(w1.dot(act(w0.dot(x))))) + # create MultilayerPerceptron mlp = MultilayerPerceptron(shape, weights, scale) # test x = np.random.random(5) assert mlp.evaluate(x) == pytest.approx(MLP(x)) -def test_evaluate_gradient(): +def test_evaluate_gradient(): ### One hidden layer, one output shape = [3, 6, 1] - w0 = np.random.random([3,6]).T - w1 = np.random.random([6,1]).T - mlp = MultilayerPerceptron(shape, [w0.flatten(),w1.flatten()], 1.2) + w0 = np.random.random([3, 6]).T + w1 = np.random.random([6, 1]).T + mlp = MultilayerPerceptron(shape, [w0.flatten(), w1.flatten()], 1.2) # test value x = np.random.random(3) f, g = mlp.evaluate_gradient(x) @@ -75,19 +82,21 @@ def test_evaluate_gradient(): for i in range(x.size): x[i] += d fp = mlp.evaluate(x)[0] - x[i] -= 2*d + x[i] -= 2 * d fm = mlp.evaluate(x)[0] x[i] += d - g_numerical[i] = (fp-fm) / (2*d) + g_numerical[i] = (fp - fm) / (2 * d) assert np.allclose(g, g_numerical, rtol=1e-4, atol=1e-6) ### Three hidden layers, one output shape = [3, 8, 8, 4, 1] - w0 = np.random.random([3,8]).T - w1 = np.random.random([8,8]).T - w2 = np.random.random([8,4]).T - w3 = np.random.random([4,1]).T - mlp = MultilayerPerceptron(shape, [w0.flatten(),w1.flatten(),w2.flatten(),w3.flatten()], 1.2) + w0 = np.random.random([3, 8]).T + w1 = np.random.random([8, 8]).T + w2 = np.random.random([8, 4]).T + w3 = np.random.random([4, 1]).T + mlp = MultilayerPerceptron( + shape, [w0.flatten(), w1.flatten(), w2.flatten(), w3.flatten()], 1.2 + ) # test value x = np.random.random(3) f, g = mlp.evaluate_gradient(x) @@ -98,66 +107,66 @@ def test_evaluate_gradient(): for i in range(x.size): x[i] += d fp = mlp.evaluate(x)[0] - x[i] -= 2*d + x[i] -= 2 * d fm = mlp.evaluate(x)[0] x[i] += d - g_numerical[i] = (fp-fm) / (2*d) + g_numerical[i] = (fp - fm) / (2 * d) assert np.allclose(g, g_numerical, rtol=1e-4, atol=1e-6) ### Two hidden layers, three outputs shape = [5, 8, 4, 3] - w0 = np.random.random([5,8]).T - w1 = np.random.random([8,4]).T - w2 = np.random.random([4,3]).T - mlp = MultilayerPerceptron(shape, [w0.flatten(),w1.flatten(),w2.flatten()], 0.9) + w0 = np.random.random([5, 8]).T + w1 = np.random.random([8, 4]).T + w2 = np.random.random([4, 3]).T + mlp = MultilayerPerceptron(shape, [w0.flatten(), w1.flatten(), w2.flatten()], 0.9) # test value x = np.random.random(5) f, g = mlp.evaluate_gradient(x) assert f == pytest.approx(mlp.evaluate(x)) # test gradient - g_numerical = np.empty(3*5) + g_numerical = np.empty(3 * 5) d = 1e-3 for i in range(3): for j in range(x.size): x[j] += d fp = mlp.evaluate(x)[i] - x[j] -= 2*d + x[j] -= 2 * d fm = mlp.evaluate(x)[i] x[j] += d - g_numerical[i*5+j] = (fp-fm) / (2*d) + g_numerical[i * 5 + j] = (fp - fm) / (2 * d) assert np.allclose(g, g_numerical, rtol=1e-4, atol=1e-6) -def test_evaluate_batch(): +def test_evaluate_batch(): ### Two hidden layers, three outputs shape = [5, 8, 4, 3] - w0 = np.random.random([5,8]).T - w1 = np.random.random([8,4]).T - w2 = np.random.random([4,3]).T + w0 = np.random.random([5, 8]).T + w1 = np.random.random([8, 4]).T + w2 = np.random.random([4, 3]).T weights = [w0.flatten(), w1.flatten(), w2.flatten()] mlp = MultilayerPerceptron(shape, weights, 0.9) # test batched evaluation - x = np.random.random([100,5]) - f = np.empty([100,3]) + x = np.random.random([100, 5]) + f = np.empty([100, 3]) for i in range(x.shape[0]): - f[i,:] = mlp.evaluate(x[i,:]) - assert f.flatten() == pytest.approx(mlp.evaluate_batch(x.flatten(),100)) + f[i, :] = mlp.evaluate(x[i, :]) + assert f.flatten() == pytest.approx(mlp.evaluate_batch(x.flatten(), 100)) -def test_evaluate_gradient_batch(): +def test_evaluate_gradient_batch(): ### Two hidden layers, three outputs shape = [5, 8, 4, 3] - w0 = np.random.random([5,8]).T - w1 = np.random.random([8,4]).T - w2 = np.random.random([4,3]).T + w0 = np.random.random([5, 8]).T + w1 = np.random.random([8, 4]).T + w2 = np.random.random([4, 3]).T weights = [w0.flatten(), w1.flatten(), w2.flatten()] mlp = MultilayerPerceptron(shape, weights, 0.9) # test batched evaluation with gradient - x = np.random.random([100,5]) - f = np.empty([100,3]) - g = np.empty([100,15]) + x = np.random.random([100, 5]) + f = np.empty([100, 3]) + g = np.empty([100, 15]) for i in range(x.shape[0]): - f[i,:], g[i,:] = mlp.evaluate_gradient(x[i,:]) - f1,g1 = mlp.evaluate_gradient_batch(x.flatten(),100) + f[i, :], g[i, :] = mlp.evaluate_gradient(x[i, :]) + f1, g1 = mlp.evaluate_gradient_batch(x.flatten(), 100) assert f.flatten() == pytest.approx(f1) assert g.flatten() == pytest.approx(g1) diff --git a/symmetrix/test/test_multilayer_perceptron_kokkos.py b/symmetrix/test/test_multilayer_perceptron_kokkos.py index 22ba463..dbf9c0d 100644 --- a/symmetrix/test/test_multilayer_perceptron_kokkos.py +++ b/symmetrix/test/test_multilayer_perceptron_kokkos.py @@ -1,7 +1,5 @@ -import os import numpy as np import pytest -import sys import symmetrix @@ -13,59 +11,65 @@ else: MultilayerPerceptron = symmetrix.MultilayerPerceptron -def test_evaluate(): +def test_evaluate(): ### Batch size: 11, input dimension: 8, hidden layers: 1 - x = np.random.random([11,8]) + x = np.random.random([11, 8]) shape = [8, 32, 1] - w0 = np.random.random([8,32]).T - w1 = np.random.random([32,1]).T + w0 = np.random.random([8, 32]).T + w1 = np.random.random([32, 1]).T weights = [w0.flatten(), w1.flatten()] scale = 0.9 # compute result mlp = MultilayerPerceptron(shape, weights, scale) f1 = np.zeros(x.shape[0]) mlp.evaluate(x, f1) + # compute reference result def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w1.dot(act(w0.dot(x))) + f2 = np.zeros(x.shape[0]) for i in range(x.shape[0]): - f2[i] = MLP(x[i,:]).item() + f2[i] = MLP(x[i, :]).item() assert f1 == pytest.approx(f2) ### Batch size: 32, input dimension: 5, hidden layers: 3 - x = np.random.random([32,5]) + x = np.random.random([32, 5]) shape = [5, 8, 8, 4, 1] - w0 = np.random.random([5,8]).T - w1 = np.random.random([8,8]).T - w2 = np.random.random([8,4]).T - w3 = np.random.random([4,1]).T + w0 = np.random.random([5, 8]).T + w1 = np.random.random([8, 8]).T + w2 = np.random.random([8, 4]).T + w3 = np.random.random([4, 1]).T weights = [w0.flatten(), w1.flatten(), w2.flatten(), w3.flatten()] scale = 1.3 # compute result mlp = MultilayerPerceptron(shape, weights, scale) f1 = np.zeros(x.shape[0]) mlp.evaluate(x, f1) + # compute reference result def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w3.dot(act(w2.dot(act(w1.dot(act(w0.dot(x))))))) + f2 = np.zeros(x.shape[0]) for i in range(x.shape[0]): - f2[i] = MLP(x[i,:]).item() + f2[i] = MLP(x[i, :]).item() assert f1 == pytest.approx(f2) -def test_evaluate_gradient_batch(): +def test_evaluate_gradient_batch(): ### Batch size: 11, input dimension: 8, hidden layers: 1 - x = np.random.random([11,8]) + x = np.random.random([11, 8]) shape = [8, 32, 1] - w0 = np.random.random([8,32]).T - w1 = np.random.random([32,1]).T + w0 = np.random.random([8, 32]).T + w1 = np.random.random([32, 1]).T weights = [w0.flatten(), w1.flatten()] scale = 0.9 # compute result with gradient @@ -73,32 +77,35 @@ def test_evaluate_gradient_batch(): f1 = np.empty(x.shape[0]) g1 = np.empty(x.shape) mlp.evaluate_gradient(x, f1, g1) + # compute reference result with gradient def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w1.dot(act(w0.dot(x))) + f2 = np.empty(x.shape[0]) g2 = np.empty(x.shape) for i in range(x.shape[0]): - f2[i] = MLP(x[i,:]).item() + f2[i] = MLP(x[i, :]).item() for j in range(x.shape[1]): - x[i,j] += 1e-3 - fp = MLP(x[i,:]).item() - x[i,j] -= 2e-3 - fm = MLP(x[i,:]).item() - x[i,j] += 1e-3 - g2[i,j] = (fp-fm) / 2e-3 + x[i, j] += 1e-3 + fp = MLP(x[i, :]).item() + x[i, j] -= 2e-3 + fm = MLP(x[i, :]).item() + x[i, j] += 1e-3 + g2[i, j] = (fp - fm) / 2e-3 assert f1 == pytest.approx(f2) assert g1 == pytest.approx(g2) ### Batch size: 32, input dimension: 5, hidden layers: 3 - x = np.random.random([32,5]) + x = np.random.random([32, 5]) shape = [5, 8, 8, 4, 1] - w0 = np.random.random([5,8]).T - w1 = np.random.random([8,8]).T - w2 = np.random.random([8,4]).T - w3 = np.random.random([4,1]).T + w0 = np.random.random([5, 8]).T + w1 = np.random.random([8, 8]).T + w2 = np.random.random([8, 4]).T + w3 = np.random.random([4, 1]).T weights = [w0.flatten(), w1.flatten(), w2.flatten(), w3.flatten()] scale = 1.3 # compute result with gradient @@ -106,21 +113,24 @@ def MLP(x): f1 = np.empty(x.shape[0]) g1 = np.empty(x.shape) mlp.evaluate_gradient(x, f1, g1) + # compute reference result with gradient def act(x): - return scale*x/(1+np.exp(-x)) + return scale * x / (1 + np.exp(-x)) + def MLP(x): return w3.dot(act(w2.dot(act(w1.dot(act(w0.dot(x))))))) + f2 = np.empty(x.shape[0]) g2 = np.empty(x.shape) for i in range(x.shape[0]): - f2[i] = MLP(x[i,:]).item() + f2[i] = MLP(x[i, :]).item() for j in range(x.shape[1]): - x[i,j] += 1e-3 - fp = MLP(x[i,:]).item() - x[i,j] -= 2e-3 - fm = MLP(x[i,:]).item() - x[i,j] += 1e-3 - g2[i,j] = (fp-fm) / 2e-3 + x[i, j] += 1e-3 + fp = MLP(x[i, :]).item() + x[i, j] -= 2e-3 + fm = MLP(x[i, :]).item() + x[i, j] += 1e-3 + g2[i, j] = (fp - fm) / 2e-3 assert f1 == pytest.approx(f2) assert g1 == pytest.approx(g2) diff --git a/symmetrix/test/test_multivariate_polynomial.py b/symmetrix/test/test_multivariate_polynomial.py index 3bc14ac..7f9216c 100644 --- a/symmetrix/test/test_multivariate_polynomial.py +++ b/symmetrix/test/test_multivariate_polynomial.py @@ -1,26 +1,24 @@ -import os import numpy as np import pytest -import sys from symmetrix import MultivariatePolynomial -def test_evaluate(): +def test_evaluate(): num_variables = 5 coefficients = np.random.rand(4) - monomials = [[0], [0,1], [1,1,2], [0,1,2,4]] + monomials = [[0], [0, 1], [1, 1, 2], [0, 1, 2, 4]] poly = MultivariatePolynomial(num_variables, coefficients, monomials) X = np.random.rand(num_variables) F0 = poly.evaluate(X) F1 = poly.evaluate_simple(X) assert F0 == pytest.approx(F1) -def test_evaluate_gradient(): +def test_evaluate_gradient(): num_variables = 8 coefficients = np.random.rand(6) - monomials = [[7], [0,1], [1,3], [1,1,2], [0,1,2,4], [4,5,6,7]] + monomials = [[7], [0, 1], [1, 3], [1, 1, 2], [0, 1, 2, 4], [4, 5, 6, 7]] poly = MultivariatePolynomial(num_variables, coefficients, monomials) X = np.random.rand(num_variables) F0, G0 = poly.evaluate_gradient(X) @@ -28,16 +26,16 @@ def test_evaluate_gradient(): assert F0 == pytest.approx(F1) assert np.allclose(G0, G1) -def test_evaluate_batch(): +def test_evaluate_batch(): num_variables = 5 coefficients = np.random.rand(4) - monomials = [[0], [0,1], [1,1,2], [0,1,2,4]] + monomials = [[0], [0, 1], [1, 1, 2], [0, 1, 2, 4]] poly = MultivariatePolynomial(num_variables, coefficients, monomials) - X = np.random.rand(3*num_variables) + X = np.random.rand(3 * num_variables) F0, G0 = poly.evaluate_gradient(X[:num_variables]) - F1, G1 = poly.evaluate_gradient(X[num_variables:2*num_variables]) - F2, G2 = poly.evaluate_gradient(X[2*num_variables:]) + F1, G1 = poly.evaluate_gradient(X[num_variables : 2 * num_variables]) + F2, G2 = poly.evaluate_gradient(X[2 * num_variables :]) F, G = poly.evaluate_batch(X, 3) assert np.allclose(F, [F0, F1, F2]) assert np.allclose(G, np.concatenate([G0, G1, G2])) diff --git a/symmetrix/test/test_symmetrix_calc.py b/symmetrix/test/test_symmetrix_calc.py index 62330c0..bf3a3cd 100755 --- a/symmetrix/test/test_symmetrix_calc.py +++ b/symmetrix/test/test_symmetrix_calc.py @@ -5,7 +5,6 @@ import os import time -from pathlib import Path import numpy as np @@ -16,20 +15,22 @@ from symmetrix import Symmetrix except ModuleNotFoundError as exc: if "No module named 'symmetrix.symmetrix'" in str(exc): - raise RuntimeError("Can't import symmetrix.symmetrix, probably need to run pytest in venv " - "and install version to be tested with " - "'(cd /path/to/repo && python3 -m pip install -e .)'") from exc + raise RuntimeError( + "Can't import symmetrix.symmetrix, probably need to run pytest in venv " + "and install version to be tested with " + "'(cd /path/to/repo && python3 -m pip install -e .)'" + ) from exc else: raise try: import mace from mace.calculators import MACECalculator - from mace.tools.utils import get_cache_dir from mace.calculators.foundations_models import download_mace_mp_checkpoint -except ImportError as exc: +except ImportError: mace = None + @pytest.fixture(scope="module") def mace_foundation_model(tmp_path_factory): if mace is None: @@ -40,7 +41,7 @@ def mace_foundation_model(tmp_path_factory): cache_dir = tmp_path_factory.mktemp("mace_cache") xdg_cache_home = os.environ.get("XDG_CACHE_HOME") os.environ["XDG_CACHE_HOME"] = str(cache_dir) - downloaded_model = download_mace_mp_checkpoint('small-omat-0') + downloaded_model = download_mace_mp_checkpoint("small-omat-0") if xdg_cache_home is None: del os.environ["XDG_CACHE_HOME"] else: @@ -51,7 +52,7 @@ def mace_foundation_model(tmp_path_factory): @pytest.mark.parametrize("use_kokkos", [True, False]) def test_calc_caching(model_cache, use_kokkos): - atoms = Atoms('O', cell=[2] * 3, pbc=[True] * 3) + atoms = Atoms("O", cell=[2] * 3, pbc=[True] * 3) atoms *= 4 rng = np.random.default_rng(5) atoms.rattle(rng=rng) @@ -60,11 +61,11 @@ def test_calc_caching(model_cache, use_kokkos): atoms.calc = calc t0 = time.time() - E = atoms.get_potential_energy() + atoms.get_potential_energy() dt_E = time.time() - t0 t0 = time.time() - E = atoms.get_forces() + atoms.get_forces() dt_F = time.time() - t0 # without perturbation, forces are from cache @@ -73,7 +74,7 @@ def test_calc_caching(model_cache, use_kokkos): atoms.positions[0, 0] += 0.1 t0 = time.time() - E = atoms.get_forces() + atoms.get_forces() dt_F_pert = time.time() - t0 # with perturbation, forces have to be recomputed @@ -82,12 +83,12 @@ def test_calc_caching(model_cache, use_kokkos): @pytest.mark.parametrize("use_kokkos", [True, False]) def test_symmetrix_calc_finite_diff(model_cache, use_kokkos): - atoms = Atoms('O', cell=[2] * 3, pbc=[True] * 3) + atoms = Atoms("O", cell=[2] * 3, pbc=[True] * 3) atoms *= 2 rng = np.random.default_rng(5) atoms.rattle(rng=rng) - F = np.eye(3) + 0.01 * rng.normal(size=(3,3)) + F = np.eye(3) + 0.01 * rng.normal(size=(3, 3)) atoms.set_cell(atoms.cell @ F, True) print("pre-converted") @@ -98,12 +99,12 @@ def test_symmetrix_calc_finite_diff(model_cache, use_kokkos): @pytest.mark.skipif(mace is None, reason="mace-torch is not available") @pytest.mark.parametrize("use_kokkos", [True, False]) def test_mace_onthefly_calc_finite_diff(mace_foundation_model, use_kokkos): - atoms = Atoms('O', cell=[2] * 3, pbc=[True] * 3) + atoms = Atoms("O", cell=[2] * 3, pbc=[True] * 3) atoms *= 2 rng = np.random.default_rng(5) atoms.rattle(rng=rng) - F = np.eye(3) + 0.01 * rng.normal(size=(3,3)) + F = np.eye(3) + 0.01 * rng.normal(size=(3, 3)) atoms.set_cell(atoms.cell @ F, True) print("converted on-the-fly") @@ -114,12 +115,12 @@ def test_mace_onthefly_calc_finite_diff(mace_foundation_model, use_kokkos): @pytest.mark.skipif(mace is None, reason="mace-torch is not available") @pytest.mark.parametrize("use_kokkos", [True, False]) def test_symmetrix_vs_pytorch(mace_foundation_model, use_kokkos): - atoms = Atoms('O', cell=[2] * 3, pbc=[True] * 3) + atoms = Atoms("O", cell=[2] * 3, pbc=[True] * 3) atoms *= 2 rng = np.random.default_rng(5) atoms.rattle(rng=rng) - F = np.eye(3) + 0.01 * rng.normal(size=(3,3)) + F = np.eye(3) + 0.01 * rng.normal(size=(3, 3)) atoms.set_cell(atoms.cell @ F, True) atoms_s = atoms.copy() @@ -132,7 +133,9 @@ def test_symmetrix_vs_pytorch(mace_foundation_model, use_kokkos): atoms_p.calc = calc_torch # are these in fact reasonable accuracies? - assert np.allclose(atoms_s.get_potential_energy(), atoms_p.get_potential_energy(), atol=0.001) + assert np.allclose( + atoms_s.get_potential_energy(), atoms_p.get_potential_energy(), atol=0.001 + ) assert np.allclose(atoms_s.get_forces(), atoms_p.get_forces(), atol=0.002) assert np.allclose(atoms_s.get_stress(), atoms_p.get_stress(), atol=0.003) @@ -153,7 +156,7 @@ def do_grad_test(atoms, calc, check, ax=None, label=None, plot_factor=1.0): passed_f = True F_scaling = None for dx_exp in np.arange(1.0, 5.1, 0.5): - dx = 0.1 ** dx_exp + dx = 0.1**dx_exp #### forces #### atoms.positions = p0 @@ -170,17 +173,19 @@ def do_grad_test(atoms, calc, check, ax=None, label=None, plot_factor=1.0): E_m = atoms.get_potential_energy() F_fd[i_a, j_a] = -(E_p - E_m) / (2 * dx) F_err = np.linalg.norm(F0 - F_fd) - print(f"F {dx:6f} {F0_norm:10.6e} {F_err:10.6e} {F_err / F0_norm:10.6e} {F_err / F0_norm / (dx ** 2):10.6e}") + print( + f"F {dx:6f} {F0_norm:10.6e} {F_err:10.6e} {F_err / F0_norm:10.6e} {F_err / F0_norm / (dx ** 2):10.6e}" + ) f_data.append([dx, F_err]) # force error only shows expected 2nd order scaling for dx = 0.1 ** 1, 0.1 ** 1.5 if F_scaling is None and dx_exp >= 1.99: # F_err / F0_norm < F_scaling * dx ** 2 - F_scaling = 2.5 * F_err / F0_norm / (dx ** 2) + F_scaling = 2.5 * F_err / F0_norm / (dx**2) if F_scaling is not None and dx_exp < 4.01: - print("test forces", dx_exp, dx, F_err / F0_norm, "= 1.99: # S_err / S0_norm < S_scaling * dx ** 2 - S_scaling = 1.5 * S_err / S0_norm / (dx ** 2) + S_scaling = 1.5 * S_err / S0_norm / (dx**2) if S_scaling is not None and dx_exp < 4.01: - print("test stress", dx_exp, dx, S_err / S0_norm, " Date: Mon, 27 Jul 2026 17:19:46 -0400 Subject: [PATCH 25/29] targeted safe_globals. --- symmetrix/source/symmetrix/extract_mace_data.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/symmetrix/source/symmetrix/extract_mace_data.py b/symmetrix/source/symmetrix/extract_mace_data.py index dfba74a..3128e87 100755 --- a/symmetrix/source/symmetrix/extract_mace_data.py +++ b/symmetrix/source/symmetrix/extract_mace_data.py @@ -6,7 +6,8 @@ import numpy as np from scipy.interpolate import CubicSpline -from e3nn.o3 import Irreps, Linear +with torch.serialization.safe_globals([slice]): + from e3nn.o3 import Irreps, Linear from mace.modules.radial import ZBLBasis from mace.tools.cg import U_matrix_real from mace.tools.scripts_utils import remove_pt_head From 18f2430dfb158fcd0ae85f06a24986de2d8cd66b Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Mon, 27 Jul 2026 17:24:36 -0400 Subject: [PATCH 26/29] prevent ruff from removing. --- symmetrix/source/symmetrix/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/symmetrix/source/symmetrix/__init__.py b/symmetrix/source/symmetrix/__init__.py index 60116d5..fbbb68c 100644 --- a/symmetrix/source/symmetrix/__init__.py +++ b/symmetrix/source/symmetrix/__init__.py @@ -6,3 +6,4 @@ _sym = importlib.import_module(".symmetrix", __name__) _sym.__all__ = [n for n in vars(_sym) if not (n.startswith("__") and n.endswith("__"))] from .symmetrix import * +from .calculator import Symmetrix as Symmetrix From 8bf1144891ac09ea359e8b3a3b9f28de2ec7db6a Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Mon, 27 Jul 2026 17:52:29 -0400 Subject: [PATCH 27/29] fix typo. --- symmetrix/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/symmetrix/README.md b/symmetrix/README.md index 7a6cbd1..9d4b6df 100644 --- a/symmetrix/README.md +++ b/symmetrix/README.md @@ -42,7 +42,7 @@ pip install ".[mace]" Then use: ``` -symmetrix_extract_mace my-mace.model --atomic-numbers 1 8 +symmetrix_extract_mace --model my-mace.model --atomic-numbers 1 8 ``` from the command line to extract a `.json` file from a Torch-based model. The result will be `my-mace-1-8.json`, and this model is only suitable From 1f4b22320ab98f6265a411112868fa11139bd150 Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Mon, 27 Jul 2026 17:54:59 -0400 Subject: [PATCH 28/29] require torch 2.5 for targeted safe_globals. --- .github/workflows/ci.yaml | 8 ++++++-- symmetrix/pyproject.toml | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 987b5dc..b6aa1ff 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -101,7 +101,11 @@ jobs: strategy: fail-fast: false matrix: - mace-torch: ["mace-torch==0.3.10", "mace-torch"] + include: + - mace-torch: "mace-torch==0.3.10" + torch: "torch==2.5.*" + - mace-torch: "mace-torch" + torch: "torch" steps: - name: Clone repo uses: actions/checkout@v6 @@ -119,7 +123,7 @@ jobs: run: | uv venv --clear source .venv/bin/activate - uv pip install ${{ matrix.mace-torch }} ./symmetrix + uv pip install "${{ matrix.torch }}" "${{ matrix.mace-torch }}" "./symmetrix[mace]" deactivate - name: Test model extraction run: | diff --git a/symmetrix/pyproject.toml b/symmetrix/pyproject.toml index c2b502a..ae69563 100644 --- a/symmetrix/pyproject.toml +++ b/symmetrix/pyproject.toml @@ -20,7 +20,7 @@ dependencies = [ [project.optional-dependencies] mace = [ - "torch", + "torch>=2.5", "mace-torch", ] test = [ From 27b2d06999f37dce743287ea6f412abb04a3c5cd Mon Sep 17 00:00:00 2001 From: Chuck Witt Date: Mon, 27 Jul 2026 17:58:55 -0400 Subject: [PATCH 29/29] restore whitespace checker. --- .github/workflows/ci.yaml | 4 +-- .pre-commit-config.yaml | 4 +-- libsymmetrix/CMakeLists.txt | 1 - libsymmetrix/source/cubic_spline_kokkos.cpp | 6 ++-- libsymmetrix/source/cubic_spline_kokkos.hpp | 2 +- .../source/cubic_spline_set_kokkos.cpp | 26 +++++++------- libsymmetrix/source/mace.cpp | 2 +- libsymmetrix/source/mace_kokkos.cpp | 34 +++++++++---------- .../source/multivariate_polynomial.cpp | 4 +-- .../source/multivariate_polynomial.hpp | 2 +- .../source/multivariate_polynomial_kokkos.cpp | 2 +- .../source/multivariate_polynomial_kokkos.hpp | 2 +- .../source/radial_function_set_kokkos.cpp | 2 +- .../source/radial_function_set_kokkos.hpp | 6 ++-- libsymmetrix/source/tools.cpp | 2 +- libsymmetrix/source/tools_kokkos.hpp | 2 +- libsymmetrix/source/zbl.hpp | 2 +- libsymmetrix/source/zbl_kokkos.cpp | 2 +- libsymmetrix/source/zbl_kokkos.hpp | 2 +- pair_symmetrix/LICENSE | 8 ++--- pair_symmetrix/pair_symmetrix_mace.cpp | 8 ++--- pair_symmetrix/pair_symmetrix_mace_kokkos.cpp | 2 +- .../test/test_pair_symmetrix_mace.py | 2 +- symmetrix/source/cpp/cubic_spline.cpp | 1 - symmetrix/source/cpp/mace.cpp | 18 +++++----- symmetrix/source/cpp/mace_kokkos.cpp | 8 ++--- .../cpp/multivariate_polynomial_kokkos.cpp | 2 +- symmetrix/source/cpp/zbl.cpp | 1 - symmetrix/source/cpp/zbl_kokkos.cpp | 1 - 29 files changed, 77 insertions(+), 81 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index b6aa1ff..8003d00 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -129,11 +129,11 @@ jobs: run: | source .venv/bin/activate # check for valid symmetrix json from atomic numbers and default filename - wget https://github.com/ACEsuit/mace-off/raw/refs/heads/main/mace_off23/MACE-OFF23_small.model + wget https://github.com/ACEsuit/mace-off/raw/refs/heads/main/mace_off23/MACE-OFF23_small.model symmetrix_extract_mace --model MACE-OFF23_small.model --atomic-numbers 1 8 python3 -c "from symmetrix import Symmetrix; calc = Symmetrix('MACE-OFF23_small-1-8.json', species=[1, 8])" # check for valid symmetrix json from checmical symbols and specified filename - wget https://github.com/ACEsuit/mace-off/raw/refs/heads/main/mace_off23/MACE-OFF23_medium.model + wget https://github.com/ACEsuit/mace-off/raw/refs/heads/main/mace_off23/MACE-OFF23_medium.model symmetrix_extract_mace --model MACE-OFF23_medium.model --chemical-symbols H O -o extract_test.json python3 -c "from symmetrix import Symmetrix; calc = Symmetrix('extract_test.json', species=[1, 8])" deactivate diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 517e2c3..154d0ba 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2,8 +2,8 @@ repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v4.5.0 hooks: - # TODO: Restore trailing-whitespace and end-of-file-fixer with the deferred - # repository-wide formatting changes. + - id: trailing-whitespace + - id: end-of-file-fixer - id: check-yaml - id: check-added-large-files diff --git a/libsymmetrix/CMakeLists.txt b/libsymmetrix/CMakeLists.txt index e7e7de7..c694cf4 100644 --- a/libsymmetrix/CMakeLists.txt +++ b/libsymmetrix/CMakeLists.txt @@ -79,4 +79,3 @@ if (SYMMETRIX_SPHERICART_SYCL) message(STATUS "Symmetrix: Will use SYCL version of sphericart.") target_compile_definitions(symmetrix PRIVATE SYMMETRIX_SPHERICART_SYCL) endif() - diff --git a/libsymmetrix/source/cubic_spline_kokkos.cpp b/libsymmetrix/source/cubic_spline_kokkos.cpp index 9226d57..640ea38 100644 --- a/libsymmetrix/source/cubic_spline_kokkos.cpp +++ b/libsymmetrix/source/cubic_spline_kokkos.cpp @@ -40,14 +40,14 @@ double CubicSplineKokkos::evaluate(double r) if (r<0 or r>h*num_coeffs/4 or std::isnan(r)) throw std::invalid_argument("Out of bounds in CubicSplineKokkos::evaluate. r=" + std::to_string(r)); const int i = std::clamp(static_cast(r / h), 0, static_cast(num_coeffs/4 - 1)); - + const double x = r - h * i; const double xx = x * x; const double xxx = xx * x; const int i4 = 4 * i; auto h_c = Kokkos::create_mirror_view(c); - + double ret = 0; const double c0 = h_c(i4); const double c1 = h_c(i4 + 1); @@ -99,7 +99,7 @@ std::tuple CubicSplineKokkos::evaluate_deriv_divided(double r) const double c1 = h_c(i4+1); const double c2=h_c(i4+2); const double c3=h_c(i4+3); - + return {c0 + c1*x + c2*xx + c3*xxx, (c1 + 2*c2*x + 3*c3*xx) / r}; } diff --git a/libsymmetrix/source/cubic_spline_kokkos.hpp b/libsymmetrix/source/cubic_spline_kokkos.hpp index 2529bef..92b9114 100644 --- a/libsymmetrix/source/cubic_spline_kokkos.hpp +++ b/libsymmetrix/source/cubic_spline_kokkos.hpp @@ -15,7 +15,7 @@ CubicSplineKokkos(double h, CubicSplineKokkos(double h, Kokkos::View nodal_values, Kokkos::View nodal_derivs); - + double evaluate(double r); std::tuple evaluate_deriv(double r); std::tuple evaluate_deriv_divided(double r); diff --git a/libsymmetrix/source/cubic_spline_set_kokkos.cpp b/libsymmetrix/source/cubic_spline_set_kokkos.cpp index 08c7484..b079397 100644 --- a/libsymmetrix/source/cubic_spline_set_kokkos.cpp +++ b/libsymmetrix/source/cubic_spline_set_kokkos.cpp @@ -14,7 +14,7 @@ CubicSplineSetKokkos::CubicSplineSetKokkos( c = Kokkos::View("coeffs", num_nodes-1, 4, num_splines); auto h_c = Kokkos::create_mirror_view(c); - + for (int i=0; i nodal_derivs; Kokkos::View c; - InitializeCoefficientsFunctor(double h_, Kokkos::View nodal_values_, + InitializeCoefficientsFunctor(double h_, Kokkos::View nodal_values_, Kokkos::View nodal_derivs_, Kokkos::View c_) : h(h_), nodal_values(nodal_values_), nodal_derivs(nodal_derivs_), c(c_) {} @@ -64,8 +64,8 @@ CubicSplineSetKokkos::CubicSplineSetKokkos( // Create and use the functor for initialization InitializeCoefficientsFunctor functor(h, nodal_values, nodal_derivs, c); - Kokkos::parallel_for("InitializeCoefficients", - Kokkos::MDRangePolicy>({0, 0}, {num_nodes - 1, num_splines}), + Kokkos::parallel_for("InitializeCoefficients", + Kokkos::MDRangePolicy>({0, 0}, {num_nodes - 1, num_splines}), functor); //Kokkos::fence(); } @@ -79,11 +79,11 @@ void CubicSplineSetKokkos::evaluate( const double x = r - h*i; const double xx = x*x; const double xxx = xx*x; - + // make a local copy //Kokkos::View local_c("local_c", c.extent(0), c.extent(1), c.extent(2)); //Kokkos::deep_copy(local_c,this->c); - + // Parallel computation of values Kokkos::parallel_for("EvaluateSpline", num_splines, KOKKOS_CLASS_LAMBDA(const int j) { values(j) = c(i, 0, j) + c(i, 1, j) * x + c(i, 2, j) * xx + c(i, 3, j) * xxx; @@ -195,23 +195,23 @@ void CubicSplineSetKokkos::evaluate( { // create test device side view Kokkos::View test_view("test_view",values.size()); - + // test view's host mirror auto h_test_view = Kokkos::create_mirror_view(test_view); - + // view to copy the values Kokkos::View> input_value_view(values.data(),values.size()); - + // copy input values to host mirror Kokkos::deep_copy(h_test_view,input_value_view); //copy it back to device side view Kokkos::deep_copy(test_view,h_test_view); - + evaluate( r, test_view); - + // copy values back to host mirror Kokkos::deep_copy(h_test_view,test_view); @@ -230,7 +230,7 @@ void CubicSplineSetKokkos::evaluate_derivs( // create test device side view Kokkos::View test_view_values("test_view_values",values.size()); Kokkos::View test_view_derivs("test_view_derivs",derivs.size()); - + // test view's host mirror auto h_test_view_values = Kokkos::create_mirror_view(test_view_values); auto h_test_view_derivs = Kokkos::create_mirror_view(test_view_derivs); @@ -239,7 +239,7 @@ void CubicSplineSetKokkos::evaluate_derivs( Kokkos::View> input_values_view(values.data(),values.size()); Kokkos::View> input_derivs_view(derivs.data(),derivs.size()); - + // copy input values to host mirror Kokkos::deep_copy(h_test_view_values,input_values_view); Kokkos::deep_copy(h_test_view_derivs,input_derivs_view); diff --git a/libsymmetrix/source/mace.cpp b/libsymmetrix/source/mace.cpp index 21c27b4..3c0644b 100644 --- a/libsymmetrix/source/mace.cpp +++ b/libsymmetrix/source/mace.cpp @@ -213,7 +213,7 @@ void MACE::reverse_A0( auto Phi0_adj_i = std::vector(num_lm*num_channels); - // [dE/dPhi0_il]_mk = \sum_k' [dE/dA0_il]_mk' [trans(W_il)]_k'k + // [dE/dPhi0_il]_mk = \sum_k' [dE/dA0_il]_mk' [trans(W_il)]_k'k for (int l=0; l<=l_max; ++l) { auto Phi0_adj_il = Phi0_adj_i.data()+l*l*num_channels; auto A0_adj_il = A0_adj.data()+(i*num_lm+l*l)*num_channels; diff --git a/libsymmetrix/source/mace_kokkos.cpp b/libsymmetrix/source/mace_kokkos.cpp index 1a4f8cd..2390be8 100644 --- a/libsymmetrix/source/mace_kokkos.cpp +++ b/libsymmetrix/source/mace_kokkos.cpp @@ -120,7 +120,7 @@ void MACEKokkos::compute_R0( Kokkos::parallel_scan("first_neigh", num_nodes, KOKKOS_LAMBDA (const int i, int& update, const bool final) { - const int num_neigh_i = num_neigh(i); + const int num_neigh_i = num_neigh(i); if (final) first_neigh(i) = update; update += num_neigh_i; @@ -166,9 +166,9 @@ void MACEKokkos::compute_R0( Kokkos::TeamVectorRange(team_member, (l_max+1)*num_channels), [&] (const int lk) { const double c0 = c(type_ij,n,0,lk); - const double c1 = c(type_ij,n,1,lk); - const double c2 = c(type_ij,n,2,lk); - const double c3 = c(type_ij,n,3,lk); + const double c1 = c(type_ij,n,1,lk); + const double c2 = c(type_ij,n,2,lk); + const double c3 = c(type_ij,n,3,lk); R0(ij,lk) = c0 + c1*x + c2*xx + c3*xxx; R0_deriv(ij,lk) = c1 + c2*two_x + c3*three_xx; }); @@ -380,7 +380,7 @@ void MACEKokkos::reverse_A0( Kokkos::parallel_scan("first_neigh", num_nodes, KOKKOS_LAMBDA (const int i, int& update, const bool final) { - const int num_neigh_i = num_neigh(i); + const int num_neigh_i = num_neigh(i); if (final) first_neigh(i) = update; update += num_neigh_i; @@ -937,7 +937,7 @@ void MACEKokkos::reverse_Phi1( bool zero_H1_adj) { if (dPhi1r.extent(0) < Phi1r.extent(0)) - Kokkos::realloc(dPhi1r, Phi1r.extent(0), Phi1r.extent(1), Phi1r.extent(2)); + Kokkos::realloc(dPhi1r, Phi1r.extent(0), Phi1r.extent(1), Phi1r.extent(2)); if (node_forces.size() < xyz.size()) Kokkos::resize(node_forces, xyz.size()); if (H1_adj.extent(0) < H1.extent(0)) @@ -986,7 +986,7 @@ void MACEKokkos::reverse_Phi1( Kokkos::parallel_scan("first_neigh", num_nodes, KOKKOS_LAMBDA (const int i, int& update, const bool final) { - const int num_neigh_i = num_neigh(i); + const int num_neigh_i = num_neigh(i); if (final) first_neigh(i) = update; update += num_neigh_i; @@ -1011,7 +1011,7 @@ void MACEKokkos::reverse_Phi1( Kokkos::parallel_reduce( Kokkos::ThreadVectorRange(team_member, num_channels), [=] (const int k, double& t1, double& t2) { - t1 += R1_deriv(ij,lel1l2*num_channels+k) * H1(neigh_indices(ij),lm2,k) * dPhi1r(i,lelm1lm2,k); + t1 += R1_deriv(ij,lel1l2*num_channels+k) * H1(neigh_indices(ij),lm2,k) * dPhi1r(i,lelm1lm2,k); t2 += R1(ij,lel1l2*num_channels+k) * H1(neigh_indices(ij),lm2,k) * dPhi1r(i,lelm1lm2,k); Kokkos::atomic_add( &H1_adj(neigh_indices(ij),lm2,k), @@ -1285,7 +1285,7 @@ void MACEKokkos::compute_M1(int num_nodes, Kokkos::View n if (M1.extent(0) < num_nodes) Kokkos::realloc(M1, num_nodes, num_channels); if (M1_poly_values.extent(0) < num_nodes) - Kokkos::realloc(M1_poly_values, num_nodes, num_lm+M1_poly_spec.extent(0), num_channels); + Kokkos::realloc(M1_poly_values, num_nodes, num_lm+M1_poly_spec.extent(0), num_channels); Kokkos::deep_copy(M1, 0.0); const auto A1 = this->A1; @@ -1373,7 +1373,7 @@ void MACEKokkos::reverse_M1(int num_nodes, Kokkos::View n Kokkos::realloc(A1_adj, A1.extent(0), A1.extent(1), A1.extent(2)); Kokkos::deep_copy(A1_adj, 0.0); if (M1_poly_adjoints.extent(0) < num_nodes) - Kokkos::realloc(M1_poly_adjoints, num_nodes, M1_poly_coeff.extent(1), num_channels); + Kokkos::realloc(M1_poly_adjoints, num_nodes, M1_poly_coeff.extent(1), num_channels); // TODO: prune const auto A1_adj = this->A1_adj; @@ -1516,7 +1516,7 @@ double MACEKokkos::compute_readouts(int num_nodes, const Kokkos::View auto H1 = this->H1; auto H1_adj = this->H1_adj; auto readout_1_weights = this->readout_1_weights; - + // atomic energies Kokkos::parallel_for("Compute Readouts 1", num_nodes, KOKKOS_LAMBDA (const int i) { node_energies(i) += atomic_energies(node_types(i)); @@ -1531,9 +1531,9 @@ double MACEKokkos::compute_readouts(int num_nodes, const Kokkos::View }); Kokkos::fence(); // second readout - auto H2 = Kokkos::subview(this->H2, make_pair(0,num_nodes), Kokkos::ALL); + auto H2 = Kokkos::subview(this->H2, make_pair(0,num_nodes), Kokkos::ALL); auto readout_2_output = Kokkos::subview(this->readout_2_output, make_pair(0,num_nodes)); - auto H2_adj = Kokkos::subview(this->H2_adj, make_pair(0,num_nodes), Kokkos::ALL); + auto H2_adj = Kokkos::subview(this->H2_adj, make_pair(0,num_nodes), Kokkos::ALL); readout_2.evaluate_gradient(H2, readout_2_output, H2_adj); Kokkos::parallel_for("Compute Readouts 2", num_nodes, KOKKOS_LAMBDA (const int i) { node_energies(i) += readout_2_output(i); @@ -1554,7 +1554,7 @@ void MACEKokkos::load_from_json(std::string filename) { std::ifstream f(filename); nlohmann::json file = nlohmann::json::parse(f); - + // Basic model information num_elements = file["num_elements"]; num_channels = file["num_channels"]; @@ -1607,9 +1607,9 @@ void MACEKokkos::load_from_json(std::string filename) for (int lk=0; lk<(l_max+1)*num_channels; ++lk) { const int k = lk % num_channels; h_c(ab,i,0,lk) *= H0_weights[b*num_channels+k]; - h_c(ab,i,1,lk) *= H0_weights[b*num_channels+k]; - h_c(ab,i,2,lk) *= H0_weights[b*num_channels+k]; - h_c(ab,i,3,lk) *= H0_weights[b*num_channels+k]; + h_c(ab,i,1,lk) *= H0_weights[b*num_channels+k]; + h_c(ab,i,2,lk) *= H0_weights[b*num_channels+k]; + h_c(ab,i,3,lk) *= H0_weights[b*num_channels+k]; } } // add A0_weights diff --git a/libsymmetrix/source/multivariate_polynomial.cpp b/libsymmetrix/source/multivariate_polynomial.cpp index 951f320..e08f1a5 100644 --- a/libsymmetrix/source/multivariate_polynomial.cpp +++ b/libsymmetrix/source/multivariate_polynomial.cpp @@ -49,7 +49,7 @@ MultivariatePolynomial::MultivariatePolynomial( node_set.insert({i}); for (auto monomial : monomials) node_set.insert(monomial); - + // add auxiliary nodes until all nodes have two upstream factors num_auxiliary_nodes = 0; auto find_parents = [](const std::vector& node, @@ -72,7 +72,7 @@ MultivariatePolynomial::MultivariatePolynomial( } } nodes = std::vector>(node_set.begin(), node_set.end()); - + // find edges for (auto node : node_set) { if (node.size() == 1) diff --git a/libsymmetrix/source/multivariate_polynomial.hpp b/libsymmetrix/source/multivariate_polynomial.hpp index d37e610..ff6b858 100644 --- a/libsymmetrix/source/multivariate_polynomial.hpp +++ b/libsymmetrix/source/multivariate_polynomial.hpp @@ -9,7 +9,7 @@ class MultivariatePolynomial public: -MultivariatePolynomial(int num_variables, +MultivariatePolynomial(int num_variables, std::vector coefficients, std::vector> monomials); diff --git a/libsymmetrix/source/multivariate_polynomial_kokkos.cpp b/libsymmetrix/source/multivariate_polynomial_kokkos.cpp index f587a28..1dcc0ec 100644 --- a/libsymmetrix/source/multivariate_polynomial_kokkos.cpp +++ b/libsymmetrix/source/multivariate_polynomial_kokkos.cpp @@ -7,7 +7,7 @@ MultivariatePolynomialKokkos::MultivariatePolynomialKokkos( - int num_variables, + int num_variables, std::vector coefficients, std::vector> monomials) { diff --git a/libsymmetrix/source/multivariate_polynomial_kokkos.hpp b/libsymmetrix/source/multivariate_polynomial_kokkos.hpp index 1a36f4b..5d0972e 100644 --- a/libsymmetrix/source/multivariate_polynomial_kokkos.hpp +++ b/libsymmetrix/source/multivariate_polynomial_kokkos.hpp @@ -8,7 +8,7 @@ class MultivariatePolynomialKokkos public: -MultivariatePolynomialKokkos(int num_variables, +MultivariatePolynomialKokkos(int num_variables, std::vector coefficients, std::vector> monomials); double evaluate(const Kokkos::View& x); diff --git a/libsymmetrix/source/radial_function_set_kokkos.cpp b/libsymmetrix/source/radial_function_set_kokkos.cpp index ae0dec3..1c4b05a 100644 --- a/libsymmetrix/source/radial_function_set_kokkos.cpp +++ b/libsymmetrix/source/radial_function_set_kokkos.cpp @@ -141,7 +141,7 @@ void RadialFunctionSetKokkos::evaluate( Kokkos::parallel_scan("first_neigh", num_nodes, KOKKOS_LAMBDA (const int i, int& update, const bool final) { - const int num_neigh_i = num_neigh(i); + const int num_neigh_i = num_neigh(i); if (final) first_neigh(i) = update; update += num_neigh_i; diff --git a/libsymmetrix/source/radial_function_set_kokkos.hpp b/libsymmetrix/source/radial_function_set_kokkos.hpp index 4859180..8d35d20 100644 --- a/libsymmetrix/source/radial_function_set_kokkos.hpp +++ b/libsymmetrix/source/radial_function_set_kokkos.hpp @@ -8,7 +8,7 @@ template class RadialFunctionSetKokkos { public: - + RadialFunctionSetKokkos(); RadialFunctionSetKokkos( double h, @@ -22,9 +22,9 @@ class RadialFunctionSetKokkos Kokkos::View r, Kokkos::View R, Kokkos::View R_deriv) const; - + private: - + double h; int num_edge_types; int num_functions; diff --git a/libsymmetrix/source/tools.cpp b/libsymmetrix/source/tools.cpp index cf053c4..6515c77 100644 --- a/libsymmetrix/source/tools.cpp +++ b/libsymmetrix/source/tools.cpp @@ -144,4 +144,4 @@ std::vector> _generate_indices(int dim, int max) { for (int r=0; r Kokkos::View toKokkosView(const char* name,const std::vector& stdVector) { - + std::string label(name); // Create a Kokkos::View with the same size as the std::vector Kokkos::View kokkosView(label, stdVector.size()); diff --git a/libsymmetrix/source/zbl.hpp b/libsymmetrix/source/zbl.hpp index c8d880e..0cc8e59 100644 --- a/libsymmetrix/source/zbl.hpp +++ b/libsymmetrix/source/zbl.hpp @@ -41,7 +41,7 @@ double a_prefactor; std::vector c; std::vector covalent_radii; int p; - + // values taken from mace/modules/radial.py static constexpr double c_exps_0 = -3.2; static constexpr double c_exps_1 = -0.9423; diff --git a/libsymmetrix/source/zbl_kokkos.cpp b/libsymmetrix/source/zbl_kokkos.cpp index 6153c2f..79769ff 100644 --- a/libsymmetrix/source/zbl_kokkos.cpp +++ b/libsymmetrix/source/zbl_kokkos.cpp @@ -117,7 +117,7 @@ void ZBLKokkos::compute_ZBL( Kokkos::parallel_scan("first_neigh", num_nodes, KOKKOS_LAMBDA (const int i, int& update, const bool final) { - const int num_neigh_i = num_neigh(i); + const int num_neigh_i = num_neigh(i); if (final) first_neigh(i) = update; update += num_neigh_i; diff --git a/libsymmetrix/source/zbl_kokkos.hpp b/libsymmetrix/source/zbl_kokkos.hpp index 46af555..474b185 100644 --- a/libsymmetrix/source/zbl_kokkos.hpp +++ b/libsymmetrix/source/zbl_kokkos.hpp @@ -46,7 +46,7 @@ double a_prefactor; Kokkos::View c; Kokkos::View covalent_radii; int p; - + // values taken from mace/modules/radial.py static constexpr double c_exps_0 = -3.2; static constexpr double c_exps_1 = -0.9423; diff --git a/pair_symmetrix/LICENSE b/pair_symmetrix/LICENSE index 202b8cd..b38a579 100644 --- a/pair_symmetrix/LICENSE +++ b/pair_symmetrix/LICENSE @@ -2,7 +2,7 @@ GNU GENERAL PUBLIC LICENSE Version 2, June 1991 -Copyright (C) 1989, 1991 Free Software Foundation, Inc. +Copyright (C) 1989, 1991 Free Software Foundation, Inc. 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA Everyone is permitted to copy and distribute verbatim copies of this @@ -313,7 +313,7 @@ You should have received a copy of the GNU General Public License along with this program; if not, write to the Free Software Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA. -Also add information on how to contact you by electronic and paper mail. +Also add information on how to contact you by electronic and paper mail. If the program is interactive, make it output a short notice like this when it starts in an interactive mode: @@ -321,7 +321,7 @@ when it starts in an interactive mode: Gnomovision version 69, Copyright (C) year name of author Gnomovision comes with ABSOLUTELY NO WARRANTY; for details type `show w'. This is free software, and you are welcome -to redistribute it under certain conditions; type `show c' +to redistribute it under certain conditions; type `show c' for details. The hypothetical commands `show w' and `show c' should show the @@ -336,7 +336,7 @@ if necessary. Here is a sample; alter the names: Yoyodyne, Inc., hereby disclaims all copyright interest in the program `Gnomovision' -(which makes passes at compilers) written +(which makes passes at compilers) written by James Hacker. signature of Ty Coon, 1 April 1989 diff --git a/pair_symmetrix/pair_symmetrix_mace.cpp b/pair_symmetrix/pair_symmetrix_mace.cpp index f16cf66..37a9eea 100644 --- a/pair_symmetrix/pair_symmetrix_mace.cpp +++ b/pair_symmetrix/pair_symmetrix_mace.cpp @@ -39,9 +39,9 @@ PairSymmetrixMACE::PairSymmetrixMACE(LAMMPS *lmp) one_coeff = 1; manybody_flag = 1; no_virial_fdotr_compute = 1; - // WARNING: for mace, these variables are model-dependent, so i + // WARNING: for mace, these variables are model-dependent, so i // reset them after the model is loaded (in coeff). - // however, i can't make them zero here, because that + // however, i can't make them zero here, because that // confusingly yields seg faults with hybrid/overlay. // so, i set them to a fairly big number here and hope. // not a great solution. @@ -159,7 +159,7 @@ void PairSymmetrixMACE::coeff(int narg, char **arg) double PairSymmetrixMACE::init_one(int i, int j) { if (setflag[i][j] == 0) error->all(FLERR, "All pair coeffs are not set"); - + return mace->r_cut; } @@ -685,7 +685,7 @@ void PairSymmetrixMACE::compute_no_mpi_message_passing(int eflag, int vflag) mace->compute_H2(num_local_nodes, node_types); mace->compute_readouts(num_local_nodes, node_types); - + mace->reverse_H2(num_local_nodes, node_types, false); mace->reverse_M1(num_local_nodes, node_types); mace->reverse_A1_scaled(num_local_nodes, node_types, num_neigh, neigh_types, xyz, r, false); diff --git a/pair_symmetrix/pair_symmetrix_mace_kokkos.cpp b/pair_symmetrix/pair_symmetrix_mace_kokkos.cpp index ae555e1..4e83705 100644 --- a/pair_symmetrix/pair_symmetrix_mace_kokkos.cpp +++ b/pair_symmetrix/pair_symmetrix_mace_kokkos.cpp @@ -1071,7 +1071,7 @@ void PairSymmetrixMACEKokkos::compute_no_mpi_message_pass mace->compute_H2(num_local_nodes, node_types); mace->compute_readouts(num_local_nodes, node_types); - + mace->reverse_H2(num_local_nodes, node_types, false); mace->reverse_M1(num_local_nodes, node_types); mace->reverse_A1_scaled(num_local_nodes, node_types, num_neigh, neigh_types, xyz, r); diff --git a/pair_symmetrix/test/test_pair_symmetrix_mace.py b/pair_symmetrix/test/test_pair_symmetrix_mace.py index 2cf90c6..b27942e 100644 --- a/pair_symmetrix/test/test_pair_symmetrix_mace.py +++ b/pair_symmetrix/test/test_pair_symmetrix_mace.py @@ -67,7 +67,7 @@ def test_h20(cmdargs, pair_style): create_atoms 2 single 0.0 -2.0 0.0 units box mass 1 1.008 mass 2 15.999 - + pair_style {} pair_coeff * * MACE-OFF23_small-1-8.json H O diff --git a/symmetrix/source/cpp/cubic_spline.cpp b/symmetrix/source/cpp/cubic_spline.cpp index 23cd408..ec5bf0f 100644 --- a/symmetrix/source/cpp/cubic_spline.cpp +++ b/symmetrix/source/cpp/cubic_spline.cpp @@ -13,4 +13,3 @@ void bind_cubic_spline(py::module_ &m) .def("evaluate_deriv", &CubicSpline::evaluate_deriv) .def("evaluate_deriv_divided", &CubicSpline::evaluate_deriv_divided); } - diff --git a/symmetrix/source/cpp/mace.cpp b/symmetrix/source/cpp/mace.cpp index 74765cd..eb3e714 100644 --- a/symmetrix/source/cpp/mace.cpp +++ b/symmetrix/source/cpp/mace.cpp @@ -41,7 +41,7 @@ void bind_mace(py::module_ &m) py::array_t xyz, py::array_t r) { self.compute_node_energies_forces( - num_nodes, + num_nodes, std::span(node_types.data(), node_types.size()), std::span(num_neigh.data(), num_neigh.size()), std::span(neigh_indices.data(), neigh_indices.size()), @@ -100,7 +100,7 @@ void bind_mace(py::module_ &m) py::array_t neigh_types, py::array_t xyz, py::array_t r) { - self.reverse_A0(num_nodes, + self.reverse_A0(num_nodes, std::span(node_types.data(), node_types.size()), std::span(num_neigh.data(), num_neigh.size()), std::span(neigh_types.data(), neigh_types.size()), @@ -140,7 +140,7 @@ void bind_mace(py::module_ &m) .def("compute_M0", [](MACE& self, const int num_nodes, py::array_t node_types) { - self.compute_M0(num_nodes, + self.compute_M0(num_nodes, std::span(node_types.data(), node_types.size())); }) .def("reverse_M0", @@ -155,7 +155,7 @@ void bind_mace(py::module_ &m) [](MACE& self, const int num_nodes, py::array_t num_neigh, py::array_t neigh_indices) { - self.compute_Phi1(num_nodes, + self.compute_Phi1(num_nodes, std::span(num_neigh.data(), num_neigh.size()), std::span(neigh_indices.data(), neigh_indices.size())); }) @@ -167,7 +167,7 @@ void bind_mace(py::module_ &m) py::array_t r, bool zero_dxyz, bool zero_H1_adj) { - self.reverse_Phi1(num_nodes, + self.reverse_Phi1(num_nodes, std::span(num_neigh.data(), num_neigh.size()), std::span(neigh_indices.data(), neigh_indices.size()), std::span(xyz.data(), xyz.size()), @@ -210,7 +210,7 @@ void bind_mace(py::module_ &m) .def("compute_M1", [](MACE& self, const int num_nodes, py::array_t node_types) { - self.compute_M1(num_nodes, + self.compute_M1(num_nodes, std::span(node_types.data(), node_types.size())); }) @@ -223,21 +223,21 @@ void bind_mace(py::module_ &m) .def("compute_H2", [](MACE& self, const int num_nodes, py::array_t node_types) { - self.compute_H2(num_nodes, + self.compute_H2(num_nodes, std::span(node_types.data(), node_types.size())); }) .def("reverse_H2", [](MACE& self, const int num_nodes, py::array_t node_types, bool zero_H1_adj) { - self.reverse_H2(num_nodes, + self.reverse_H2(num_nodes, std::span(node_types.data(), node_types.size()), zero_H1_adj); }) .def("compute_readouts", [](MACE& self, const int num_nodes, py::array_t node_types) { - self.compute_readouts(num_nodes, + self.compute_readouts(num_nodes, std::span(node_types.data(), node_types.size())); }); } diff --git a/symmetrix/source/cpp/mace_kokkos.cpp b/symmetrix/source/cpp/mace_kokkos.cpp index b7d28d1..61bd226 100644 --- a/symmetrix/source/cpp/mace_kokkos.cpp +++ b/symmetrix/source/cpp/mace_kokkos.cpp @@ -45,7 +45,7 @@ void bind_mace_kokkos(py::module_ &m, const char* class_name) py::array_t xyz, py::array_t r) { self.compute_node_energies_forces( - num_nodes, + num_nodes, create_kokkos_view("node_types", node_types), create_kokkos_view("num_neigh", num_neigh), create_kokkos_view("neigh_indices", neigh_indices), @@ -62,7 +62,7 @@ void bind_mace_kokkos(py::module_ &m, const char* class_name) const int total_num_neigh = R0.size()/((self.l_max+1)*self.num_channels); set_kokkos_view(self.R0, R0, total_num_neigh, (self.l_max+1)*self.num_channels); }) - .def("compute_R0", + .def("compute_R0", [](MACEKokkos& self, const int num_nodes, py::array_t node_types, @@ -86,7 +86,7 @@ void bind_mace_kokkos(py::module_ &m, const char* class_name) const int total_num_neigh = R1.size()/(num_le*self.num_channels); set_kokkos_view(self.R1, R1, total_num_neigh, num_le*self.num_channels); }) - .def("compute_R1", + .def("compute_R1", [](MACEKokkos& self, const int num_nodes, py::array_t node_types, @@ -271,7 +271,7 @@ void bind_mace_kokkos(py::module_ &m, const char* class_name) bool zero_dxyz, bool zero_H1_adj) { self.reverse_Phi1( - num_nodes, + num_nodes, create_kokkos_view("num_neigh", num_neigh), create_kokkos_view("neigh_indices", neigh_indices), create_kokkos_view("xyz", xyz), diff --git a/symmetrix/source/cpp/multivariate_polynomial_kokkos.cpp b/symmetrix/source/cpp/multivariate_polynomial_kokkos.cpp index dc8ff8d..4da97ae 100644 --- a/symmetrix/source/cpp/multivariate_polynomial_kokkos.cpp +++ b/symmetrix/source/cpp/multivariate_polynomial_kokkos.cpp @@ -3,7 +3,7 @@ #include #include -#include "utilities_kokkos.hpp" +#include "utilities_kokkos.hpp" #include "multivariate_polynomial_kokkos.hpp" diff --git a/symmetrix/source/cpp/zbl.cpp b/symmetrix/source/cpp/zbl.cpp index 072648b..63c66e8 100644 --- a/symmetrix/source/cpp/zbl.cpp +++ b/symmetrix/source/cpp/zbl.cpp @@ -14,4 +14,3 @@ void bind_zbl(py::module_ &m) .def("compute_envelope", &ZBL::compute_envelope) .def("compute_envelope_gradient", &ZBL::compute_envelope_gradient); } - diff --git a/symmetrix/source/cpp/zbl_kokkos.cpp b/symmetrix/source/cpp/zbl_kokkos.cpp index a4dadfd..3532d9a 100644 --- a/symmetrix/source/cpp/zbl_kokkos.cpp +++ b/symmetrix/source/cpp/zbl_kokkos.cpp @@ -14,4 +14,3 @@ void bind_zbl_kokkos(py::module_ &m) .def("compute_envelope", &ZBLKokkos::compute_envelope) .def("compute_envelope_gradient", &ZBLKokkos::compute_envelope_gradient); } -