From d201661b5463c9ea3b584024260290b35e81518e Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 11:36:38 +0200
Subject: [PATCH 01/34] Add calculation of SNAP descriptors

---
 structuretoolkit/analyse/__init__.py |   5 +
 structuretoolkit/analyse/snap.py     | 359 +++++++++++++++++++++++++++
 tests/test_snap.py                   |  71 ++++++
 3 files changed, 435 insertions(+)
 create mode 100644 structuretoolkit/analyse/snap.py
 create mode 100644 tests/test_snap.py

diff --git a/structuretoolkit/analyse/__init__.py b/structuretoolkit/analyse/__init__.py
index 34d4a8115..023c33a52 100644
--- a/structuretoolkit/analyse/__init__.py
+++ b/structuretoolkit/analyse/__init__.py
@@ -22,6 +22,11 @@
     get_voronoi_vertices,
 )
 from structuretoolkit.analyse.strain import get_strain
+from structuretoolkit.analyse.snap import (
+    calc_bispectrum_names,
+    calc_bisepctrum_lmp,
+    calc_a_matrix_snappy,
+)
 
 
 def get_symmetry(
diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
new file mode 100644
index 000000000..68c5a8a72
--- /dev/null
+++ b/structuretoolkit/analyse/snap.py
@@ -0,0 +1,359 @@
+from ctypes import c_double, c_int, cast, POINTER
+import numpy as np
+from scipy.constants import physical_constants
+
+eV_div_A3_to_bar = (
+        1e25 / physical_constants["joule-electron volt relationship"][0]
+)
+
+
+def convert_mat(mat):
+    mat[np.diag_indices_from(mat)] /= 2
+    return mat[np.triu_indices(len(mat))]
+
+
+def calc_per_atom_quad(linear_per_atom):
+    return np.array([
+        np.concatenate((atom, convert_mat(mat=np.dot(atom.reshape(len(atom), 1), atom.reshape(len(atom), 1).T))))
+        for atom in linear_per_atom
+    ])
+
+
+def calc_sum_quad(linear_sum):
+    return np.concatenate((
+        linear_sum,
+        convert_mat(mat=np.dot(linear_sum.reshape(len(linear_sum), 1), linear_sum.reshape(len(linear_sum), 1).T))
+    ))
+
+
+def get_apre(cell):
+    a, b, c = cell
+    an, bn, cn = [np.linalg.norm(v) for v in cell]
+
+    alpha = np.arccos(np.dot(b, c) / (bn * cn))
+    beta = np.arccos(np.dot(a, c) / (an * cn))
+    gamma = np.arccos(np.dot(a, b) / (an * bn))
+
+    xhi = an
+    xyp = np.cos(gamma) * bn
+    yhi = np.sin(gamma) * bn
+    xzp = np.cos(beta) * cn
+    yzp = (bn * cn * np.cos(alpha) - xyp * xzp) / yhi
+    zhi = np.sqrt(cn ** 2 - xzp ** 2 - yzp ** 2)
+
+    return np.array(((xhi, 0, 0), (xyp, yhi, 0), (xzp, yzp, zhi)))
+
+
+def write_ase_structure(lmp, structure):
+    number_species = len(set(structure.get_chemical_symbols()))
+
+    apre = get_apre(cell=structure.cell)
+    (
+        (xhi, xy, xz),
+        (_, yhi, yz),
+        (_, _, zhi),
+    ) = apre.T
+    lmp.command(
+        "region 1 prism"
+        + " 0.0 "
+        + str(xhi)
+        + " 0.0 "
+        + str(yhi)
+        + " 0.0 "
+        + str(zhi)
+        + " "
+        + str(xy)
+        + " "
+        + str(xz)
+        + " "
+        + str(yz)
+        + " units box"
+    )
+    lmp.command("create_box " + str(number_species) + " 1")
+
+    el_dict = {el: i for i, el in enumerate(set(structure.get_chemical_symbols()))}
+
+    R = np.dot(np.linalg.inv(structure.cell), apre)
+
+    positions = structure.positions.flatten()
+    if np.matrix.trace(R) != 3:
+        positions = np.array(positions).reshape(-1, 3)
+        positions = np.matmul(positions, R)
+    positions = positions.flatten()
+    elem_all = np.array([el_dict[el] + 1 for el in structure.get_chemical_symbols()])
+    lmp.create_atoms(
+        n=len(structure),
+        id=None,
+        type=(len(elem_all) * c_int)(*elem_all),
+        x=(len(positions) * c_double)(*positions),
+        v=None,
+        image=None,
+        shrinkexceed=False
+    )
+
+
+def extract_compute_np(lmp, name, compute_type, result_type, array_shape):
+    """
+    Convert a lammps compute to a numpy array.
+    Assumes the compute returns a floating point numbers.
+    Note that the result is a view into the original memory.
+    If the result type is 0 (scalar) then conversion to numpy is skipped and a python float is returned.
+    """
+    ptr = lmp.extract_compute(
+        name, compute_type, result_type
+    )  # 1,2: Style (1) is per-atom compute, returns array type (2).
+    if result_type == 0:
+        return ptr  # No casting needed, lammps.py already works
+    if result_type == 2:
+        ptr = ptr.contents
+    total_size = int(np.prod(array_shape))
+    buffer_ptr = cast(ptr, POINTER(c_double * total_size))
+    array_np = np.frombuffer(buffer_ptr.contents, dtype=float)
+    array_np.shape = array_shape
+    return array_np.copy()
+
+
+def reset_lmp(lmp):
+    lmp.command("clear")
+    lmp.command("units metal")
+    lmp.command("dimension 3")
+    lmp.command("boundary p p p")
+    lmp.command("atom_style charge")
+    lmp.command("atom_modify map array sort 0 2.0")
+
+
+def set_potential_lmp(lmp, cutoff=10.0):
+    lmp.command("pair_style zero " + str(cutoff))
+    lmp.command("pair_coeff * *")
+    lmp.command("mass * 1.0e-20")
+    lmp.command("neighbor 1.0e-20 nsq")
+    lmp.command("neigh_modify one 10000")
+
+
+def calc_bispectrum_names(twojmax):
+    lst = []
+    for j1 in range(0, twojmax + 1):
+        for j2 in range(0, j1 + 1):
+            for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
+                if j >= j1:
+                    lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
+    return lst
+
+
+def set_compute_lammps(lmp, bispec_options, numtypes):
+    lmp_variable_args = {
+        k: bispec_options[k] for k in ["rcutfac", "rfac0", "rmin0", "twojmax"]
+    }
+    lmp_variable_args.update(
+        {
+            (k + str(i + 1)): bispec_options[k][i]
+            for k in ["wj", "radelem"]
+            for i, v in enumerate(bispec_options[k])
+        }
+    )
+
+    for k, v in lmp_variable_args.items():
+        lmp.command(f"variable {k} equal {v}")
+
+    base_b = "compute b all sna/atom ${rcutfac} ${rfac0} ${twojmax}"
+    radelem = " ".join([f"${{radelem{i}}}" for i in range(1, numtypes + 1)])
+    wj = " ".join([f"${{wj{i}}}" for i in range(1, numtypes + 1)])
+    kw_options = {
+        k: bispec_options[v]
+        for k, v in {
+            "rmin0": "rmin0",
+            "bzeroflag": "bzeroflag",
+            "quadraticflag": "quadraticflag",
+            "switchflag": "switchflag",
+            "chem": "chemflag",
+            "bnormflag": "bnormflag",
+            "wselfallflag": "wselfallflag",
+            "bikflag": "bikflag",
+            "switchinnerflag": "switchinnerflag",
+            "sinner": "sinner",
+            "dinner": "dinner",
+        }.items()
+        if v in bispec_options
+    }
+    kw_substrings = [f"{k} {v}" for k, v in kw_options.items()]
+    kwargs = " ".join(kw_substrings)
+    lmp.command(f"{base_b} {radelem} {wj} {kwargs}")
+
+
+def calc_bisepctrum_lmp(lmp, structure, bispec_options, cutoff=10.0):
+    number_coef = len(calc_bispectrum_names(twojmax=bispec_options["twojmax"]))
+    reset_lmp(lmp=lmp)
+    write_ase_structure(lmp=lmp, structure=structure)
+    set_potential_lmp(lmp=lmp, cutoff=cutoff)
+    set_compute_lammps(
+        lmp=lmp,
+        bispec_options=bispec_options,
+        numtypes=len(set(structure.get_chemical_symbols())),
+    )
+    try:
+        lmp.command("run 0")
+    except:
+        return np.array([])
+    else:
+        if "quadraticflag" in bispec_options.keys() and int(bispec_options["quadraticflag"]) == 1:
+            return extract_compute_np(
+                lmp=lmp,
+                name="b",
+                compute_type=1,
+                result_type=2,
+                # basically we only care about the off diagonal elements and from those we need only half
+                # plus the linear terms:   n + sum_{i: 1->n} i
+                array_shape=(len(structure), int(number_coef * (number_coef * (1 - 1 / 2) + 3 / 2)))
+            )
+        else:
+            return extract_compute_np(
+                lmp=lmp,
+                name="b",
+                compute_type=1,
+                result_type=2,
+                array_shape=(len(structure), number_coef)
+            )
+
+
+def lammps_variables(bispec_options):
+    d = {k: bispec_options[k] for k in
+         ["rcutfac",
+          "rfac0",
+          "rmin0",
+          # "diagonalstyle",
+          # "zblcutinner",
+          # "zblcutouter",
+          "twojmax"]}
+    d.update(
+        {
+            (k + str(i + 1)): bispec_options[k][i]
+            for k in ["wj", "radelem"]  # ["zblz", "wj", "radelem"]
+            for i, v in enumerate(bispec_options[k])
+        })
+    return d
+
+
+def set_variables(lmp, **lmp_variable_args):
+    for k, v in lmp_variable_args.items():
+        lmp.command(f"variable {k} equal {v}")
+
+
+def set_computes_snappy(lmp, bispec_options):
+    # # Bispectrum coefficient computes
+    base_b = "compute b all sna/atom ${rcutfac} ${rfac0} ${twojmax}"
+    base_db = "compute db all snad/atom ${rcutfac} ${rfac0} ${twojmax}"
+    base_vb = "compute vb all snav/atom ${rcutfac} ${rfac0} ${twojmax}"
+
+    numtypes = bispec_options["numtypes"]
+    radelem = " ".join([f"${{radelem{i}}}" for i in range(1, numtypes + 1)])
+    wj = " ".join([f"${{wj{i}}}" for i in range(1, numtypes + 1)])
+
+    kw_options = {
+        k: bispec_options[v]
+        for k, v in
+        {
+            # "diagonal": "diagonalstyle", For legacy diagonalstyle option
+            "rmin0": "rmin0",
+            "bzeroflag": "bzeroflag",
+            "quadraticflag": "quadraticflag",
+            "switchflag": "switchflag",
+        }.items()
+        if v in bispec_options
+    }
+    kw_substrings = [f"{k} {v}" for k, v in kw_options.items()]
+    kwargs = " ".join(kw_substrings)
+
+    for op, base in zip(("b", "db", "vb"), (base_b, base_db, base_vb)):
+        # print("Setting up compute",op)
+        command = f"{base} {radelem} {wj} {kwargs}"
+        # print(command)
+        lmp.command(command)
+
+    for cname in ["b", "db", "vb"]:
+        lmp.command(f"compute {cname}_sum all reduce sum c_{cname}[*]")
+
+
+def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
+    lmp_atom_ids = lmp.numpy.extract_atom_iarray("id", num_atoms).flatten()
+    assert np.all(lmp_atom_ids == 1 + np.arange(num_atoms)), "LAMMPS seems to have lost atoms"
+
+    # Extract types
+    lmp_types = lmp.numpy.extract_atom_iarray(name="type", nelem=num_atoms).flatten()
+    lmp_volume = lmp.get_thermo("vol")
+
+    # Extract Bsum
+    lmp_bsum = extract_compute_np(lmp, "b_sum", 0, 1, (n_coeff))
+
+    # Extract B
+    lmp_barr = extract_compute_np(lmp, "b", 1, 2, (num_atoms, n_coeff))
+
+    type_onehot = np.eye(num_types)[lmp_types - 1]  # lammps types are 1-indexed.
+    # has shape n_atoms, n_types, num_coeffs.
+    # Constructing so it has the similar form to db and vb arrays. This adds some memory usage,
+    # but not nearly as much as vb or db (which are factor of 6 and n_atoms*3 larger, respectively)
+
+    b_atom = type_onehot[:, :, np.newaxis] * lmp_barr[:, np.newaxis, :]
+    b_sum = b_atom.sum(axis=0) / num_atoms
+
+    try:
+        assert np.allclose(lmp_bsum, lmp_barr.sum(axis=0), rtol=1e-12, atol=1e-12), \
+            "b_sum doesn't match sum of b"
+    except AssertionError:
+        print(lmp_bsum)
+        print(lmp_barr.sum(axis=0))
+
+    lmp_dbarr = extract_compute_np(lmp, "db", 1, 2, (num_atoms, num_types, 3, n_coeff))
+    lmp_dbsum = extract_compute_np(lmp, "db_sum", 0, 1, (num_types, 3, n_coeff))
+    assert np.allclose(lmp_dbsum, lmp_dbarr.sum(axis=0), rtol=1e-12, atol=1e-12), \
+        "db_sum doesn't match sum of db"
+    db_atom = np.transpose(lmp_dbarr, (0, 2, 1, 3))
+
+    lmp_vbarr = extract_compute_np(lmp, "vb", 1, 2, (num_atoms, num_types, 6, n_coeff))
+    lmp_vbsum = extract_compute_np(lmp, "vb_sum", 0, 1, (num_types, 6, n_coeff))
+    assert np.allclose(lmp_vbsum, lmp_vbarr.sum(axis=0), rtol=1e-12, atol=1e-12), \
+        "vb_sum doesn't match sum of vb"
+    vb_sum = np.transpose(lmp_vbsum, (1, 0, 2)) / lmp_volume * eV_div_A3_to_bar
+
+    dbatom_shape = db_atom.shape
+    vbsum_shape = vb_sum.shape
+    a_fit = np.concatenate(
+        (
+            b_sum,
+            db_atom.reshape(dbatom_shape[0] * dbatom_shape[1], dbatom_shape[2] * dbatom_shape[3]),
+            vb_sum.reshape(vbsum_shape[0] * vbsum_shape[1], vbsum_shape[2])
+        ),
+        axis=0
+    )
+    return np.concatenate((np.array([np.eye(a_fit.shape[0])[0]]).T, a_fit), axis=1).copy()
+
+
+def calc_a_matrix_snappy(lmp, structure, bispec_options, cutoff=10.0):
+    number_coef = len(calc_bispectrum_names(twojmax=bispec_options["twojmax"]))
+    number_species = len(set(structure.get_chemical_symbols()))
+    reset_lmp(lmp=lmp)
+    write_ase_structure(lmp=lmp, structure=structure)
+    set_potential_lmp(lmp=lmp, cutoff=cutoff)
+    set_variables(lmp, **lammps_variables(bispec_options))
+    set_computes_snappy(
+        lmp=lmp,
+        bispec_options=bispec_options,
+    )
+    try:
+        lmp.command("run 0")
+    except:
+        return np.array([])
+    else:
+        if "quadraticflag" in bispec_options.keys() and int(bispec_options["quadraticflag"]) == 1:
+            return extract_computes_snappy(
+                lmp=lmp,
+                num_atoms=len(structure),
+                n_coeff=int(number_coef * (number_coef * (1 - 1 / 2) + 3 / 2)),
+                num_types=number_species
+            )
+        else:
+            return extract_computes_snappy(
+                lmp=lmp,
+                num_atoms=len(structure),
+                n_coeff=number_coef,
+                num_types=number_species,
+            )
diff --git a/tests/test_snap.py b/tests/test_snap.py
new file mode 100644
index 000000000..d4ff27d87
--- /dev/null
+++ b/tests/test_snap.py
@@ -0,0 +1,71 @@
+from ase.build.bulk import bulk
+import structuretoolkit as stk
+import unittest
+
+
+try:
+    from lammps import lammps
+
+    skip_snap_test = False
+except ImportError:
+    skip_snap_test = True
+
+
+@unittest.skipIf(
+    skip_snap_test, "LAMMPS is not installed, so the SNAP tests are skipped."
+)
+class TestCu(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
+        cls.structure = bulk("Cu", cubic=True)
+        bispec_options = {
+            "numtypes": 1,
+            "twojmax": 6,
+            "rcutfac": 1.0,
+            "rfac0": 0.99363,
+            "rmin0": 0.0,
+            "bzeroflag": 0,
+            "radelem": [4.0],
+            "type": ['Cu'],
+            "wj": [1.0],
+        }
+        cls.bispec_options = bispec_options
+
+    def test_calc_bispectrum_lmp(self):
+        n_coeff = len(stk.analyse.calc_bispectrum_names(
+            twojmax=self.bispec_options["twojmax"]
+        ))
+        coeff = stk.analyse.calc_bisepctrum_lmp(
+            lmp=self.lmp,
+            structure=self.structure,
+            bispec_options=self.bispec_options,
+            cutoff=10.0
+        )
+        self.assertTrue(coeff.shape, (len(self.structure), n_coeff))
+
+    def test_calc_bispectrum_lmp_quad(self):
+        n_coeff = len(stk.analyse.calc_bispectrum_names(
+            twojmax=self.bispec_options["twojmax"]
+        ))
+        bispec_options = self.bispec_options.copy()
+        bispec_options["quadraticflag"] = 1
+        coeff = stk.analyse.calc_bisepctrum_lmp(
+            lmp=self.lmp,
+            structure=self.structure,
+            bispec_options=bispec_options,
+            cutoff=10.0
+        )
+        self.assertTrue(coeff.shape, (len(self.structure), n_coeff * 31))
+
+    def test_calc_a_matrix_snappy(self):
+        n_coeff = len(stk.analyse.calc_bispectrum_names(
+            twojmax=self.bispec_options["twojmax"]
+        ))
+        mat_a = stk.analyse.calc_a_matrix_snappy(
+            lmp=self.lmp,
+            structure=self.structure,
+            bispec_options=self.bispec_options,
+            cutoff=10.0
+        )
+        self.assertTrue(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff))
\ No newline at end of file

From 4014f0c171d7490f31ee75119c17f238707619fd Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 11:40:28 +0200
Subject: [PATCH 02/34] renaming

---
 structuretoolkit/analyse/__init__.py |  6 +++---
 structuretoolkit/analyse/snap.py     | 10 +++++-----
 tests/test_snap.py                   | 12 ++++++------
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/structuretoolkit/analyse/__init__.py b/structuretoolkit/analyse/__init__.py
index 023c33a52..0415aa33c 100644
--- a/structuretoolkit/analyse/__init__.py
+++ b/structuretoolkit/analyse/__init__.py
@@ -23,9 +23,9 @@
 )
 from structuretoolkit.analyse.strain import get_strain
 from structuretoolkit.analyse.snap import (
-    calc_bispectrum_names,
-    calc_bisepctrum_lmp,
-    calc_a_matrix_snappy,
+    get_snap_descriptor_names,
+    calc_snap_descriptors_per_atom,
+    calc_snap_descriptor_derivatives,
 )
 
 
diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 68c5a8a72..5bedcfee6 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -130,7 +130,7 @@ def set_potential_lmp(lmp, cutoff=10.0):
     lmp.command("neigh_modify one 10000")
 
 
-def calc_bispectrum_names(twojmax):
+def get_snap_descriptor_names(twojmax):
     lst = []
     for j1 in range(0, twojmax + 1):
         for j2 in range(0, j1 + 1):
@@ -180,8 +180,8 @@ def set_compute_lammps(lmp, bispec_options, numtypes):
     lmp.command(f"{base_b} {radelem} {wj} {kwargs}")
 
 
-def calc_bisepctrum_lmp(lmp, structure, bispec_options, cutoff=10.0):
-    number_coef = len(calc_bispectrum_names(twojmax=bispec_options["twojmax"]))
+def calc_snap_descriptors_per_atom(lmp, structure, bispec_options, cutoff=10.0):
+    number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     reset_lmp(lmp=lmp)
     write_ase_structure(lmp=lmp, structure=structure)
     set_potential_lmp(lmp=lmp, cutoff=cutoff)
@@ -327,8 +327,8 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
     return np.concatenate((np.array([np.eye(a_fit.shape[0])[0]]).T, a_fit), axis=1).copy()
 
 
-def calc_a_matrix_snappy(lmp, structure, bispec_options, cutoff=10.0):
-    number_coef = len(calc_bispectrum_names(twojmax=bispec_options["twojmax"]))
+def calc_snap_descriptor_derivatives(lmp, structure, bispec_options, cutoff=10.0):
+    number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     number_species = len(set(structure.get_chemical_symbols()))
     reset_lmp(lmp=lmp)
     write_ase_structure(lmp=lmp, structure=structure)
diff --git a/tests/test_snap.py b/tests/test_snap.py
index d4ff27d87..dcc69e5f6 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -33,10 +33,10 @@ def setUpClass(cls):
         cls.bispec_options = bispec_options
 
     def test_calc_bispectrum_lmp(self):
-        n_coeff = len(stk.analyse.calc_bispectrum_names(
+        n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
-        coeff = stk.analyse.calc_bisepctrum_lmp(
+        coeff = stk.analyse.calc_snap_descriptors_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=self.bispec_options,
@@ -45,12 +45,12 @@ def test_calc_bispectrum_lmp(self):
         self.assertTrue(coeff.shape, (len(self.structure), n_coeff))
 
     def test_calc_bispectrum_lmp_quad(self):
-        n_coeff = len(stk.analyse.calc_bispectrum_names(
+        n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
         bispec_options = self.bispec_options.copy()
         bispec_options["quadraticflag"] = 1
-        coeff = stk.analyse.calc_bisepctrum_lmp(
+        coeff = stk.analyse.calc_snap_descriptors_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=bispec_options,
@@ -59,10 +59,10 @@ def test_calc_bispectrum_lmp_quad(self):
         self.assertTrue(coeff.shape, (len(self.structure), n_coeff * 31))
 
     def test_calc_a_matrix_snappy(self):
-        n_coeff = len(stk.analyse.calc_bispectrum_names(
+        n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
-        mat_a = stk.analyse.calc_a_matrix_snappy(
+        mat_a = stk.analyse.calc_snap_descriptor_derivatives(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=self.bispec_options,

From 0619bddcc931792d14f551bbe867778bb87dc1b3 Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 11:45:25 +0200
Subject: [PATCH 03/34] add lammps as optional dependency - only unix

---
 .ci_support/environment.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.ci_support/environment.yml b/.ci_support/environment.yml
index d91b4e386..baccbe2bc 100644
--- a/.ci_support/environment.yml
+++ b/.ci_support/environment.yml
@@ -18,3 +18,4 @@ dependencies:
 - spglib =2.0.2
 - sqsgenerator =0.2
 - pyxtal =0.6.0
+- lammps =2023.08.02  # [unix]
\ No newline at end of file

From a2aec83cee722562becc64e2b2f1439def91db66 Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 11:48:21 +0200
Subject: [PATCH 04/34] black formatting

---
 structuretoolkit/analyse/snap.py | 118 ++++++++++++++++++++-----------
 1 file changed, 78 insertions(+), 40 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 5bedcfee6..528c2103b 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -2,9 +2,7 @@
 import numpy as np
 from scipy.constants import physical_constants
 
-eV_div_A3_to_bar = (
-        1e25 / physical_constants["joule-electron volt relationship"][0]
-)
+eV_div_A3_to_bar = 1e25 / physical_constants["joule-electron volt relationship"][0]
 
 
 def convert_mat(mat):
@@ -13,17 +11,35 @@ def convert_mat(mat):
 
 
 def calc_per_atom_quad(linear_per_atom):
-    return np.array([
-        np.concatenate((atom, convert_mat(mat=np.dot(atom.reshape(len(atom), 1), atom.reshape(len(atom), 1).T))))
-        for atom in linear_per_atom
-    ])
+    return np.array(
+        [
+            np.concatenate(
+                (
+                    atom,
+                    convert_mat(
+                        mat=np.dot(
+                            atom.reshape(len(atom), 1), atom.reshape(len(atom), 1).T
+                        )
+                    ),
+                )
+            )
+            for atom in linear_per_atom
+        ]
+    )
 
 
 def calc_sum_quad(linear_sum):
-    return np.concatenate((
-        linear_sum,
-        convert_mat(mat=np.dot(linear_sum.reshape(len(linear_sum), 1), linear_sum.reshape(len(linear_sum), 1).T))
-    ))
+    return np.concatenate(
+        (
+            linear_sum,
+            convert_mat(
+                mat=np.dot(
+                    linear_sum.reshape(len(linear_sum), 1),
+                    linear_sum.reshape(len(linear_sum), 1).T,
+                )
+            ),
+        )
+    )
 
 
 def get_apre(cell):
@@ -39,7 +55,7 @@ def get_apre(cell):
     yhi = np.sin(gamma) * bn
     xzp = np.cos(beta) * cn
     yzp = (bn * cn * np.cos(alpha) - xyp * xzp) / yhi
-    zhi = np.sqrt(cn ** 2 - xzp ** 2 - yzp ** 2)
+    zhi = np.sqrt(cn**2 - xzp**2 - yzp**2)
 
     return np.array(((xhi, 0, 0), (xyp, yhi, 0), (xzp, yzp, zhi)))
 
@@ -88,7 +104,7 @@ def write_ase_structure(lmp, structure):
         x=(len(positions) * c_double)(*positions),
         v=None,
         image=None,
-        shrinkexceed=False
+        shrinkexceed=False,
     )
 
 
@@ -195,7 +211,10 @@ def calc_snap_descriptors_per_atom(lmp, structure, bispec_options, cutoff=10.0):
     except:
         return np.array([])
     else:
-        if "quadraticflag" in bispec_options.keys() and int(bispec_options["quadraticflag"]) == 1:
+        if (
+            "quadraticflag" in bispec_options.keys()
+            and int(bispec_options["quadraticflag"]) == 1
+        ):
             return extract_compute_np(
                 lmp=lmp,
                 name="b",
@@ -203,7 +222,10 @@ def calc_snap_descriptors_per_atom(lmp, structure, bispec_options, cutoff=10.0):
                 result_type=2,
                 # basically we only care about the off diagonal elements and from those we need only half
                 # plus the linear terms:   n + sum_{i: 1->n} i
-                array_shape=(len(structure), int(number_coef * (number_coef * (1 - 1 / 2) + 3 / 2)))
+                array_shape=(
+                    len(structure),
+                    int(number_coef * (number_coef * (1 - 1 / 2) + 3 / 2)),
+                ),
             )
         else:
             return extract_compute_np(
@@ -211,25 +233,30 @@ def calc_snap_descriptors_per_atom(lmp, structure, bispec_options, cutoff=10.0):
                 name="b",
                 compute_type=1,
                 result_type=2,
-                array_shape=(len(structure), number_coef)
+                array_shape=(len(structure), number_coef),
             )
 
 
 def lammps_variables(bispec_options):
-    d = {k: bispec_options[k] for k in
-         ["rcutfac",
-          "rfac0",
-          "rmin0",
-          # "diagonalstyle",
-          # "zblcutinner",
-          # "zblcutouter",
-          "twojmax"]}
+    d = {
+        k: bispec_options[k]
+        for k in [
+            "rcutfac",
+            "rfac0",
+            "rmin0",
+            # "diagonalstyle",
+            # "zblcutinner",
+            # "zblcutouter",
+            "twojmax",
+        ]
+    }
     d.update(
         {
             (k + str(i + 1)): bispec_options[k][i]
             for k in ["wj", "radelem"]  # ["zblz", "wj", "radelem"]
             for i, v in enumerate(bispec_options[k])
-        })
+        }
+    )
     return d
 
 
@@ -250,8 +277,7 @@ def set_computes_snappy(lmp, bispec_options):
 
     kw_options = {
         k: bispec_options[v]
-        for k, v in
-        {
+        for k, v in {
             # "diagonal": "diagonalstyle", For legacy diagonalstyle option
             "rmin0": "rmin0",
             "bzeroflag": "bzeroflag",
@@ -275,7 +301,9 @@ def set_computes_snappy(lmp, bispec_options):
 
 def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
     lmp_atom_ids = lmp.numpy.extract_atom_iarray("id", num_atoms).flatten()
-    assert np.all(lmp_atom_ids == 1 + np.arange(num_atoms)), "LAMMPS seems to have lost atoms"
+    assert np.all(
+        lmp_atom_ids == 1 + np.arange(num_atoms)
+    ), "LAMMPS seems to have lost atoms"
 
     # Extract types
     lmp_types = lmp.numpy.extract_atom_iarray(name="type", nelem=num_atoms).flatten()
@@ -296,22 +324,25 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
     b_sum = b_atom.sum(axis=0) / num_atoms
 
     try:
-        assert np.allclose(lmp_bsum, lmp_barr.sum(axis=0), rtol=1e-12, atol=1e-12), \
-            "b_sum doesn't match sum of b"
+        assert np.allclose(
+            lmp_bsum, lmp_barr.sum(axis=0), rtol=1e-12, atol=1e-12
+        ), "b_sum doesn't match sum of b"
     except AssertionError:
         print(lmp_bsum)
         print(lmp_barr.sum(axis=0))
 
     lmp_dbarr = extract_compute_np(lmp, "db", 1, 2, (num_atoms, num_types, 3, n_coeff))
     lmp_dbsum = extract_compute_np(lmp, "db_sum", 0, 1, (num_types, 3, n_coeff))
-    assert np.allclose(lmp_dbsum, lmp_dbarr.sum(axis=0), rtol=1e-12, atol=1e-12), \
-        "db_sum doesn't match sum of db"
+    assert np.allclose(
+        lmp_dbsum, lmp_dbarr.sum(axis=0), rtol=1e-12, atol=1e-12
+    ), "db_sum doesn't match sum of db"
     db_atom = np.transpose(lmp_dbarr, (0, 2, 1, 3))
 
     lmp_vbarr = extract_compute_np(lmp, "vb", 1, 2, (num_atoms, num_types, 6, n_coeff))
     lmp_vbsum = extract_compute_np(lmp, "vb_sum", 0, 1, (num_types, 6, n_coeff))
-    assert np.allclose(lmp_vbsum, lmp_vbarr.sum(axis=0), rtol=1e-12, atol=1e-12), \
-        "vb_sum doesn't match sum of vb"
+    assert np.allclose(
+        lmp_vbsum, lmp_vbarr.sum(axis=0), rtol=1e-12, atol=1e-12
+    ), "vb_sum doesn't match sum of vb"
     vb_sum = np.transpose(lmp_vbsum, (1, 0, 2)) / lmp_volume * eV_div_A3_to_bar
 
     dbatom_shape = db_atom.shape
@@ -319,12 +350,16 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
     a_fit = np.concatenate(
         (
             b_sum,
-            db_atom.reshape(dbatom_shape[0] * dbatom_shape[1], dbatom_shape[2] * dbatom_shape[3]),
-            vb_sum.reshape(vbsum_shape[0] * vbsum_shape[1], vbsum_shape[2])
+            db_atom.reshape(
+                dbatom_shape[0] * dbatom_shape[1], dbatom_shape[2] * dbatom_shape[3]
+            ),
+            vb_sum.reshape(vbsum_shape[0] * vbsum_shape[1], vbsum_shape[2]),
         ),
-        axis=0
+        axis=0,
     )
-    return np.concatenate((np.array([np.eye(a_fit.shape[0])[0]]).T, a_fit), axis=1).copy()
+    return np.concatenate(
+        (np.array([np.eye(a_fit.shape[0])[0]]).T, a_fit), axis=1
+    ).copy()
 
 
 def calc_snap_descriptor_derivatives(lmp, structure, bispec_options, cutoff=10.0):
@@ -343,12 +378,15 @@ def calc_snap_descriptor_derivatives(lmp, structure, bispec_options, cutoff=10.0
     except:
         return np.array([])
     else:
-        if "quadraticflag" in bispec_options.keys() and int(bispec_options["quadraticflag"]) == 1:
+        if (
+            "quadraticflag" in bispec_options.keys()
+            and int(bispec_options["quadraticflag"]) == 1
+        ):
             return extract_computes_snappy(
                 lmp=lmp,
                 num_atoms=len(structure),
                 n_coeff=int(number_coef * (number_coef * (1 - 1 / 2) + 3 / 2)),
-                num_types=number_species
+                num_types=number_species,
             )
         else:
             return extract_computes_snappy(

From 8ab9a86312fa9c45fe86a525217ead1fa60314f8 Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 12:06:16 +0200
Subject: [PATCH 05/34] Add LAMMPS dependency only for unix

---
 .ci_support/condamerge.py          | 57 ++++++++++++++++++++++++++++++
 .ci_support/environment-lammps.yml |  4 +++
 .ci_support/environment.yml        |  3 +-
 .github/workflows/unittests.yml    |  6 +++-
 4 files changed, 67 insertions(+), 3 deletions(-)
 create mode 100644 .ci_support/condamerge.py
 create mode 100644 .ci_support/environment-lammps.yml

diff --git a/.ci_support/condamerge.py b/.ci_support/condamerge.py
new file mode 100644
index 000000000..fb2df69eb
--- /dev/null
+++ b/.ci_support/condamerge.py
@@ -0,0 +1,57 @@
+import argparse
+import sys
+import yaml
+
+
+def read_file(path):
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+def parse_args(argv=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--base', dest='base', help='base environment.yml file')
+    parser.add_argument('--add', dest='add', help='addon environment.yml file')
+    return parser.parse_args(argv)
+
+
+def merge_dependencies(env_base, env_add):
+    base_dict = {f.split()[0]: f for f in env_base}
+    add_dict = {f.split()[0]: f for f in env_add}
+    for k,v in add_dict.items():
+        if k not in base_dict.keys():
+            base_dict[k] = v
+    return list(base_dict.values())
+
+
+def merge_channels(env_base, env_add):
+    for c in env_add:
+        if c not in env_base:
+            env_base.append(c)
+    return env_base
+
+
+def merge_env(env_base, env_add):
+    return {
+        "channels": merge_channels(
+            env_base=env_base['channels'],
+            env_add=env_add['channels']
+        ),
+        'dependencies': merge_dependencies(
+            env_base=env_base['dependencies'],
+            env_add=env_add['dependencies']
+        )
+    }
+
+
+if __name__ == '__main__':
+    arguments = parse_args(argv=None)
+    yaml.dump(
+        merge_env(
+            env_base=read_file(arguments.base),
+            env_add=read_file(arguments.add)
+        ),
+        sys.stdout,
+        indent=2,
+        default_flow_style=False
+    )
diff --git a/.ci_support/environment-lammps.yml b/.ci_support/environment-lammps.yml
new file mode 100644
index 000000000..675d3a4df
--- /dev/null
+++ b/.ci_support/environment-lammps.yml
@@ -0,0 +1,4 @@
+channels:
+- conda-forge
+dependencies:
+- lammps =2023.08.02  # [unix]
\ No newline at end of file
diff --git a/.ci_support/environment.yml b/.ci_support/environment.yml
index baccbe2bc..fa7675bd5 100644
--- a/.ci_support/environment.yml
+++ b/.ci_support/environment.yml
@@ -17,5 +17,4 @@ dependencies:
 - scipy =1.11.2
 - spglib =2.0.2
 - sqsgenerator =0.2
-- pyxtal =0.6.0
-- lammps =2023.08.02  # [unix]
\ No newline at end of file
+- pyxtal =0.6.0
\ No newline at end of file
diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index 6d96f23b1..72ab3161f 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -46,8 +46,12 @@ jobs:
 
     steps:
     - uses: actions/checkout@v2
-    - name: Setup environment
+    - name: Setup environment (windows)
+      if: matrix.operating-system == 'windows-latest'
       run: cp .ci_support/environment.yml environment.yml
+    - name: Setup environment (unix)
+      if: matrix.operating-system != 'windows-latest'
+      run: python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2
       with:

From 1f93d78e3ff20e6e659287ab3790b2ae9ca5ebd1 Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 12:08:36 +0200
Subject: [PATCH 06/34] Add shell

---
 .github/workflows/unittests.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index 72ab3161f..a3f8bcd28 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -51,6 +51,7 @@ jobs:
       run: cp .ci_support/environment.yml environment.yml
     - name: Setup environment (unix)
       if: matrix.operating-system != 'windows-latest'
+      shell: bash -l {0}
       run: python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2

From 204779e593c74030d3ca2e14c59d5f25006ac2ce Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 12:11:07 +0200
Subject: [PATCH 07/34] Add yaml

---
 .github/workflows/unittests.yml | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index a3f8bcd28..a52a85d76 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -52,7 +52,9 @@ jobs:
     - name: Setup environment (unix)
       if: matrix.operating-system != 'windows-latest'
       shell: bash -l {0}
-      run: python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
+      run: |
+        mamba install -y yaml
+        python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2
       with:

From 8fc6b2c372b3cc94839c9fa3dcac9edf45d084fd Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 13:07:36 +0200
Subject: [PATCH 08/34] add pyyaml

---
 .github/workflows/unittests.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index a52a85d76..fe71da99d 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -53,7 +53,7 @@ jobs:
       if: matrix.operating-system != 'windows-latest'
       shell: bash -l {0}
       run: |
-        mamba install -y yaml
+        mamba install -y pyyaml
         python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2

From f1130ed4d683720eb545e2f2d94f2d51bba9a92d Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 13:11:16 +0200
Subject: [PATCH 09/34] fix order

---
 .github/workflows/unittests.yml | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index fe71da99d..293072530 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -46,15 +46,6 @@ jobs:
 
     steps:
     - uses: actions/checkout@v2
-    - name: Setup environment (windows)
-      if: matrix.operating-system == 'windows-latest'
-      run: cp .ci_support/environment.yml environment.yml
-    - name: Setup environment (unix)
-      if: matrix.operating-system != 'windows-latest'
-      shell: bash -l {0}
-      run: |
-        mamba install -y pyyaml
-        python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2
       with:
@@ -64,6 +55,14 @@ jobs:
         channel-priority: strict
         activate-environment: my-env
         use-mamba: true
+    - name: Setup environment (windows)
+      if: matrix.operating-system == 'windows-latest'
+      run: cp .ci_support/environment.yml environment.yml
+    - name: Setup environment (unix)
+      if: matrix.operating-system != 'windows-latest'
+      shell: bash -l {0}
+      run: |
+        python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Update environment
       run: mamba env update -n my-env -f environment.yml
     - name: Setup

From da7df818dda66c125487f43ae6839fb8c322c378 Mon Sep 17 00:00:00 2001
From: Jan Janssen <jan.janssen@outlook.com>
Date: Mon, 11 Sep 2023 13:14:23 +0200
Subject: [PATCH 10/34] install pyyaml

---
 .github/workflows/unittests.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index 293072530..51baf441f 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -62,6 +62,7 @@ jobs:
       if: matrix.operating-system != 'windows-latest'
       shell: bash -l {0}
       run: |
+        mamba install -y pyyaml
         python .ci_support/condamerge.py --base .ci_support/environment.yml --add .ci_support/environment-lammps.yml > environment.yml
     - name: Update environment
       run: mamba env update -n my-env -f environment.yml

From b26a795d51c43bd958d21cfcf56065c8391918a8 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Sun, 14 Jan 2024 19:27:41 +0100
Subject: [PATCH 11/34] simplify the API

---
 structuretoolkit/analyse/snap.py | 69 +++++++++++++++++++++++++++++++-
 tests/test_snap.py               |  6 +--
 2 files changed, 70 insertions(+), 5 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 528c2103b..d552c31a2 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -1,4 +1,5 @@
 from ctypes import c_double, c_int, cast, POINTER
+from lammps import lammps
 import numpy as np
 from scipy.constants import physical_constants
 
@@ -196,7 +197,7 @@ def set_compute_lammps(lmp, bispec_options, numtypes):
     lmp.command(f"{base_b} {radelem} {wj} {kwargs}")
 
 
-def calc_snap_descriptors_per_atom(lmp, structure, bispec_options, cutoff=10.0):
+def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
     number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     reset_lmp(lmp=lmp)
     write_ase_structure(lmp=lmp, structure=structure)
@@ -362,7 +363,7 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
     ).copy()
 
 
-def calc_snap_descriptor_derivatives(lmp, structure, bispec_options, cutoff=10.0):
+def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
     number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     number_species = len(set(structure.get_chemical_symbols()))
     reset_lmp(lmp=lmp)
@@ -395,3 +396,67 @@ def calc_snap_descriptor_derivatives(lmp, structure, bispec_options, cutoff=10.0
                 n_coeff=number_coef,
                 num_types=number_species,
             )
+
+
+def get_default_parameters(atom_types, twojmax=6, element_radius=4.0, rcutfac=1.0, rfac0=0.99363, rmin0=0.0, bzeroflag=0, weights=None, cutoff=10.0):
+    if weights is None:
+        wj = [1.0] * len(atom_types)
+    else:
+        wj = weights
+    if isinstance(element_radius, float):
+        radelem = [element_radius] * len(atom_types)
+    else:
+        radelem = element_radius
+    bispec_options = {
+        "numtypes": len(atom_types),
+        "twojmax": twojmax,
+        "rcutfac": rcutfac,
+        "rfac0": rfac0,
+        "rmin0": rmin0,
+        "bzeroflag": bzeroflag,
+        "radelem": radelem,
+        "type": atom_types,
+        "wj": wj,
+    }
+    lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
+    return lmp, bispec_options, cutoff
+
+
+def calc_snap_descriptors_per_atom(structure, atom_types, twojmax=6, element_radius=4.0, rcutfac=1.0, rfac0=0.99363, rmin0=0.0, bzeroflag=0, weights=None, cutoff=10.0):
+    lmp, bispec_options, cutoff = get_default_parameters(
+        atom_types=atom_types,
+        twojmax=twojmax,
+        element_radius=element_radius,
+        rcutfac=rcutfac,
+        rfac0=rfac0,
+        rmin0=rmin0,
+        bzeroflag=bzeroflag,
+        weights=weights,
+        cutoff=cutoff,
+    )
+    return _calc_snap_per_atom(
+        lmp=lmp,
+        structure=structure,
+        bispec_options=bispec_options,
+        cutoff=cutoff
+    )
+
+
+def calc_snap_descriptor_derivatives(structure, atom_types, twojmax=6, element_radius=4.0, rcutfac=1.0, rfac0=0.99363, rmin0=0.0, bzeroflag=0, weights=None, cutoff=10.0):
+    lmp, bispec_options, cutoff = get_default_parameters(
+        atom_types=atom_types,
+        twojmax=twojmax,
+        element_radius=element_radius,
+        rcutfac=rcutfac,
+        rfac0=rfac0,
+        rmin0=rmin0,
+        bzeroflag=bzeroflag,
+        weights=weights,
+        cutoff=cutoff,
+    )
+    return _calc_snap_derivatives(
+        lmp=lmp,
+        structure=structure,
+        bispec_options=bispec_options,
+        cutoff=cutoff
+    )
diff --git a/tests/test_snap.py b/tests/test_snap.py
index dcc69e5f6..eb7076647 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -36,7 +36,7 @@ def test_calc_bispectrum_lmp(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
-        coeff = stk.analyse.calc_snap_descriptors_per_atom(
+        coeff = stk.analyse._calc_snap_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=self.bispec_options,
@@ -50,7 +50,7 @@ def test_calc_bispectrum_lmp_quad(self):
         ))
         bispec_options = self.bispec_options.copy()
         bispec_options["quadraticflag"] = 1
-        coeff = stk.analyse.calc_snap_descriptors_per_atom(
+        coeff = stk.analyse._calc_snap_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=bispec_options,
@@ -62,7 +62,7 @@ def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
-        mat_a = stk.analyse.calc_snap_descriptor_derivatives(
+        mat_a = stk.analyse._calc_snap_derivatives(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=self.bispec_options,

From 4f994197afedb2179740b0edb5bc23a9406dc88f Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 13:18:27 -0600
Subject: [PATCH 12/34] Add more docstrings, refactor the order of the
 functions and separate internal and external functions.

---
 structuretoolkit/analyse/snap.py | 297 ++++++++++++++++++++++---------
 1 file changed, 210 insertions(+), 87 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index d552c31a2..1d1ae825d 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -6,18 +6,23 @@
 eV_div_A3_to_bar = 1e25 / physical_constants["joule-electron volt relationship"][0]
 
 
-def convert_mat(mat):
-    mat[np.diag_indices_from(mat)] /= 2
-    return mat[np.triu_indices(len(mat))]
+def calc_per_atom_quad(linear_per_atom):
+    """
+    Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
+    the SNAP descriptors.
 
+    Args:
+        linear_per_atom (np.ndarray): Numpy array of the linear per atom SNAP descriptors
 
-def calc_per_atom_quad(linear_per_atom):
+    Returns:
+        np.ndarray: Numpy array of the quadratic per atom SNAP descriptors
+    """
     return np.array(
         [
             np.concatenate(
                 (
                     atom,
-                    convert_mat(
+                    _convert_mat(
                         mat=np.dot(
                             atom.reshape(len(atom), 1), atom.reshape(len(atom), 1).T
                         )
@@ -33,7 +38,7 @@ def calc_sum_quad(linear_sum):
     return np.concatenate(
         (
             linear_sum,
-            convert_mat(
+            _convert_mat(
                 mat=np.dot(
                     linear_sum.reshape(len(linear_sum), 1),
                     linear_sum.reshape(len(linear_sum), 1).T,
@@ -43,7 +48,114 @@ def calc_sum_quad(linear_sum):
     )
 
 
+def calc_snap_descriptors_per_atom(
+    structure,
+    atom_types,
+    twojmax=6,
+    element_radius=4.0,
+    rcutfac=1.0,
+    rfac0=0.99363,
+    rmin0=0.0,
+    bzeroflag=0,
+    weights=None,
+    cutoff=10.0,
+):
+    """
+    Calculate per atom SNAP descriptors
+
+    Args:
+        structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
+        atom_types:
+        twojmax (int):
+        element_radius (float):
+        rcutfac (float):
+        rfac0 (float):
+        rmin0 (float):
+        bzeroflag (bool):
+        weights (list/np.ndarry/None):
+        cutoff (float):
+
+    Returns:
+        np.ndarray: Numpy array with the calculated descriptor derivatives
+    """
+    lmp, bispec_options, cutoff = _get_default_parameters(
+        atom_types=atom_types,
+        twojmax=twojmax,
+        element_radius=element_radius,
+        rcutfac=rcutfac,
+        rfac0=rfac0,
+        rmin0=rmin0,
+        bzeroflag=bzeroflag,
+        weights=weights,
+        cutoff=cutoff,
+    )
+    return _calc_snap_per_atom(
+        lmp=lmp,
+        structure=structure,
+        bispec_options=bispec_options,
+        cutoff=cutoff
+    )
+
+
+def calc_snap_descriptor_derivatives(
+    structure,
+    atom_types,
+    twojmax=6,
+    element_radius=4.0,
+    rcutfac=1.0,
+    rfac0=0.99363,
+    rmin0=0.0,
+    bzeroflag=0,
+    weights=None,
+    cutoff=10.0,
+):
+    """
+    Calculate per atom derivatives of the SNAP descriptors.
+
+    Args:
+        structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
+        atom_types:
+        twojmax (int):
+        element_radius (float):
+        rcutfac (float):
+        rfac0 (float):
+        rmin0 (float):
+        bzeroflag (bool):
+        weights (list/np.ndarry/None):
+        cutoff (float):
+
+    Returns:
+        np.ndarray: Numpy array with the calculated descriptor derivatives
+    """
+    lmp, bispec_options, cutoff = _get_default_parameters(
+        atom_types=atom_types,
+        twojmax=twojmax,
+        element_radius=element_radius,
+        rcutfac=rcutfac,
+        rfac0=rfac0,
+        rmin0=rmin0,
+        bzeroflag=bzeroflag,
+        weights=weights,
+        cutoff=cutoff,
+    )
+    return _calc_snap_derivatives(
+        lmp=lmp,
+        structure=structure,
+        bispec_options=bispec_options,
+        cutoff=cutoff
+    )
+
+
 def get_apre(cell):
+    """
+    Convert ASE cell to LAMMPS cell - LAMMPS required the upper triangle to be zero
+
+    Args:
+        cell (np.ndarray): ASE cell as 3x3 matrix
+
+    Returns:
+        np.ndarray: LAMMPS cell as 3x3 matrix
+    """
     a, b, c = cell
     an, bn, cn = [np.linalg.norm(v) for v in cell]
 
@@ -56,12 +168,36 @@ def get_apre(cell):
     yhi = np.sin(gamma) * bn
     xzp = np.cos(beta) * cn
     yzp = (bn * cn * np.cos(alpha) - xyp * xzp) / yhi
-    zhi = np.sqrt(cn**2 - xzp**2 - yzp**2)
+    zhi = np.sqrt(cn ** 2 - xzp ** 2 - yzp ** 2)
 
     return np.array(((xhi, 0, 0), (xyp, yhi, 0), (xzp, yzp, zhi)))
 
 
-def write_ase_structure(lmp, structure):
+def get_snap_descriptor_names(twojmax):
+    """
+    Get names of the SNAP descriptors
+
+    Args:
+        twojmax (int):
+
+    Returns:
+        list: List of SNAP descriptor names
+    """
+    lst = []
+    for j1 in range(0, twojmax + 1):
+        for j2 in range(0, j1 + 1):
+            for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
+                if j >= j1:
+                    lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
+    return lst
+
+
+def _convert_mat(mat):
+    mat[np.diag_indices_from(mat)] /= 2
+    return mat[np.triu_indices(len(mat))]
+
+
+def _write_ase_structure(lmp, structure):
     number_species = len(set(structure.get_chemical_symbols()))
 
     apre = get_apre(cell=structure.cell)
@@ -109,7 +245,7 @@ def write_ase_structure(lmp, structure):
     )
 
 
-def extract_compute_np(lmp, name, compute_type, result_type, array_shape):
+def _extract_compute_np(lmp, name, compute_type, result_type, array_shape):
     """
     Convert a lammps compute to a numpy array.
     Assumes the compute returns a floating point numbers.
@@ -130,7 +266,13 @@ def extract_compute_np(lmp, name, compute_type, result_type, array_shape):
     return array_np.copy()
 
 
-def reset_lmp(lmp):
+def _reset_lmp(lmp):
+    """
+    Reset the LAMMPS library instance
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+    """
     lmp.command("clear")
     lmp.command("units metal")
     lmp.command("dimension 3")
@@ -139,7 +281,14 @@ def reset_lmp(lmp):
     lmp.command("atom_modify map array sort 0 2.0")
 
 
-def set_potential_lmp(lmp, cutoff=10.0):
+def _set_potential_lmp(lmp, cutoff=10.0):
+    """
+    Set interatomic potential parameters to LAMMPS library instance
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        cutoff (float): cutoff radius for the construction of the neighbor list
+    """
     lmp.command("pair_style zero " + str(cutoff))
     lmp.command("pair_coeff * *")
     lmp.command("mass * 1.0e-20")
@@ -147,17 +296,7 @@ def set_potential_lmp(lmp, cutoff=10.0):
     lmp.command("neigh_modify one 10000")
 
 
-def get_snap_descriptor_names(twojmax):
-    lst = []
-    for j1 in range(0, twojmax + 1):
-        for j2 in range(0, j1 + 1):
-            for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
-                if j >= j1:
-                    lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
-    return lst
-
-
-def set_compute_lammps(lmp, bispec_options, numtypes):
+def _set_compute_lammps(lmp, bispec_options, numtypes):
     lmp_variable_args = {
         k: bispec_options[k] for k in ["rcutfac", "rfac0", "rmin0", "twojmax"]
     }
@@ -199,10 +338,10 @@ def set_compute_lammps(lmp, bispec_options, numtypes):
 
 def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
     number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
-    reset_lmp(lmp=lmp)
-    write_ase_structure(lmp=lmp, structure=structure)
-    set_potential_lmp(lmp=lmp, cutoff=cutoff)
-    set_compute_lammps(
+    _reset_lmp(lmp=lmp)
+    _write_ase_structure(lmp=lmp, structure=structure)
+    _set_potential_lmp(lmp=lmp, cutoff=cutoff)
+    _set_compute_lammps(
         lmp=lmp,
         bispec_options=bispec_options,
         numtypes=len(set(structure.get_chemical_symbols())),
@@ -216,7 +355,7 @@ def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
             "quadraticflag" in bispec_options.keys()
             and int(bispec_options["quadraticflag"]) == 1
         ):
-            return extract_compute_np(
+            return _extract_compute_np(
                 lmp=lmp,
                 name="b",
                 compute_type=1,
@@ -229,7 +368,7 @@ def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
                 ),
             )
         else:
-            return extract_compute_np(
+            return _extract_compute_np(
                 lmp=lmp,
                 name="b",
                 compute_type=1,
@@ -238,7 +377,7 @@ def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
             )
 
 
-def lammps_variables(bispec_options):
+def _lammps_variables(bispec_options):
     d = {
         k: bispec_options[k]
         for k in [
@@ -261,12 +400,26 @@ def lammps_variables(bispec_options):
     return d
 
 
-def set_variables(lmp, **lmp_variable_args):
+def _set_variables(lmp, **lmp_variable_args):
+    """
+    Internal helper function to set LAMMPS variables
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        **lmp_variable_args (dict): key value pairs of LAMMPS variables to set
+    """
     for k, v in lmp_variable_args.items():
         lmp.command(f"variable {k} equal {v}")
 
 
-def set_computes_snappy(lmp, bispec_options):
+def _set_computes_snap(lmp, bispec_options):
+    """
+    Set LAMMPS computes to calculate SNAP descriptors
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        bispec_options (dict): bi-spectrum component settings
+    """
     # # Bispectrum coefficient computes
     base_b = "compute b all sna/atom ${rcutfac} ${rfac0} ${twojmax}"
     base_db = "compute db all snad/atom ${rcutfac} ${rfac0} ${twojmax}"
@@ -300,7 +453,7 @@ def set_computes_snappy(lmp, bispec_options):
         lmp.command(f"compute {cname}_sum all reduce sum c_{cname}[*]")
 
 
-def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
+def _extract_computes_snap(lmp, num_atoms, n_coeff, num_types):
     lmp_atom_ids = lmp.numpy.extract_atom_iarray("id", num_atoms).flatten()
     assert np.all(
         lmp_atom_ids == 1 + np.arange(num_atoms)
@@ -311,10 +464,10 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
     lmp_volume = lmp.get_thermo("vol")
 
     # Extract Bsum
-    lmp_bsum = extract_compute_np(lmp, "b_sum", 0, 1, (n_coeff))
+    lmp_bsum = _extract_compute_np(lmp, "b_sum", 0, 1, (n_coeff))
 
     # Extract B
-    lmp_barr = extract_compute_np(lmp, "b", 1, 2, (num_atoms, n_coeff))
+    lmp_barr = _extract_compute_np(lmp, "b", 1, 2, (num_atoms, n_coeff))
 
     type_onehot = np.eye(num_types)[lmp_types - 1]  # lammps types are 1-indexed.
     # has shape n_atoms, n_types, num_coeffs.
@@ -332,15 +485,15 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
         print(lmp_bsum)
         print(lmp_barr.sum(axis=0))
 
-    lmp_dbarr = extract_compute_np(lmp, "db", 1, 2, (num_atoms, num_types, 3, n_coeff))
-    lmp_dbsum = extract_compute_np(lmp, "db_sum", 0, 1, (num_types, 3, n_coeff))
+    lmp_dbarr = _extract_compute_np(lmp, "db", 1, 2, (num_atoms, num_types, 3, n_coeff))
+    lmp_dbsum = _extract_compute_np(lmp, "db_sum", 0, 1, (num_types, 3, n_coeff))
     assert np.allclose(
         lmp_dbsum, lmp_dbarr.sum(axis=0), rtol=1e-12, atol=1e-12
     ), "db_sum doesn't match sum of db"
     db_atom = np.transpose(lmp_dbarr, (0, 2, 1, 3))
 
-    lmp_vbarr = extract_compute_np(lmp, "vb", 1, 2, (num_atoms, num_types, 6, n_coeff))
-    lmp_vbsum = extract_compute_np(lmp, "vb_sum", 0, 1, (num_types, 6, n_coeff))
+    lmp_vbarr = _extract_compute_np(lmp, "vb", 1, 2, (num_atoms, num_types, 6, n_coeff))
+    lmp_vbsum = _extract_compute_np(lmp, "vb_sum", 0, 1, (num_types, 6, n_coeff))
     assert np.allclose(
         lmp_vbsum, lmp_vbarr.sum(axis=0), rtol=1e-12, atol=1e-12
     ), "vb_sum doesn't match sum of vb"
@@ -366,11 +519,11 @@ def extract_computes_snappy(lmp, num_atoms, n_coeff, num_types):
 def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
     number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     number_species = len(set(structure.get_chemical_symbols()))
-    reset_lmp(lmp=lmp)
-    write_ase_structure(lmp=lmp, structure=structure)
-    set_potential_lmp(lmp=lmp, cutoff=cutoff)
-    set_variables(lmp, **lammps_variables(bispec_options))
-    set_computes_snappy(
+    _reset_lmp(lmp=lmp)
+    _write_ase_structure(lmp=lmp, structure=structure)
+    _set_potential_lmp(lmp=lmp, cutoff=cutoff)
+    _set_variables(lmp, **_lammps_variables(bispec_options))
+    _set_computes_snap(
         lmp=lmp,
         bispec_options=bispec_options,
     )
@@ -383,14 +536,14 @@ def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
             "quadraticflag" in bispec_options.keys()
             and int(bispec_options["quadraticflag"]) == 1
         ):
-            return extract_computes_snappy(
+            return _extract_computes_snap(
                 lmp=lmp,
                 num_atoms=len(structure),
                 n_coeff=int(number_coef * (number_coef * (1 - 1 / 2) + 3 / 2)),
                 num_types=number_species,
             )
         else:
-            return extract_computes_snappy(
+            return _extract_computes_snap(
                 lmp=lmp,
                 num_atoms=len(structure),
                 n_coeff=number_coef,
@@ -398,7 +551,17 @@ def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
             )
 
 
-def get_default_parameters(atom_types, twojmax=6, element_radius=4.0, rcutfac=1.0, rfac0=0.99363, rmin0=0.0, bzeroflag=0, weights=None, cutoff=10.0):
+def _get_default_parameters(
+    atom_types,
+    twojmax=6,
+    element_radius=4.0,
+    rcutfac=1.0,
+    rfac0=0.99363,
+    rmin0=0.0,
+    bzeroflag=0,
+    weights=None,
+    cutoff=10.0
+):
     if weights is None:
         wj = [1.0] * len(atom_types)
     else:
@@ -419,44 +582,4 @@ def get_default_parameters(atom_types, twojmax=6, element_radius=4.0, rcutfac=1.
         "wj": wj,
     }
     lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
-    return lmp, bispec_options, cutoff
-
-
-def calc_snap_descriptors_per_atom(structure, atom_types, twojmax=6, element_radius=4.0, rcutfac=1.0, rfac0=0.99363, rmin0=0.0, bzeroflag=0, weights=None, cutoff=10.0):
-    lmp, bispec_options, cutoff = get_default_parameters(
-        atom_types=atom_types,
-        twojmax=twojmax,
-        element_radius=element_radius,
-        rcutfac=rcutfac,
-        rfac0=rfac0,
-        rmin0=rmin0,
-        bzeroflag=bzeroflag,
-        weights=weights,
-        cutoff=cutoff,
-    )
-    return _calc_snap_per_atom(
-        lmp=lmp,
-        structure=structure,
-        bispec_options=bispec_options,
-        cutoff=cutoff
-    )
-
-
-def calc_snap_descriptor_derivatives(structure, atom_types, twojmax=6, element_radius=4.0, rcutfac=1.0, rfac0=0.99363, rmin0=0.0, bzeroflag=0, weights=None, cutoff=10.0):
-    lmp, bispec_options, cutoff = get_default_parameters(
-        atom_types=atom_types,
-        twojmax=twojmax,
-        element_radius=element_radius,
-        rcutfac=rcutfac,
-        rfac0=rfac0,
-        rmin0=rmin0,
-        bzeroflag=bzeroflag,
-        weights=weights,
-        cutoff=cutoff,
-    )
-    return _calc_snap_derivatives(
-        lmp=lmp,
-        structure=structure,
-        bispec_options=bispec_options,
-        cutoff=cutoff
-    )
+    return lmp, bispec_options, cutoff
\ No newline at end of file

From 416a9cce738ba9861a9a914696962e9591ea57f7 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 13:19:12 -0600
Subject: [PATCH 13/34] black formatting

---
 structuretoolkit/analyse/snap.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 1d1ae825d..0d77ba428 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -90,10 +90,7 @@ def calc_snap_descriptors_per_atom(
         cutoff=cutoff,
     )
     return _calc_snap_per_atom(
-        lmp=lmp,
-        structure=structure,
-        bispec_options=bispec_options,
-        cutoff=cutoff
+        lmp=lmp, structure=structure, bispec_options=bispec_options, cutoff=cutoff
     )
 
 
@@ -139,10 +136,7 @@ def calc_snap_descriptor_derivatives(
         cutoff=cutoff,
     )
     return _calc_snap_derivatives(
-        lmp=lmp,
-        structure=structure,
-        bispec_options=bispec_options,
-        cutoff=cutoff
+        lmp=lmp, structure=structure, bispec_options=bispec_options, cutoff=cutoff
     )
 
 
@@ -168,7 +162,7 @@ def get_apre(cell):
     yhi = np.sin(gamma) * bn
     xzp = np.cos(beta) * cn
     yzp = (bn * cn * np.cos(alpha) - xyp * xzp) / yhi
-    zhi = np.sqrt(cn ** 2 - xzp ** 2 - yzp ** 2)
+    zhi = np.sqrt(cn**2 - xzp**2 - yzp**2)
 
     return np.array(((xhi, 0, 0), (xyp, yhi, 0), (xzp, yzp, zhi)))
 
@@ -560,7 +554,7 @@ def _get_default_parameters(
     rmin0=0.0,
     bzeroflag=0,
     weights=None,
-    cutoff=10.0
+    cutoff=10.0,
 ):
     if weights is None:
         wj = [1.0] * len(atom_types)
@@ -582,4 +576,4 @@ def _get_default_parameters(
         "wj": wj,
     }
     lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
-    return lmp, bispec_options, cutoff
\ No newline at end of file
+    return lmp, bispec_options, cutoff

From 2c9f6315f7411ee0fe102c029dab379a8b523dfc Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 13:21:55 -0600
Subject: [PATCH 14/34] remove print functions

---
 structuretoolkit/analyse/snap.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 0d77ba428..1a666c779 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -438,9 +438,7 @@ def _set_computes_snap(lmp, bispec_options):
     kwargs = " ".join(kw_substrings)
 
     for op, base in zip(("b", "db", "vb"), (base_b, base_db, base_vb)):
-        # print("Setting up compute",op)
         command = f"{base} {radelem} {wj} {kwargs}"
-        # print(command)
         lmp.command(command)
 
     for cname in ["b", "db", "vb"]:
@@ -471,14 +469,6 @@ def _extract_computes_snap(lmp, num_atoms, n_coeff, num_types):
     b_atom = type_onehot[:, :, np.newaxis] * lmp_barr[:, np.newaxis, :]
     b_sum = b_atom.sum(axis=0) / num_atoms
 
-    try:
-        assert np.allclose(
-            lmp_bsum, lmp_barr.sum(axis=0), rtol=1e-12, atol=1e-12
-        ), "b_sum doesn't match sum of b"
-    except AssertionError:
-        print(lmp_bsum)
-        print(lmp_barr.sum(axis=0))
-
     lmp_dbarr = _extract_compute_np(lmp, "db", 1, 2, (num_atoms, num_types, 3, n_coeff))
     lmp_dbsum = _extract_compute_np(lmp, "db_sum", 0, 1, (num_types, 3, n_coeff))
     assert np.allclose(

From 0f881a46772c070362edcceb315c8980b1a46c0f Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 14:37:32 -0600
Subject: [PATCH 15/34] fix tests

---
 structuretoolkit/analyse/snap.py |   4 +-
 tests/test_snap.py               | 122 +++++++++++++++++++++++++++++--
 2 files changed, 119 insertions(+), 7 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 1a666c779..6a1e9a8d4 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -67,7 +67,7 @@ def calc_snap_descriptors_per_atom(
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
         atom_types:
         twojmax (int):
-        element_radius (float):
+        element_radius (list):
         rcutfac (float):
         rfac0 (float):
         rmin0 (float):
@@ -113,7 +113,7 @@ def calc_snap_descriptor_derivatives(
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
         atom_types:
         twojmax (int):
-        element_radius (float):
+        element_radius (list):
         rcutfac (float):
         rfac0 (float):
         rmin0 (float):
diff --git a/tests/test_snap.py b/tests/test_snap.py
index eb7076647..5ad4781ab 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -1,5 +1,6 @@
 from ase.build.bulk import bulk
 import structuretoolkit as stk
+from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives
 import unittest
 
 
@@ -14,7 +15,118 @@
 @unittest.skipIf(
     skip_snap_test, "LAMMPS is not installed, so the SNAP tests are skipped."
 )
-class TestCu(unittest.TestCase):
+class TestSNAP(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
+        cls.structure = bulk("Cu", cubic=True)
+        cls.numtypes = 1
+        cls.twojmax = 6
+        cls.rcutfac = 1.0
+        cls.rfac0 = 0.99363
+        cls.rmin0 = 0.0
+        cls.bzeroflag = 0
+        cls.radelem = [4.0]
+        cls.type = ['Cu']
+        cls.wj = [1.0]
+
+    def test_get_descriptor_name(self):
+        descriptors = [
+            [0.0, 0.0, 0.0],
+            [0.5, 0.0, 0.5],
+            [0.5, 0.5, 1.0],
+            [1.0, 0.0, 1.0],
+            [1.0, 0.5, 1.5],
+            [1.0, 1.0, 1.0],
+            [1.0, 1.0, 2.0],
+            [1.5, 0.0, 1.5],
+            [1.5, 0.5, 2.0],
+            [1.5, 1.0, 1.5],
+            [1.5, 1.0, 2.5],
+            [1.5, 1.5, 2.0],
+            [1.5, 1.5, 3.0],
+            [2.0, 0.0, 2.0],
+            [2.0, 0.5, 2.5],
+            [2.0, 1.0, 2.0],
+            [2.0, 1.0, 3.0],
+            [2.0, 1.5, 2.5],
+            [2.0, 2.0, 2.0],
+            [2.0, 2.0, 3.0],
+            [2.5, 0.0, 2.5],
+            [2.5, 0.5, 3.0],
+            [2.5, 1.0, 2.5],
+            [2.5, 1.5, 3.0],
+            [2.5, 2.0, 2.5],
+            [2.5, 2.5, 3.0],
+            [3.0, 0.0, 3.0],
+            [3.0, 1.0, 3.0],
+            [3.0, 2.0, 3.0],
+            [3.0, 3.0, 3.0]
+        ]
+        names = stk.analyse.get_snap_descriptor_names(
+            twojmax=self.twojmax
+        )
+        self.assertEqual(descriptors, names)
+
+    def test_calc_bispectrum_lmp(self):
+        n_coeff = len(stk.analyse.get_snap_descriptor_names(
+            twojmax=self.twojmax
+        ))
+        coeff = stk.analyse.calc_snap_descriptors_per_atom(
+            structure=self.structure,
+            atom_types=self.type,
+            twojmax=self.twojmax,
+            element_radius=self.radelem,
+            rcutfac=self.rcutfac,
+            rfac0=self.rfac0,
+            rmin0=self.rmin0,
+            bzeroflag=self.bzeroflag,
+            weights=self.wj,
+            cutoff=10.0,
+        )
+        self.assertTrue(coeff.shape, (len(self.structure), n_coeff))
+
+    def test_calc_bispectrum_lmp_quad(self):
+        n_coeff = len(stk.analyse.get_snap_descriptor_names(
+            twojmax=self.twojmax
+        ))
+        coeff = stk.analyse.calc_snap_descriptors_per_atom(
+                structure=self.structure,
+                atom_types=self.type,
+                twojmax=self.twojmax,
+                element_radius=self.radelem,
+                rcutfac=self.rcutfac,
+                rfac0=self.rfac0,
+                rmin0=self.rmin0,
+                bzeroflag=self.bzeroflag,
+                weights=self.wj,
+                cutoff=10.0,
+            )
+        self.assertTrue(coeff.shape, (len(self.structure), n_coeff * 31))
+
+    def test_calc_a_matrix_snappy(self):
+        n_coeff = len(stk.analyse.get_snap_descriptor_names(
+            twojmax=self.twojmax
+        ))
+        mat_a = stk.analyse.calc_snap_descriptor_derivatives(
+            structure=self.structure,
+            atom_types=self.type,
+            twojmax=self.twojmax,
+            element_radius=self.radelem,
+            rcutfac=self.rcutfac,
+            rfac0=self.rfac0,
+            rmin0=self.rmin0,
+            bzeroflag=self.bzeroflag,
+            weights=self.wj,
+            cutoff=10.0,
+        )
+        self.assertTrue(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff))
+
+
+@unittest.skipIf(
+    skip_snap_test, "LAMMPS is not installed, so the SNAP tests are skipped."
+)
+class TestSNAPInternal(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         cls.lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
@@ -36,7 +148,7 @@ def test_calc_bispectrum_lmp(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
-        coeff = stk.analyse._calc_snap_per_atom(
+        coeff = _calc_snap_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=self.bispec_options,
@@ -50,7 +162,7 @@ def test_calc_bispectrum_lmp_quad(self):
         ))
         bispec_options = self.bispec_options.copy()
         bispec_options["quadraticflag"] = 1
-        coeff = stk.analyse._calc_snap_per_atom(
+        coeff = _calc_snap_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=bispec_options,
@@ -62,10 +174,10 @@ def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
-        mat_a = stk.analyse._calc_snap_derivatives(
+        mat_a = _calc_snap_derivatives(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=self.bispec_options,
             cutoff=10.0
         )
-        self.assertTrue(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff))
\ No newline at end of file
+        self.assertTrue(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff))

From 37b253eb2d6c8bda94334b9ca9016b9f2a083d7d Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 14:45:42 -0600
Subject: [PATCH 16/34] Move LAMMPS import to _get_default_parameters()
 function

---
 structuretoolkit/analyse/snap.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 6a1e9a8d4..7077945fa 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -1,5 +1,4 @@
 from ctypes import c_double, c_int, cast, POINTER
-from lammps import lammps
 import numpy as np
 from scipy.constants import physical_constants
 
@@ -546,6 +545,8 @@ def _get_default_parameters(
     weights=None,
     cutoff=10.0,
 ):
+    from lammps import lammps
+
     if weights is None:
         wj = [1.0] * len(atom_types)
     else:

From 61c15bd9b6ca82ddcb5b0f455dc504f43b298eb4 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 15:15:13 -0600
Subject: [PATCH 17/34] add quadratic tests

---
 structuretoolkit/analyse/snap.py | 17 +++++++++--
 tests/test_snap.py               | 52 ++++++++++++++++++++------------
 2 files changed, 47 insertions(+), 22 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 7077945fa..9c4dc2770 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -55,7 +55,8 @@ def calc_snap_descriptors_per_atom(
     rcutfac=1.0,
     rfac0=0.99363,
     rmin0=0.0,
-    bzeroflag=0,
+    bzeroflag=False,
+    quadraticflag=False,
     weights=None,
     cutoff=10.0,
 ):
@@ -71,6 +72,7 @@ def calc_snap_descriptors_per_atom(
         rfac0 (float):
         rmin0 (float):
         bzeroflag (bool):
+        quadraticflag (bool):
         weights (list/np.ndarry/None):
         cutoff (float):
 
@@ -85,6 +87,7 @@ def calc_snap_descriptors_per_atom(
         rfac0=rfac0,
         rmin0=rmin0,
         bzeroflag=bzeroflag,
+        quadraticflag=quadraticflag,
         weights=weights,
         cutoff=cutoff,
     )
@@ -102,6 +105,7 @@ def calc_snap_descriptor_derivatives(
     rfac0=0.99363,
     rmin0=0.0,
     bzeroflag=0,
+    quadraticflag=0,
     weights=None,
     cutoff=10.0,
 ):
@@ -131,6 +135,7 @@ def calc_snap_descriptor_derivatives(
         rfac0=rfac0,
         rmin0=rmin0,
         bzeroflag=bzeroflag,
+        quadraticflag=quadraticflag,
         weights=weights,
         cutoff=cutoff,
     )
@@ -544,6 +549,7 @@ def _get_default_parameters(
     bzeroflag=0,
     weights=None,
     cutoff=10.0,
+    quadraticflag=0,
 ):
     from lammps import lammps
 
@@ -561,10 +567,17 @@ def _get_default_parameters(
         "rcutfac": rcutfac,
         "rfac0": rfac0,
         "rmin0": rmin0,
-        "bzeroflag": bzeroflag,
         "radelem": radelem,
         "type": atom_types,
         "wj": wj,
     }
+    if bzeroflag:
+        bispec_options["bzeroflag"] = 1
+    else:
+        bispec_options["bzeroflag"] = 0
+    if quadraticflag:
+        bispec_options["quadraticflag"] = 1
+    else:
+        bispec_options["quadraticflag"] = 0
     lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
     return lmp, bispec_options, cutoff
diff --git a/tests/test_snap.py b/tests/test_snap.py
index 5ad4781ab..522c490ff 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -1,6 +1,6 @@
 from ase.build.bulk import bulk
 import structuretoolkit as stk
-from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives
+from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives, calc_per_atom_quad, calc_sum_quad
 import unittest
 
 
@@ -25,7 +25,8 @@ def setUpClass(cls):
         cls.rcutfac = 1.0
         cls.rfac0 = 0.99363
         cls.rmin0 = 0.0
-        cls.bzeroflag = 0
+        cls.bzeroflag = False
+        cls.quadraticflag = False
         cls.radelem = [4.0]
         cls.type = ['Cu']
         cls.wj = [1.0]
@@ -81,28 +82,30 @@ def test_calc_bispectrum_lmp(self):
             rfac0=self.rfac0,
             rmin0=self.rmin0,
             bzeroflag=self.bzeroflag,
+            quadraticflag=self.quadraticflag,
             weights=self.wj,
             cutoff=10.0,
         )
-        self.assertTrue(coeff.shape, (len(self.structure), n_coeff))
+        self.assertEqual(coeff.shape, (len(self.structure), n_coeff))
 
     def test_calc_bispectrum_lmp_quad(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.twojmax
         ))
         coeff = stk.analyse.calc_snap_descriptors_per_atom(
-                structure=self.structure,
-                atom_types=self.type,
-                twojmax=self.twojmax,
-                element_radius=self.radelem,
-                rcutfac=self.rcutfac,
-                rfac0=self.rfac0,
-                rmin0=self.rmin0,
-                bzeroflag=self.bzeroflag,
-                weights=self.wj,
-                cutoff=10.0,
-            )
-        self.assertTrue(coeff.shape, (len(self.structure), n_coeff * 31))
+            structure=self.structure,
+            atom_types=self.type,
+            twojmax=self.twojmax,
+            element_radius=self.radelem,
+            rcutfac=self.rcutfac,
+            rfac0=self.rfac0,
+            rmin0=self.rmin0,
+            bzeroflag=self.bzeroflag,
+            quadraticflag=True,
+            weights=self.wj,
+            cutoff=10.0,
+        )
+        self.assertEqual(coeff.shape, (len(self.structure), ((n_coeff + 1) * n_coeff) / 2 + 30))
 
     def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
@@ -117,10 +120,11 @@ def test_calc_a_matrix_snappy(self):
             rfac0=self.rfac0,
             rmin0=self.rmin0,
             bzeroflag=self.bzeroflag,
+            quadraticflag=self.quadraticflag,
             weights=self.wj,
             cutoff=10.0,
         )
-        self.assertTrue(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff))
+        self.assertEqual(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff + 1))
 
 
 @unittest.skipIf(
@@ -154,21 +158,29 @@ def test_calc_bispectrum_lmp(self):
             bispec_options=self.bispec_options,
             cutoff=10.0
         )
-        self.assertTrue(coeff.shape, (len(self.structure), n_coeff))
+        self.assertEqual(coeff.shape, (len(self.structure), n_coeff))
 
     def test_calc_bispectrum_lmp_quad(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.bispec_options["twojmax"]
         ))
+        coeff_lin = _calc_snap_per_atom(
+            lmp=self.lmp,
+            structure=self.structure,
+            bispec_options=self.bispec_options,
+            cutoff=10.0
+        )
         bispec_options = self.bispec_options.copy()
         bispec_options["quadraticflag"] = 1
-        coeff = _calc_snap_per_atom(
+        coeff_quad = _calc_snap_per_atom(
             lmp=self.lmp,
             structure=self.structure,
             bispec_options=bispec_options,
             cutoff=10.0
         )
-        self.assertTrue(coeff.shape, (len(self.structure), n_coeff * 31))
+        self.assertEqual(coeff_quad.shape, (len(self.structure), ((n_coeff+1) * n_coeff)/2 + 30))
+        coeff_quad_py = calc_per_atom_quad(coeff_lin)
+        self.assertEqual(coeff_quad.shape, coeff_quad_py.shape)
 
     def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
@@ -180,4 +192,4 @@ def test_calc_a_matrix_snappy(self):
             bispec_options=self.bispec_options,
             cutoff=10.0
         )
-        self.assertTrue(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff))
+        self.assertEqual(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff + 1))

From b4f0038e9c5a89aebec06ad27ad2d1e53a60aa92 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 15:31:51 -0600
Subject: [PATCH 18/34] add more tests

---
 structuretoolkit/analyse/snap.py | 61 +++++++++++++++++++-------------
 tests/test_snap.py               | 17 +++++++--
 2 files changed, 50 insertions(+), 28 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 9c4dc2770..66366c8b0 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -7,8 +7,8 @@
 
 def calc_per_atom_quad(linear_per_atom):
     """
-    Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
-    the SNAP descriptors.
+    Calculate quadratic par-atom SNAP descriptors from the linear SNAP descriptors, by multiplying the individual
+    components of the SNAP descriptors.
 
     Args:
         linear_per_atom (np.ndarray): Numpy array of the linear per atom SNAP descriptors
@@ -34,6 +34,16 @@ def calc_per_atom_quad(linear_per_atom):
 
 
 def calc_sum_quad(linear_sum):
+    """
+   Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
+   the SNAP descriptors.
+
+   Args:
+       linear_sum (np.ndarray): Numpy array of the linear SNAP descriptors
+
+   Returns:
+       np.ndarray: Numpy array of the quadratic SNAP descriptors
+   """
     return np.concatenate(
         (
             linear_sum,
@@ -65,16 +75,16 @@ def calc_snap_descriptors_per_atom(
 
     Args:
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
-        atom_types:
-        twojmax (int):
-        element_radius (list):
-        rcutfac (float):
-        rfac0 (float):
-        rmin0 (float):
-        bzeroflag (bool):
-        quadraticflag (bool):
-        weights (list/np.ndarry/None):
-        cutoff (float):
+        atom_types (list): list of element types
+        twojmax (int): band limit for bispectrum components (non-negative integer)
+        element_radius (list): list of radii for the individual elements
+        rcutfac (float): scale factor applied to all cutoff radii (positive real)
+        rfac0 (float): parameter in distance to angle conversion (0 < rcutfac < 1)
+        rmin0 (float): parameter in distance to angle conversion (distance units)
+        bzeroflag (bool): subtract B0
+        quadraticflag (bool): generate quadratic terms
+        weights (list/np.ndarry/None): list of neighbor weights, one for each type
+        cutoff (float): cutoff radius for the construction of the neighbor list
 
     Returns:
         np.ndarray: Numpy array with the calculated descriptor derivatives
@@ -104,8 +114,8 @@ def calc_snap_descriptor_derivatives(
     rcutfac=1.0,
     rfac0=0.99363,
     rmin0=0.0,
-    bzeroflag=0,
-    quadraticflag=0,
+    bzeroflag=False,
+    quadraticflag=False,
     weights=None,
     cutoff=10.0,
 ):
@@ -114,15 +124,16 @@ def calc_snap_descriptor_derivatives(
 
     Args:
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
-        atom_types:
-        twojmax (int):
-        element_radius (list):
-        rcutfac (float):
-        rfac0 (float):
-        rmin0 (float):
-        bzeroflag (bool):
-        weights (list/np.ndarry/None):
-        cutoff (float):
+        atom_types (list): list of element types
+        twojmax (int): band limit for bispectrum components (non-negative integer)
+        element_radius (list): list of radii for the individual elements
+        rcutfac (float): scale factor applied to all cutoff radii (positive real)
+        rfac0 (float): parameter in distance to angle conversion (0 < rcutfac < 1)
+        rmin0 (float): parameter in distance to angle conversion (distance units)
+        bzeroflag (bool): subtract B0
+        quadraticflag (bool): generate quadratic terms
+        weights (list/np.ndarry/None): list of neighbor weights, one for each type
+        cutoff (float): cutoff radius for the construction of the neighbor list
 
     Returns:
         np.ndarray: Numpy array with the calculated descriptor derivatives
@@ -146,7 +157,7 @@ def calc_snap_descriptor_derivatives(
 
 def get_apre(cell):
     """
-    Convert ASE cell to LAMMPS cell - LAMMPS required the upper triangle to be zero
+    Convert ASE cell to LAMMPS cell - LAMMPS requires the upper triangle to be zero
 
     Args:
         cell (np.ndarray): ASE cell as 3x3 matrix
@@ -176,7 +187,7 @@ def get_snap_descriptor_names(twojmax):
     Get names of the SNAP descriptors
 
     Args:
-        twojmax (int):
+        twojmax (int): band limit for bispectrum components (non-negative integer)
 
     Returns:
         list: List of SNAP descriptor names
diff --git a/tests/test_snap.py b/tests/test_snap.py
index 522c490ff..de7233a17 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -1,6 +1,7 @@
 from ase.build.bulk import bulk
+import numpy as np
 import structuretoolkit as stk
-from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives, calc_per_atom_quad, calc_sum_quad
+from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives, calc_per_atom_quad, calc_sum_quad, get_apre
 import unittest
 
 
@@ -179,8 +180,10 @@ def test_calc_bispectrum_lmp_quad(self):
             cutoff=10.0
         )
         self.assertEqual(coeff_quad.shape, (len(self.structure), ((n_coeff+1) * n_coeff)/2 + 30))
-        coeff_quad_py = calc_per_atom_quad(coeff_lin)
-        self.assertEqual(coeff_quad.shape, coeff_quad_py.shape)
+        coeff_quad_per_atom = calc_per_atom_quad(coeff_lin)
+        coeff_quad_sum = calc_sum_quad(np.sum(coeff_lin, axis=0))
+        self.assertEqual(coeff_quad.shape, coeff_quad_per_atom.shape)
+        self.assertEqual(np.sum(coeff_quad, axis=0).shape, coeff_quad_sum.shape)
 
     def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
@@ -193,3 +196,11 @@ def test_calc_a_matrix_snappy(self):
             cutoff=10.0
         )
         self.assertEqual(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff + 1))
+
+    def test_get_apre(self):
+        cell = bulk("Cu").cell
+        lmp_cell = get_apre(cell=cell)
+        self.assertEqual(lmp_cell[0, 1], 0.0)
+        self.assertEqual(lmp_cell[0, 2], 0.0)
+        self.assertEqual(lmp_cell[1, 2], 0.0)
+        self.assertEqual(cell.shape, lmp_cell.shape)

From acb4e28b5f24d9801cbf129466f6a7bb968afe9d Mon Sep 17 00:00:00 2001
From: pyiron-runner <pyiron@mpie.de>
Date: Fri, 8 Mar 2024 21:32:54 +0000
Subject: [PATCH 19/34] Format black

---
 structuretoolkit/analyse/snap.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 66366c8b0..53ab19063 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -35,15 +35,15 @@ def calc_per_atom_quad(linear_per_atom):
 
 def calc_sum_quad(linear_sum):
     """
-   Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
-   the SNAP descriptors.
+    Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
+    the SNAP descriptors.
 
-   Args:
-       linear_sum (np.ndarray): Numpy array of the linear SNAP descriptors
+    Args:
+        linear_sum (np.ndarray): Numpy array of the linear SNAP descriptors
 
-   Returns:
-       np.ndarray: Numpy array of the quadratic SNAP descriptors
-   """
+    Returns:
+        np.ndarray: Numpy array of the quadratic SNAP descriptors
+    """
     return np.concatenate(
         (
             linear_sum,

From 6d47420727f52ac9428dba7d4e2edd5f1eb589bc Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 15:42:50 -0600
Subject: [PATCH 20/34] execute lammps commands with for loop

---
 structuretoolkit/analyse/snap.py | 42 ++++++++++++++++++--------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 66366c8b0..629c40901 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -35,15 +35,15 @@ def calc_per_atom_quad(linear_per_atom):
 
 def calc_sum_quad(linear_sum):
     """
-   Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
-   the SNAP descriptors.
+    Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
+    the SNAP descriptors.
 
-   Args:
-       linear_sum (np.ndarray): Numpy array of the linear SNAP descriptors
+    Args:
+        linear_sum (np.ndarray): Numpy array of the linear SNAP descriptors
 
-   Returns:
-       np.ndarray: Numpy array of the quadratic SNAP descriptors
-   """
+    Returns:
+        np.ndarray: Numpy array of the quadratic SNAP descriptors
+    """
     return np.concatenate(
         (
             linear_sum,
@@ -282,12 +282,15 @@ def _reset_lmp(lmp):
     Args:
         lmp (lammps.Lammps): Lammps library instance
     """
-    lmp.command("clear")
-    lmp.command("units metal")
-    lmp.command("dimension 3")
-    lmp.command("boundary p p p")
-    lmp.command("atom_style charge")
-    lmp.command("atom_modify map array sort 0 2.0")
+    for c in [
+        "clear",
+        "units metal",
+        "dimension 3",
+        "boundary p p p",
+        "atom_style charge",
+        "atom_modify map array sort 0 2.0",
+    ]:
+        lmp.command(c)
 
 
 def _set_potential_lmp(lmp, cutoff=10.0):
@@ -298,11 +301,14 @@ def _set_potential_lmp(lmp, cutoff=10.0):
         lmp (lammps.Lammps): Lammps library instance
         cutoff (float): cutoff radius for the construction of the neighbor list
     """
-    lmp.command("pair_style zero " + str(cutoff))
-    lmp.command("pair_coeff * *")
-    lmp.command("mass * 1.0e-20")
-    lmp.command("neighbor 1.0e-20 nsq")
-    lmp.command("neigh_modify one 10000")
+    for c in [
+        "pair_style zero " + str(cutoff),
+        "pair_coeff * *",
+        "mass * 1.0e-20",
+        "neighbor 1.0e-20 nsq",
+        "neigh_modify one 10000",
+    ]:
+        lmp.command(c)
 
 
 def _set_compute_lammps(lmp, bispec_options, numtypes):

From ad4da416fdbde2fdf91ac40ed5dd35156eb6ab0b Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 15:47:01 -0600
Subject: [PATCH 21/34] remove outdated conda_merge.py script

---
 .ci_support/condamerge.py       | 57 ---------------------------------
 .github/workflows/unittests.yml |  7 +++-
 2 files changed, 6 insertions(+), 58 deletions(-)
 delete mode 100644 .ci_support/condamerge.py

diff --git a/.ci_support/condamerge.py b/.ci_support/condamerge.py
deleted file mode 100644
index fb2df69eb..000000000
--- a/.ci_support/condamerge.py
+++ /dev/null
@@ -1,57 +0,0 @@
-import argparse
-import sys
-import yaml
-
-
-def read_file(path):
-    with open(path) as f:
-        return yaml.safe_load(f)
-
-
-def parse_args(argv=None):
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--base', dest='base', help='base environment.yml file')
-    parser.add_argument('--add', dest='add', help='addon environment.yml file')
-    return parser.parse_args(argv)
-
-
-def merge_dependencies(env_base, env_add):
-    base_dict = {f.split()[0]: f for f in env_base}
-    add_dict = {f.split()[0]: f for f in env_add}
-    for k,v in add_dict.items():
-        if k not in base_dict.keys():
-            base_dict[k] = v
-    return list(base_dict.values())
-
-
-def merge_channels(env_base, env_add):
-    for c in env_add:
-        if c not in env_base:
-            env_base.append(c)
-    return env_base
-
-
-def merge_env(env_base, env_add):
-    return {
-        "channels": merge_channels(
-            env_base=env_base['channels'],
-            env_add=env_add['channels']
-        ),
-        'dependencies': merge_dependencies(
-            env_base=env_base['dependencies'],
-            env_add=env_add['dependencies']
-        )
-    }
-
-
-if __name__ == '__main__':
-    arguments = parse_args(argv=None)
-    yaml.dump(
-        merge_env(
-            env_base=read_file(arguments.base),
-            env_add=read_file(arguments.add)
-        ),
-        sys.stdout,
-        indent=2,
-        default_flow_style=False
-    )
diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index b433e6dc5..b9b6185a5 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -46,6 +46,11 @@ jobs:
 
     steps:
     - uses: actions/checkout@v4
+    - name: Merge conda environment
+      if:  matrix.operating-system != 'windows-latest'
+      run: |
+        cp .ci_support/environment.yml environment.yml
+        tail --lines=+4 .ci_support/environment-lammps.yml >> environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2
       with:
@@ -54,7 +59,7 @@ jobs:
         channels: conda-forge
         channel-priority: strict
         activate-environment: my-env
-        environment-file: .ci_support/environment.yml
+        environment-file: environment.yml
         use-mamba: true
     - name: Test
       shell: bash -l {0}

From 23682404ead5333f6b04f466c2129956e71921eb Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 15:47:47 -0600
Subject: [PATCH 22/34] update lammps version

---
 .ci_support/environment-lammps.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.ci_support/environment-lammps.yml b/.ci_support/environment-lammps.yml
index 675d3a4df..4a201417c 100644
--- a/.ci_support/environment-lammps.yml
+++ b/.ci_support/environment-lammps.yml
@@ -1,4 +1,4 @@
 channels:
 - conda-forge
 dependencies:
-- lammps =2023.08.02  # [unix]
\ No newline at end of file
+- lammps =2023.11.21
\ No newline at end of file

From 62907501013f882cbe69c8566c59cb59360ab670 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 15:52:32 -0600
Subject: [PATCH 23/34] fix windows tests

---
 .github/workflows/unittests.yml | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml
index b9b6185a5..ac2f3e232 100644
--- a/.github/workflows/unittests.yml
+++ b/.github/workflows/unittests.yml
@@ -48,9 +48,7 @@ jobs:
     - uses: actions/checkout@v4
     - name: Merge conda environment
       if:  matrix.operating-system != 'windows-latest'
-      run: |
-        cp .ci_support/environment.yml environment.yml
-        tail --lines=+4 .ci_support/environment-lammps.yml >> environment.yml
+      run: tail --lines=+4 .ci_support/environment-lammps.yml >> .ci_support/environment.yml
     - name: Setup Mambaforge
       uses: conda-incubator/setup-miniconda@v2
       with:
@@ -59,7 +57,7 @@ jobs:
         channels: conda-forge
         channel-priority: strict
         activate-environment: my-env
-        environment-file: environment.yml
+        environment-file: .ci_support/environment.yml
         use-mamba: true
     - name: Test
       shell: bash -l {0}

From 3f1a05b90bcdd0dcdc149aa548fa7e6accdab031 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 16:03:56 -0600
Subject: [PATCH 24/34] Add links to LAMMPS documentation

---
 structuretoolkit/analyse/snap.py | 33 ++++++++++++++++----------------
 1 file changed, 16 insertions(+), 17 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 629c40901..0f84072cf 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -71,7 +71,7 @@ def calc_snap_descriptors_per_atom(
     cutoff=10.0,
 ):
     """
-    Calculate per atom SNAP descriptors
+    Calculate per atom SNAP descriptors using LAMMPS https://docs.lammps.org/compute_sna_atom.html
 
     Args:
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
@@ -120,7 +120,7 @@ def calc_snap_descriptor_derivatives(
     cutoff=10.0,
 ):
     """
-    Calculate per atom derivatives of the SNAP descriptors.
+    Calculate per atom derivatives of the SNAP descriptors using LAMMPS https://docs.lammps.org/compute_sna_atom.html
 
     Args:
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
@@ -312,6 +312,19 @@ def _set_potential_lmp(lmp, cutoff=10.0):
 
 
 def _set_compute_lammps(lmp, bispec_options, numtypes):
+    compute_parameter = [
+        "rmin0",
+        "bzeroflag",
+        "quadraticflag",
+        "switchflag",
+        "chem",
+        "bnormflag",
+        "wselfallflag",
+        "bikflag",
+        "switchinnerflag",
+        "sinner",
+        "dinner",
+    ]
     lmp_variable_args = {
         k: bispec_options[k] for k in ["rcutfac", "rfac0", "rmin0", "twojmax"]
     }
@@ -330,21 +343,7 @@ def _set_compute_lammps(lmp, bispec_options, numtypes):
     radelem = " ".join([f"${{radelem{i}}}" for i in range(1, numtypes + 1)])
     wj = " ".join([f"${{wj{i}}}" for i in range(1, numtypes + 1)])
     kw_options = {
-        k: bispec_options[v]
-        for k, v in {
-            "rmin0": "rmin0",
-            "bzeroflag": "bzeroflag",
-            "quadraticflag": "quadraticflag",
-            "switchflag": "switchflag",
-            "chem": "chemflag",
-            "bnormflag": "bnormflag",
-            "wselfallflag": "wselfallflag",
-            "bikflag": "bikflag",
-            "switchinnerflag": "switchinnerflag",
-            "sinner": "sinner",
-            "dinner": "dinner",
-        }.items()
-        if v in bispec_options
+        k: bispec_options[k] for k in compute_parameter if k in bispec_options
     }
     kw_substrings = [f"{k} {v}" for k, v in kw_options.items()]
     kwargs = " ".join(kw_substrings)

From 5d3f80479af403b39e034f5d78ebaf143890ea90 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 16:13:58 -0600
Subject: [PATCH 25/34] rename get_apre()

---
 structuretoolkit/analyse/snap.py | 42 ++++++++++++++++----------------
 tests/test_snap.py               |  6 ++---
 2 files changed, 24 insertions(+), 24 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 0f84072cf..9a430f7d9 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -155,7 +155,26 @@ def calc_snap_descriptor_derivatives(
     )
 
 
-def get_apre(cell):
+def get_snap_descriptor_names(twojmax):
+    """
+    Get names of the SNAP descriptors
+
+    Args:
+        twojmax (int): band limit for bispectrum components (non-negative integer)
+
+    Returns:
+        list: List of SNAP descriptor names
+    """
+    lst = []
+    for j1 in range(0, twojmax + 1):
+        for j2 in range(0, j1 + 1):
+            for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
+                if j >= j1:
+                    lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
+    return lst
+
+
+def _get_lammps_compatible_cell(cell):
     """
     Convert ASE cell to LAMMPS cell - LAMMPS requires the upper triangle to be zero
 
@@ -182,25 +201,6 @@ def get_apre(cell):
     return np.array(((xhi, 0, 0), (xyp, yhi, 0), (xzp, yzp, zhi)))
 
 
-def get_snap_descriptor_names(twojmax):
-    """
-    Get names of the SNAP descriptors
-
-    Args:
-        twojmax (int): band limit for bispectrum components (non-negative integer)
-
-    Returns:
-        list: List of SNAP descriptor names
-    """
-    lst = []
-    for j1 in range(0, twojmax + 1):
-        for j2 in range(0, j1 + 1):
-            for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
-                if j >= j1:
-                    lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
-    return lst
-
-
 def _convert_mat(mat):
     mat[np.diag_indices_from(mat)] /= 2
     return mat[np.triu_indices(len(mat))]
@@ -209,7 +209,7 @@ def _convert_mat(mat):
 def _write_ase_structure(lmp, structure):
     number_species = len(set(structure.get_chemical_symbols()))
 
-    apre = get_apre(cell=structure.cell)
+    apre = _get_lammps_compatible_cell(cell=structure.cell)
     (
         (xhi, xy, xz),
         (_, yhi, yz),
diff --git a/tests/test_snap.py b/tests/test_snap.py
index de7233a17..a0d92a2e2 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -1,7 +1,7 @@
 from ase.build.bulk import bulk
 import numpy as np
 import structuretoolkit as stk
-from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives, calc_per_atom_quad, calc_sum_quad, get_apre
+from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives, calc_per_atom_quad, calc_sum_quad, _get_lammps_compatible_cell
 import unittest
 
 
@@ -197,9 +197,9 @@ def test_calc_a_matrix_snappy(self):
         )
         self.assertEqual(mat_a.shape, (len(self.structure) * 3 + 7, n_coeff + 1))
 
-    def test_get_apre(self):
+    def test_get_lammps_compatible_cell(self):
         cell = bulk("Cu").cell
-        lmp_cell = get_apre(cell=cell)
+        lmp_cell = _get_lammps_compatible_cell(cell=cell)
         self.assertEqual(lmp_cell[0, 1], 0.0)
         self.assertEqual(lmp_cell[0, 2], 0.0)
         self.assertEqual(lmp_cell[1, 2], 0.0)

From 7ad13a85fd50b1e789f69389be4562ccb5f83cb7 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 16:19:20 -0600
Subject: [PATCH 26/34] rename variables

---
 structuretoolkit/analyse/snap.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 9a430f7d9..48c4c793f 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -209,12 +209,12 @@ def _convert_mat(mat):
 def _write_ase_structure(lmp, structure):
     number_species = len(set(structure.get_chemical_symbols()))
 
-    apre = _get_lammps_compatible_cell(cell=structure.cell)
+    lammps_cell = _get_lammps_compatible_cell(cell=structure.cell)
     (
         (xhi, xy, xz),
         (_, yhi, yz),
         (_, _, zhi),
-    ) = apre.T
+    ) = lammps_cell.T
     lmp.command(
         "region 1 prism"
         + " 0.0 "
@@ -235,12 +235,12 @@ def _write_ase_structure(lmp, structure):
 
     el_dict = {el: i for i, el in enumerate(set(structure.get_chemical_symbols()))}
 
-    R = np.dot(np.linalg.inv(structure.cell), apre)
+    rotation_mat = np.dot(np.linalg.inv(structure.cell), lammps_cell)
 
     positions = structure.positions.flatten()
-    if np.matrix.trace(R) != 3:
+    if np.matrix.trace(rotation_mat) != 3:
         positions = np.array(positions).reshape(-1, 3)
-        positions = np.matmul(positions, R)
+        positions = np.matmul(positions, rotation_mat)
     positions = positions.flatten()
     elem_all = np.array([el_dict[el] + 1 for el in structure.get_chemical_symbols()])
     lmp.create_atoms(
@@ -261,16 +261,16 @@ def _extract_compute_np(lmp, name, compute_type, result_type, array_shape):
     Note that the result is a view into the original memory.
     If the result type is 0 (scalar) then conversion to numpy is skipped and a python float is returned.
     """
-    ptr = lmp.extract_compute(
+    pointer = lmp.extract_compute(
         name, compute_type, result_type
     )  # 1,2: Style (1) is per-atom compute, returns array type (2).
     if result_type == 0:
-        return ptr  # No casting needed, lammps.py already works
+        return pointer  # No casting needed, lammps.py already works
     if result_type == 2:
-        ptr = ptr.contents
+        pointer = pointer.contents
     total_size = int(np.prod(array_shape))
-    buffer_ptr = cast(ptr, POINTER(c_double * total_size))
-    array_np = np.frombuffer(buffer_ptr.contents, dtype=float)
+    buffer_pointer = cast(pointer, POINTER(c_double * total_size))
+    array_np = np.frombuffer(buffer_pointer.contents, dtype=float)
     array_np.shape = array_shape
     return array_np.copy()
 

From 6934cc288cebe547f46fb23a52f6d38a1f0e0275 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 16:34:18 -0600
Subject: [PATCH 27/34] fix imports in test

---
 tests/test_snap.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tests/test_snap.py b/tests/test_snap.py
index a0d92a2e2..363d8190c 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -1,7 +1,13 @@
 from ase.build.bulk import bulk
 import numpy as np
 import structuretoolkit as stk
-from structuretoolkit.analyse.snap import _calc_snap_per_atom, _calc_snap_derivatives, calc_per_atom_quad, calc_sum_quad, _get_lammps_compatible_cell
+from structuretoolkit.analyse.snap import (
+    calc_per_atom_quad,
+    calc_sum_quad,
+    _calc_snap_per_atom,
+    _calc_snap_derivatives,
+    _get_lammps_compatible_cell
+)
 import unittest
 
 

From a5cde35c63e536566d35ce1abb152d4946c0283d Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Fri, 8 Mar 2024 16:49:21 -0600
Subject: [PATCH 28/34] add sorting on the set

---
 structuretoolkit/analyse/snap.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 48c4c793f..fb6f11a36 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -233,7 +233,9 @@ def _write_ase_structure(lmp, structure):
     )
     lmp.command("create_box " + str(number_species) + " 1")
 
-    el_dict = {el: i for i, el in enumerate(set(structure.get_chemical_symbols()))}
+    el_dict = {
+        el: i for i, el in enumerate(sorted(set(structure.get_chemical_symbols())))
+    }
 
     rotation_mat = np.dot(np.linalg.inv(structure.cell), lammps_cell)
 

From 4bb4e4a41b704e7451c1c2628e7a3a4eab8feb02 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Sun, 10 Mar 2024 13:40:29 -0500
Subject: [PATCH 29/34] compare with reference

---
 tests/test_snap.py | 155 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 155 insertions(+)

diff --git a/tests/test_snap.py b/tests/test_snap.py
index 363d8190c..4b9bd0318 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -19,6 +19,145 @@
     skip_snap_test = True
 
 
+coeff_lin_atom_0 = np.array([
+    4.57047509e+04, 4.85993797e+02, 9.67664872e+00, 2.41595356e+01,
+    6.79985845e-01, 3.20694545e-01, 1.02500964e-01, 2.69980900e+00,
+    1.63003719e-01, 3.58373618e-02, 8.17819735e-02, 8.59080127e-03,
+    1.84391799e-01, 1.49927925e+00, 3.03343563e-01, 1.94938837e-02,
+    4.30230643e-01, 1.53366482e-02, 3.80620803e-03, 7.60529037e-02,
+    1.02198715e+01, 5.81245462e+00, 1.22813679e-01, 2.67620628e-01,
+    2.32501387e-02, 3.96561016e-01, 5.76813472e+02, 6.63806331e+00,
+    1.17222436e+00, 1.62662275e+01
+])
+
+coeff_quad_atom_0 = np.array([
+    4.57047509e+04, 4.85993797e+02, 9.67664872e+00, 2.41595356e+01,
+    6.79985845e-01, 3.20694545e-01, 1.02500964e-01, 2.69980900e+00,
+    1.63003719e-01, 3.58373618e-02, 8.17819735e-02, 8.59080127e-03,
+    1.84391799e-01, 1.49927925e+00, 3.03343563e-01, 1.94938837e-02,
+    4.30230643e-01, 1.53366482e-02, 3.80620803e-03, 7.60529037e-02,
+    1.02198715e+01, 5.81245462e+00, 1.22813679e-01, 2.67620628e-01,
+    2.32501387e-02, 3.96561016e-01, 5.76813472e+02, 6.63806331e+00,
+    1.17222436e+00, 1.62662275e+01, 1.04446213e+09, 2.22122254e+07,
+    4.42268820e+05, 1.10420556e+06, 3.10785837e+04, 1.46572643e+04,
+    4.68478104e+03, 1.23394098e+05, 7.45004438e+03, 1.63793770e+03,
+    3.73782473e+03, 3.92640432e+02, 8.42758125e+03, 6.85241845e+04,
+    1.38642420e+04, 8.90963101e+02, 1.96635844e+04, 7.00957686e+02,
+    1.73961790e+02, 3.47597902e+03, 4.67096682e+05, 2.65656791e+05,
+    5.61316860e+03, 1.22315342e+04, 1.06264180e+03, 1.81247224e+04,
+    2.63631161e+07, 3.03391030e+05, 5.35762225e+04, 7.43443878e+05,
+    1.18094985e+05, 4.70279126e+03, 1.17413844e+04, 3.30468903e+02,
+    1.55855560e+02, 4.98148328e+01, 1.31209043e+03, 7.92187963e+01,
+    1.74167356e+01, 3.97455318e+01, 4.17507613e+00, 8.96132706e+01,
+    7.28640413e+02, 1.47423090e+02, 9.47390658e+00, 2.09089424e+02,
+    7.45351589e+00, 1.84979349e+00, 3.69612395e+01, 4.96679417e+03,
+    2.82481689e+03, 5.96866861e+01, 1.30061965e+02, 1.12994232e+01,
+    1.92726194e+02, 2.80327769e+05, 3.22605759e+03, 5.69693769e+02,
+    7.90528568e+03, 4.68187653e+01, 2.33783339e+02, 6.57998416e+00,
+    3.10324846e+00, 9.91865824e-01, 2.61251033e+01, 1.57732973e+00,
+    3.46785562e-01, 7.91375429e-01, 8.31301662e-02, 1.78429467e+00,
+    1.45079986e+01, 2.93534910e+00, 1.88635465e-01, 4.16319080e+00,
+    1.48407357e-01, 3.68313381e-02, 7.35937234e-01, 9.88941068e+01,
+    5.62450816e+01, 1.18842483e+00, 2.58967081e+00, 2.24983425e-01,
+    3.83738165e+00, 5.58162135e+03, 6.42342068e+01, 1.13432034e+01,
+    1.57402570e+02, 2.91841581e+02, 1.64281422e+01, 7.74783128e+00,
+    2.47637569e+00, 6.52261317e+01, 3.93809416e+00, 8.65814020e-01,
+    1.97581450e+00, 2.07549769e-01, 4.45482024e+00, 3.62218903e+01,
+    7.32863961e+00, 4.70963179e-01, 1.03941725e+01, 3.70526298e-01,
+    9.19562186e-02, 1.83740284e+00, 2.46907350e+02, 1.40426205e+02,
+    2.96712145e+00, 6.46559010e+00, 5.61712554e-01, 9.58072999e+00,
+    1.39355456e+04, 1.60372527e+02, 2.83203962e+01, 3.92984503e+02,
+    2.31190374e-01, 2.18067751e-01, 6.96992047e-02, 1.83583190e+00,
+    1.10840222e-01, 2.43688988e-02, 5.56105843e-02, 5.84162326e-03,
+    1.25383813e-01, 1.01948866e+00, 2.06269329e-01, 1.32555650e-02,
+    2.92550747e-01, 1.04287037e-02, 2.58816759e-03, 5.17148980e-02,
+    6.94936797e+00, 3.95238687e+00, 8.35115631e-02, 1.81978239e-01,
+    1.58097652e-02, 2.69655877e-01, 3.92224996e+02, 4.51378909e+00,
+    7.97095973e-01, 1.10608045e+01, 5.14224956e-02, 3.28715001e-02,
+    8.65814020e-01, 5.22744035e-02, 1.14928465e-02, 2.62270328e-02,
+    2.75502311e-03, 5.91334442e-02, 4.80810676e-01, 9.72806258e-02,
+    6.25158218e-03, 1.37972620e-01, 4.91837942e-03, 1.22063015e-03,
+    2.43897514e-02, 3.27745705e+00, 1.86402249e+00, 3.93856768e-02,
+    8.58244756e-02, 7.45619266e-03, 1.27174955e-01, 1.84980934e+02,
+    2.12879069e+00, 3.75925959e-01, 5.21649044e+00, 5.25322383e-03,
+    2.76733026e-01, 1.67080384e-02, 3.67336414e-03, 8.38273113e-03,
+    8.80565413e-04, 1.89003372e-02, 1.53677568e-01, 3.10930076e-02,
+    1.99814188e-03, 4.40990557e-02, 1.57202123e-03, 3.90139993e-04,
+    7.79549596e-03, 1.04754669e+00, 5.95782203e-01, 1.25885205e-02,
+    2.74313724e-02, 2.38316163e-03, 4.06478865e-02, 5.91239370e+01,
+    6.80407889e-01, 1.20154127e-01, 1.66730401e+00, 3.64448432e+00,
+    4.40078908e-01, 9.67540321e-02, 2.20795708e-01, 2.31935226e-02,
+    4.97822639e-01, 4.04776760e+00, 8.18969681e-01, 5.26297628e-02,
+    1.16154056e+00, 4.14060209e-02, 1.02760347e-02, 2.05328314e-01,
+    2.75917011e+01, 1.56925173e+01, 3.31573476e-01, 7.22524581e-01,
+    6.27709338e-02, 1.07063900e+00, 1.55728620e+03, 1.79215031e+01,
+    3.16478189e+00, 4.39157075e+01, 1.32851062e-02, 5.84162326e-03,
+    1.33307658e-02, 1.40033256e-03, 3.00565490e-02, 2.44388093e-01,
+    4.94461289e-02, 3.17757555e-03, 7.01291949e-02, 2.49993069e-03,
+    6.20426065e-04, 1.23969061e-02, 1.66587707e+00, 9.47451720e-01,
+    2.00190864e-02, 4.36231577e-02, 3.78985908e-03, 6.46409204e-02,
+    9.40227411e+01, 1.08202901e+00, 1.91076931e-01, 2.65145558e+00,
+    6.42158252e-04, 2.93085018e-03, 3.07871654e-04, 6.60811563e-03,
+    5.37302128e-02, 1.08710330e-02, 6.98609366e-04, 1.54183312e-02,
+    5.49625011e-04, 1.36404455e-04, 2.72553543e-03, 3.66253234e-01,
+    2.08303040e-01, 4.40131825e-03, 9.59081729e-03, 8.33223634e-04,
+    1.42117006e-02, 2.06714731e+01, 2.37890677e-01, 4.20094286e-02,
+    5.82938682e-01, 3.34414559e-03, 7.02572682e-04, 1.50799252e-02,
+    1.22614015e-01, 2.48080352e-02, 1.59424828e-03, 3.51851111e-02,
+    1.25426136e-03, 3.11279204e-04, 6.21975656e-03, 8.35801262e-01,
+    4.75354010e-01, 1.00439450e-02, 2.18865431e-02, 1.90144223e-03,
+    3.24315425e-02, 4.71729441e+01, 5.42873918e-01, 9.58668217e-02,
+    1.33028419e+00, 3.69009333e-05, 1.58407330e-03, 1.28800100e-02,
+    2.60596426e-03, 1.67468081e-04, 3.69602596e-03, 1.31754097e-04,
+    3.26983768e-05, 6.53355382e-04, 8.77968853e-02, 4.99336426e-02,
+    1.05506791e-03, 2.29907563e-03, 1.99737321e-04, 3.40677688e-03,
+    4.95528991e+00, 5.70262827e-02, 1.00703465e-02, 1.39739928e-01,
+    1.70001678e-02, 2.76454798e-01, 5.59340653e-02, 3.59451230e-03,
+    7.93310024e-02, 2.82795216e-03, 7.01833548e-04, 1.40235318e-02,
+    1.88446050e+00, 1.07176897e+00, 2.26458352e-02, 4.93470492e-02,
+    4.28713491e-03, 7.31225992e-02, 1.06359674e+02, 1.22400444e+00,
+    2.16148559e-01, 2.99935896e+00, 1.12391913e+00, 4.54796708e-01,
+    2.92267753e-02, 6.45035874e-01, 2.29939183e-02, 5.70656871e-03,
+    1.14024540e-01, 1.53224413e+01, 8.71449258e+00, 1.84132000e-01,
+    4.01238054e-01, 3.48584504e-02, 5.94555701e-01, 8.64804467e+02,
+    9.95231055e+00, 1.75749166e+00, 2.43876173e+01, 4.60086585e-02,
+    5.91334415e-03, 1.30507696e-01, 4.65227350e-03, 1.15458871e-03,
+    2.30701588e-02, 3.10013224e+00, 1.76317069e+00, 3.72547389e-02,
+    8.11809948e-02, 7.05277991e-03, 1.20294231e-01, 1.74972654e+02,
+    2.01361377e+00, 3.55586714e-01, 4.93425541e+00, 1.90005752e-04,
+    8.38686614e-03, 2.98970837e-04, 7.41977769e-05, 1.48256646e-03,
+    1.99224987e-01, 1.13307315e-01, 2.39411558e-03, 5.21696542e-03,
+    4.53235501e-04, 7.73051434e-03, 1.12443348e+01, 1.29401634e-01,
+    2.28512054e-02, 3.17091949e-01, 9.25492032e-02, 6.59829602e-03,
+    1.63754733e-03, 3.27202897e-02, 4.39690190e+00, 2.50069609e+00,
+    5.28382080e-02, 1.15138595e-01, 1.00029221e-02, 1.70612701e-01,
+    2.48162831e+02, 2.85589825e+00, 5.04326841e-01, 6.99822953e+00,
+    1.17606389e-04, 5.83744736e-05, 1.16639663e-03, 1.56738574e-01,
+    8.91435718e-02, 1.88355019e-03, 4.10440343e-03, 3.56579198e-04,
+    6.08191679e-03, 8.84638530e+00, 1.01805642e-01, 1.79779927e-02,
+    2.49469409e-01, 7.24360980e-06, 2.89473173e-04, 3.88989571e-02,
+    2.21234115e-02, 4.67454411e-04, 1.01861979e-03, 8.84948648e-05,
+    1.50939372e-03, 2.19547207e+00, 2.52658499e-02, 4.46172979e-03,
+    6.19126459e-02, 2.89202208e-03, 7.77250905e-01, 4.42054052e-01,
+    9.34033689e-03, 2.03533259e-02, 1.76824056e-03, 3.01596168e-02,
+    4.38683395e+01, 5.04843990e-01, 8.91510666e-02, 1.23709384e+00,
+    5.22228870e+01, 5.94025395e+01, 1.25514002e+00, 2.73504844e+00,
+    2.37613431e-01, 4.05280263e+00, 5.89495958e+03, 6.78401542e+01,
+    1.19799824e+01, 1.66238756e+02, 1.68923144e+01, 7.13848935e-01,
+    1.55553276e+00, 1.35140376e-01, 2.30499291e+00, 3.35270213e+03,
+    3.85834418e+01, 6.81350092e+00, 9.45467094e+01, 7.54159985e-03,
+    3.28674739e-02, 2.85543507e-03, 4.87031172e-02, 7.08405845e+01,
+    8.15244975e-01, 1.43965186e-01, 1.99771524e+00, 3.58104003e-02,
+    6.22221673e-03, 1.06127908e-01, 1.54367184e+02, 1.77648267e+00,
+    3.13711420e-01, 4.35317803e+00, 2.70284475e-04, 9.22009862e-03,
+    1.34109932e+01, 1.54335893e-01, 2.72543790e-02, 3.78192046e-01,
+    7.86303196e-02, 2.28741736e+02, 2.63239713e+00, 4.64858484e-01,
+    6.45055171e+00, 1.66356891e+05, 3.82892435e+03, 6.76154805e+02,
+    9.38257918e+03, 2.20319423e+01, 7.78129953e+00, 1.07976248e+02,
+    6.87054978e-01, 1.90676682e+01, 1.32295079e+02
+])
+
+
 @unittest.skipIf(
     skip_snap_test, "LAMMPS is not installed, so the SNAP tests are skipped."
 )
@@ -94,6 +233,10 @@ def test_calc_bispectrum_lmp(self):
             cutoff=10.0,
         )
         self.assertEqual(coeff.shape, (len(self.structure), n_coeff))
+        self.assertTrue(np.isclose(
+            np.array(coeff),
+            np.array([coeff_lin_atom_0, coeff_lin_atom_0, coeff_lin_atom_0, coeff_lin_atom_0])
+        ).all())
 
     def test_calc_bispectrum_lmp_quad(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
@@ -113,6 +256,10 @@ def test_calc_bispectrum_lmp_quad(self):
             cutoff=10.0,
         )
         self.assertEqual(coeff.shape, (len(self.structure), ((n_coeff + 1) * n_coeff) / 2 + 30))
+        self.assertTrue(np.isclose(
+            np.array(coeff),
+            np.array([coeff_quad_atom_0, coeff_quad_atom_0, coeff_quad_atom_0, coeff_quad_atom_0])
+        ).all())
 
     def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
@@ -166,6 +313,10 @@ def test_calc_bispectrum_lmp(self):
             cutoff=10.0
         )
         self.assertEqual(coeff.shape, (len(self.structure), n_coeff))
+        self.assertTrue(np.isclose(
+            np.array(coeff),
+            np.array([coeff_lin_atom_0, coeff_lin_atom_0, coeff_lin_atom_0, coeff_lin_atom_0])
+        ).all())
 
     def test_calc_bispectrum_lmp_quad(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
@@ -190,6 +341,10 @@ def test_calc_bispectrum_lmp_quad(self):
         coeff_quad_sum = calc_sum_quad(np.sum(coeff_lin, axis=0))
         self.assertEqual(coeff_quad.shape, coeff_quad_per_atom.shape)
         self.assertEqual(np.sum(coeff_quad, axis=0).shape, coeff_quad_sum.shape)
+        self.assertTrue(np.isclose(
+            np.array(coeff_quad_per_atom),
+            np.array([coeff_quad_atom_0, coeff_quad_atom_0, coeff_quad_atom_0, coeff_quad_atom_0])
+        ).all())
 
     def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(

From f6d6cd514e065ffd76dcb4438185d19580199929 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Thu, 14 Mar 2024 08:39:22 -0500
Subject: [PATCH 30/34] Add comment about ignoring LAMMPS crashes

---
 structuretoolkit/analyse/snap.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index fb6f11a36..fcf27253b 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -536,6 +536,8 @@ def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
     try:
         lmp.command("run 0")
     except:
+        # When LAMMPS crashes return an empty array - this can be the case when atoms are overlapping.
+        # To iterate over large structure databases, the code does not raise an exception here.
         return np.array([])
     else:
         if (

From 21dc2a34fb8b26cbfdead9b73bd1b13e08d4b0b8 Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Thu, 14 Mar 2024 09:18:35 -0500
Subject: [PATCH 31/34] rename calc_*() to get_*() for consistency with other
 descriptors

---
 structuretoolkit/analyse/__init__.py | 4 ++--
 structuretoolkit/analyse/snap.py     | 4 ++--
 tests/test_snap.py                   | 6 +++---
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/structuretoolkit/analyse/__init__.py b/structuretoolkit/analyse/__init__.py
index e52be6fc9..dd13bacc0 100644
--- a/structuretoolkit/analyse/__init__.py
+++ b/structuretoolkit/analyse/__init__.py
@@ -25,8 +25,8 @@
 from structuretoolkit.analyse.strain import get_strain
 from structuretoolkit.analyse.snap import (
     get_snap_descriptor_names,
-    calc_snap_descriptors_per_atom,
-    calc_snap_descriptor_derivatives,
+    get_snap_descriptors_per_atom,
+    get_snap_descriptor_derivatives,
 )
 
 
diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index fcf27253b..6dc60c863 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -57,7 +57,7 @@ def calc_sum_quad(linear_sum):
     )
 
 
-def calc_snap_descriptors_per_atom(
+def get_snap_descriptors_per_atom(
     structure,
     atom_types,
     twojmax=6,
@@ -106,7 +106,7 @@ def calc_snap_descriptors_per_atom(
     )
 
 
-def calc_snap_descriptor_derivatives(
+def get_snap_descriptor_derivatives(
     structure,
     atom_types,
     twojmax=6,
diff --git a/tests/test_snap.py b/tests/test_snap.py
index 4b9bd0318..43b5d9a22 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -219,7 +219,7 @@ def test_calc_bispectrum_lmp(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.twojmax
         ))
-        coeff = stk.analyse.calc_snap_descriptors_per_atom(
+        coeff = stk.analyse.get_snap_descriptors_per_atom(
             structure=self.structure,
             atom_types=self.type,
             twojmax=self.twojmax,
@@ -242,7 +242,7 @@ def test_calc_bispectrum_lmp_quad(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.twojmax
         ))
-        coeff = stk.analyse.calc_snap_descriptors_per_atom(
+        coeff = stk.analyse.get_snap_descriptors_per_atom(
             structure=self.structure,
             atom_types=self.type,
             twojmax=self.twojmax,
@@ -265,7 +265,7 @@ def test_calc_a_matrix_snappy(self):
         n_coeff = len(stk.analyse.get_snap_descriptor_names(
             twojmax=self.twojmax
         ))
-        mat_a = stk.analyse.calc_snap_descriptor_derivatives(
+        mat_a = stk.analyse.get_snap_descriptor_derivatives(
             structure=self.structure,
             atom_types=self.type,
             twojmax=self.twojmax,

From 054c53f67698a4318b6bece27afd72f4ab99f24f Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Thu, 14 Mar 2024 09:21:00 -0500
Subject: [PATCH 32/34] remove if statement

---
 structuretoolkit/analyse/snap.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 6dc60c863..97d190389 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -168,9 +168,8 @@ def get_snap_descriptor_names(twojmax):
     lst = []
     for j1 in range(0, twojmax + 1):
         for j2 in range(0, j1 + 1):
-            for j in range(j1 - j2, min(twojmax, j1 + j2) + 1, 2):
-                if j >= j1:
-                    lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
+            for j in range(j1 + j2 % 2, min(twojmax, j1 + j2) + 1, 2):
+                lst.append([j1 / 2.0, j2 / 2.0, j / 2.0])
     return lst
 
 

From f5f191bfb5d8679e695d96231a19d801b1185b4e Mon Sep 17 00:00:00 2001
From: Jan Janssen <janssen@mpie.de>
Date: Thu, 14 Mar 2024 10:24:45 -0500
Subject: [PATCH 33/34] Add more docstrings and type hints

---
 structuretoolkit/analyse/snap.py | 197 ++++++++++++++++++++++---------
 tests/test_snap.py               |   8 +-
 2 files changed, 148 insertions(+), 57 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index 97d190389..c9dc760c3 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -1,3 +1,5 @@
+from ase.atoms import Atoms
+from typing import Optional, Union
 from ctypes import c_double, c_int, cast, POINTER
 import numpy as np
 from scipy.constants import physical_constants
@@ -5,7 +7,7 @@
 eV_div_A3_to_bar = 1e25 / physical_constants["joule-electron volt relationship"][0]
 
 
-def calc_per_atom_quad(linear_per_atom):
+def get_per_atom_quad(linear_per_atom: np.ndarray) -> np.ndarray:
     """
     Calculate quadratic par-atom SNAP descriptors from the linear SNAP descriptors, by multiplying the individual
     components of the SNAP descriptors.
@@ -33,7 +35,7 @@ def calc_per_atom_quad(linear_per_atom):
     )
 
 
-def calc_sum_quad(linear_sum):
+def get_sum_quad(linear_sum: np.ndarray) -> np.ndarray:
     """
     Calculate quadratic SNAP descriptors from the linear SNAP descriptors, by multiplying the individual components of
     the SNAP descriptors.
@@ -58,26 +60,26 @@ def calc_sum_quad(linear_sum):
 
 
 def get_snap_descriptors_per_atom(
-    structure,
-    atom_types,
-    twojmax=6,
-    element_radius=4.0,
-    rcutfac=1.0,
-    rfac0=0.99363,
-    rmin0=0.0,
-    bzeroflag=False,
-    quadraticflag=False,
-    weights=None,
-    cutoff=10.0,
-):
+    structure: Atoms,
+    atom_types: list[str],
+    twojmax: int = 6,
+    element_radius: list[int] = [4.0],
+    rcutfac: float = 1.0,
+    rfac0: float = 0.99363,
+    rmin0: float = 0.0,
+    bzeroflag: bool = False,
+    quadraticflag: bool = False,
+    weights: Optional[Union[list, np.ndarray]] = None,
+    cutoff: float = 10.0,
+) -> np.ndarray:
     """
     Calculate per atom SNAP descriptors using LAMMPS https://docs.lammps.org/compute_sna_atom.html
 
     Args:
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
-        atom_types (list): list of element types
+        atom_types (list[str]): list of element types
         twojmax (int): band limit for bispectrum components (non-negative integer)
-        element_radius (list): list of radii for the individual elements
+        element_radius (list[int]): list of radii for the individual elements
         rcutfac (float): scale factor applied to all cutoff radii (positive real)
         rfac0 (float): parameter in distance to angle conversion (0 < rcutfac < 1)
         rmin0 (float): parameter in distance to angle conversion (distance units)
@@ -107,26 +109,26 @@ def get_snap_descriptors_per_atom(
 
 
 def get_snap_descriptor_derivatives(
-    structure,
-    atom_types,
-    twojmax=6,
-    element_radius=4.0,
-    rcutfac=1.0,
-    rfac0=0.99363,
-    rmin0=0.0,
-    bzeroflag=False,
-    quadraticflag=False,
-    weights=None,
-    cutoff=10.0,
+    structure: Atoms,
+    atom_types: list[str],
+    twojmax: int = 6,
+    element_radius: list[int] = [4.0],
+    rcutfac: float = 1.0,
+    rfac0: float = 0.99363,
+    rmin0: float = 0.0,
+    bzeroflag: bool = False,
+    quadraticflag: bool = False,
+    weights: Optional[Union[list, np.ndarray]] = None,
+    cutoff: float = 10.0,
 ):
     """
     Calculate per atom derivatives of the SNAP descriptors using LAMMPS https://docs.lammps.org/compute_sna_atom.html
 
     Args:
         structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
-        atom_types (list): list of element types
+        atom_types (list[str]): list of element types
         twojmax (int): band limit for bispectrum components (non-negative integer)
-        element_radius (list): list of radii for the individual elements
+        element_radius (list[int]): list of radii for the individual elements
         rcutfac (float): scale factor applied to all cutoff radii (positive real)
         rfac0 (float): parameter in distance to angle conversion (0 < rcutfac < 1)
         rmin0 (float): parameter in distance to angle conversion (distance units)
@@ -155,7 +157,7 @@ def get_snap_descriptor_derivatives(
     )
 
 
-def get_snap_descriptor_names(twojmax):
+def get_snap_descriptor_names(twojmax: int) -> np.ndarray:
     """
     Get names of the SNAP descriptors
 
@@ -173,7 +175,7 @@ def get_snap_descriptor_names(twojmax):
     return lst
 
 
-def _get_lammps_compatible_cell(cell):
+def _get_lammps_compatible_cell(cell: np.ndarray) -> np.ndarray:
     """
     Convert ASE cell to LAMMPS cell - LAMMPS requires the upper triangle to be zero
 
@@ -200,12 +202,19 @@ def _get_lammps_compatible_cell(cell):
     return np.array(((xhi, 0, 0), (xyp, yhi, 0), (xzp, yzp, zhi)))
 
 
-def _convert_mat(mat):
+def _convert_mat(mat: np.ndarray) -> np.ndarray:
     mat[np.diag_indices_from(mat)] /= 2
     return mat[np.triu_indices(len(mat))]
 
 
-def _write_ase_structure(lmp, structure):
+def _set_ase_structure(lmp, structure: Atoms):
+    """
+    Use ASE structure object to initialize the structure in the LAMMPS calculator
+
+    Args:
+        lmp (lammps.Lammps): LAMMPS library instance
+        structure (ase.atoms.Atoms): ASE structure object
+    """
     number_species = len(set(structure.get_chemical_symbols()))
 
     lammps_cell = _get_lammps_compatible_cell(cell=structure.cell)
@@ -255,12 +264,22 @@ def _write_ase_structure(lmp, structure):
     )
 
 
-def _extract_compute_np(lmp, name, compute_type, result_type, array_shape):
+def _extract_compute_np(lmp, name: str, compute_type: int, result_type: int, array_shape: tuple) -> np.ndarray:
     """
     Convert a lammps compute to a numpy array.
     Assumes the compute returns a floating point numbers.
     Note that the result is a view into the original memory.
     If the result type is 0 (scalar) then conversion to numpy is skipped and a python float is returned.
+
+    Args:
+        lmp (lammps.Lammps): LAMMPS library instance
+        name (str): LAMMPS internal compute ID
+        compute_type (int): style of the data retrieve (global, atom, or local)
+        result_type (int): type or size of the returned data (scalar, vector, or array)
+        array_shape (tuple[int]): shape of the array as tuple of integers
+
+    Returns:
+        np.ndarray: output of the LAMMPS compute
     """
     pointer = lmp.extract_compute(
         name, compute_type, result_type
@@ -294,7 +313,7 @@ def _reset_lmp(lmp):
         lmp.command(c)
 
 
-def _set_potential_lmp(lmp, cutoff=10.0):
+def _set_potential_lmp(lmp, cutoff: float = 10.0):
     """
     Set interatomic potential parameters to LAMMPS library instance
 
@@ -312,7 +331,15 @@ def _set_potential_lmp(lmp, cutoff=10.0):
         lmp.command(c)
 
 
-def _set_compute_lammps(lmp, bispec_options, numtypes):
+def _set_compute_lammps(lmp, bispec_options: dict, numtypes: int):
+    """
+    Internal function to set the LAMMPS compute
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        bispec_options (dict): dictionary of the parameters to compute the bi-spectral components
+        numtypes (int): number of atomic types
+    """
     compute_parameter = [
         "rmin0",
         "bzeroflag",
@@ -351,10 +378,22 @@ def _set_compute_lammps(lmp, bispec_options, numtypes):
     lmp.command(f"{base_b} {radelem} {wj} {kwargs}")
 
 
-def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
+def _calc_snap_per_atom(lmp, structure: Atoms, bispec_options: dict, cutoff: float = 10.0) -> np.ndarray:
+    """
+    Internal function to calculate the per-atom SNAP descriptors
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        structure (ase.atoms.Atoms):
+        bispec_options (dict): dictionary of the parameters to compute the bi-spectral components
+        cutoff (float): cutoff radius for the construction of the neighbor list
+
+    Returns:
+        np.ndarray: Output of the LAMMPS compute command
+    """
     number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     _reset_lmp(lmp=lmp)
-    _write_ase_structure(lmp=lmp, structure=structure)
+    _set_ase_structure(lmp=lmp, structure=structure)
     _set_potential_lmp(lmp=lmp, cutoff=cutoff)
     _set_compute_lammps(
         lmp=lmp,
@@ -392,7 +431,17 @@ def _calc_snap_per_atom(lmp, structure, bispec_options, cutoff=10.0):
             )
 
 
-def _lammps_variables(bispec_options):
+def _lammps_variables(bispec_options: dict):
+    """
+    Identify which of the parameters to compute the bi-spectral components can be set as LAMMPS variables rather than
+    as inputs to the LAMMPS compute command.
+
+    Args:
+        bispec_options (dict): dictionary of the parameters to compute the bi-spectral components
+
+    Returns:
+        dict: dictionary of LAMMPS variables and their values to be set
+    """
     d = {
         k: bispec_options[k]
         for k in [
@@ -427,13 +476,13 @@ def _set_variables(lmp, **lmp_variable_args):
         lmp.command(f"variable {k} equal {v}")
 
 
-def _set_computes_snap(lmp, bispec_options):
+def _set_computes_snap(lmp, bispec_options: dict):
     """
     Set LAMMPS computes to calculate SNAP descriptors
 
     Args:
         lmp (lammps.Lammps): Lammps library instance
-        bispec_options (dict): bi-spectrum component settings
+        bispec_options (dict): dictionary of the parameters to compute the bi-spectral components
     """
     # # Bispectrum coefficient computes
     base_b = "compute b all sna/atom ${rcutfac} ${rfac0} ${twojmax}"
@@ -466,7 +515,19 @@ def _set_computes_snap(lmp, bispec_options):
         lmp.command(f"compute {cname}_sum all reduce sum c_{cname}[*]")
 
 
-def _extract_computes_snap(lmp, num_atoms, n_coeff, num_types):
+def _extract_computes_snap(lmp, num_atoms: int, n_coeff: int, num_types: int) -> np.ndarray:
+    """
+    Internal function to extract the compute from the LAMMPS instance
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        num_atoms (int): Number of atoms
+        n_coeff (int): Number of SNAP coefficients to compure
+        num_types (int): Number of chemical types
+
+    Returns:
+        np.ndarray: Output of the LAMMPS compute command
+    """
     lmp_atom_ids = lmp.numpy.extract_atom_iarray("id", num_atoms).flatten()
     assert np.all(
         lmp_atom_ids == 1 + np.arange(num_atoms)
@@ -521,11 +582,23 @@ def _extract_computes_snap(lmp, num_atoms, n_coeff, num_types):
     ).copy()
 
 
-def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
+def _calc_snap_derivatives(lmp, structure: Atoms, bispec_options: dict, cutoff=10.0) -> np.ndarray:
+    """
+    Internal function to calculate per-atom derivatives of the SNAP descriptors
+
+    Args:
+        lmp (lammps.Lammps): Lammps library instance
+        structure (ase.atoms.Atoms): atomistic structure as ASE atoms object
+        bispec_options (dict): dictionary of the parameters to compute the bi-spectral components
+        cutoff (float): cutoff radius for the construction of the neighbor list
+
+    Returns:
+        np.ndarray: Output of the LAMMPS compute command
+    """
     number_coef = len(get_snap_descriptor_names(twojmax=bispec_options["twojmax"]))
     number_species = len(set(structure.get_chemical_symbols()))
     _reset_lmp(lmp=lmp)
-    _write_ase_structure(lmp=lmp, structure=structure)
+    _set_ase_structure(lmp=lmp, structure=structure)
     _set_potential_lmp(lmp=lmp, cutoff=cutoff)
     _set_variables(lmp, **_lammps_variables(bispec_options))
     _set_computes_snap(
@@ -559,17 +632,35 @@ def _calc_snap_derivatives(lmp, structure, bispec_options, cutoff=10.0):
 
 
 def _get_default_parameters(
-    atom_types,
-    twojmax=6,
+    atom_types: list[str],
+    twojmax: int = 6,
     element_radius=4.0,
-    rcutfac=1.0,
-    rfac0=0.99363,
-    rmin0=0.0,
-    bzeroflag=0,
-    weights=None,
-    cutoff=10.0,
-    quadraticflag=0,
+    rcutfac: float = 1.0,
+    rfac0: float = 0.99363,
+    rmin0: float = 0.0,
+    bzeroflag: bool = False,
+    quadraticflag: bool = False,
+    weights: Optional[Union[list, np.ndarray]] = None,
+    cutoff: float = 10.0,
 ):
+    """
+    Convert input parameters to the internal dictonary format
+
+    Args:
+        atom_types (list[str]): list of element types
+        twojmax (int): band limit for bispectrum components (non-negative integer)
+        element_radius (list[int]): list of radii for the individual elements
+        rcutfac (float): scale factor applied to all cutoff radii (positive real)
+        rfac0 (float): parameter in distance to angle conversion (0 < rcutfac < 1)
+        rmin0 (float): parameter in distance to angle conversion (distance units)
+        bzeroflag (bool): subtract B0
+        quadraticflag (bool): generate quadratic terms
+        weights (list/np.ndarry/None): list of neighbor weights, one for each type
+        cutoff (float): cutoff radius for the construction of the neighbor list
+
+    Returns:
+        (lammps.Lammps, dict, flaot): LAMMPS instance, dictionary of bi-spectral component parameters, cut-off radius
+    """
     from lammps import lammps
 
     if weights is None:
diff --git a/tests/test_snap.py b/tests/test_snap.py
index 43b5d9a22..c5f0f5e9b 100644
--- a/tests/test_snap.py
+++ b/tests/test_snap.py
@@ -2,8 +2,8 @@
 import numpy as np
 import structuretoolkit as stk
 from structuretoolkit.analyse.snap import (
-    calc_per_atom_quad,
-    calc_sum_quad,
+    get_per_atom_quad,
+    get_sum_quad,
     _calc_snap_per_atom,
     _calc_snap_derivatives,
     _get_lammps_compatible_cell
@@ -337,8 +337,8 @@ def test_calc_bispectrum_lmp_quad(self):
             cutoff=10.0
         )
         self.assertEqual(coeff_quad.shape, (len(self.structure), ((n_coeff+1) * n_coeff)/2 + 30))
-        coeff_quad_per_atom = calc_per_atom_quad(coeff_lin)
-        coeff_quad_sum = calc_sum_quad(np.sum(coeff_lin, axis=0))
+        coeff_quad_per_atom = get_per_atom_quad(coeff_lin)
+        coeff_quad_sum = get_sum_quad(np.sum(coeff_lin, axis=0))
         self.assertEqual(coeff_quad.shape, coeff_quad_per_atom.shape)
         self.assertEqual(np.sum(coeff_quad, axis=0).shape, coeff_quad_sum.shape)
         self.assertTrue(np.isclose(

From 65ece1a30548eba3e5455e74af483a7e1d9a131d Mon Sep 17 00:00:00 2001
From: pyiron-runner <pyiron@mpie.de>
Date: Thu, 14 Mar 2024 15:27:47 +0000
Subject: [PATCH 34/34] Format black

---
 structuretoolkit/analyse/snap.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/structuretoolkit/analyse/snap.py b/structuretoolkit/analyse/snap.py
index c9dc760c3..9a9b6629e 100644
--- a/structuretoolkit/analyse/snap.py
+++ b/structuretoolkit/analyse/snap.py
@@ -264,7 +264,9 @@ def _set_ase_structure(lmp, structure: Atoms):
     )
 
 
-def _extract_compute_np(lmp, name: str, compute_type: int, result_type: int, array_shape: tuple) -> np.ndarray:
+def _extract_compute_np(
+    lmp, name: str, compute_type: int, result_type: int, array_shape: tuple
+) -> np.ndarray:
     """
     Convert a lammps compute to a numpy array.
     Assumes the compute returns a floating point numbers.
@@ -378,7 +380,9 @@ def _set_compute_lammps(lmp, bispec_options: dict, numtypes: int):
     lmp.command(f"{base_b} {radelem} {wj} {kwargs}")
 
 
-def _calc_snap_per_atom(lmp, structure: Atoms, bispec_options: dict, cutoff: float = 10.0) -> np.ndarray:
+def _calc_snap_per_atom(
+    lmp, structure: Atoms, bispec_options: dict, cutoff: float = 10.0
+) -> np.ndarray:
     """
     Internal function to calculate the per-atom SNAP descriptors
 
@@ -515,7 +519,9 @@ def _set_computes_snap(lmp, bispec_options: dict):
         lmp.command(f"compute {cname}_sum all reduce sum c_{cname}[*]")
 
 
-def _extract_computes_snap(lmp, num_atoms: int, n_coeff: int, num_types: int) -> np.ndarray:
+def _extract_computes_snap(
+    lmp, num_atoms: int, n_coeff: int, num_types: int
+) -> np.ndarray:
     """
     Internal function to extract the compute from the LAMMPS instance
 
@@ -582,7 +588,9 @@ def _extract_computes_snap(lmp, num_atoms: int, n_coeff: int, num_types: int) ->
     ).copy()
 
 
-def _calc_snap_derivatives(lmp, structure: Atoms, bispec_options: dict, cutoff=10.0) -> np.ndarray:
+def _calc_snap_derivatives(
+    lmp, structure: Atoms, bispec_options: dict, cutoff=10.0
+) -> np.ndarray:
     """
     Internal function to calculate per-atom derivatives of the SNAP descriptors