Repository: cesmix-mit/LAMMPS.jl
Pull request #74: Add Support for fix_external
Author: login: Joroks
Created: 2025-04-22T21:17:58Z | Merged: 2025-05-09T13:25:43Z

@vchuravy I've reworked the progress you've made in your own fix_external branch. It's still very much work in progress, however, and needs some cleaning up and proper documentation.

I've changed PairExternal such that you only provide the potential function. I then use AD (provided through Tensors/ForwardDiff) to calculate the Force. I've also made sure that the potentials, the forces, and the virials (both global and per-atom) are set correcly.


Diff:
diff --git a/Project.toml b/Project.toml
index 5068d5f..6e3571b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -4,19 +4,27 @@ authors = ["Valentin Churavy <v.churavy@gmail.com>"]
 version = "0.8.0"
 
 [deps]
+Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LAMMPS_jll = "5b3ab26d-9607-527c-88ea-8fe5ba57cafe"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 OpenBLAS32_jll = "656ef2d0-ae68-5445-9ca0-591084a874a2"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [compat]
+Bumper = "0.7.1"
 CEnum = "0.4, 0.5"
+DifferentiationInterface = "0.6.52"
+ForwardDiff = "1.0.1"
 LAMMPS_jll = "2.8.0"
 MPI = "0.20"
 Preferences = "1"
+StaticArrays = "1.9.13"
 UnsafeArrays = "1.0.8"
 julia = "1.10"
diff --git a/examples/fitting_snap.jl b/examples/fitting_snap.jl
index f87f7aa..2226f90 100644
--- a/examples/fitting_snap.jl
+++ b/examples/fitting_snap.jl
@@ -25,6 +25,7 @@
 #
 # We start off by importing the necessary packages
 
+using MPI; MPI.Init()
 using LAMMPS
 using LinearAlgebra: norm, pinv
 using Printf
diff --git a/examples/fix_external.jl b/examples/fix_external.jl
new file mode 100644
index 0000000..ab4ff13
--- /dev/null
+++ b/examples/fix_external.jl
@@ -0,0 +1,65 @@
+using MPI; MPI.Init()
+using LAMMPS
+
+lmp = LMP()
+
+command(lmp, "units lj")
+command(lmp, "atom_style atomic")
+command(lmp, "atom_modify map array sort 0 0")
+command(lmp, "box tilt large")
+
+# Setup box
+x_hi = 10.0
+y_hi = 10.0
+z_hi = 10.0
+command(lmp, "boundary p p p")
+command(lmp, "region cell block 0 $x_hi 0 $y_hi 0 $z_hi units box")
+command(lmp, "create_box 1 cell")
+
+# Use `pair_style zero` to create neighbor list for `julia_lj`
+cutoff = 2.5
+command(lmp, "pair_style zero $cutoff")
+command(lmp, "pair_coeff * *")
+command(lmp, "fix julia_lj all external pf/callback 1 1")
+
+const coefficients = Base.ImmutableDict(
+    1 => Base.ImmutableDict(
+        1 => [48.0, 24.0, 4.0, 4.0]
+    )
+)
+
+@inline function compute_force(rsq, itype, jtype)
+    coeff = coefficients[itype][jtype]
+    r2inv  = 1.0/rsq
+    r6inv  = r2inv^3
+    lj1 = coeff[1]
+    lj2 = coeff[2]
+    return (r6inv * (lj1*r6inv - lj2))*r2inv
+end
+
+@inline function compute_energy(rsq, itype, jtype)
+    coeff = coefficients[itype][jtype]
+    r2inv  = 1.0/rsq
+    r6inv  = r2inv^3
+    lj3 = coeff[3]
+    lj4 = coeff[4]
+    return (r6inv * (lj3*r6inv - lj4))
+end
+
+# Register external fix
+lj = LAMMPS.PairExternal(lmp, "julia_lj", "zero", compute_force, compute_energy, cutoff)
+
+# Setup atoms
+natoms = 10
+command(lmp, "create_atoms 1 random $natoms 1 NULL")
+command(lmp, "mass 1 1.0")
+
+# (x,y,z), natoms
+positions = rand(3, 10) .* 5
+
+LAMMPS.API.lammps_scatter_atoms(lmp, "x", 1, 3, positions)
+
+command(lmp, "run 0")
+
+# extract forces
+forces = extract_atom(lmp, "f")
diff --git a/examples/lj_forces.jl b/examples/lj_forces.jl
index 1374552..aec35ae 100644
--- a/examples/lj_forces.jl
+++ b/examples/lj_forces.jl
@@ -1,3 +1,4 @@
+using MPI; MPI.Init()
 using LAMMPS
 
 lmp = LMP()
diff --git a/examples/snap.jl b/examples/snap.jl
index 46c086a..662a9da 100644
--- a/examples/snap.jl
+++ b/examples/snap.jl
@@ -16,6 +16,7 @@
 #
 # We start off by importing `LAMMPS`.
 
+using MPI; MPI.Init()
 using LAMMPS
 
 const DATA = joinpath(dirname(pathof(LAMMPS)), "..", "examples", "example_GaN_data")
diff --git a/src/LAMMPS.jl b/src/LAMMPS.jl
index 560db8f..00adf89 100644
--- a/src/LAMMPS.jl
+++ b/src/LAMMPS.jl
@@ -1,15 +1,20 @@
 module LAMMPS
 import MPI
-import LinearAlgebra
+using LinearAlgebra
 import OpenBLAS32_jll
+import StaticArrays: SVector, SMatrix, MVector, SA
+import DifferentiationInterface: AbstractADType, Constant, value_and_derivative
+import Bumper: @no_escape, @alloc
 import UnsafeArrays: UnsafeArray
 
 include("api.jl")
 
 export LMP, command, create_atoms, get_natoms, extract_atom, extract_compute, extract_global,
        extract_setting, extract_box, reset_box, gather, gather_bonds, gather_angles, gather_dihedrals,
-       gather_impropers, scatter!, group_to_atom_ids, get_category_ids, extract_variable, LAMMPSError,
+       gather_impropers, scatter!, group_to_atom_ids, get_category_ids, extract_variable, LAMMPSError, FixExternal,
+       PairExternal, set_energy!, set_virial!, InteractionConfig,
        encode_image_flags, decode_image_flags, compute_neighborlist, fix_neighborlist, pair_neighborlist,
+       get_mpi_comm,
 
        # _LMP_DATATYPE
        LAMMPS_NONE,
@@ -156,7 +161,7 @@ function set_library!(path)
 end
 
 """
-    LMP(args::Vector{String}=String[], comm::Union{Nothing, MPI.Comm}=nothing)
+    LMP(args::Vector{String}=String[], comm::MPI.Comm=MPI.COMM_WORLD)
 
 Create a new LAMMPS instance while passing in a list of strings as if they were command-line arguments for the LAMMPS executable.
 
@@ -164,19 +169,20 @@ A full ist of command-line options can be found in the [lammps documentation](ht
 """
 mutable struct LMP
     @atomic handle::Ptr{Cvoid}
+    external_fixes::Dict{String, Any}
 
-    function LMP(args::Vector{String}=String[], comm::Union{Nothing, MPI.Comm}=nothing)
+    function LMP(args::Vector{String}=String[], comm::MPI.Comm=MPI.COMM_WORLD)
         args = copy(args)
         pushfirst!(args, "lammps")
 
         GC.@preserve args begin
-            if comm !== nothing
+            if API.lammps_config_has_mpi_support() == 0
+                handle = API.lammps_open_no_mpi(length(args), args, C_NULL)
+            else
                 if !MPI.Initialized()
                     error("MPI has not been initialized. Make sure to first call `MPI.Init()`")
                 end
                 handle = API.lammps_open(length(args), args, comm, C_NULL)
-            else
-                handle = API.lammps_open_no_mpi(length(args), args, C_NULL)
             end
         end
 
@@ -187,7 +193,7 @@ mutable struct LMP
             throw(LAMMPSError(msg))
         end
 
-        this = new(handle)
+        this = new(handle, Dict{String, Any}())
         finalizer(close!, this)
 
         ver = version(this)
@@ -214,16 +220,18 @@ Shutdown a LAMMPS instance.
 function close!(lmp::LMP)
     handle = @atomicswap lmp.handle = C_NULL
     if handle !== C_NULL 
-       API.lammps_close(handle)
+        empty!(lmp.external_fixes)
+        API.lammps_close(handle)
     end
+    return nothing
 end
 
 """
-    LMP(f::Function, args=String[], comm=nothing)
+    LMP(f::Function, args=String[], comm=MPI.COMM_WORLD)
 
 Create a new LAMMPS instance and call `f` on that instance while returning the result from `f`.
 """
-function LMP(f::Function, args=String[], comm=nothing)
+function LMP(f::Function, args=String[], comm=MPI.COMM_WORLD)
     lmp = LMP(args, comm)
     return f(lmp)
     # `close!` is registered as a finalizer for LMP, no need to close it here.
@@ -233,6 +241,18 @@ function version(lmp::LMP)
     API.lammps_version(lmp)
 end
 
+"""
+    get_mpi_comm(lmp::LMP)::Union{Nothing, MPI.Comm}
+
+Return the MPI communicator used by the lammps instance or `nothing` if the lammps build doesn't support MPI.
+"""
+function get_mpi_comm(lmp::LMP)::Union{Nothing, MPI.Comm}
+    comm_f = API.lammps_get_mpi_comm(lmp)
+    comm_f == -1 && return nothing
+    comm_c = MPI.API.MPI_Comm_f2c(comm_f)
+    return MPI.Comm(comm_c)
+end
+
 struct LAMMPSError <: Exception
     msg::String
 end
@@ -417,7 +437,9 @@ A full list of settings can be found in the [lammps documentation](https://docs.
 ```
 """
 function extract_setting(lmp::LMP, name::String)::Int32
-    return API.lammps_extract_setting(lmp, name)
+    val = API.lammps_extract_setting(lmp, name)
+    val == -1 && error("Could not find setting $name")
+    return val
 end
 
 """
@@ -874,7 +896,6 @@ function _get_T(lmp::LMP, name::String)
     else
         error("Unkown per atom property $name")
     end
-
 end
 
 """
@@ -1116,4 +1137,6 @@ function pair_neighborlist(lmp::LMP, style::String; exact=false, nsub=0, request
 end
 
 
+include("external.jl")
+
 end # module
diff --git a/src/api.jl b/src/api.jl
index 54773b1..b5b06f2 100644
--- a/src/api.jl
+++ b/src/api.jl
@@ -448,6 +448,10 @@ function lammps_flush_buffers(ptr)
     ccall((:lammps_flush_buffers, liblammps), Cvoid, (Ptr{Cvoid},), ptr)
 end
 
+# function lammps_fix_external_set_energy_peratom(handle, id, eng)
+#     ccall((:lammps_fix_external_set_energy_peratom, liblammps), Cvoid, (Ptr{Cvoid}, Ptr{Cchar}, Ptr{Cdouble}), handle, id, eng)
+# end
+
 function lammps_free(ptr)
     ccall((:lammps_free, liblammps), Cvoid, (Ptr{Cvoid},), ptr)
 end
diff --git a/src/external.jl b/src/external.jl
new file mode 100644
index 0000000..c6d5891
--- /dev/null
+++ b/src/external.jl
@@ -0,0 +1,416 @@
+function fix_external_callback end
+
+"""
+    FixExternal(callback, lmp::LMP, name::String, group::String, ncall::Int, napply::Int)
+
+Creates a fix in lammps that calls a Julia function `callback` every `ncall` during the simulation.
+
+!!! info
+    The following command is executed in LAMMPS when `FixExternal` is called in order to setup the fix:
+    ```lammps
+    fix <name> <group> external pf/callback <ncall> <napply>
+    ```
+
+The `FixExternal` object gets passed to the `callback` function, it contains the parameters passed to `FixExternal` as fields:
+- `lmp::LMP`: The LAMMPS object.
+- `name::String`: The name of the fix.
+- `group::String`: The group of atoms to which the fix is applied.
+- `ncall::Int`: The number of timesteps between calls to the callback function.
+- `napply::Int`: The number of times the callback function is applied to the atoms in the group.
+Additionally, the following fields are set by the callback function:
+- `timestep::Int`: The current timestep of the simulation.
+- `nlocal::Int`: The number of local atoms on the MPI rank.
+- `ids::UnsafeArray{Int32, 1}`: The IDs of the local atoms in an `nall`-element Vector.
+- `x::UnsafeArray{Float64, 2}`: The positions of the local atoms in a (3 × `nall`) Matrix.
+- `f::UnsafeArray{Float64, 2}`: The forces on the local atoms in a (3 × `nlocal`) Matrix.
+These values are only valid during the callback execution and should *not* be used outside of it, as they refer directly to LAMMPS internal memory.
+
+Here, the intention is for the callback to write forces to the `f` field, which will be applied to the atoms in the group every `napply` timesteps.
+`f` is *not* zeroed before the callback is called, so the forces from previous calls are preserved.
+
+Contributions to the per-atom energies or virials can be set with [`set_energy!`](@ref) and [`set_virial!`](@ref), respectively.
+
+"""
+mutable struct FixExternal
+    const callback
+    const lmp::LMP
+    const name::String
+    const group::String
+    const ncall::Int
+    const napply::Int
+    # The following fields are set by the callback function
+    timestep::Int
+    nlocal::Int
+    ids::UnsafeArray{Int32, 1}
+    x::UnsafeArray{Float64, 2}
+    f::UnsafeArray{Float64, 2}
+
+    function FixExternal(callback, lmp::LMP, name::String, group::String, ncall::Int, napply::Int)
+        command(lmp, "fix $name $group external pf/callback $ncall $napply")
+        this = new(callback, lmp, name, group, ncall, napply,  0, 0,
+            UnsafeArray(Ptr{Int32}(C_NULL), (0, )),
+            UnsafeArray(Ptr{Float64}(C_NULL), (3, 0)), 
+            UnsafeArray(Ptr{Float64}(C_NULL), (3, 0))
+        )
+        lmp.external_fixes[name] = this # preserves pair globally
+
+        ctx = Base.pointer_from_objref(this)
+        callback = @cfunction(fix_external_callback, Cvoid, (Ptr{Cvoid}, Int64, Cint, Ptr{Cint}, Ptr{Ptr{Float64}}, Ptr{Ptr{Float64}}))
+        API.lammps_set_fix_external_callback(lmp, name, callback, ctx)
+
+        # Ensure function is compiled before timestep 0
+        if !precompile(this.callback, (FixExternal,))
+            @warn "Failed to precompile the callback" this.callback
+        end
+        return this
+    end
+end
+
+function fix_external_callback(ctx::Ptr{Cvoid}, timestep::Int64, nlocal::Cint, ids::Ptr{Cint}, x::Ptr{Ptr{Float64}}, fexternal::Ptr{Ptr{Float64}})
+    fix = Base.unsafe_pointer_to_objref(ctx)::FixExternal
+    fix.timestep = timestep
+    fix.nlocal = nlocal
+
+    nall = extract_setting(fix.lmp, "nall")
+    fix.x = _extract(x, (Int32(3), nall))
+    fix.f = _extract(fexternal, (Int32(3), nlocal))
+    fix.ids = _extract(ids, nall)
+
+    # necessary dynamic
+    fix.callback(fix)
+
+    fix.timestep = 0
+    fix.nlocal = 0
+
+    fix.ids = UnsafeArray(Ptr{Int32}(C_NULL), (0, ))
+    fix.x = UnsafeArray(Ptr{Float64}(C_NULL), (3, 0))
+    fix.f = UnsafeArray(Ptr{Float64}(C_NULL), (3, 0))
+
+    return nothing
+end
+
+"""
+    set_virial!(fix::FixExternal, virial_peratom::AbstractVector{SVector{6, Float64}})
+    set_virial!(fix::FixExternal, virial_peratom::AbstractMatrix{Float64})
+
+Sets the contribution to the local per-atom and total global virial stress tensors resulting from the `FixExternal`.
+This function should only be used within the callback of the `FixExternal`
+
+The virial tensor is represented as a 6-component vector with the order: 
+  `[xx, yy, zz, xy, xz, yz]`.
+"""
+function set_virial!(fix::FixExternal, virial_peratom::AbstractVector{SVector{6, Float64}})
+    @assert length(virial_peratom) >= fix.nlocal
+
+    @no_escape begin
+        ptrs = @alloc(Ptr{Ptr{Float64}}, fix.nlocal)
+        for i in eachindex(ptrs)
+            @inbounds ptrs[i] = pointer(virial_peratom, i)
+        end
+        API.lammps_fix_external_set_virial_peratom(fix.lmp, fix.name, ptrs)
+    end
+
+    virial_global = MVector(sum(view(virial_peratom, 1:fix.nlocal)))
+    comm = get_mpi_comm(fix.lmp)
+    GC.@preserve virial_global begin
+        if comm !== nothing
+            buffer = MPI.RBuffer(MPI.IN_PLACE, pointer(virial_global), 6, MPI.Datatype(Float64))
+            @inline MPI.Allreduce!(buffer, +, comm)
+        end
+        @inline API.lammps_fix_external_set_virial_global(fix.lmp, fix.name, virial_global)
+    end
+
+    check(fix.lmp)
+end
+
+function set_virial(fix::FixExternal, virial::AbstractMatrix{Float64})
+    virial_vec = reinterpret(reshape, SVector{6, Float64}, virial)
+    set_virial_peratom(fix, virial_vec; set_global)
+end
+
+"""
+    set_energy!(fix::FixExternal, energy::AbstractVector{Float64}; set_global=false)
+
+Sets the contribution to the local per-atom and total global energy resulting from the `FixExternal`.
+This function should only be used within the callback of the `FixExternal`
+"""
+function set_energy!(fix::FixExternal, energy::AbstractVector{Float64})
+    @assert length(energy) >= fix.nlocal
+
+    API.lammps_fix_external_set_energy_peratom(fix.lmp, fix.name, energy)
+
+    energy_global = sum(view(energy, 1:fix.nlocal))
+    comm = get_mpi_comm(fix.lmp)
+    if comm !== nothing
+        energy_global = MPI.Allreduce(energy_global, +, comm)
+    end
+    API.lammps_fix_external_set_energy_global(fix.lmp, fix.name, energy_global)
+
+    check(fix.lmp)
+end
+
+function _check_atom_datatype(lmp, name, T)
+    lammps_type = API.lammps_extract_atom_datatype(lmp, name)
+    rows::Int = API.lammps_extract_atom_size(lmp, name, API.LMP_SIZE_COLS)
+
+    if lammps_type == API.LAMMPS_INT64
+        T !== Int64 && throw(ArgumentError("Expected type Int64 for per-atom property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_INT
+        T !== Int32 && throw(ArgumentError("Expected type Int32 for per-atom property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_DOUBLE
+        T !== Float64 && throw(ArgumentError("Expected type Float64 for per-atom property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_INT64_2D
+        (T !== NTuple{rows, Int64} && T !== SVector{rows, Int64}) && throw(ArgumentError("Expected type NTuple{$rows, Int64} or SVector{$rows, Int64} for per-atom property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_INT_2D
+        (T !== NTuple{rows, Int32} && T !== SVector{rows, Int32}) && throw(ArgumentError("Expected type NTuple{$rows, Int32} or SVector{$rows, Int32} for per-atom property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_DOUBLE_2D
+        (T !== NTuple{rows, Float64} && T !== SVector{rows, Float64}) && throw(ArgumentError("Expected type NTuple{$rows, Float64} or SVector{$rows, Float64} for per-atom property :$name, got $T instead"))
+    else
+        throw(KeyError(name))
+    end
+end
+
+function _check_atom_ghost(lmp, name, nall)
+    cols = API.lammps_extract_atom_size(lmp, name, API.LMP_SIZE_ROWS)
+    cols == nall || throw(ArgumentError("Per-atom property :$name doesn't communicate ghost data"))
+    nothing
+end
+
+struct ExtractAtomMultiple{T<:NamedTuple, P} <: AbstractVector{T}
+    ptrs::P
+    length::Int
+    
+    function ExtractAtomMultiple{T}(lmp) where T
+        nall = extract_setting(lmp, "nall")
+        ptrs = map(fieldnames(T)) do name
+            Base.@constprop :aggressive
+            type = fieldtype(T, name)
+            void_ptr = API.lammps_extract_atom(lmp, name)
+            void_ptr == C_NULL && throw(KeyError(name))
+            API.lammps_extract_atom_size(lmp, name, API.LMP_SIZE_ROWS)
+            _check_atom_datatype(lmp, name, type)
+            _check_atom_ghost(lmp, name, nall)
+            if name === :mass
+                type_ptr::Ptr{Int32} = API.lammps_extract_atom(lmp, :type)
+                (type_ptr, Ptr{type}(void_ptr))
+            elseif type<:Real
+                Ptr{type}(void_ptr)
+            elseif nall == 0
+                Ptr{type}(C_NULL)
+            else
+                unsafe_load(Ptr{Ptr{type}}(void_ptr))
+            end
+        end
+        return new{T, typeof(ptrs)}(ptrs, nall)
+    end
+end
+
+Base.size(self::ExtractAtomMultiple) = (self.length,)
+function Base.getindex(self::ExtractAtomMultiple{T}, i::Integer) where T
+    @boundscheck checkbounds(self, i)
+    map(self.ptrs) do ptr
+        if ptr isa Tuple
+            type = unsafe_load(ptr[1])
+            return unsafe_load(ptr[2], type+1)
+        end
+        unsafe_load(ptr, i)
+    end |> T
+end
+
+function _check_global_datatype(lmp, name, T)
+    lammps_type = API.lammps_extract_global_datatype(lmp, name)
+
+    if name in (:sublo, :subhi, :sublo_lambda, :subhi_lambda)
+        (T !== NTuple{3, Float64} && T !== SVector{3, Float64}) && throw(ArgumentError("Expected type NTuple{3, Float64} or SVector{3, Float64} for global property :$name, got $T instead"))
+    elseif name == :procgrid
+        (T !== NTuple{3, Int32} && T !== SVector{3, Int32}) && throw(ArgumentError("Expected type NTuple{3, Int32} or SVector{3, Int32} for global property :$name, got $T instead"))
+    elseif name in (:special_lj, :special_coul)
+        (T !== NTuple{4, Float64} && T !== SVector{4, Float64}) && throw(ArgumentError("Expected type NTuple{4, Float64} or SVector{4, Float64} for global property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_INT64
+        T !== Int64 && throw(ArgumentError("Expected type Int64 for global property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_INT
+        T !== Int32 && throw(ArgumentError("Expected type Int32 for global property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_DOUBLE
+        T !== Float64 && throw(ArgumentError("Expected type Float64 for global property :$name, got $T instead"))
+    elseif lammps_type == API.LAMMPS_STRING
+        T !== String && throw(ArgumentError("Expected type String for global property :$name, got $T instead"))
+    else
+        throw(KeyError(name))
+    end
+end
+
+struct ExtractGlobalMultiple{T<:NamedTuple, P}
+    ptrs::P
+
+    function ExtractGlobalMultiple{T}(lmp) where T
+        ptrs = map(fieldnames(T)) do name
+            Base.@constprop :aggressive
+            type = fieldtype(T, name)
+            void_ptr = API.lammps_extract_global(lmp, name)
+            void_ptr == C_NULL && throw(KeyError(name))
+            _check_global_datatype(lmp, name, type)
+            if type === String
+                Ptr{Cchar}(void_ptr)
+            else
+                Ptr{type}(void_ptr)
+            end
+        end
+        return new{T, typeof(ptrs)}(ptrs)
+    end
+end
+
+function Base.getindex(self::ExtractGlobalMultiple{T}) where T
+    map(self.ptrs) do ptr
+        ptr isa Ptr{Cchar} ? unsafe_string(ptr) : unsafe_load(ptr)
+    end |> T
+end
+
+"""
+    InteractionConfig(;
+        system::Type{<:NamedTuple} = @NamedTuple{},
+        atom::Type{<:NamedTuple} = @NamedTuple{type::Int32},
+        backend::Union{Nothing, AbstractADType} = nothing
+    )
+
+A configuration struct that encapsulates the system properties, atom properties, and an optional differentiation backend for use in external pair interactions.
+
+# Fields
+- `system`: A `NamedTuple` type that defines the global system properties to be extracted from LAMMPS. Defaults to an empty `NamedTuple`.
+    A list of the available global properties can be found in the [LAMMPS documentation](https://docs.lammps.org/Library_properties.html#_CPPv421lammps_extract_globalPvPKc)
+- `atom`: A `NamedTuple` type that defines the per-atom properties to be extracted from LAMMPS. Defaults to `@NamedTuple{type::Int32}`.
+    A list of the available atom properties can be found in the [LAMMPS documentation](https://docs.lammps.org/Classes_atom.html#_CPPv4N9LAMMPS_NS4Atom7extractEPKc)
+- `backend`: An optional automatic differentiation backend used to automatically calculate interaction forces from a potential function. Defaults to `nothing`.
+    The backends are defined through the [`ADTypes` package](https://docs.sciml.ai/ADTypes/stable/).
+
+# Example
+```julia
+config = InteractionConfig(
+    system = @NamedTuple{qqrd2e::Float64, boltz::Float64},
+    atom = @NamedTuple{type::Int32, q::Float64},
+    backend = nothing
+)
+```
+"""
+struct InteractionConfig{T<:NamedTuple, U<:NamedTuple, B<:Union{Nothing, AbstractADType}}
+    backend::B
+
+    function InteractionConfig(;
+            system::Type{T} = @NamedTuple{},
+            atom::Type{U} = @NamedTuple{type::Int32},
+            backend::B = nothing,
+        ) where {T<:NamedTuple, U<:NamedTuple, B<:Union{Nothing, AbstractADType}}
+        return new{T, U, B}(backend)
+    end
+end
+
+_dott(v) = SA[v.x*v.x, v.y*v.y, v.z*v.z, v.x*v.y, v.x*v.z, v.y*v.z]
+
+"""
+    PairExternal(compute_potential, lmp::LMP, config::InteractionConfig, cutoff::Float64)
+
+Defines a custom pair style in LAMMPS using a user-provided potential function.
+
+!!! info
+    The following commands are executed in LAMMPS during the setup process of `PairExternal`:
+    ```lammps
+    fix pair_julia all external pf/callback 1 1
+    pair_style zero <cutoff> nocoeff
+    pair_coeff * *
+    ```
+
+`compute_potential` should have the following signature:
+```julia
+function compute_potential(r::Real, system::config.system, iatom::config.atom, jatom::config.atom)
+```
+If no differentiation backend is provided through `config`, the function should return a tuple of the pair energy and force magnitude.
+Otherwise it should return only the energy.
+
+## Examples
+```julia
+# without automatic differentiation
+config = InteractionConfig(
+    system = @NamedTuple{qqrd2e::Float64},
+    atom = @NamedTuple{q::Float64},
+)
+
+PairExternal(lmp, config, 2.5) do r, system, iatom, jatom
+    energy = system.qqrd2e * (iatom.q * jatom.q) / r
+    force = system.qqrd2e * (iatom.q * jatom.q) / r^2
+    return energy, force
+end
+
+# with automatic differentiation
+import ADTypes: AutoEnzyme()
+import Enzyme
+
+config = InteractionConfig(
+    system = @NamedTuple{qqrd2e::Float64},
+    atom = @NamedTuple{q::Float64},
+    backend = AutoEnzyme(),
+)
+
+PairExternal(lmp, config, 2.5) do r, system, iatom, jatom
+    system.qqrd2e * (iatom.q * jatom.q) / r
+end
+```
+
+"""
+function PairExternal(compute_potential::F, lmp::LMP, config::InteractionConfig{T, U}, cutoff::Float64) where {F, T, U}
+    command(lmp, """
+        pair_style zero $cutoff nocoeff
+        pair_coeff * *
+    """)
+
+    system_ptrs = ExtractGlobalMultiple{T}(lmp) # persistent in memory
+
+    FixExternal(lmp, "pair_julia", "all", 1, 1) do fix::FixExternal
+        system = system_ptrs[]
+        atom = ExtractAtomMultiple{U}(lmp)
+
+        @no_escape begin
+            x = reinterpret(reshape, SVector{3, Float64}, fix.x)
+            e = @alloc(Float64, fix.nlocal)
+            f = reinterpret(reshape, SVector{3, Float64}, fix.f)
+            v = @alloc(SVector{6, Float64}, fix.nlocal)
+
+            fill!(e, zero(eltype(e)))
+            fill!(f, zero(eltype(f)))
+            fill!(v, zero(eltype(v)))
+
+            @inbounds for (i, neigh) in pair_neighborlist(fix.lmp, "zero")
+                iatom = atom[i]
+                ipos = x[i]
+
+                for j in neigh
+                    jatom = atom[j]
+                    diff = ipos - x[j]
+                    r = norm(diff)
+                    r > cutoff && continue
+
+                    if config.backend === nothing
+                        energy, force_magnitude = compute_potential(r, system, iatom, jatom)
+                    else
+                        energy, derivative = value_and_derivative(compute_potential, config.backend, r, Constant(system), Constant(iatom), Constant(jatom))
+                        force_magnitude = -derivative
+                    end
+
+                    force = diff * (force_magnitude / r)
+                    virial = _dott(diff) * (force_magnitude / r)
+
+                    e[i] += 0.5 * energy
+                    f[i] += force
+                    v[i] += 0.5 * virial
+
+                    if j <= fix.nlocal
+                        e[j] += 0.5 * energy
+                        f[j] -= force
+                        v[j] += 0.5 * virial
+                    end
+                end
+            end
+
+            set_energy!(fix, e)
+            set_virial!(fix, v)
+        end
+    end
+end
\ No newline at end of file
diff --git a/test/Project.toml b/test/Project.toml
index cf4df89..5f90a3c 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,4 +1,7 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
diff --git a/test/external_pair.jl b/test/external_pair.jl
new file mode 100644
index 0000000..f7bdea0
--- /dev/null
+++ b/test/external_pair.jl
@@ -0,0 +1,160 @@
+using LAMMPS
+using Test
+using ADTypes
+import ForwardDiff
+import Enzyme
+
+const coefficients = Base.ImmutableDict(
+    1 => Base.ImmutableDict(
+        1 => (1.0, 1.0)
+    )
+)
+
+@testset "external_pair_lj" begin
+    for units in ("lj", "si"),  backend in (nothing, AutoForwardDiff(), AutoEnzyme())
+        lmp_native = LMP(["-screen", "none"])
+        lmp_julia = LMP(["-screen", "none"])
+
+        for lmp in (lmp_native, lmp_julia)
+            command(lmp, """
+                units $units
+                atom_style atomic
+                box tilt large
+                atom_modify map array sort 0 0
+
+                boundary p p p
+                region cell block 0 10.0 0 10.0 0 10.0 units box
+                create_box 1 cell
+            """)
+        end
+
+        cutoff = 2.5
+        command(lmp_native, "pair_style lj/cut $cutoff")
+        command(lmp_native, "pair_coeff * * 1 1")
+
+        # Register external fix
+        config = InteractionConfig(
+            atom = @NamedTuple{type::Int32},
+            backend = backend,
+        )
+
+        PairExternal(lmp_julia, config, cutoff) do r, system, iatom, jatom
+            ε, σ = coefficients[iatom.type][jatom.type]
+            r6inv  = (σ/r)^6
+            energy = 4ε * (r6inv * (r6inv - 1))
+            force = 24ε * (r6inv * (2r6inv - 1)) / r
+            return backend === nothing ? (energy, force) : energy
+        end
+
+        # Setup atoms
+        natoms = 10
+        positions = rand(3, 10) .* 5
+        for lmp in (lmp_native, lmp_julia)
+            command(lmp, """
+                create_atoms 1 random $natoms 1 NULL
+                mass 1 1.0
+                compute potential all pe/atom
+                compute virial all stress/atom NULL
+                compute virial_tot all pressure thermo_temp
+            """)
+
+            scatter!(lmp, "x", positions)
+
+            command(lmp, "run 0")
+        end
+
+        # extract forces
+        potential_native = gather(lmp_native, "c_potential", Float64)
+        potential_julia = gather(lmp_julia, "c_potential", Float64)
+        forces_native = gather(lmp_native, "f", Float64)
+        forces_julia = gather(lmp_julia, "f", Float64)
+        virial_native = gather(lmp_native, "c_virial", Float64)
+        virial_julia = gather(lmp_julia, "c_virial", Float64)
+        virial_tot_native = extract_compute(lmp_native, "virial_tot", STYLE_GLOBAL, TYPE_VECTOR)
+        virial_tot_julia = extract_compute(lmp_julia, "virial_tot", STYLE_GLOBAL, TYPE_VECTOR)
+        energy_native = LAMMPS.API.lammps_get_thermo(lmp_native, "etotal")
+        energy_julia = LAMMPS.API.lammps_get_thermo(lmp_julia, "etotal")
+
+        @testset "$units $backend" begin
+            @test potential_native ≈ potential_julia
+            @test forces_native ≈ forces_julia
+            @test virial_native ≈ virial_julia
+            @test virial_tot_native ≈ virial_tot_julia
+            @test energy_native ≈ energy_julia
+        end
+    end
+end
+
+@testset "external_pair_coul" begin
+    for units in ("lj", "si"),  backend in (nothing, AutoForwardDiff(), AutoEnzyme())
+        lmp_native = LMP(["-screen", "none"])
+        lmp_julia = LMP(["-screen", "none"])
+
+        for lmp in (lmp_native, lmp_julia)
+            command(lmp, """
+                units $units
+                atom_style charge
+                box tilt large
+                atom_modify map array sort 0 0
+
+                boundary p p p
+                region cell block 0 10.0 0 10.0 0 10.0 units box
+                create_box 1 cell
+            """)
+        end
+
+        cutoff = 2.5
+        command(lmp_native, "pair_style coul/cut $cutoff")
+        command(lmp_native, "pair_coeff * *")
+
+        config = InteractionConfig(
+            system = @NamedTuple{qqrd2e::Float64},
+            atom = @NamedTuple{q::Float64},
+            backend = backend,
+        )
+
+        PairExternal(lmp_julia, config, cutoff) do r, system, iatom, jatom
+            energy = system.qqrd2e * (iatom.q * jatom.q) / r
+            force = system.qqrd2e * (iatom.q * jatom.q) / r^2
+            return backend === nothing ? (energy, force) : energy
+        end
+
+        # Setup atoms
+        natoms = 10
+        positions = rand(3, 10) .* 5
+        for lmp in (lmp_native, lmp_julia)
+            command(lmp, """
+                create_atoms 1 random $natoms 1 NULL
+                set type 1 charge 2.0
+                mass 1 1.0
+                compute potential all pe/atom
+                compute virial all stress/atom NULL
+                compute virial_tot all pressure thermo_temp
+            """)
+
+            scatter!(lmp, "x", positions)
+
+            command(lmp, "run 0")
+        end
+
+        # extract forces
+        potential_native = gather(lmp_native, "c_potential", Float64)
+        potential_julia = gather(lmp_julia, "c_potential", Float64)
+        forces_native = gather(lmp_native, "f", Float64)
+        forces_julia = gather(lmp_julia, "f", Float64)
+        virial_native = gather(lmp_native, "c_virial", Float64)
+        virial_julia = gather(lmp_julia, "c_virial", Float64)
+        virial_tot_native = extract_compute(lmp_native, "virial_tot", STYLE_GLOBAL, TYPE_VECTOR)
+        virial_tot_julia = extract_compute(lmp_julia, "virial_tot", STYLE_GLOBAL, TYPE_VECTOR)
+        energy_native = LAMMPS.API.lammps_get_thermo(lmp_native, "etotal")
+        energy_julia = LAMMPS.API.lammps_get_thermo(lmp_julia, "etotal")
+
+        @testset "$units $backend" begin
+            @test potential_native ≈ potential_julia
+            @test forces_native ≈ forces_julia
+            @test virial_native ≈ virial_julia
+            @test virial_tot_native ≈ virial_tot_julia
+            @test energy_native ≈ energy_julia
+        end
+    end
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index dc347eb..7c062a9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,6 @@
 using Test
 using LAMMPS
-using MPI
+using MPI; MPI.Init()
 using UnsafeArrays
 
 @test_logs (:warn,"LAMMPS library path changed, you will need to restart Julia for the change to take effect") LAMMPS.set_library!(LAMMPS.locate())
@@ -71,6 +71,16 @@ end
     end
 end
 
+
+function f()
+    lmp = LMP(["-screen", "none"])
+    @test LAMMPS.version(lmp) >= 0
+    command(lmp, "clear")
+    @test_throws ErrorException command(lmp, "nonsense")
+    LAMMPS.close!(lmp)
+end
+
+
 @testset "Variables" begin
     LMP(["-screen", "none"]) do lmp
         command(lmp, """
@@ -436,6 +446,21 @@ end
     end
 end
 
+LMP(["-screen", "none"]) do lmp
+    called = Ref(false)
+    command(lmp, "boundary p p p")
+    command(lmp, "region cell block 0 1 0 1 0 1 units box")
+    command(lmp, "create_box 1 cell")
+    LAMMPS.FixExternal(lmp, "julia", "all", 1, 1) do fix
+        called[] = true
+    end
+    command(lmp, "mass 1 1.0")
+    command(lmp, "run 0")
+    @test called[] == true
+end
+
+include("external_pair.jl")
+
 if !Sys.iswindows()
     @testset "MPI" begin
          @test success(pipeline(`$(MPI.mpiexec()) -n 2 $(Base.julia_cmd()) mpitest.jl`, stderr=stderr, stdout=stdout))


Reviews:
user: login: vchuravy
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-05T07:21:46Z
user: login: Joroks
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-05T11:00:21Z
user: login: vchuravy
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-05T11:04:26Z
user: login: Joroks
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-05T11:09:42Z
user: login: vchuravy
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-08T15:10:20Z
user: login: Joroks
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T16:09:20Z
user: login: Joroks
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-08T16:14:58Z

Review comments:
diff_hunk: @@ -163,19 +167,21 @@ A full ist of command-line options can be found in the [lammps documentation](ht
 """
 mutable struct LMP
     @atomic handle::Ptr{Cvoid}
+    external_fixes::Dict{String, Any}
 
     function LMP(args::Vector{String}=String[], comm::Union{Nothing, MPI.Comm}=nothing)
         args = copy(args)
         pushfirst!(args, "lammps")
 
+        if !MPI.Initialized()
+            MPI.Init()
+        end
path: src/LAMMPS.jl
user: login: vchuravy
body: Why make that unconditional? 
created_at: 2025-05-05T07:02:00Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,152 @@
+function fix_external_callback end
+
+mutable struct FixExternal
+    const lmp::LMP
+    const name::String
+    const callback
+    timestep::Int
+    nlocal::Int
+    ids::Union{Nothing, Vector{Cint}}
+    x::Union{Nothing, Matrix{Float64}}
+    f::Union{Nothing, Matrix{Float64}}
path: src/external.jl
user: login: vchuravy
body: Why store the fields here? As far as I can tell the state is only used during the execution of the callback?
created_at: 2025-05-05T07:06:25Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,152 @@
+function fix_external_callback end
+
+mutable struct FixExternal
+    const lmp::LMP
+    const name::String
+    const callback
+    timestep::Int
+    nlocal::Int
+    ids::Union{Nothing, Vector{Cint}}
+    x::Union{Nothing, Matrix{Float64}}
+    f::Union{Nothing, Matrix{Float64}}
+
+    function FixExternal(callback, lmp::LMP, name::String, group::String, ncall::Int, napply::Int)
+        command(lmp, "fix $name $group external pf/callback $ncall $napply")
+        this = new(lmp, name, callback, 0, 0, nothing, nothing, nothing)
+        lmp.external_fixes[name] = this # preserves pair globally
+
+        ctx = Base.pointer_from_objref(this)
+        callback = @cfunction(fix_external_callback, Cvoid, (Ptr{Cvoid}, Int64, Cint, Ptr{Cint}, Ptr{Ptr{Float64}}, Ptr{Ptr{Float64}}))
+        API.lammps_set_fix_external_callback(lmp, name, callback, ctx)
+
+        # Ensure function is compiled before timestep 0
+        if !precompile(this.callback, (FixExternal,))
+            @warn "Failed to precompile the callback" this.callback
+        end
+        return this
+    end
+end
+
+function fix_external_callback(ctx::Ptr{Cvoid}, timestep::Int64, nlocal::Cint, ids::Ptr{Cint}, x::Ptr{Ptr{Float64}}, fexternal::Ptr{Ptr{Float64}})
+    fix = Base.unsafe_pointer_to_objref(ctx)::FixExternal
+    fix.timestep = timestep
+    fix.nlocal = nlocal
+
+    nall = extract_setting(fix.lmp, "nall")
+    fix.x = _extract(x, (Int32(3), nall))
+    fix.f = _extract(fexternal, (Int32(3), nall))
+    fix.ids = _extract(ids, nall)
+
+    # necessary dynamic
+    fix.callback(fix)
+
+    fix.timestep = 0
+    fix.nlocal = 0
+    fix.x = nothing
+    fix.f = nothing
+    fix.ids = nothing
+
+    return nothing
+end
+
+function set_virial_peratom(fix::FixExternal, virial::AbstractVector{SVector{6, Float64}}; set_global=false)
+    @assert length(virial) >= fix.nlocal
+
+    @no_escape begin
+        ptrs = @alloc(Ptr{Ptr{Float64}}, fix.nlocal)
+        for i in eachindex(ptrs)
+            @inbounds ptrs[i] = pointer(virial, i)
+        end
+        API.lammps_fix_external_set_virial_peratom(fix.lmp, fix.name, ptrs)
path: src/external.jl
user: login: vchuravy
body: **sigh** that's an unfortunate API design. 

I would much prefer if we could access `vatom` directly.  
https://github.com/lammps/lammps/blob/develop/src/fix_external.cpp#L257C7-L257C12

In this context I am okay with Bumper, and UnsafeArray since it is in a self-contained function and the details of how this works is not user facing.
created_at: 2025-05-05T07:12:02Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,152 @@
+function fix_external_callback end
+
+mutable struct FixExternal
+    const lmp::LMP
+    const name::String
+    const callback
+    timestep::Int
+    nlocal::Int
+    ids::Union{Nothing, Vector{Cint}}
+    x::Union{Nothing, Matrix{Float64}}
+    f::Union{Nothing, Matrix{Float64}}
+
+    function FixExternal(callback, lmp::LMP, name::String, group::String, ncall::Int, napply::Int)
+        command(lmp, "fix $name $group external pf/callback $ncall $napply")
+        this = new(lmp, name, callback, 0, 0, nothing, nothing, nothing)
+        lmp.external_fixes[name] = this # preserves pair globally
+
+        ctx = Base.pointer_from_objref(this)
+        callback = @cfunction(fix_external_callback, Cvoid, (Ptr{Cvoid}, Int64, Cint, Ptr{Cint}, Ptr{Ptr{Float64}}, Ptr{Ptr{Float64}}))
+        API.lammps_set_fix_external_callback(lmp, name, callback, ctx)
+
+        # Ensure function is compiled before timestep 0
+        if !precompile(this.callback, (FixExternal,))
+            @warn "Failed to precompile the callback" this.callback
+        end
+        return this
+    end
+end
+
+function fix_external_callback(ctx::Ptr{Cvoid}, timestep::Int64, nlocal::Cint, ids::Ptr{Cint}, x::Ptr{Ptr{Float64}}, fexternal::Ptr{Ptr{Float64}})
+    fix = Base.unsafe_pointer_to_objref(ctx)::FixExternal
+    fix.timestep = timestep
+    fix.nlocal = nlocal
+
+    nall = extract_setting(fix.lmp, "nall")
+    fix.x = _extract(x, (Int32(3), nall))
+    fix.f = _extract(fexternal, (Int32(3), nall))
+    fix.ids = _extract(ids, nall)
+
+    # necessary dynamic
+    fix.callback(fix)
+
+    fix.timestep = 0
+    fix.nlocal = 0
+    fix.x = nothing
+    fix.f = nothing
+    fix.ids = nothing
+
+    return nothing
+end
+
+function set_virial_peratom(fix::FixExternal, virial::AbstractVector{SVector{6, Float64}}; set_global=false)
+    @assert length(virial) >= fix.nlocal
+
+    @no_escape begin
+        ptrs = @alloc(Ptr{Ptr{Float64}}, fix.nlocal)
+        for i in eachindex(ptrs)
+            @inbounds ptrs[i] = pointer(virial, i)
+        end
+        API.lammps_fix_external_set_virial_peratom(fix.lmp, fix.name, ptrs)
+    end
+
+    if set_global
+        virial_global = MVector(sum(virial))
+        comm = API.lammps_get_mpi_comm(fix.lmp)
path: src/external.jl
user: login: vchuravy
body: FFS...

I think on OpenMPI & co you will need `MPI.API.MPI_Comm_f2c(comm)` since we use the C API
created_at: 2025-05-05T07:17:07Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,152 @@
+function fix_external_callback end
+
+mutable struct FixExternal
+    const lmp::LMP
+    const name::String
+    const callback
+    timestep::Int
+    nlocal::Int
+    ids::Union{Nothing, Vector{Cint}}
+    x::Union{Nothing, Matrix{Float64}}
+    f::Union{Nothing, Matrix{Float64}}
+
+    function FixExternal(callback, lmp::LMP, name::String, group::String, ncall::Int, napply::Int)
+        command(lmp, "fix $name $group external pf/callback $ncall $napply")
+        this = new(lmp, name, callback, 0, 0, nothing, nothing, nothing)
+        lmp.external_fixes[name] = this # preserves pair globally
+
+        ctx = Base.pointer_from_objref(this)
+        callback = @cfunction(fix_external_callback, Cvoid, (Ptr{Cvoid}, Int64, Cint, Ptr{Cint}, Ptr{Ptr{Float64}}, Ptr{Ptr{Float64}}))
+        API.lammps_set_fix_external_callback(lmp, name, callback, ctx)
+
+        # Ensure function is compiled before timestep 0
+        if !precompile(this.callback, (FixExternal,))
+            @warn "Failed to precompile the callback" this.callback
+        end
+        return this
+    end
+end
+
+function fix_external_callback(ctx::Ptr{Cvoid}, timestep::Int64, nlocal::Cint, ids::Ptr{Cint}, x::Ptr{Ptr{Float64}}, fexternal::Ptr{Ptr{Float64}})
+    fix = Base.unsafe_pointer_to_objref(ctx)::FixExternal
+    fix.timestep = timestep
+    fix.nlocal = nlocal
+
+    nall = extract_setting(fix.lmp, "nall")
+    fix.x = _extract(x, (Int32(3), nall))
+    fix.f = _extract(fexternal, (Int32(3), nall))
+    fix.ids = _extract(ids, nall)
+
+    # necessary dynamic
+    fix.callback(fix)
+
+    fix.timestep = 0
+    fix.nlocal = 0
+    fix.x = nothing
+    fix.f = nothing
+    fix.ids = nothing
+
+    return nothing
+end
+
+function set_virial_peratom(fix::FixExternal, virial::AbstractVector{SVector{6, Float64}}; set_global=false)
+    @assert length(virial) >= fix.nlocal
+
+    @no_escape begin
+        ptrs = @alloc(Ptr{Ptr{Float64}}, fix.nlocal)
+        for i in eachindex(ptrs)
+            @inbounds ptrs[i] = pointer(virial, i)
+        end
+        API.lammps_fix_external_set_virial_peratom(fix.lmp, fix.name, ptrs)
+    end
+
+    if set_global
+        virial_global = MVector(sum(virial))
+        comm = API.lammps_get_mpi_comm(fix.lmp)
+        GC.@preserve virial_global begin
+            if comm != -1
+                buffer = MPI.RBuffer(MPI.IN_PLACE, pointer(virial_global), 6, MPI.Datatype(Float64))
+                @inline MPI.Allreduce!(buffer, +, MPI.Comm(comm))
+            end
+            @inline API.lammps_fix_external_set_virial_global(fix.lmp, fix.name, virial_global)
+        end
+    end
+
+    check(fix.lmp)
+end
+
+function set_virial_peratom(fix::FixExternal, virial::AbstractMatrix{Float64}; set_global=false)
+    virial_vec = reinterpret(reshape, SVector{6, Float64}, virial)
+    set_virial_peratom(fix, virial_vec; set_global)
+end
+
+function set_energy_peratom(fix::FixExternal, energy; set_global=false)
+    @assert length(energy) >= fix.nlocal
+
+    API.lammps_fix_external_set_energy_peratom(fix.lmp, fix.name, energy)
+
+    if set_global
+        energy_global = sum(view(energy, 1:fix.nlocal))
+        comm = API.lammps_get_mpi_comm(fix.lmp)
+        if comm != -1
+            energy_global = MPI.Allreduce(energy_global, +, MPI.Comm(comm))
+        end
+        API.lammps_fix_external_set_energy_global(fix.lmp, fix.name, energy_global)
+    end
+
+    check(fix.lmp)
+end
+
+_dott(v) = SA[v.x*v.x, v.y*v.y, v.z*v.z, v.x*v.y, v.x*v.z, v.y*v.z]
+
+function PairExternal(compute_potential::F, lmp::LMP, name::String, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where F
+    command(lmp, """
+        pair_style zero $cutoff nocoeff full
+        pair_coeff * *
+    """)
+
+    FixExternal(lmp, name, "all", 1, 1) do fix::FixExternal
path: src/external.jl
user: login: vchuravy
body: ah, I see you prefer this as an interface.
created_at: 2025-05-05T07:19:53Z
author_association: MEMBER
diff_hunk: @@ -163,19 +167,21 @@ A full ist of command-line options can be found in the [lammps documentation](ht
 """
 mutable struct LMP
     @atomic handle::Ptr{Cvoid}
+    external_fixes::Dict{String, Any}
 
     function LMP(args::Vector{String}=String[], comm::Union{Nothing, MPI.Comm}=nothing)
         args = copy(args)
         pushfirst!(args, "lammps")
 
+        if !MPI.Initialized()
+            MPI.Init()
+        end
path: src/LAMMPS.jl
user: login: Joroks
body: Lammps will internally initialize MPI even when calling `lammps_open_no_mpi`. It will then just assume the global communicator.
created_at: 2025-05-05T11:00:21Z
author_association: COLLABORATOR
diff_hunk: @@ -163,19 +167,21 @@ A full ist of command-line options can be found in the [lammps documentation](ht
 """
 mutable struct LMP
     @atomic handle::Ptr{Cvoid}
+    external_fixes::Dict{String, Any}
 
     function LMP(args::Vector{String}=String[], comm::Union{Nothing, MPI.Comm}=nothing)
         args = copy(args)
         pushfirst!(args, "lammps")
 
+        if !MPI.Initialized()
+            MPI.Init()
+        end
path: src/LAMMPS.jl
user: login: vchuravy
body: Should we then instead of passing in `comm=nothing` pass in `MPI.COMM_WORLD`?

I generally prefer errors over magical behavior, so the `MPI.Initialized` check should not cause us to `init`.
created_at: 2025-05-05T11:04:26Z
author_association: MEMBER
diff_hunk: @@ -163,19 +167,21 @@ A full ist of command-line options can be found in the [lammps documentation](ht
 """
 mutable struct LMP
     @atomic handle::Ptr{Cvoid}
+    external_fixes::Dict{String, Any}
 
     function LMP(args::Vector{String}=String[], comm::Union{Nothing, MPI.Comm}=nothing)
         args = copy(args)
         pushfirst!(args, "lammps")
 
+        if !MPI.Initialized()
+            MPI.Init()
+        end
path: src/LAMMPS.jl
user: login: Joroks
body: And then allways error if `MPI.Initialized()` is false?
created_at: 2025-05-05T11:09:42Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +220,15 @@ end
 
 _dott(v) = SA[v.x*v.x, v.y*v.y, v.z*v.z, v.x*v.y, v.x*v.z, v.y*v.z]
 
-function PairExternal(compute_potential::F, lmp::LMP, name::String, system_properties::Type{T}, atom_properties::Type{U}, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where {F, T<:NamedTuple, U<:NamedTuple}
+function PairExternal(compute_potential::F, lmp::LMP, system_properties::Type{T}, atom_properties::Type{U}, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where {F, T<:NamedTuple, U<:NamedTuple}
     command(lmp, """
         pair_style zero $cutoff nocoeff
         pair_coeff * *
     """)
 
     system_ptrs = ExtractGlobalMultiple{system_properties}(lmp) # persistent in memory
 
-    FixExternal(lmp, name, "all", 1, 1) do fix::FixExternal
+    FixExternal(lmp, "pair_julia", "all", 1, 1) do fix::FixExternal
path: src/external.jl
user: login: vchuravy
body: Can I then still have multiple pair styles? (If that is a sensible thing to have)
created_at: 2025-05-08T15:10:19Z
author_association: MEMBER
diff_hunk: @@ -220,15 +220,15 @@ end
 
 _dott(v) = SA[v.x*v.x, v.y*v.y, v.z*v.z, v.x*v.y, v.x*v.z, v.y*v.z]
 
-function PairExternal(compute_potential::F, lmp::LMP, name::String, system_properties::Type{T}, atom_properties::Type{U}, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where {F, T<:NamedTuple, U<:NamedTuple}
+function PairExternal(compute_potential::F, lmp::LMP, system_properties::Type{T}, atom_properties::Type{U}, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where {F, T<:NamedTuple, U<:NamedTuple}
     command(lmp, """
         pair_style zero $cutoff nocoeff
         pair_coeff * *
     """)
 
     system_ptrs = ExtractGlobalMultiple{system_properties}(lmp) # persistent in memory
 
-    FixExternal(lmp, name, "all", 1, 1) do fix::FixExternal
+    FixExternal(lmp, "pair_julia", "all", 1, 1) do fix::FixExternal
path: src/external.jl
user: login: Joroks
body: Indirectly, yes. You can allways just sum up the potentials in`compute_potential` for two different pair styles.
Defining two `PairExternals` is a bit tricky, however. In order to accomodate two different cutoff distances, the pair style would have to be  set up as follow:
```
pair_style hybrid zero $cutoff1 zero $cutoff2
```

and then for the respective neigborlist
```julia
neigh1 = pair_neighborlist(lmp, "zero"; nsub=0)
neigh2 = pair_neighborlist(lmp, "zero"; nsub=1)
```
Which wasn't accounted for previously either. I want to keep the interface relatively simple and have therefore decided to have the setup within lammps be handled directly by `FixExternal` and `PairExternal` but this is a limitation of this approach. Even without doing this, it's a bit of a hassle to account for two pair styles.

We encounter a similar issue when trying to overlay lammps and julia pair styles.

created_at: 2025-05-08T16:09:20Z
author_association: COLLABORATOR
diff_hunk: @@ -220,15 +220,15 @@ end
 
 _dott(v) = SA[v.x*v.x, v.y*v.y, v.z*v.z, v.x*v.y, v.x*v.z, v.y*v.z]
 
-function PairExternal(compute_potential::F, lmp::LMP, name::String, system_properties::Type{T}, atom_properties::Type{U}, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where {F, T<:NamedTuple, U<:NamedTuple}
+function PairExternal(compute_potential::F, lmp::LMP, system_properties::Type{T}, atom_properties::Type{U}, cutoff::Float64; backend::Union{Nothing, AbstractADType} = nothing) where {F, T<:NamedTuple, U<:NamedTuple}
     command(lmp, """
         pair_style zero $cutoff nocoeff
         pair_coeff * *
     """)
 
     system_ptrs = ExtractGlobalMultiple{system_properties}(lmp) # persistent in memory
 
-    FixExternal(lmp, name, "all", 1, 1) do fix::FixExternal
+    FixExternal(lmp, "pair_julia", "all", 1, 1) do fix::FixExternal
path: src/external.jl
user: login: Joroks
body: Allways using "pair_julia" is actually more consistent with what lammps does since creating a new `PairExternal` simply overwrites the old pair. Which apparently is also what lammps does when you call `pair_style` multiple times. Otherwise creating two `PairExternals` would quietly change the cutoff radius (at least for neighborlist creation) of the first `PairExternal` to the cutoff radius of the second `PairExternal`
created_at: 2025-05-08T16:14:58Z
author_association: COLLABORATOR

PR comments:
user: login: vchuravy
created_at: 2025-05-03T16:32:35Z
body: Thanks! This is the first time I encounter Tensors.jl

Regarding forward diff, I am wondering if we could use ADtypes and DifferentiationInterface to allow the user to use Enzyme or ForwardDiff? 

Would be good if one could also provide an explicit implementation (for those chasing peak performance)
author_association: MEMBER
user: login: Joroks
created_at: 2025-05-08T22:01:30Z
body: @vchuravy for some reason, my recent commits haven't triggered the automatic tests. I can't find a way to start them manually either. Do you know what's causing this?
author_association: COLLABORATOR
user: login: vchuravy
created_at: 2025-05-09T14:16:12Z
body: Very nice! Thank you 
author_association: MEMBER