Skip to content

Commit

Permalink
Merge pull request #22 from numericalEFT/uniformbzmesh
Browse files Browse the repository at this point in the history 
Uniformbzmesh
  • Loading branch information
@iintSjds
iintSjds authored Oct 19, 2022
2 parents 5279d5e + e2e6eca commit 4c0bc1c
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 26 deletions.
119 changes: 93 additions & 26 deletions src/Basemesh.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ function _compute_inverse_lattice(lattice)
return inv(lattice)
end

function _compute_recip_lattice(lattice)
function _compute_recip_lattice(lattice::Matrix{T}) where {T}
return 2T(π) * _compute_inverse_lattice(lattice)
end

Expand All @@ -25,22 +25,21 @@ struct Brillouin{T,DIM}
unit_cell_volume::T
recip_cell_volume::T
G_vector::Vector{SVector{DIM,Int}}
end

function Brillouin(; lattice::Matrix{T}, G_vector=nothing) where {T}
DIM = size(lattice, 1)
recip_lattice = _compute_recip_lattice(lattice)
inv_lattice = _compute_inverse_lattice(lattice)
inv_recip_lattice = _compute_inverse_lattice(recip_lattice)
unit_cell_volume = abs(det(lattice))
recip_cell_volume = abs(det(recip_lattice))
if isnothing(G_vector)
G_vector = [SVector{DIM,Int}(zeros(DIM)),]
end

return new{T,DIM}(lattice, recip_lattice, inv_lattice, inv_recip_lattice, unit_cell_volume, recip_cell_volume, G_vector)
function Brillouin(; lattice::Matrix{T}, G_vector=nothing) where {T}
DIM = size(lattice, 1)
recip_lattice = _compute_recip_lattice(lattice)
inv_lattice = _compute_inverse_lattice(lattice)
inv_recip_lattice = _compute_inverse_lattice(recip_lattice)
unit_cell_volume = abs(det(lattice))
recip_cell_volume = abs(det(recip_lattice))
if isnothing(G_vector)
G_vector = [SVector{DIM,Int}(zeros(DIM)),]
end
end

return Brillouin{T,DIM}(lattice, recip_lattice, inv_lattice, inv_recip_lattice, unit_cell_volume, recip_cell_volume, G_vector)
end

struct UniformBZMesh{T,DIM} <: AbstractMesh{T,DIM}
br::Brillouin{T,DIM}
Expand All @@ -49,6 +48,11 @@ struct UniformBZMesh{T,DIM} <: AbstractMesh{T,DIM}
shift::SVector{DIM,Rational}
end

# default shift is 1/2, result in Monkhorst-Pack mesh
# with shift = 0, result in Gamma-centered
# can also customize with shift::SVector by calling default constructor
UniformBZMesh(; br::Brillouin{T,DIM}, size, shift::Number=1 // 2) where {T,DIM} = UniformBZMesh{T,DIM}(br, size, SVector{DIM,Rational}(shift .* ones(Int, DIM)))

Base.length(mesh::UniformBZMesh) = prod(mesh.size)
Base.size(mesh::UniformBZMesh) = mesh.size
Base.size(mesh::UniformBZMesh, I) = mesh.size[I]
Expand All @@ -57,6 +61,81 @@ function Base.show(io::IO, mesh::UniformBZMesh)
println("UniformBZMesh with $(length(mesh)) mesh points")
end

@generated function _inds2ind(size::NTuple{DIM,Int}, I) where {DIM}
ex = :(I[DIM] - 1)
for i = (DIM-1):-1:1
ex = :(I[$i] - 1 + size[$i] * $ex)
end
return :($ex + 1)
end

@generated function _ind2inds(size::NTuple{DIM,Int}, I::Int) where {DIM}
inds, quotient = :((I - 1) % size[1] + 1), :((I - 1) ÷ size[1])
for i = 2:DIM-1
inds, quotient = :($inds..., $quotient % size[$i] + 1), :($quotient ÷ size[$i])
end
inds = :($inds..., $quotient + 1)
return :(SVector{DIM,Int}($inds))
end

function Base.getindex(mesh::UniformBZMesh{T,DIM}, inds...) where {T,DIM}
n = SVector{DIM,Int}(inds)
return mesh.br.recip_lattice * ((n .- 1 .+ mesh.shift) ./ mesh.size)
end

function Base.getindex(mesh::UniformBZMesh, I::Int)
return Base.getindex(mesh, _ind2inds(mesh.size, I)...)
end

Base.firstindex(mesh::UniformBZMesh) = 1
Base.lastindex(mesh::UniformBZMesh) = length(mesh)
# # iterator
Base.iterate(mesh::UniformBZMesh) = (mesh[1], 1)
Base.iterate(mesh::UniformBZMesh, state) = (state >= length(mesh)) ? nothing : (mesh[state+1], state + 1)

function _indfloor(x, N; edgeshift=1)
# edgeshift = 1 by default in floor function so that end point return N-1
# edgeshift = 0 in locate function
if x < 1
return 1
elseif x >= N
return N - edgeshift
else
return floor(Int, x)
end
end

function locate(mesh::UniformBZMesh{T,DIM}, x) where {T,DIM}
# find index of nearest grid point to the point
displacement = SVector{DIM,T}(x)
inds = (mesh.br.inv_recip_lattice * displacement) .* mesh.size .+ 1.5 .- mesh.shift .+ 2 .* eps.(T.(mesh.size))
indexall = 1
# println((mesh.invlatvec * displacement))
# println(inds)
factor = 1
indexall += (_indfloor(inds[1], mesh.size[1]; edgeshift=0) - 1) * factor
for i in 2:DIM
factor *= mesh.size[i-1]
indexall += (_indfloor(inds[i], mesh.size[i]; edgeshift=0) - 1) * factor
end

return indexall
end

volume(mesh::UniformBZMesh) = mesh.br.recip_cell_volume
function volume(mesh::UniformBZMesh{T,DIM}, i) where {T,DIM}
inds = _ind2inds(mesh.size, i)
cellarea = T(1.0)
for j in 1:DIM
if inds[j] == 1
cellarea *= T(0.5) + mesh.shift[j]
elseif inds[j] == mesh.size[j]
cellarea *= T(1.5) - mesh.shift[j]
end
# else cellarea *= 1.0 so nothing
end
return cellarea / length(mesh) * volume(mesh)
end

# c.f. DFTK.jl/src/Model.jl
# UniformBZMesh iterate on 1st Brillouin Zone
Expand Down Expand Up @@ -159,18 +238,6 @@ end
return :(SVector{DIM,Int}($inds))
end

function _indfloor(x, N; edgeshift=1)
# edgeshift = 1 by default in floor function so that end point return N-1
# edgeshift = 0 in locate function
if x < 1
return 1
elseif x >= N
return N - edgeshift
else
return floor(Int, x)
end
end

function Base.floor(mesh::UniformMesh{DIM,N}, x) where {DIM,N}
# find index of nearest grid point to the point
displacement = SVector{DIM,Float64}(x) - mesh.origin
Expand Down
70 changes: 70 additions & 0 deletions test/BaseMesh.jl
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,76 @@
@testset "Base Mesh" begin
rng = MersenneTwister(1234)

@testset "Brillouin" begin

# square lattice
DIM = 2
lattice = Matrix([1.0 0; 0 1]')
br = BaseMesh.Brillouin(lattice=lattice)
@test br.inv_lattice .* 2π ≈ br.recip_lattice
@test br.unit_cell_volume ≈ abs(det(lattice))
@test br.recip_cell_volume ≈ 1 / abs(det(lattice)) * (2π)^DIM

# triagular lattice
DIM = 2
lattice = Matrix([2.0 0; 1 sqrt(3)]')
br = BaseMesh.Brillouin(lattice=lattice)
@test br.inv_lattice .* 2π ≈ br.recip_lattice
@test br.unit_cell_volume ≈ abs(det(lattice))
@test br.recip_cell_volume ≈ 1 / abs(det(lattice)) * (2π)^DIM

# 3d testing lattice
DIM = 3
lattice = Matrix([2.0 0 0; 1 sqrt(3) 0; 7 11 19]')
br = BaseMesh.Brillouin(lattice=lattice)
@test br.inv_lattice .* 2π ≈ br.recip_lattice
@test br.unit_cell_volume ≈ abs(det(lattice))
@test br.recip_cell_volume ≈ 1 / abs(det(lattice)) * (2π)^DIM
end

@testset "UniformBZMesh" begin
@testset "Indexing" begin
size = (3, 4, 5)
for i in 1:prod(size)
@test i == BaseMesh._inds2ind(size, BaseMesh._ind2inds(size, i))
end
end

@testset "Array Interface" begin
N1, N2 = 3, 5
lattice = Matrix([1/N1/2 0; 0 1.0/N2/2]') .* 2π
# so that bzmesh[i,j] = (2i-1,2j-1)
br = BaseMesh.Brillouin(lattice=lattice)
bzmesh = BaseMesh.UniformBZMesh(br=br, size=(N1, N2))
for (pi, p) in enumerate(bzmesh)
@test bzmesh[pi] ≈ p # linear index
inds = BaseMesh._ind2inds(bzmesh.size, pi)
@test p ≈ inds .* 2.0 .- 1.0
@test bzmesh[inds...] ≈ p # cartesian index
end
end

@testset "locate and volume" begin
size = (3, 5, 7)
lattice = Matrix([2.0 0 0; 1 sqrt(3) 0; 7 11 19]')
# size = (3, 5)
# lattice = Matrix([2.3 0; 0 7.0]')
br = BaseMesh.Brillouin(lattice=lattice)
bzmesh = BaseMesh.UniformBZMesh(br=br, size=size)
vol = 0.0
for (pi, p) in enumerate(bzmesh)
@test bzmesh[pi] ≈ p # linear index
inds = BaseMesh._ind2inds(bzmesh.size, pi)
@test bzmesh[inds...] ≈ p # cartesian index

@test BaseMesh.locate(bzmesh, p) == pi
vol += BaseMesh.volume(bzmesh, pi)
end
@test vol ≈ BaseMesh.volume(bzmesh)
end
end


function dispersion(k)
me = 0.5
return dot(k, k) / 2me
Expand Down

0 comments on commit 4c0bc1c

Please # to comment.