Bravais lattice with a gaussian at every site

[ash]

I am trying to generate a numerical initial condition for a PDE that looks like this: it is a Bravais lattice, with a gaussian at each site (in 2D). This is not difficult to do by hand for a square lattice with two atoms per unit cell, for example:

using Plots

const N = 256
dx = 1/143
dy = dx

x = collect(-N/2*dx:dx:(N/2-1)*dx)
y = collect(-N/2*dy:dy:(N/2-1)*dy)
Nₚ = 4  # Has to be power of two, number of partitions. You get Nₚ^2 sites
xₗ = collect(-N/Nₚ/2*dx:dx:(N/Nₚ/2-1)*dx) # Local coordinate
yₗ = collect(-N/Nₚ/2*dy:dy:(N/Nₚ/2-1)*dy) # Local Coordinate
# hack fix for rect lattices
if length(yₗ) < length(xₗ)
  append!(yₗ,yₗ[end]+dy)
end
α = 500.0
u01 = kron(ones(Nₚ, Nₚ), (1 .- exp.(-α*((xₗ).^2 .+ (yₗ').^2)) .- exp.(-α*((xₗ.-xₗ[end]).^2 .+ (yₗ'.-yₗ[end]).^2)) .- exp.(-α*((xₗ.-xₗ[1]).^2 .+ (yₗ'.-yₗ[end]).^2)) .-  exp.(-α*((xₗ.-xₗ[end]).^2 .+ (yₗ'.-yₗ[1]).^2)) .- exp.(-α*((xₗ.-xₗ[1]).^2 .+ (yₗ'.-yₗ[1]).^2))))

heatmap(u01)

Or even a triangular (hexagonal) lattice if we expand our PDE model to handle rectangular grids:

using Plots

const N = 256
dx = 1/143
dy = dx*2/sqrt(3)
x = collect(-N/2*dx:dx:(N/2-1)*dx)
y = collect(-N/2*dy:dy:(N/2-1)*dy)
Nₚ = 4  # Has to be power of two, number of partitions. You get Nₚ^2 sites
xₗ = collect(-N/Nₚ/2*dx:dx:(N/Nₚ/2-1)*dx) # Local coordinate
yₗ = collect(-N/Nₚ/2*dy:dy:(N/Nₚ/2-1)*dy) # Local Coordinate
# hack fix for rect lattices
if length(yₗ) < length(xₗ)
  append!(yₗ,yₗ[end]+dy)
end
α = 500.0
u01_piece = (1 .- exp.(-α*((xₗ).^2 .+ (yₗ'.- yₗ[end]).^2)) .- exp.(-α*((xₗ).^2 .+ (yₗ'.- yₗ[1]).^2)) .-
                exp.(-α*((xₗ.+sqrt(3)/2*yₗ[1]).^2 .+ (yₗ'.+yₗ[1]/2).^2)) .-
                exp.(-α*((xₗ.+sqrt(3)/2*yₗ[1]).^2 .+ (yₗ'.-yₗ[1]/2).^2)) .-
                exp.(-α*((xₗ.-sqrt(3)/2*yₗ[1]).^2 .+ (yₗ'.-yₗ[1]/2).^2)) .-
                exp.(-α*((xₗ.-sqrt(3)/2*yₗ[1]).^2 .+ (yₗ'.+yₗ[1]/2).^2)) .-
                exp.(-α*((xₗ).^2 .+ (yₗ').^2)))

u01 = kron(ones(Nₚ, Nₚ), u01_piece)
heatmap(u01)

However:

1. This code is very clunky and ugly
2. This code is not generalizable to more complex patterns, such as the one I’m actually after shown below:
   
   

I tried using this Bravais.jl package but it is a WIP and doesn’t exactly work with rectangular grids.

using Bravais
using Plots

N = 5
M = 5
c = 6
Lattice = TriangularLattice([N,M])
x = zeros(1, N*M)
y = zeros(1, N*M)
for i in 1:N*M
    x[i], y[i] = realspace(lattice, lattice[i])
end
x = [-x -x x x] # Need to remove duplicates but no big deal for now
y = [-y y -y y]
xₘ = maximum(x)
yₘ =  maximum(y)


Ng = 256
dx = 1/143
dy = dx*yₘ/xₘ
xg = collect(-Ng/2*dx:dx:(Ng/2-1)*dx)
yg = collect(-Ng/2*dy:dy:(Ng/2-1)*dy)
x = x./xₘ*maximum(xg)
y = y./xₘ*maximum(xg)

α = 500.0
u01_piece = ones(Ng,Ng)
for i in 1:length(x)
    u01_piece .-= exp.(-α*((xg .- x[i]).^2 .+ (yg'.- y[i]).^2))
end
#u01 = kron(ones(Nₚ, Nₚ), u01_piece)


heatmap(xg, yg, u01_piece', aspect_ratio=yₘ/xₘ)

Any thoughts on this problem? Is there a package that makes writing code like this more elegant? Or a different approach than what I’m doing?

[lmiq]

In CellListMap.jl I have a function that replicates the a periodic system according to the unit cell matrix, it works like this:

julia> using CellListMap

julia> x = SVector{2,Float64}[ [0,0], [0,1], [0.5,0.5] ];

julia> CellListMap.replicate_system!(x,[ 1 0 ; 0.5 1 ],(0:5,0:5)); # replicate (updates x)

julia> using Plots

julia> scatter(Tuple.(x))

where the second argument is the unit cell matrix, and the third is a tuple that defines how many times the system is replicated in each direction. With that you get:

Not sure if that helps, I was intending to split some of this functionality into another package, but didn’t have the time, for now it is an internal function.

[rafael.guerra]

@ash

Bravais lattice

FWIW, the packages indicated in this thread handle Bravais lattices, and there is also the PWDFT.jl package which deals with them.

[stevengj]

@ash

This code is very clunky and ugly

Just writing a loop over unit cells, accumulating one or more gaussians per unit cell, seems a lot easier than fiddling with broadcasting.

[ash]

Thank you all for the suggestions. I moved over to accumulating gaussians per unit cell instead of broadcasting as @stevengj suggested and used LightLattices.jl. Here’s a script I spent a few hours hacking together and testing (most of the time went to figuring out the geometry and making sure the pattern fills a NxN grid with constant dx and dy = dy0*dx regardless of the size of the pattern).

using LightLattices
using Plots
using StaticArrays


function init_cond(type, N; dx=1/143, M = 3, dy = 1/143, α=500.0)
    # Use odd sized pattern
    @assert M ÷ 2 + 1 == ceil(M/2)
    # Done
    if type==:Square
        @assert M ≥ 3
        x = collect(-N/2*dx:dx:(N/2-1)*dx) # Global coordinate
        y = collect(-N/2*dx:dx:(N/2-1)*dx) # Global Coordinate
        len = x[end]/(M-1)*2 # The factor here is to make the whole MxM grid fit into the lattice
        basis = len*hcat([0,1],[1,0]) |> SMatrix{2,2}
        cell = HomogeneousCell([[0.0,0.0]])
        lattice = RegularLattice((M,M), basis, cell)
        u01 = zeros(N,N)
        for i in 1:M
            for j in 1:M
                ii, jj= (lattice[i,j,1] .- len*(M+1)/2)
                @show ii,jj
                u01 .+= exp.(-α*((x .- ii).^2 .+ (y' .- jj).^2))
            end
        end
    elseif type==:ErMnO3
        @assert M ≥ 5
        # Grid Stuff
        x = collect(-N/2*dx:dx:(N/2-1)*dx) # Global coordinate
        dy = dx*sqrt(3)/2 # factor of 6/5??
        y = collect(-N/2*dy:dy:(N/2-1)*dy) # Global Coordinate
        #y = [y; y[end]*(1+dx)] # Hack Fix, idk why y sometimes has N-1 elements
        u01 = zeros(N,N)
        len = x[end]/(M-3)*2
        # Declare Basis
        fbasis = len*hcat([0.5, 0.5*sqrt(3)],
                      [0.5, -0.5*sqrt(3)],
                     ) |> SMatrix{2,2}
        # Atom positions
        cell_vectors_raw1 = [[0.0, 0.0], [1/3, 2/3], [2/3, 1/3]]
        cell_vectors_raw2 = [[0, 1/3], [0, 2/3], [1/3, 1/3], [1/3, 0], [2/3, 0], [2/3, 2/3]]

        fcell = InhomogeneousCell([fbasis*vec for vec in cell_vectors_raw1], [fbasis*vec for vec in cell_vectors_raw2]; label = :ermno3)
        lattice = RegularLattice((M,M), fbasis, fcell; label = :hexagonal)

        u01 = zeros(N,N)
        for i in 1:M
            for j in 1:M
                for ic in 1:length(cell_vectors_raw1)
                    ii, jj= lattice[i,j,ic,1]
                    ii -= len*(M ÷ 2 + 1)
                    @show ii,jj
                    u01 .+= exp.(-α*((x .- ii).^2 .+ (y' .- jj).^2))
                end
                for ic in 1:length(cell_vectors_raw2)
                    ii, jj= lattice[i,j,ic,2]
                    ii -= len*(M ÷ 2 + 1)
                    u01 .-= exp.(-α*((x .- ii).^2 .+ (y' .- jj).^2))
                end
            end
        end
    elseif type==:HoneyComb
        @assert M ≥ 5
        # Grid Stuff
        x = collect(-N/2*dx:dx:(N/2-1)*dx) # Global coordinate
        dy = dx*sqrt(3)/2*6/5 # factor of 6/5
        y = collect(-N/2*dy:dy:(N/2-1)*dy) # Global Coordinate
        #y = [y; y[end]*(1+dx)] # Hack Fix, idk why y sometimes has N-1 elements
        u01 = zeros(N,N)
        len = x[end]/(M-3)*2
        # Declare Basis
        fbasis = len*hcat([0.5, 0.5*sqrt(3)],
                      [0.5, -0.5*sqrt(3)],
                     ) |> SMatrix{2,2}
        # Atom positions
        cell_vectors_raw = [[2/3, 1/3], [1/3, 2/3]]
        pos = [fbasis*vec for vec in cell_vectors_raw]
        fcell = HomogeneousCell(pos)
        lattice = RegularLattice((M,M), fbasis, fcell; label = :hexagonal)

        u01 = zeros(N,N)
        for i in 1:M
            for j in 1:M
                for ic in 1:length(cell_vectors_raw)
                    ii, jj= lattice[i,j,ic]
                    ii -= len*(M ÷ 2 + 1)
                    @show ii,jj
                    u01 .+= exp.(-α*((x .- ii).^2 .+ (y' .- jj).^2))
                end
            end
        end
    elseif type==:Hexagonal
        @assert M ≥ 5
        # Grid Stuff
        x = collect(-N/2*dx:dx:(N/2-1)*dx) # Global coordinate
        dy = dx*sqrt(3)/2
        y = collect(-N/2*dy:dy:(N/2-1)*dy) # Global Coordinate
        #y = [y; y[end]*(1+dx)] # Hack Fix, idk why y sometimes has N-1 elements
        u01 = zeros(N,N)
        len = x[end]/(M-3)*2
        # Declare Basis
        fbasis = len*hcat([0.5, 0.5*sqrt(3)],
                      [0.5, -0.5*sqrt(3)],
                     ) |> SMatrix{2,2}
        # Atom positions
        cell_vectors_raw = [[0,0]]
        pos = [fbasis*vec for vec in cell_vectors_raw]
        fcell = HomogeneousCell(pos)
        lattice = RegularLattice((M,M), fbasis, fcell; label = :hexagonal)

        for i in 1:M
            for j in 1:M
                for ic in 1:length(cell_vectors_raw)
                    ii, jj= lattice[i,j, ic]
                    ii -= len*(M ÷ 2 + 1)
                    @show ii
                    u01 .+= exp.(-α*((x .- ii).^2 .+ (y' .- jj).^2))
                end
            end
        end
    end
    return u01, x, y, dx, dy
end

Let’s go for a square lattice:

N = 256
u, x, y, dx, dy = init_cond(:Square, N, M=5, α = 500.0)
heatmap(x, y, u', aspect_ratio=dx/dy, axis=([], false), colorbar=false, title="Square Grid, M = 5")

Or a hexagonal (triangular) lattice:

N = 256
u, x, y, dx, dy = init_cond(:Hexagonal, N, M=11, α = 500.0)
heatmap(x, y, u', aspect_ratio=dx/dy, axis=([], false), colorbar=false, title="Hexagonal Lattice, M = 11")

And finally, an Erbium Manganate (ErMnO3) lattice, the one shown above as an example of the complicated lattices I am interested in.

N = 256
u, x, y, dx, dy = init_cond(:ErMnO3, N, M=5, α = 500.0)
heatmap(x, y, u', aspect_ratio=dx/dy, axis=([], false), colorbar=false, title="ErMnO3 Lattice, M = 5")

However, I haven’t got the honeycomb lattice to work.

It might look fine initially, but pay attention to the boundary. Let’s try using periodic boundary conditions:

heatmap(kron(ones(3,3), u), aspect_ratio=dx/dy, axis=([], false), colorbar=false, title="Honeycomb \"Lattice\", M = 7")

Any improvements, generalizations, fixes, etc. to the code are more than welcome!

Overall, I have (mostly) achieved what I wanted with LightLattices.jl, nifty package!

P.S. I know the code isn’t very well written, I will optimize it when I clean up my entire PDE solver.

[Oscar_Smith]

I’d you just remind the bottom and right bits, it would tile properly.

[stevengj]

@ash

if type==:Square

I don’t understand why you need a different loop for each lattice type. Can’t you just specify the Bravais lattice by two primitive lattice vectors, and then have an array of displacements per unit cell for each “atom”? (e.g. honeycomb is a hexagonal lattice with two “atoms” per unit cell.)

[antoine-levitt]

Annoyances with boundaries are so much nicer to deal with in Fourier space, so I’d recommend doing that (you’re probably going to do FFTs later anyway…). Also seconding @stevengj 's comment on the lattice construction. Here’s an example that does that, piggybacking on the routine in our code DFTK (https://github.com/JuliaMolSim/DFTK.jl/) that builds an initial input density as a superposition of gaussians (the relevant code is at https://github.com/JuliaMolSim/DFTK.jl/blob/master/src/guess_density.jl#L100: we build the function in Fourier space and take an FFT) and using pymatgen to build the geometry (see https://docs.dftk.org/dev/examples/pymatgen/, but of course you can do that yourself):

using DFTK

kgrid = [1, 1, 1]
lattice = [4.66 -2.33 0.00;
           0.00  4.04 0.00
           0.00  0.00 .1] # that .1 is arbitrary
C = ElementPsp(:C, psp=load_psp("hgh/pbe/c-q4"))
atoms = [C => [[0.0, 0.0, 0.0], [1//3, 2//3, 0.0]]]

# Next we make a `[2, 2, 2]` supercell using pymatgen
pystruct = pymatgen_structure(lattice, atoms)
pystruct.make_supercell([2, 2, 1])
lattice = load_lattice(pystruct)
atoms = [C => [s.frac_coords for s in pystruct.sites]]

model = Model(lattice; atoms)
N = 100
basis = PlaneWaveBasis(model; Ecut=Inf, kgrid, fft_size=[N, N, 1], variational=false)

ρ = guess_density(basis)

using Plots
heatmap(ρ[:, :, 1])

[antoine-levitt]

Oh I see you wanted the plot in cartesian coordinates; then the above will not do what you want, because it’ll plot it in the crystal coordinate system (the change of variables is x_cart=lattice*x_crystal): you’d have to either do the plot in these skewed coordinates, or map back to real space and interpolate. In general if you’re solving a PDE in a lattice you’re probably much better off doing the change of variable to crystal coordinates at the outset, though.

[Gregstrq]

By the way, if you are dealing with a Bravais Lattice (single atom per unit cell), then there is a special type of unit cell called TrivialCell, which emulates a single atom with zero displacement.

So, if you want to specify a Bravais lattice, you can simply drop unit cell from Regular Lattice constructor. Then, TrivialCell is picked automatically.