From aa836e351d6d9b91e7503c709ea1dae341aa73a1 Mon Sep 17 00:00:00 2001
From: Spencer Wyant <17836774+swyant@users.noreply.github.com>
Date: Mon, 18 Nov 2024 08:44:47 -0500
Subject: [PATCH 1/2] enable users to use a special kwarg indicating size of 2D
 array for TYPE_ARRAY in extract_compute

---
 src/LAMMPS.jl | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/LAMMPS.jl b/src/LAMMPS.jl
index 7516d9d..5e755db 100644
--- a/src/LAMMPS.jl
+++ b/src/LAMMPS.jl
@@ -582,7 +582,10 @@ LMP(["-screen", "none"]) do lmp
 end
 ```
 """
-function extract_compute(lmp::LMP, name::String, style::_LMP_STYLE_CONST, lmp_type::_LMP_TYPE; copy::Bool=false)
+function extract_compute(lmp::LMP, name::String, style::_LMP_STYLE_CONST, lmp_type::_LMP_TYPE; 
+                         copy::Bool=false,
+                         size_2d::Union{Tuple{Int64,Int64}, Nothing}=nothing)
+
     API.lammps_has_id(lmp, "compute", name) != 1 && throw(KeyError("Unknown compute $name"))
 
     void_ptr = API.lammps_extract_compute(lmp, name, style, get_enum(lmp_type))
@@ -608,6 +611,12 @@ function extract_compute(lmp::LMP, name::String, style::_LMP_STYLE_CONST, lmp_ty
         return  _extract(ptr, ndata; copy=copy)
     end
 
+    # if you know the size of the 2D array, skip getting the row size, column size
+    if lmp_type == TYPE_ARRAY && !isnothing(size_2d)
+        ptr = _reinterpret(LAMMPS_DOUBLE_2D, void_ptr)
+        return _extract(ptr, size_2d; copy=copy)
+    end
+
     ndata = (style == STYLE_ATOM) ?
         extract_setting(lmp, "nlocal") :
         extract_compute(lmp, name, style, SIZE_ROWS)[]

From 3080ee590829df3c9d00281464fe1c028adb63b4 Mon Sep 17 00:00:00 2001
From: Spencer Wyant <17836774+swyant@users.noreply.github.com>
Date: Thu, 21 Nov 2024 18:38:00 -0500
Subject: [PATCH 2/2] add documentation for size_2d option, tests for
 TYPE_ARRAY extract_compute with the size_2d option

---
 src/LAMMPS.jl    | 10 +++++++++-
 test/runtests.jl | 23 +++++++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/src/LAMMPS.jl b/src/LAMMPS.jl
index 5e755db..8db4d2d 100644
--- a/src/LAMMPS.jl
+++ b/src/LAMMPS.jl
@@ -541,7 +541,7 @@ function extract_atom_datatype(lmp::LMP, name)
 end
 
 """
-    extract_compute(lmp::LMP, name::String, style::_LMP_STYLE_CONST, lmp_type::_LMP_TYPE; copy::Bool=true)
+    extract_compute(lmp::LMP, name::String, style::_LMP_STYLE_CONST, lmp_type::_LMP_TYPE; copy::Bool=true, size_2d::Union{Tuple{Int64,Int64}, Nothing}=nothing)
 
 Extract data provided by a compute command identified by the compute-ID.
 Computes may provide global, per-atom, or local data, and those may be a scalar, a vector or an array.
@@ -565,6 +565,14 @@ Since computes may provide multiple kinds of data, it is required to set style a
 Scalar values get returned as a vector with a single element. This way it's possible to
 modify the internal state of the LAMMPS instance even if the data is scalar.
 
+The `size_2d` argument allows users to pre-specify the size of the 2D array when `style=STYLE_GLOBAL` 
+and `lmp_type=TYPE_ARRAY`, thus bypassing the need to perform multiple extracts to determine row and column size. 
+
+!!! warning
+    The `size_2d` option is fundamentally memory-unsafe and will lead to out-of-bounds access if the specified 
+    size is larger than the array returned by the LAMMPS instance. It should only be used if the user confidently 
+    knows what size the returned 2D array will be. Otherwise, do not use this keyword argument!
+
 !!! info
     The returned data may become invalid as soon as another LAMMPS command has been issued at any point after calling this method.
     If this has happened, trying to read from this data will likely cause julia to crash.
diff --git a/test/runtests.jl b/test/runtests.jl
index 221646d..971ae6c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -267,6 +267,29 @@ end
     end
 end
 
+@testset "Extract Compute Global Array" begin
+    LMP(["-screen", "none"]) do lmp
+        command(lmp, """
+            atom_modify map yes
+            region cell block 0 6 0 6 0 6
+            create_box 1 cell
+            lattice sc 1
+            create_atoms 1 region cell
+            mass 1 1
+
+            compute bin3d all chunk/atom bin/3d x lower 3.0 y lower 3.0 z lower 3.0 compress yes units box
+            compute prop3 all property/chunk bin3d id count
+        """)
+
+        prop_arr1 = extract_compute(lmp, "prop3", STYLE_GLOBAL, TYPE_ARRAY)
+        @test size(prop_arr1) == (2,8)
+
+        prop_arr2 = extract_compute(lmp, "prop3", STYLE_GLOBAL, TYPE_ARRAY; size_2d=(2,8))
+        @test prop_arr1 == prop_arr2
+    end
+end
+
+
 @testset "Utilities" begin
     LMP(["-screen", "none"]) do lmp
         # setting up example data