diff --git a/src/LAMMPS.jl b/src/LAMMPS.jl index 7516d9d..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. @@ -582,7 +590,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 +619,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)[] 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