Export PeriodicDistance2 to unify API instead of periodic_distance! & co

docs/make.jl

@@ -2,7 +2,8 @@ using Documenter
 using PeriodicGraphEmbeddings, PeriodicGraphs, Graphs
 
 DocMeta.setdocmeta!(PeriodicGraphEmbeddings, :DocTestSetup, quote
-    using PeriodicGraphEmbeddings, PeriodicGraphs, Graphs
+    using PeriodicGraphEmbeddings, PeriodicGraphs, Graphs, StaticArrays
+    using LinearAlgebra
 end; recursive=true)
 
 makedocs(

docs/src/utilities.md

@@ -2,13 +2,10 @@
 
 ## Periodic distance
 
-The [`periodic_distance`](@ref) function is useful to compute the shortest distance between
-two vertices of a graph of dimension 3 or less,
-*i.e.* for which a [`Cell`](@ref) has been provided.
-It assumes that the unit cell is not too much skewed.
+The [`PeriodicDistance2`](@ref) setup is useful to compute the shortest distance between two vertices of a graph of dimension 3 or less, *i.e.* for which a [`Cell`](@ref) has been provided. It assumes that the unit cell is not too much skewed.
 
 ```@docs
-periodic_distance
+PeriodicDistance2
 ```
 
 ## Other

src/symmetries.jl

@@ -264,15 +264,23 @@ function _find_closest_point_after_symop(pge::PeriodicGraphEmbedding{D,T}, t::SV
     if T <: Rational
         notencountered, buffer = rest
     else
-        notencountered, buffer, mat, ortho, safemin = rest
+        notencountered, pd2 = rest
     end
     i = notencountered isa Int ? 1 : findfirst(notencountered)
     j = notencountered isa Int ? 2 : findnext(notencountered, i+1)
-    buffer .= pge.pos[i] .- x
-    mindist = T <: Rational ? norm(pge.pos[i] .- x) : periodic_distance!(buffer, mat, ortho, safemin)
+    mindist = if T <: Rational
+        buffer .= pge.pos[i] .- x
+        norm(buffer)
+    else
+        pd2(pge.pos[i], x)
+    end
     while notencountered isa Int ? j ≤ notencountered : j isa Int
-        buffer .= pge.pos[j] .- x
-        newdist = T <: Rational ? norm(buffer) : periodic_distance!(buffer, mat, ortho, safemin)
+        newdist = if T <: Rational
+            buffer .= pge.pos[j] .- x
+            norm(buffer)
+        else
+            pd2(pge.pos[j], x)
+        end
         if newdist < mindist
             mindist = newdist
             i = j
@@ -308,9 +316,7 @@ function check_valid_symmetry(pge::PeriodicGraphEmbedding{D,T}, t::SVector{D,T},
         if T <: Rational
             rest = (notencountered, buffer)
         else
-            mat = Float64.(pge.cell.mat)
-            _, ortho, safemin = prepare_periodic_distance_computations(mat)
-            rest = (notencountered, buffer, mat, ortho, safemin)
+            rest = (notencountered, PeriodicDistance2(Float64.(pge.cell.mat)))
         end
     end
     for k in 1:n

src/utils.jl

@@ -1,4 +1,5 @@
-export prepare_periodic_distance_computations, periodic_distance!, periodic_distance
+export PeriodicDistance2
+using LinearAlgebra: mul!
 
 @static if VERSION < v"1.8-"
     # copy-pasted from the implementation of @lazy_str in Base
@@ -17,61 +18,177 @@ export prepare_periodic_distance_computations, periodic_distance!, periodic_dist
     end
 end
 
+function norm2(u::T) where {T}
+    r2 = zero(eltype(T))^2
+    @simd for x in u
+        r2 += x^2
+    end
+    r2
+end
+
+"""
+    (pd2::PeriodicDistance2)(x, y=nothing, ofsx=nothing, ofsy=nothing; fromcartesian=false, ofs=nothing)
+
+Squared periodic distance between points `x` and `y`, optionally with respective offsets
+`ofsx` and `ofsy`.
+The periodic distance is the shortest distance between all the periodic images of the
+inputs. For `x` and `y` given as triplets of fractional coordinates, putting integer
+offsets in `ofsx` and `ofsy` does not have any effect thus.
+
+If `y` is not set, compute the periodic distance between `x` and the origin.
+
+If `fromcartesian` is set, the inputs must be given as cartesian coordinates. Otherwise, it
+is assumed that the input is given in fractional coordinates.
+
+If `ofs` is set to a mutable vector of integers (`MVector{3,Int}` from StaticArrays.jl is
+suggested), then the offset of the translation of `y .+ ofsy` with respect to `x .+ ofsx`
+is stored in `ofs`.
+This means that the returned periodic distance is equal to the non-periodic distance
+between `y .+ ofsy .+ ofs` and `x .+ ofsx` (assuming fractional coordinates).
+The offset is always returned as a triplet of integers, though the input can be given in
+cartesian coordinates as long as `fromcartesian` is set.
+
+!!! warning
+    This function modifies an internal  state of `pd2`, and is thus not thread-safe.
 
+## Example
+```jldoctest
+julia> mat = [26.04 -7.71 -8.32; 0.0 32.72 -3.38; 0.0 0.0 26.66];
 
-function prepare_periodic_distance_computations(mat)
-    (a, b, c), (α, β, γ) = cell_parameters(mat)
+julia> pd2 = PeriodicDistance2(mat)
+PeriodicDistance2([26.04 -7.71 -8.32; 0.0 32.72 -3.38; 0.0 0.0 26.66])
+
+julia> sqrt(pd2([0.5, 0, 0])) # distance to the middle of the a axis is simply a/2
+13.02
+
+julia> ofs = MVector{3,Int}(undef);
+
+julia> vec1 = [0.9, 0.6, 0.5]; vec2 = [0.0, 0.5, 0.01];
+
+julia> d2 = pd2(vec1, vec2; ofs)
+210.57932044000003
+
+julia> println(ofs) # the offset of vec2 to have the closest image to vec1
+[1, 0, 1]
+
+julia> norm(mat*(vec2 .+ ofs .- vec1))^2 ≈ d2
+true
+
+julia> pd2(mat*vec1, mat*vec2; fromcartesian=true) ≈ d2
+true
+```
+"""
+struct PeriodicDistance2{T,Ti,T2}
+    buffer::MVector{3,T}
+    buffer2::MVector{3,Float64}
+    mat::SMatrix{3,3,T,9}
+    invmat::SMatrix{3,3,Ti,9}
+    ortho::Bool
+    safemin2::T2
+end
+Base.show(io::IO, pd2::PeriodicDistance2) = println(io, PeriodicDistance2, '(', pd2.mat, ')')
+
+"""
+    PeriodicDistance2(mat::AbstractMatrix)
+
+Build a [`PeriodicDistance2`](@ref) object which can be called to compute squared periodic
+distances in a unit cell of given matrix `mat`.
+
+!!! warning
+    It is assumed that the unit cell is in standard settings: its angles must be above 60°
+    or the returned distance may not be the shortest.
+"""
+function PeriodicDistance2(mat::AbstractMatrix{T}) where {T}
+    smat = SMatrix{3,3,T,9}(mat)
+    (a, b, c), (α, β, γ) = cell_parameters(smat)
     ortho = all(x -> isapprox(Float16(x), 90; rtol=0.02), (α, β, γ))
-    _a, _b, _c = eachcol(mat)
+    _a, _b, _c = eachcol(smat)
     safemin = min(Float64(dot(cross(_b, _c), _a)/(b*c)),
                   Float64(dot(cross(_c, _a), _b)/(a*c)),
                   Float64(dot(cross(_a, _b), _c)/(a*b)))/2
     # safemin is the half-distance between opposite planes of the unit cell
-    return MVector{3,Float64}(undef), ortho, safemin
+    buffer = MVector{3,T}(undef)
+    buffer2 = MVector{3,Float64}(undef)
+    T1 = oneunit(first(smat))
+    invmat = inv(smat./T1)./T1 # to handle Unitful
+    PeriodicDistance2(buffer, buffer2, smat, invmat, ortho, safemin^2)
 end
 
-function periodic_distance!(buffer, u, mat, ortho, safemin)
-    @simd for i in 1:3
-        diff = u[i] + 0.5
-        buffer[i] = diff - floor(diff) - 0.5
+@inline function _set_if_ofs!(buffer, x, y, ofsx, ofsy)
+    if y isa Nothing
+        if ofsx isa Nothing
+            if ofsy isa Nothing
+                buffer .= x
+            else
+                buffer .= x .- ofsy
+            end
+        else
+            if ofsy isa Nothing
+                buffer .= x .+ ofsx
+            else
+                buffer .= x .+ ofsx .- ofsy
+            end
+        end
+    else
+        if ofsx isa Nothing
+            if ofsy isa Nothing
+                buffer .= x .- y
+            else
+                buffer .= x .- y .- ofsy
+            end
+        else
+            if ofsy isa Nothing
+                buffer .= x .+ ofsx .- y
+            else
+                buffer .= x .+ ofsx .- y .- ofsy
+            end
+        end
     end
-    ref = norm(mat*buffer)
-    (ortho || ref ≤ safemin) && return ref
+end
+
+function _periodic_distance2!(pd2::PeriodicDistance2, ofs)
+    if ofs isa Nothing
+        @simd for i in 1:3
+            diff = pd2.buffer2[i] + 0.5
+            pd2.buffer2[i] = diff - floor(diff) - 0.5
+        end
+    else
+        @simd for i in 1:3
+            diff = pd2.buffer2[i] + 0.5
+            tmp = floor(Int, diff)
+            ofs[i] = tmp
+            pd2.buffer2[i] = diff - tmp - 0.5
+        end
+    end
+    mul!(pd2.buffer, pd2.mat, pd2.buffer2)
+    ref2 = norm2(pd2.buffer)
+    (pd2.ortho || ref2 ≤ pd2.safemin2) && return ref2
     @inbounds for i in 1:3
-        buffer[i] += 1
-        newnorm = norm(mat*buffer)
-        newnorm < ref && return newnorm # in a reduced lattice, there should be at most one
-        buffer[i] -= 2
-        newnorm = norm(mat*buffer)
-        newnorm < ref && return newnorm
-        buffer[i] += 1
+        pd2.buffer2[i] += 1
+        mul!(pd2.buffer, pd2.mat, pd2.buffer2)
+        newnorm2 = norm2(pd2.buffer)
+        if newnorm2 < ref2
+            ofs isa Nothing || (ofs[i] -= 1)
+            return newnorm2 # in a reduced lattice, there should be at most one
+        end
+        pd2.buffer2[i] -= 2
+        mul!(pd2.buffer, pd2.mat, pd2.buffer2)
+        newnorm2 = norm2(pd2.buffer)
+        if newnorm2 < ref2
+            ofs isa Nothing || (ofs[i] += 1)
+            return newnorm2
+        end
+        pd2.buffer2[i] += 1
     end
-    return ref
+    return ref2
 end
-periodic_distance!(u, mat, ortho, safemin) = periodic_distance!(u, u, mat, ortho, safemin)
 
-"""
-    periodic_distance(u, mat, ortho=nothing, safemin=nothing)
-
-Distance between point `u` and the origin, given as a triplet of fractional coordinates, in
-a repeating unit cell of matrix `mat`.
-The distance is the shortest between all equivalents of `u` and the origin.
-If `ortho` is set to `true`, the angles α, β and γ of the cell are assumed right, which
-accelerates the computation by up to 7 times.
-If a distance lower than `safemin` is computed, stop trying to find a periodic image of `u`
-closer to the origin.
-If unspecified, both `ortho` and `safemin` are automatically determined from `mat`.
-
-This implementation assumes that the cell corresponds to a reduced lattice. It may be
-invalid for some edge cases otherwise.
-
-For optimal performance, use `periodic_distance!` with `buffer`, `ortho` and `safemin`
-obtained from `prepare_periodic_distance_computations`.
-"""
-function periodic_distance(u, mat, ortho=nothing, safemin=nothing)
-    if ortho === nothing || safemin === nothing
-        _, ortho, safemin = prepare_periodic_distance_computations(mat)
+function (pd2::PeriodicDistance2)(x, y=nothing, ofsx=nothing, ofsy=nothing; fromcartesian=false, ofs=nothing)
+    if fromcartesian
+        _set_if_ofs!(pd2.buffer, x, y, ofsx, ofsy)
+        mul!(pd2.buffer2, pd2.invmat, pd2.buffer)
+    else
+        _set_if_ofs!(pd2.buffer2, x, y, ofsx, ofsy)
     end
-    periodic_distance!(similar(u), u, mat, ortho::Bool, safemin::Float64)
+    _periodic_distance2!(pd2, ofs)
 end
-