relax_cell_geometry.lua
local flos = require "flos"
local geom = {}
geom[1] = flos.LBFGS{H0 = 1. / 75.}
geom[2] = flos.LBFGS{H0 = 1. / 50.}
local lattice = {}
lattice[1] = flos.LBFGS{H0 = 1. / 75.}
lattice[2] = flos.LBFGS{H0 = 1. / 50.}
local Unit = siesta.Units
local strain = flos.Array.zeros(6)
local stress_mask = flos.Array.ones(6)
stress_mask[4] = 0.
stress_mask[5] = 0.
stress_mask[6] = 0.
local cell_first
local relax_geom = true
function siesta_comm()
local ret_tbl = {}
if siesta.state == siesta.INITIALIZE then
siesta.receive({"geom.cell",
"MD.Relax.Cell",
"MD.MaxDispl",
"MD.MaxForceTol",
"MD.MaxStressTol"})
if not siesta.MD.Relax.Cell then
siesta.MD.Relax.Cell = true
ret_tbl = {"MD.Relax.Cell"}
end
IOprint("\nLUA convergence information for the LBFGS algorithms:")
cell_first = flos.Array.from(siesta.geom.cell) / Unit.Ang
IOprint("Lattice optimization:")
for i = 1, #lattice do
lattice[i].tolerance = siesta.MD.MaxStressTol * Unit.Ang ^ 3 / Unit.eV
lattice[i].max_dF = siesta.MD.MaxDispl / Unit.Ang
if siesta.IONode then
lattice[i]:info()
end
end
IOprint("\nGeometry optimization:")
for i = 1, #geom do
geom[i].tolerance = siesta.MD.MaxForceTol * Unit.Ang / Unit.eV
geom[i].max_dF = siesta.MD.MaxDispl / Unit.Ang
if siesta.IONode then
geom[i]:info()
end
end
if relax_geom then
IOprint("\nLUA: Starting with geometry relaxation!\n")
else
IOprint("\nLUA: Starting with cell relaxation!\n")
end
end
if siesta.state == siesta.MOVE then
siesta.receive({"geom.cell",
"geom.xa",
"geom.fa",
"geom.stress",
"MD.Relaxed"})
ret_tbl = siesta_move(siesta)
end
siesta.send(ret_tbl)
end
function siesta_move(siesta)
siesta.geom.cell = flos.Array.from(siesta.geom.cell) / Unit.Ang
siesta.geom.xa = flos.Array.from(siesta.geom.xa) / Unit.Ang
siesta.geom.fa = flos.Array.from(siesta.geom.fa) * Unit.Ang / Unit.eV
siesta.geom.stress = -flos.Array.from(siesta.geom.stress) * Unit.Ang ^ 3 / Unit.eV
local voigt = stress_to_voigt(siesta.geom.stress)
voigt = voigt * stress_mask
local conv_lattice = lattice[1]:optimized(voigt)
voigt = nil
local conv_geom = geom[1]:optimized(siesta.geom.fa)
if conv_lattice and conv_geom then
siesta.MD.Relaxed = true
return {'MD.Relaxed'}
end
if relax_geom and conv_geom then
relax_geom = false
for i = 1, #geom do
geom[i]:reset()
end
cell_first = siesta.geom.cell:copy()
strain = flos.Array.zeros(6)
IOprint("\nLUA: switching to cell relaxation!\n")
elseif (not relax_geom) and conv_lattice then
relax_geom = true
for i = 1, #lattice do
lattice[i]:reset()
end
IOprint("\nLUA: switching to geometry relaxation!\n")
end
if relax_geom then
return siesta_geometry(siesta)
else
return siesta_cell(siesta)
end
end
function stress_to_voigt(stress)
local voigt = flos.Array.empty(6)
voigt[1] = stress[1][1]
voigt[2] = stress[2][2]
voigt[3] = stress[3][3]
voigt[4] = (stress[2][3] + stress[3][2]) * 0.5
voigt[5] = (stress[1][3] + stress[3][1]) * 0.5
voigt[6] = (stress[1][2] + stress[2][1]) * 0.5
return voigt
end
function stress_from_voigt(voigt)
local stress = flos.Array.empty(3, 3)
stress[1][1] = voigt[1]
stress[1][2] = voigt[6]
stress[1][3] = voigt[5]
stress[2][1] = voigt[6]
stress[2][2] = voigt[2]
stress[2][3] = voigt[4]
stress[3][1] = voigt[5]
stress[3][2] = voigt[4]
stress[3][3] = voigt[3]
return stress
end
function siesta_geometry(siesta)
local xa = siesta.geom.xa
local fa = siesta.geom.fa
local all_xa = {}
local weight = flos.Array.empty(#geom)
for i = 1, #geom do
all_xa[i] = geom[i]:optimize(xa, fa)
weight[i] = geom[i].weight
end
weight = weight / weight:sum()
if #geom > 1 then
IOprint("\nGeometry weighted average: ", weight)
end
local out_xa = all_xa[1] * weight[1]
for i = 2, #geom do
out_xa = out_xa + all_xa[i] * weight[i]
end
all_xa = nil
siesta.geom.xa = out_xa * Unit.Ang
return {"geom.xa"}
end
function siesta_cell(siesta)
local cell = siesta.geom.cell
local xa = siesta.geom.xa
local stress = stress_to_voigt(siesta.geom.stress)
stress = stress * stress_mask
local vol = cell[1]:cross(cell[2]):dot(cell[3])
local all_strain = {}
local weight = flos.Array.empty(#lattice)
for i = 1, #lattice do
all_strain[i] = lattice[i]:optimize(strain, stress * vol)
lattice[i]:optimized(stress)
weight[i] = lattice[i].weight
end
weight = weight / weight:sum()
if #lattice > 1 then
IOprint("\nLattice weighted average: ", weight)
end
local out_strain = all_strain[1] * weight[1]
for i = 2, #lattice do
out_strain = out_strain + all_strain[i] * weight[i]
end
all_strain = nil
strain = out_strain * stress_mask
out_strain = nil
local dcell = flos.Array( cell.shape )
dcell[1][1] = 1.0 + strain[1]
dcell[1][2] = 0.5 * strain[6]
dcell[1][3] = 0.5 * strain[5]
dcell[2][1] = 0.5 * strain[6]
dcell[2][2] = 1.0 + strain[2]
dcell[2][3] = 0.5 * strain[4]
dcell[3][1] = 0.5 * strain[5]
dcell[3][2] = 0.5 * strain[4]
dcell[3][3] = 1.0 + strain[3]
local out_cell = cell_first:dot(dcell)
dcell = nil
local lat = flos.Lattice:new(cell)
local fxa = lat:fractional(xa)
xa = fxa:dot(out_cell)
lat = nil
fxa = nil
siesta.geom.cell = out_cell * Unit.Ang
siesta.geom.xa = xa * Unit.Ang
return {"geom.cell",
"geom.xa"}
end