From 24910439ac3eeec406d3a8cf1dec782b478d05bb Mon Sep 17 00:00:00 2001
From: Joroks <johannes@roksvaag.com>
Date: Mon, 15 Jul 2024 15:05:52 +0200
Subject: [PATCH] support gather[bonds|angles|dihedrals|impropers]

---
 src/LAMMPS.jl                                 | 77 ++++++++++++++++++-
 test/runtests.jl                              | 37 ++++++++-
 .../bonds_angles_dihedrals_impropers.data     | 48 ++++++++++++
 3 files changed, 158 insertions(+), 4 deletions(-)
 create mode 100644 test/test_files/bonds_angles_dihedrals_impropers.data

diff --git a/src/LAMMPS.jl b/src/LAMMPS.jl
index 9c0199f..66ce65a 100644
--- a/src/LAMMPS.jl
+++ b/src/LAMMPS.jl
@@ -3,8 +3,8 @@ import MPI
 include("api.jl")
 
 export LMP, command, get_natoms, extract_atom, extract_compute, extract_global,
-       extract_setting, gather, scatter!, group_to_atom_ids, get_category_ids,
-       extract_variable,
+       extract_setting, gather, gather_bonds, gather_angles, gather_dihedrals, gather_impropers,
+       scatter!, group_to_atom_ids, get_category_ids, extract_variable,
 
        # _LMP_DATATYPE
        LAMMPS_NONE,
@@ -740,6 +740,79 @@ function _get_T(lmp::LMP, name::String)
 
 end
 
+"""
+    gather_bonds(lmp::LMP)
+
+Gather the list of all bonds into a 3 x nbonds Matrix:
+```
+row1 -> bond type
+row2 -> atom 1
+row3 -> atom 2
+```
+"""
+function gather_bonds(lmp::LMP)
+    ndata = extract_global(lmp, "nbonds", LAMMPS_INT64)[]
+    data = Matrix{Int32}(undef, 3, ndata)
+    API.lammps_gather_bonds(lmp, data)
+    return data
+end
+
+"""
+    gather_angles(lmp::LMP)
+
+Gather the list of all angles into a 4 x nangles Matrix:
+```
+row1 -> angle type
+row2 -> atom 1
+row3 -> atom 2
+row4 -> atom 3
+```
+"""
+function gather_angles(lmp::LMP)
+    ndata = extract_global(lmp, "nangles", LAMMPS_INT64)[]
+    data = Matrix{Int32}(undef, 4, ndata)
+    API.lammps_gather_angles(lmp, data)
+    return data
+end
+
+"""
+    gather_dihedrals(lmp::LMP)
+
+Gather the list of all dihedrals into a 5 x ndihedrals Matrix:
+```
+row1 -> dihedral type
+row2 -> atom 1
+row3 -> atom 2
+row4 -> atom 3
+row5 -> atom 4
+```
+"""
+function gather_dihedrals(lmp::LMP)
+    ndata = extract_global(lmp, "ndihedrals", LAMMPS_INT64)[]
+    data = Matrix{Int32}(undef, 5, ndata)
+    API.lammps_gather_dihedrals(lmp, data)
+    return data
+end
+
+"""
+    gather_impropers(lmp::LMP)
+
+Gather the list of all impropers into a 5 x nimpropers Matrix:
+```
+row1 -> improper type
+row2 -> atom 1
+row3 -> atom 2
+row4 -> atom 3
+row5 -> atom 4
+```
+"""
+function gather_impropers(lmp::LMP)
+    ndata = extract_global(lmp, "nimpropers", LAMMPS_INT64)[]
+    data = Matrix{Int32}(undef, 5, ndata)
+    API.lammps_gather_impropers(lmp, data)
+    return data
+end
+
 """
     group_to_atom_ids(lmp::LMP, group::String)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 267cfcc..e9812f7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -188,8 +188,41 @@ end
 
         @test gather(lmp, "x", Float64, subset) == gather(lmp, "c_pos", Float64, subset) == gather(lmp, "f_pos", Float64, subset) == data_subset
 
-    # verify that no errors were missed
-    @test LAMMPS.API.lammps_has_error(lmp) == 0
+        # verify that no errors were missed
+        @test LAMMPS.API.lammps_has_error(lmp) == 0
+    end
+end
+
+@testset "Gather bonds/angles/dihedrals/impropers" begin
+    LMP(["-screen", "none"]) do lmp
+        file = joinpath(@__DIR__, "test_files/bonds_angles_dihedrals_impropers.data")
+
+        command(lmp,  """
+            atom_style molecular
+            read_data $file
+            """)
+
+        @test gather_bonds(lmp) == transpose([
+            1 1 2
+            1 2 3
+            1 3 4
+            1 4 1
+        ])
+        @test gather_angles(lmp) == transpose([
+            1 1 2 3
+            1 2 3 4
+        ])
+        @test gather_angles(lmp) == transpose([
+            1 1 2 3
+            1 2 3 4
+        ])
+        @test gather_dihedrals(lmp) == transpose([
+            1 1 2 3 4
+        ])
+        @test gather_impropers(lmp) == transpose([
+            1 4 3 2 1
+        ])
+        @test LAMMPS.API.lammps_has_error(lmp) == 0
     end
 end
 
diff --git a/test/test_files/bonds_angles_dihedrals_impropers.data b/test/test_files/bonds_angles_dihedrals_impropers.data
new file mode 100644
index 0000000..7fa7d1d
--- /dev/null
+++ b/test/test_files/bonds_angles_dihedrals_impropers.data
@@ -0,0 +1,48 @@
+#
+
+4 atoms
+4 bonds
+2 angles
+1 dihedrals
+1 impropers
+
+1 atom types
+1 bond types
+1 angle types
+1 dihedral types
+1 improper types
+
+-1 1 xlo xhi
+-1 1 ylo yhi
+-1 1 zlo zhi
+
+Masses
+
+1 1
+
+Atoms # molecular
+
+1 1 1 0 0 0
+2 1 1 0 0 0
+3 1 1 0 0 0
+4 1 1 0 0 0
+
+Bonds
+
+1 1 1 2
+2 1 2 3
+3 1 3 4
+4 1 4 1
+
+Angles
+
+1 1 1 2 3
+2 1 2 3 4
+
+Dihedrals
+
+1 1 1 2 3 4
+
+Impropers
+
+1 1 4 3 2 1
\ No newline at end of file