add energy logger

src/ExTinyMD.jl

@@ -3,10 +3,11 @@ module ExTinyMD
 using LinearAlgebra, Random, Distributions, CellListMap, StaticArrays, DelimitedFiles
 
 export Point, Atom, Boundary, Q2dBoundary, CubicBoundary, MDSys, position_check3D, position_checkQ2D, BoundaryCheck!, SimulationInfo, thermostat_update!, update_acceleration!, update_finder!, NoInteraction, AllNeighborFinder, NoNeighborFinder, NoThermoStat, dist2, random_position, random_velocity, create_atoms
+export energy
 export VerletProcess, simulate!
 export AndersenThermoStat, BerendsenThermoStat, NHVerletProcess
 export SubNeighborFinder, CellList3D, CellList2D, CellListDir3D, CellListDirQ2D, CellListQ2D
-export TemperatureLogger, TrajectoryLogger
+export TemperatureLogger, TrajectoryLogger, EnergyLogger
 
 export SubLennardJones, LennardJones, ExternalField
 
@@ -32,6 +33,7 @@ include("MD_core/neighbor_finder/substrate_finder.jl")
 
 include("MD_core/recorder/temperatue_logger.jl")
 include("MD_core/recorder/trajectory_logger.jl")
+include("MD_core/recorder/energy_logger.jl")
 
 # this part will be about the interactions
 include("interactions/lennard_jones.jl")

src/MD_core/neighbor_finder/cell_list.jl

@@ -4,6 +4,7 @@ mutable struct CellList3D{T, TI} <: AbstractNeighborFinder
     update_steps::TI
 end
 
+Base.show(io::IO, cell_list::CellList3D) = print(io, "CellList3D")
 
 function CellList3D(info::SimulationInfo{T}, cutoff::T, boundary::Boundary{T}, update_steps::TI) where {T<:Number, TI<:Integer}
     coords = [SVector{3, T}(p_info.position[1], p_info.position[2], p_info.position[3]) for p_info in info.particle_info]
@@ -32,6 +33,8 @@ mutable struct CellListQ2D{T, TI} <: AbstractNeighborFinder
     update_steps::TI
 end
 
+Base.show(io::IO, cell_list::CellListQ2D) = print(io, "CellListQ2D")
+
 function CellListQ2D(info::SimulationInfo{T}, cutoff::T, boundary::Boundary{T}, update_steps::TI) where {T<:Number, TI<:Integer}
     coords = [SVector{2, T}(p_info.position[1], p_info.position[2]) for p_info in info.particle_info]
 

src/MD_core/recorder/energy_logger.jl

@@ -0,0 +1,46 @@
+struct EnergyLogger{TI} <: AbstractLogger 
+    step::TI 
+    energy_file::String
+    output::Bool
+    interactions::Vector{Tuple{AbstractInteraction, AbstractNeighborFinder}}
+end
+
+Base.show(io::IO, logger::EnergyLogger) = print(io, "EnergyLogger")
+
+function EnergyLogger(step::TI; interactions::Vector{Tuple{T_interaction, T_neighbor}}, energy_names::Vector{String}, energy_file::String = "energy.txt", output::Bool = true) where{TI<:Integer, T_interaction<:AbstractInteraction, T_neighbor<:AbstractNeighborFinder}
+    @assert length(interactions) == length(energy_names)
+    if output
+        # if there is already a file used to store the result, refresh them
+        f = open(energy_file, "w")
+        names = "step, Ek, " * join(energy_names, ", ")
+        write(f, names, "\n")
+        close(f)
+    end
+    return EnergyLogger{TI}(step, energy_file, output, interactions)
+end
+
+@inbounds function kinetic_energy(sys::MDSys{T}, info::SimulationInfo{T}) where {T<:Number}
+    Ek = zero(T)
+    for p_info in info.particle_info
+        id = p_info.id
+        Ek += 0.5 * sys.atoms[id].mass * sum(abs2, p_info.velocity)
+    end
+    return Ek
+end
+
+function record!(logger::EnergyLogger{TI}, sys::MDSys{T}, info::SimulationInfo{T}) where {T<:Number, TI<:Integer}
+    if iszero(info.running_step % logger.step) && logger.output
+        IO_energy = open(logger.energy_file, "a")
+        step = info.running_step
+        Ek = kinetic_energy(sys, info)
+        write(IO_energy, "$step, $Ek")
+        for i in 1:length(logger.interactions)
+            interaction, neighborfinder = logger.interactions[i]
+            E = energy(interaction, neighborfinder, sys, info)
+            write(IO_energy, ", $E")
+        end
+        write(IO_energy, "\n")
+        close(IO_energy)
+    end
+    return nothing
+end

src/MD_core/recorder/temperatue_logger.jl

@@ -4,6 +4,8 @@ mutable struct TemperatureLogger{T, TI} <: AbstractLogger
     output::Bool
 end
 
+Base.show(io::IO, logger::TemperatureLogger) = print(io, "TemperatureLogger")
+
 mutable struct PressureLogger{T, TI} <: AbstractLogger
     step::TI
     data::Vector{T}

src/MD_core/recorder/trajectory_logger.jl

@@ -4,10 +4,12 @@ struct TrajectoryLogger{TI} <: AbstractLogger
     output::Bool
 end
 
+Base.show(io::IO, logger::TrajectoryLogger) = print(io, "TrajectoryLogger")
+
 function TrajectoryLogger(;step::TI, trajectory_file::String = "trajectory.txt", output::Bool = true) where{TI<:Integer}
     if output
         # if there is already a file used to store the result, refresh them
-        f = open(trajectory_file, "a")
+        f = open(trajectory_file, "w")
         close(f)
     end
     return TrajectoryLogger{TI}(step, trajectory_file, output)

src/interactions/external_field.jl

@@ -2,6 +2,8 @@ struct ExternalField{T} <: AbstractInteraction
     E::Point{3, T}
 end
 
+Base.show(io::IO, interaction::ExternalField) = print(io, "ExternalField with E = $(interaction.E)")
+
 function ExternalField(;E::Point{3, T} = Point((0.0, 0.0, 0.0))) where {T} 
     return ExternalField{T}(E)
 end

src/interactions/lennard_jones.jl

@@ -4,6 +4,8 @@ struct LennardJones{T} <: AbstractInteraction
     σ::T
 end
 
+Base.show(io::IO, interaction::LennardJones) = print(io, "LennardJones with ϵ = $(interaction.ϵ), cutoff = $(interaction.cutoff), σ = $(interaction.σ)")
+
 LennardJones(;ϵ::T = 1.0, cutoff::T = 3.5, σ::T = 1.0) where T = LennardJones(ϵ, cutoff, σ)
 
 function update_acceleration!(interaction::LennardJones{T}, neighborfinder::T_NIEGHBER, sys::MDSys{T}, info::SimulationInfo{T}) where {T<:Number, T_NIEGHBER<:AbstractNeighborFinder}

src/interactions/substrate_lennard_jones.jl

@@ -6,6 +6,8 @@ struct SubLennardJones{T} <: AbstractInteraction
     sub_up::T
 end
 
+Base.show(io::IO, interaction::SubLennardJones) = print(io, "SubLennardJones with ϵ = $(interaction.ϵ), cutoff = $(interaction.cutoff), σ = $(interaction.σ), sub_down = $(interaction.sub_down), sub_up = $(interaction.sub_up)")
+
 SubLennardJones(sub_down::T, sub_up::T;ϵ::T = 1.0, cutoff::T = 3.5, σ::T = 1.0) where T<:Number = SubLennardJones(ϵ, cutoff, σ, sub_down, sub_up)
 
 function update_acceleration!(interaction::SubLennardJones{T}, neighborfinder::SubNeighborFinder{T, TI}, sys::MDSys{T}, info::SimulationInfo{T}) where {T<:Number, TI<:Integer}

src/types.jl

@@ -44,6 +44,8 @@ struct MDSys{T_NUM, T_INTERACTION, T_LOGGER, T_SIMULATOR}
     simulator::T_SIMULATOR
 end
 
+Base.show(io::IO, sys::MDSys) = print(io, " MDSys with $(sys.n_atoms) atoms \n boundary: $(sys.boundary) \n simulator: $(sys.simulator) \n interactions: $(sys.interactions) \n loggers: $(sys.loggers)")
+
 function MDSys(;
     n_atoms::Int64,
     atoms::Vector{Atom{T_NUM}},
@@ -68,6 +70,8 @@ mutable struct SimulationInfo{T}
     id_dict::Dict{Int, Int}
 end
 
+Base.show(io::IO, info::SimulationInfo) = print(io, "SimulationInfo: $(info.running_step) steps, $(length(info.particle_info)) particles")
+
 struct NoThermoStat <: AbstractThermoStat
     nostat::Bool
 end
@@ -83,6 +87,8 @@ struct AllNeighborFinder{T} <: AbstractNeighborFinder
 end
 AllNeighborFinder(n_atoms::TI, T::Type = Float64) where {TI <: Integer} = AllNeighborFinder{T}([(i, j, zero(T)) for i in 1:n_atoms - 1 for j in i+1:n_atoms])
 
+Base.show(io::IO, neighborfinder::AllNeighborFinder) = print(io, "AllNeighborFinder with $(length(neighborfinder.neighborlist)) pairs")
+
 function update_finder!(neighborfinder::T_NIEGHBOR, info::SimulationInfo{T}) where {T<:Number, T_NIEGHBOR <: AllNeighborFinder}
     return nothing
 end
@@ -92,6 +98,8 @@ struct NoNeighborFinder{T} <: AbstractNeighborFinder
 end
 NoNeighborFinder(T::Type = Float64) = NoNeighborFinder{T}([(0, 0, zero(T))])
 
+Base.show(io::IO, neighborfinder::NoNeighborFinder) = print(io, "NoNeighborFinder")
+
 function update_finder!(neighborfinder::T_NIEGHBOR, info::SimulationInfo{T}) where {T<:Number, T_NIEGHBOR <: NoNeighborFinder}
     return nothing
 end
@@ -101,6 +109,8 @@ struct NoInteraction <: AbstractInteraction
 end
 NoInteraction() = NoInteraction(true)
 
+Base.show(io::IO, interaction::NoInteraction) = print(io, "NoInteraction")
+
 function update_acceleration!(interaction::NoInteraction, neighborfinder::T_NEIGHBOR, sys::MDSys{T}, info::SimulationInfo{T}) where {T<:Number, T_NEIGHBOR<:AbstractNeighborFinder}
     return nothing
 end

test/simulation.jl

@@ -11,7 +11,7 @@
 
     interactions = [(LennardJones(), CellList3D(info, 4.5, boundary, 100)), (SubLennardJones(0.0, L; cutoff = 1.0, σ = 0.5), SubNeighborFinder(1.5, info, 0.0, L)), (ExternalField(E = Point((0.0, 0.0, 1.0))), NoNeighborFinder())]
 
-    loggers = [TemperatureLogger(100, output = false), TrajectoryLogger(step = 100, output = false)]
+    loggers = [TemperatureLogger(100, output = false), TrajectoryLogger(step = 100, output = false), EnergyLogger(100, output = false, interactions = [(LennardJones(), CellList3D(info, 4.5, boundary, 100))], energy_names = ["E_lj"])]
     simulator = VerletProcess(dt = 0.001, thermostat = AndersenThermoStat(1.0, 0.05))
 
     sys = MDSys(