From 6ba7280b665b574773af2751d861e363e24e5e78 Mon Sep 17 00:00:00 2001
From: bpuchala <bpuchala@users.noreply.github.com>
Date: Mon, 22 Apr 2024 13:35:20 -0400
Subject: [PATCH] Add PrimNeighborList.weight_matrix accessor

---
 CHANGELOG.md                                  |  7 +++++++
 python/src/clexulator.cpp                     | 19 +++++++++++++++++++
 python/tests/clexulator/test_neighbor_list.py |  1 +
 3 files changed, 27 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a7b40d5..f83d011 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,13 @@ All notable changes to `libcasm-clexulator` will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- Added `libcasm.clexulator.PrimNeighborList.weight_matrix()` to access the weight matrix
+
+
 ## [2.0a4] - 2024-03-14
 
 ### Fixed
diff --git a/python/src/clexulator.cpp b/python/src/clexulator.cpp
index c47627e..361148a 100644
--- a/python/src/clexulator.cpp
+++ b/python/src/clexulator.cpp
@@ -632,6 +632,25 @@ PYBIND11_MODULE(_clexulator, m) {
               (DoF) or 1 or more continuous DoF.
           )pbdoc",
           py::arg("xtal_prim"))
+      .def(
+          "weight_matrix",
+          [](clexulator::PrimNeighborListWrapper const &x) {
+            return x.prim_neighbor_list->weight_matrix();
+          },
+          R"pbdoc(
+          Get the weight matrix that defines the shape of neighborhood and \
+          orders neighbors.
+
+          Returns
+          -------
+          weight_matrix: np.ndarray, shape=(3,3), dtype=int
+              Weight matrix, :math:`W`, that defines the shape of neighborhood
+              that orders neighbors. In the
+              :class:`~libcasm.clexulator.PrimNeighborList`,
+              unit cells are ordered by the integer distance,
+              :math:`r = x^{\mathsf{T}} W x`, where :math:`x` is the integer
+              unit cell coordinates.
+          )pbdoc")
       .def(
           "sublattice_indices",
           [](clexulator::PrimNeighborListWrapper const &x) {
diff --git a/python/tests/clexulator/test_neighbor_list.py b/python/tests/clexulator/test_neighbor_list.py
index 23938ab..414dbf2 100644
--- a/python/tests/clexulator/test_neighbor_list.py
+++ b/python/tests/clexulator/test_neighbor_list.py
@@ -34,6 +34,7 @@ def test_prim_neighbor_list_2():
     W = PrimNeighborList.make_weight_matrix(lattice.column_vector_matrix())
     prim_neighbor_list = PrimNeighborList(W, [0], 1)
     assert isinstance(prim_neighbor_list, PrimNeighborList)
+    assert (prim_neighbor_list.weight_matrix() == W).all()
 
 
 def test_prim_neighbor_list_3():