From 41a609631d0cdbcfaefe6437faea02a6daf255cd Mon Sep 17 00:00:00 2001
From: Kipton Barros <kbarros@gmail.com>
Date: Sun, 14 Jul 2024 13:00:11 -0600
Subject: [PATCH] Refactor reshape_supercell_aux

---
 src/Reshaping.jl                     | 18 +++++++-----------
 src/SpinWaveTheory/SpinWaveTheory.jl | 22 ++++++++++++++++++----
 src/Symmetry/Crystal.jl              | 20 +++++++++++++++-----
 3 files changed, 40 insertions(+), 20 deletions(-)

diff --git a/src/Reshaping.jl b/src/Reshaping.jl
index c6701c767..101eac4cf 100644
--- a/src/Reshaping.jl
+++ b/src/Reshaping.jl
@@ -19,10 +19,11 @@ function reshape_supercell(sys::System{N}, shape) where N
     @assert prim_shape′ ≈ prim_shape
 
     # Unit cell for new system, in units of original unit cell.
-    enlarged_latsize = NTuple{3, Int}(gcd.(eachcol(prim_shape′)))
-    reduced_shape = Mat3(shape * diagm(collect(inv.(enlarged_latsize))))
+    new_latsize = NTuple{3, Int}(gcd.(eachcol(prim_shape′)))
+    new_shape = Mat3(shape * diagm(collect(inv.(new_latsize))))
+    new_cryst = reshape_crystal(orig_crystal(sys), new_shape)
 
-    return reshape_supercell_aux(sys, enlarged_latsize, reduced_shape)
+    return reshape_supercell_aux(sys, new_cryst, new_latsize)
 end
 
 
@@ -59,12 +60,9 @@ function transfer_interactions!(sys::System{N}, src::System{N}) where N
 end
 
 
-function reshape_supercell_aux(sys::System{N}, new_latsize::NTuple{3, Int}, new_cell_shape::Mat3) where N
-    # Reshape the crystal
-    new_cryst            = reshape_crystal(orig_crystal(sys), new_cell_shape)
-    new_na               = natoms(new_cryst)
-
+function reshape_supercell_aux(sys::System{N}, new_cryst::Crystal, new_latsize::NTuple{3, Int}) where N
     # Allocate data for new system, but with an empty list of interactions
+    new_na               = natoms(new_cryst)
     new_Ns               = zeros(Int, new_latsize..., new_na)
     new_κs               = zeros(Float64, new_latsize..., new_na)
     new_gs               = zeros(Mat3, new_latsize..., new_na)
@@ -143,7 +141,5 @@ See also [`reshape_supercell`](@ref).
 """
 function repeat_periodically(sys::System{N}, counts::NTuple{3,Int}) where N
     all(>=(1), counts) || error("Require at least one count in each direction.")
-
-    # Scale each column by `counts` and reshape
-    return reshape_supercell_aux(sys, counts .* sys.latsize, cell_shape(sys))
+    return reshape_supercell_aux(sys, sys.crystal, counts .* sys.latsize)
 end
diff --git a/src/SpinWaveTheory/SpinWaveTheory.jl b/src/SpinWaveTheory/SpinWaveTheory.jl
index 60f4f3b5d..49395d94b 100644
--- a/src/SpinWaveTheory/SpinWaveTheory.jl
+++ b/src/SpinWaveTheory/SpinWaveTheory.jl
@@ -29,10 +29,24 @@ struct SpinWaveTheory
 end
 
 function SpinWaveTheory(sys::System{N}; energy_ϵ::Float64=1e-8, observables=nothing, correlations=nothing, apply_g=true) where N
-    # Reshape into single unit cell. A clone will always be performed, even if
-    # no reshaping happens.
-    cellsize_mag = cell_shape(sys) * diagm(Vec3(sys.latsize))
-    sys = reshape_supercell_aux(sys, (1,1,1), cellsize_mag)
+    # Create single chemical cell that matches the full system size.
+    new_shape = cell_shape(sys) * diagm(Vec3(sys.latsize))
+    new_cryst = reshape_crystal(orig_crystal(sys), new_shape)
+
+    # Sort crystal positions so that their order matches sites in sys. Quadratic
+    # scaling in system size.
+    global_positions = global_position.(Ref(sys), vec(eachsite(sys)))
+    p = map(new_cryst.positions) do r
+        pos = new_cryst.latvecs * r
+        findfirst(global_positions) do refpos
+            isapprox(pos, refpos, atol=new_cryst.symprec)
+        end
+    end
+    @assert allunique(p)
+    permute_sites!(new_cryst, p)
+
+    # Create a new system with latsize (1,1,1). A clone happens in all cases.
+    sys = reshape_supercell_aux(sys, new_cryst, (1,1,1))
 
     # Rotate local operators to quantization axis
     if sys.mode == :SUN
diff --git a/src/Symmetry/Crystal.jl b/src/Symmetry/Crystal.jl
index 7ea9bafc4..cc07302c9 100644
--- a/src/Symmetry/Crystal.jl
+++ b/src/Symmetry/Crystal.jl
@@ -152,7 +152,19 @@ function spacegroup_name(hall_number::Int)
 end
 
 
-# Sort the sites according to class and fractional coordinates.
+# Crystal is immutable in its public API. Therefore, this is an internal
+# function should not overload Base.permute!
+function permute_sites!(cryst::Crystal, p)
+    cryst.positions .= cryst.positions[p]
+    cryst.classes .= cryst.classes[p]
+    cryst.types .= cryst.types[p]
+    if !isnothing(cryst.sitesyms)
+        cryst.sitesyms .= cryst.sitesyms[p]
+    end
+end
+
+# Sort the sites according to class and fractional coordinates. This is an
+# internal function.
 function sort_sites!(cryst::Crystal)
     function less_than(i, j)
         ci = cryst.classes[i]
@@ -171,10 +183,8 @@ function sort_sites!(cryst::Crystal)
         error("""Detected two very close atoms ($str1 and $str2).
                  If positions inferred from spacegroup, try increasing `symprec` parameter to `Crystal`.""")
     end
-    perm = sort(eachindex(cryst.positions), lt=less_than)
-    cryst.positions .= cryst.positions[perm]
-    cryst.classes .= cryst.classes[perm]
-    cryst.types .= cryst.types[perm]
+    p = sort(eachindex(cryst.positions), lt=less_than)
+    permute_sites!(cryst, p)
 end