Fix symmetries & symmetry-breaking with ForwardDiff (#1082)

---------

Co-authored-by: Bruno Ploumhans <[email protected]>

Author: niklasschmitz

src/SymOp.jl

@@ -66,18 +66,49 @@ function Base.:*(op1::SymOp, op2::SymOp)
 end
 Base.inv(op::SymOp) = SymOp(inv(op.W), -op.W\op.w)
 
+is_approx_in(symop, group; kwargs...) = any(s -> isapprox(s, symop; kwargs...), group)
 function check_group(symops::Vector; kwargs...)
-    is_approx_in_symops(s1) = any(s -> isapprox(s, s1; kwargs...), symops)
+    is_approx_in_symops(s1) = is_approx_in(s1, symops; kwargs...)
     is_approx_in_symops(one(SymOp)) || error("check_group: no identity element")
     for s in symops
         if !is_approx_in_symops(inv(s))
             error("check_group: symop $s with inverse $(inv(s)) is not in the group")
         end
         for s2 in symops
-            if !is_approx_in_symops(s*s2) || !is_approx_in_symops(s2*s)
-                error("check_group: product is not stable")
+            if !is_approx_in_symops(s*s2)
+                error("check_group: product is not stable: $(s*s2) is not in the group")
             end
         end
     end
     symops
 end
+
+function complete_symop_group(symops; maxiter=10, kwargs...)
+    completed_group = Vector(symops)
+
+    function add_to_group(to_add, s1)
+        if !is_approx_in(s1, completed_group; kwargs...) && !is_approx_in(s1, to_add; kwargs...)
+            push!(to_add, s1)
+        end
+    end
+
+    for it = 1:maxiter
+        if it == maxiter
+            error("Could not complete group in $maxiter iterations")
+        end
+        to_add = []
+        # Identity always needs to be there!
+        add_to_group(to_add, one(SymOp))
+        for s in completed_group
+            add_to_group(to_add, inv(s))
+            for t in completed_group
+                add_to_group(to_add, s*t)
+            end
+        end
+        if isempty(to_add)
+            return completed_group
+        end
+        append!(completed_group, to_add)
+    end
+    DFTK.check_group(completed_group) # returns the completed group
+end

src/symmetry.jl

@@ -364,31 +364,40 @@ function symmetrize_stresses(basis::PlaneWaveBasis, stresses)
     symmetrize_stresses(basis.model, stresses; basis.symmetries)
 end
 
+"""
+Find the symmetry preimage of `position`, returning the corresponding index in `positions_group`.
+"""
+function find_symmetry_preimage(positions_group, position, symop;
+                                tol_symmetry=SYMMETRY_TOLERANCE)
+    # see (A.27) of https://arxiv.org/pdf/0906.2569.pdf
+    # (but careful that our symmetries are r -> Wr+w, not R(r+f))
+    other_at = symop.W \ (position - symop.w)
+
+    # Find the index of the atom to which idx is mapped to by the symmetry operation.
+    # To avoid issues due to numerical noise we compute the deviations from being
+    # an integer shift (thus equivalent by translational symmetry) for all atoms in
+    # the group and pick the smallest one.
+    smallest_deviation, i_other_at = findmin(positions_group) do at
+        δat = at - other_at
+        maximum(abs, δat - round.(δat))
+    end
+    # Note, that without a fudging factor this occasionally fails:
+    @assert smallest_deviation < 10tol_symmetry
+
+    i_other_at
+end
+
 function symmetrize_forces(positions::AbstractVector, atom_groups::AbstractVector, forces;
                            symmetries, tol_symmetry=SYMMETRY_TOLERANCE)
     symmetrized_forces = zero(forces)
     for group in atom_groups, symop in symmetries
         positions_group = positions[group]
-        W, w = symop.W, symop.w
         for (idx, position) in enumerate(positions_group)
-            # see (A.27) of https://arxiv.org/pdf/0906.2569.pdf
-            # (but careful that our symmetries are r -> Wr+w, not R(r+f))
-            other_at = W \ (position - w)
-
-            # Find the index of the atom to which idx is mapped to by the symmetry operation.
-            # To avoid issues due to numerical noise we compute the deviations from being
-            # an integer shift (thus equivalent by translational symmetry) for all atoms in
-            # the group and pick the smallest one.
-            smallest_deviation, i_other_at = findmin(positions_group) do at
-                δat = at - other_at
-                maximum(abs, δat - round.(δat))
-            end
-            # Note, that without a fudging factor this occasionally fails:
-            @assert smallest_deviation < 10tol_symmetry
+            i_other_at = find_symmetry_preimage(positions_group, position, symop; tol_symmetry)
 
             # (A.27) is in Cartesian coordinates, and since Wcart is orthogonal,
             # Fsymcart = Wcart * Fcart <=> Fsymred = inv(Wred') Fred
-            symmetrized_forces[group[idx]] += inv(W') * forces[group[i_other_at]]
+            symmetrized_forces[group[idx]] += inv(symop.W') * forces[group[i_other_at]]
         end
     end
     symmetrized_forces / length(symmetries)

src/workarounds/forwarddiff_rules.jl

@@ -71,12 +71,76 @@ end
 
 next_working_fft_size(::Type{<:Dual}, size::Int) = size
 
-# determine symmetry operations only from primal lattice values
+# determine symmetry operations only from primal values, then filter out symmetries broken by dual part
 function symmetry_operations(lattice::AbstractMatrix{<:Dual},
-                             atoms, positions, magnetic_moments=[]; kwargs...)
+                             atoms, positions, magnetic_moments=[];
+                             tol_symmetry=SYMMETRY_TOLERANCE, kwargs...)
     positions_value = [ForwardDiff.value.(pos) for pos in positions]
-    symmetry_operations(ForwardDiff.value.(lattice), atoms, positions_value,
-                        magnetic_moments; kwargs...)
+    symmetries = symmetry_operations(ForwardDiff.value.(lattice), atoms,
+                                     positions_value, magnetic_moments;
+                                     tol_symmetry, kwargs...)
+    remove_dual_broken_symmetries(lattice, atoms, positions, symmetries; tol_symmetry)
+end
+
+function remove_dual_broken_symmetries(lattice, atoms, positions,
+                                       symmetries; tol_symmetry=SYMMETRY_TOLERANCE)
+    filter(symmetries) do symmetry
+        !is_symmetry_broken_by_dual(lattice, atoms, positions, symmetry; tol_symmetry)
+    end
+end
+
+"""
+Return `true` if a symmetry that holds for the primal part is broken by
+a perturbation in the lattice or in the positions, `false` otherwise.
+"""
+function is_symmetry_broken_by_dual(lattice, atoms, positions, symmetry::SymOp; tol_symmetry)
+    # For any lattice atom at position x, W*x + w should be in the lattice.
+    # In cartesian coordinates, with a perturbed lattice A = A₀ + εA₁,
+    # this means that for any atom position xcart in the unit cell and any 3 integers u,
+    # there should be an atom at position ycart and 3 integers v such that:
+    # Wcart * (xcart + A*u) + wcart = ycart + A*v
+    # where
+    #     Wcart = A₀ * W * A₀⁻¹; note that W is an integer matrix
+    #     wcart = A₀ * w.
+    #
+    # In relative coordinates this gives:
+    # A⁻¹ * Wcart * (A*x + A*u) + A⁻¹ * wcart = y + v   (*)
+    #
+    # The strategy is then to check that:
+    # 1. A⁻¹ * Wcart * A is still an integer matrix (i.e. no dual part),
+    #    such that any change in u is easily compensated for in v.
+    # 2. The primal component of (*), i.e. with ε=0, is already known to hold.
+    #    Since v does not have a dual component, it is enough to check that
+    #    the dual part of the following is 0:
+    #      A⁻¹ * Wcart * A*x + A⁻¹ * wcart - y 
+
+    lattice_primal = ForwardDiff.value.(lattice)
+    W = (compute_inverse_lattice(lattice) * lattice_primal
+        * symmetry.W * compute_inverse_lattice(lattice_primal) * lattice)
+    w = compute_inverse_lattice(lattice) * lattice_primal * symmetry.w
+
+    is_dual_nonzero(x::AbstractArray) = any(x) do xi
+        maximum(abs, ForwardDiff.partials(xi)) >= tol_symmetry
+    end
+    # Check 1.
+    if is_dual_nonzero(W)
+        return true
+    end
+
+    atom_groups = [findall(Ref(pot) .== atoms) for pot in Set(atoms)]
+    for group in atom_groups
+        positions_group = positions[group]
+        for position in positions_group
+            i_other_at = find_symmetry_preimage(positions_group, position, symmetry; tol_symmetry)
+
+            # Check 2. with x = positions_group[i_other_at] and y = position
+            if is_dual_nonzero(positions_group[i_other_at] + inv(W) * (w - position))
+                return true
+            end
+        end
+    end
+
+    false
 end
 
 function _is_well_conditioned(A::AbstractArray{<:Dual}; kwargs...)

test/forwarddiff.jl

@@ -260,3 +260,117 @@ end
     derivative_fd = ForwardDiff.derivative(compute_force, 0.0)
     @test norm(derivative_ε - derivative_fd) < 1e-4
 end
+
+@testitem "Symmetries broken by perturbation are filtered out" tags=[:dont_test_mpi] begin
+    using DFTK
+    using ForwardDiff
+    using LinearAlgebra
+
+    lattice = [2. 0. 0.; 0. 1. 0.; 0. 0. 1.]
+    positions = [[0., 0., 0.], [0.5, 0., 0.]]
+    gauss = ElementGaussian(1.0, 0.5)
+    atoms = [gauss, gauss]
+    atom_groups = [findall(Ref(pot) .== atoms) for pot in Set(atoms)]
+
+    # Select some "interesting" subset of the symmetries
+    # Rotation in the yz plane by 90 degrees
+    rotyz = SymOp([1 0 0; 0 0 1; 0 -1 0], [0., 0., 0.])
+    mirroryz = rotyz * rotyz
+    # Mirror y
+    mirrory = SymOp([1 0 0; 0 -1 0; 0 0 1], [0., 0., 0.])
+    # Translation by 0.5 in the x direction
+    transx = SymOp(diagm([1, 1, 1]), [0.5, 0., 0.])
+
+    # Generate the full group
+    symmetries_full = DFTK.complete_symop_group([rotyz, mirrory, transx])
+    @test length(symmetries_full) == 16
+
+    DFTK._check_symmetries(symmetries_full, lattice, atom_groups, positions)
+
+    function check_symmetries(filtered_symmetries, expected_symmetries)
+        expected_symmetries = DFTK.complete_symop_group(expected_symmetries)
+
+        @test length(filtered_symmetries) == length(expected_symmetries)
+        for s in expected_symmetries
+            @test DFTK.is_approx_in(s, filtered_symmetries)
+        end
+        for s in filtered_symmetries
+            @test DFTK.is_approx_in(s, expected_symmetries)
+        end
+    end
+
+    # Instantiate Dual to test with perturbations
+    ε = ForwardDiff.Dual{ForwardDiff.Tag{Nothing, Float64}}(0.0, 1.0)
+
+    @testset "Atom movement" begin
+        # Moving the second atom should break the transx symmetry, but not the others
+        positions_modified = [[0., 0., 0.], [0.5 + ε, 0., 0.]]
+        symmetries_filtered = DFTK.remove_dual_broken_symmetries(lattice, atoms, positions_modified, symmetries_full)
+
+        @test length(symmetries_filtered) == 8
+        check_symmetries(symmetries_filtered, [rotyz, mirrory])
+    end
+
+    @testset "Lattice strain" begin
+        # Straining the lattice along z should break the rotyz symmetry, but not the others
+        # In particular it should not break mirroryz which is normally generated by rotyz.
+        lattice_modified = diagm([2., 1., 1. + ε])
+        symmetries_filtered = DFTK.remove_dual_broken_symmetries(lattice_modified, atoms, positions, symmetries_full)
+
+        @test length(symmetries_filtered) == 8
+        check_symmetries(symmetries_filtered, [mirrory, transx, mirroryz])
+    end
+
+    @testset "Atom movement + lattice strain" begin
+        # Only the mirrory and mirroryz symmetries should be left
+        positions_modified = [[0., 0., 0.], [0.5 + ε, 0., 0.]]
+        lattice_modified = diagm([2., 1., 1. + ε])
+        symmetries_filtered = DFTK.remove_dual_broken_symmetries(lattice_modified, atoms, positions_modified, symmetries_full)
+
+        @test length(symmetries_filtered) == 4
+        check_symmetries(symmetries_filtered, [mirrory, mirroryz])
+    end
+end
+
+@testitem "Symmetry-breaking perturbation using ForwardDiff" #=
+    =#    tags=[:dont_test_mpi] setup=[TestCases] begin
+    using DFTK
+    using ForwardDiff
+    using LinearAlgebra
+    aluminium = TestCases.aluminium
+
+    @testset for perturbation in (:lattice, :positions)
+        function run_scf(ε)
+            lattice = if perturbation == :lattice
+                v = ε * [0., 0., 0., 0., 0., 1.]
+                DFTK.voigt_strain_to_full(v) * aluminium.lattice
+            else
+                # Lattice has to be a dual for position perturbations to work
+                aluminium.lattice * one(typeof(ε))
+            end
+            pos = if perturbation == :lattice
+                aluminium.positions
+            else
+                map(enumerate(aluminium.positions)) do (i, x)
+                    i == 1 ? x + ε * [1., 0, 0] : x
+                end
+            end
+
+            model = model_DFT(lattice, aluminium.atoms, pos;
+                              functionals=LDA(), temperature=1e-2,
+                              smearing=Smearing.Gaussian())
+            basis = PlaneWaveBasis(model; Ecut=5, kgrid=[2, 2, 2])
+            self_consistent_field(basis; tol=1e-10)
+        end
+
+        δρ = ForwardDiff.derivative(ε -> run_scf(ε).ρ, 0.)
+
+        h = 1e-6
+        scfres1 = run_scf(-h)
+        scfres2 = run_scf(+h)
+        δρ_finitediff = (scfres2.ρ - scfres1.ρ) / 2h
+
+        rtol = 1e-4
+        @test norm(δρ - δρ_finitediff, 1) < rtol * norm(δρ, 1)
+    end
+end