Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Provide size_2d argument to extract_compute #69

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions src/LAMMPS.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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))
Expand All @@ -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)[]
Expand Down
23 changes: 23 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading