Fix ShiftedLattice & ShiftedCell & shift

Project.toml

@@ -9,6 +9,7 @@ CrystallographyCore = "80545937-1184-4bc9-b283-396e91386b5c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+StructEquality = "6ec83bb0-ed9f-11e9-3b4c-2b04cb4e219c"
 
 [weakdeps]
 Brillouin = "23470ee3-d0df-4052-8b1a-8cbd6363e7f0"

src/CrystallographyBase.jl

@@ -16,12 +16,12 @@ include("systems.jl")
 include("lattice.jl")
 include("cell.jl")
 include("reciprocal.jl")
-include("geometry.jl")
 include("metric.jl")
 include("volume.jl")
 include("transform.jl")
 include("eachatomgroup.jl")
 include("shifted.jl")
+include("geometry.jl")
 include("show.jl")
 
 end

src/cell.jl

@@ -32,14 +32,3 @@ function super(cell::Cell, factors::AbstractMatrix{<:Integer})
     new_lattice = super(cell.lattice, factors)
     return Cell(new_lattice, new_positions, new_atoms)
 end
-
-function shift(cell::Cell, 𝐱::AbstractVector)
-    new_lattice = shift(Lattice(cell), 𝐱)
-    new_positions = Ref(ShiftedLattice(Lattice(cell), 𝐱)) .* cell.positions
-    new_positions = Ref(new_lattice) .\ new_positions
-    return Cell(new_lattice, new_positions, cell.atoms)
-end
-function shift(cell::Cell, x::Integer, y::Integer, z::Integer)
-    𝐚, 𝐛, 𝐜 = basisvectors(Lattice(cell))
-    return shift(cell, x * 𝐚 + y * 𝐛 + z * 𝐜)
-end

src/eachatomgroup.jl

@@ -1,21 +1,21 @@
-using CrystallographyCore: EachAtom
+using CrystallographyCore: AbstractCell, EachAtom
 
 import CrystallographyCore: eachatom
 
 export eachatomgroup
 
-struct AtomGroup{A,B}
+struct AtomGroup{N,A,B}
     atom::A
-    positions::Vector{B}
+    positions::NTuple{N,B}
 end
 
 eachatom(group::AtomGroup) =
-    EachAtom(fill(group.atom, length(group.positions)), group.positions)
+    EachAtom(ntuple(Returns(group.atom), length(group.positions)), group.positions)
 
-function eachatomgroup(cell::Cell)
+function eachatomgroup(cell::AbstractCell)
     types = atomtypes(cell)
     return Iterators.map(types) do type
         indices = findall(==(type), cell.atoms)
-        AtomGroup(type, cell.positions[indices])
+        AtomGroup(type, Tuple(cell.positions[indices]))
     end
 end

src/geometry.jl

@@ -4,16 +4,23 @@ const FACES = Base.vect(
     (1, 2, 3, 4), (5, 6, 7, 8), (1, 2, 6, 5), (3, 4, 8, 7), (2, 3, 7, 6), (5, 8, 4, 1)
 )
 
-function vertices(lattice::Lattice, O⃗=zeros(eltype(lattice), 3))
+function vertices(lattice::Lattice)
+    O⃗ = zeros(eltype(lattice), 3)
     A⃗, B⃗, C⃗ = basisvectors(lattice)
-    A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗ = A⃗ + B⃗ - O⃗, A⃗ + C⃗ - O⃗, B⃗ + C⃗ - O⃗, A⃗ + B⃗ + C⃗ - 2O⃗
-    return (O⃗, A⃗, B⃗, C⃗, A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗)
+    A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗ = A⃗ + B⃗, A⃗ + C⃗, B⃗ + C⃗, A⃗ + B⃗ + C⃗
+    return O⃗, A⃗, B⃗, C⃗, A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗
+end
+function vertices(lattice::ShiftedLattice)
+    O⃗ = lattice.by
+    A⃗, B⃗, C⃗ = basisvectors(lattice)
+    A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗ = A⃗ + B⃗, A⃗ + C⃗, B⃗ + C⃗, A⃗ + B⃗ + C⃗
+    return (O⃗, (A⃗, B⃗, C⃗, A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗) .+ Ref(O⃗)...)
 end
 
 edge(A⃗, B⃗) = hcat(([Aᵢ, Bᵢ] for (Aᵢ, Bᵢ) in zip(A⃗, B⃗))...)
 
-function edges(lattice::Lattice, O⃗=zeros(eltype(lattice), 3))
-    O⃗, A⃗, B⃗, C⃗, A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗ = vertices(lattice, O⃗)
+function edges(lattice::AbstractLattice)
+    O⃗, A⃗, B⃗, C⃗, A⃗B⃗, A⃗C⃗, B⃗C⃗, A⃗B⃗C⃗ = vertices(lattice)
     return (
         edge(O⃗, A⃗),
         edge(O⃗, C⃗),
@@ -30,7 +37,7 @@ function edges(lattice::Lattice, O⃗=zeros(eltype(lattice), 3))
     )
 end
 
-function faces(lattice::Lattice, O⃗=zeros(eltype(lattice), 3))
-    verts = vertices(lattice, O⃗)
-    return map(face -> [verts[i] for i in face], FACES)
+function faces(lattice::AbstractLattice)
+    verts = vertices(lattice)
+    return Tuple(Tuple(verts[i] for i in face) for face in FACES)
 end

src/lattice.jl

@@ -203,9 +203,3 @@ super(lattice_or_cell, factors::AbstractVector{<:Integer}) =
     super(lattice_or_cell, Diagonal(factors))
 # See https://stackoverflow.com/a/57270841
 super(lattice_or_cell, factor::Integer) = super(lattice_or_cell, fill(factor, 3))
-
-shift(lattice::Lattice, 𝐱::AbstractVector) = Lattice(lattice .+ 𝐱)
-function shift(lattice::Lattice, x::Integer, y::Integer, z::Integer)
-    𝐚, 𝐛, 𝐜 = basisvectors(lattice)
-    return shift(lattice, x * 𝐚 + y * 𝐛 + z * 𝐜)
-end

src/shifted.jl

@@ -1,26 +1,70 @@
-import CrystallographyCore: basisvectors
+using CrystallographyCore: AbstractCell, Inverted
+using StructEquality: @struct_hash_equal_isequal, @struct_hash_equal_isequal_isapprox
 
-export ShiftedLattice
+import CrystallographyCore: basisvectors, natoms, atomtypes
 
-struct ShiftedLattice{T,S} <: AbstractLattice{T}
+export ShiftedLattice, ShiftedCell, shift
+
+@struct_hash_equal_isequal_isapprox struct ShiftedLattice{T} <: AbstractLattice{T}
     original::Lattice{T}
-    by::SVector{3,S}
+    by::SVector{3,T}
+end
+
+basisvectors(shifted::ShiftedLattice) = basisvectors(parent(shifted))
+
+function shift(lattice::Lattice, 𝐱::AbstractVector)
+    T = Base.promote_eltype(lattice, 𝐱)
+    return ShiftedLattice(convert(Lattice{T}, lattice), SVector{3,T}(𝐱))
+end
+shift(shifted::ShiftedLattice, 𝐱::AbstractVector) = shift(parent(shifted), shifted.by .+ 𝐱)
+function shift(lattice::Lattice, x::Integer, y::Integer, z::Integer)
+    𝐚, 𝐛, 𝐜 = basisvectors(lattice)
+    return shift(lattice, x * 𝐚 + y * 𝐛 + z * 𝐜)
+end
+function shift(shifted::ShiftedLattice, x::Integer, y::Integer, z::Integer)
+    𝐚, 𝐛, 𝐜 = basisvectors(shifted)
+    return shift(shifted, x * 𝐚 + y * 𝐛 + z * 𝐜)
 end
-ShiftedLattice(original::AbstractMatrix, by::AbstractVector) =
-    ShiftedLattice(Lattice(original), SVector{3}(by))
 
-# basisvectors(lattice::ShiftedLattice) = basisvectors(lattice.original) .+ Ref(lattice.by)
+Base.parent(shifted::ShiftedLattice) = shifted.original
 
-Base.parent(lattice::ShiftedLattice) = lattice.original
+# See https://github.com/MineralsCloud/CrystallographyCore.jl/blob/d9b808c/src/transform.jl#L10C80-L10C80
+(shifted::ShiftedLattice)(reduced::AbstractVector) = parent(shifted)(reduced) .+ shifted.by
 
-Base.size(::ShiftedLattice) = (3, 3)
+(inverted::Inverted{<:ShiftedLattice})(cartesian::AbstractVector) =
+    inv(parent(inverted.lattice))(cartesian .- inverted.lattice.by)
 
-Base.getindex(lattice::ShiftedLattice, i::Int) = getindex(parent(lattice), i)
+@struct_hash_equal_isequal mutable struct ShiftedCell{L,P,T} <: AbstractCell
+    original::Cell{L,P,T}
+    by::SVector{3,L}
+end
+
+Base.parent(shifted::ShiftedCell) = shifted.original
+
+natoms(shifted::ShiftedCell) = length(parent(shifted).atoms)
 
-Base.setindex!(lattice::ShiftedLattice, v, i::Int) = setindex!(parent(lattice), v, i)
+atomtypes(shifted::ShiftedCell) = unique(parent(shifted).atoms)
 
-Base.IndexStyle(::Type{<:ShiftedLattice}) = IndexLinear()
+ShiftedLattice(shifted::ShiftedCell) = shift(Lattice(parent(shifted)), shifted.by)
+
+function shift(cell::Cell{L}, 𝐱::AbstractVector{X}) where {L,X}
+    T = Base.promote_type(L, X)
+    return ShiftedCell(
+        Cell(convert(Lattice{T}, Lattice(cell)), cell.positions, cell.atoms),
+        SVector{3,T}(𝐱),
+    )
+end
+function shift(cell::Cell, x::Integer, y::Integer, z::Integer)
+    𝐚, 𝐛, 𝐜 = basisvectors(Lattice(cell))
+    return shift(cell, x * 𝐚 + y * 𝐛 + z * 𝐜)
+end
 
-function Base.:*(lattice::ShiftedLattice, x::AbstractVector)
-    return parent(lattice) * x + lattice.by
+function Base.getproperty(shifted::ShiftedCell, name::Symbol)
+    if name in (:positions, :atoms)
+        return getproperty(parent(shifted), name)
+    elseif name == :lattice
+        return ShiftedLattice(shifted)
+    else
+        return getfield(shifted, name)
+    end
 end