Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Gadget2 read and units #8

Merged
merged 11 commits into from
Oct 29, 2021
5 changes: 5 additions & 0 deletions src/AstroIO.jl
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ export

GadgetKeys,
GadgetTypes,
uGadget2,

# RAMSES
read_ramses,
Expand Down Expand Up @@ -83,4 +84,8 @@ include("Houdini.jl")
include("PrettyPrint.jl")
include("Tools.jl")

function __init__()
Unitful.register(AstroIO)
end

end # module
4 changes: 2 additions & 2 deletions src/CSV.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ function write_csv(filename::AbstractString, particles::Union{Array{T,N}, Struct
uAcc = getuAcc(units)
uMass = getuMass(units)
uTime = getuTime(units)
uPotential = getuEnergy(units)
uPotential = getuEnergyUnit(units)
uEntropy = getuEntropy(units)
uDensity2D = getuDensity2D(units)
uPressure = getuPressure(units)
Expand Down Expand Up @@ -66,7 +66,7 @@ function write_csv(filename::AbstractString, particles::Union{Array{T,N}, Struct
uAcc = getuAcc(units)
uMass = getuMass(units)
uTime = getuTime(units)
uPotential = getuEnergy(units)
uPotential = getuEnergyUnit(units)
uEntropy = getuEntropy(units)
uDensity = getuDensity(units)
uPressure = getuPressure(units)
Expand Down
85 changes: 58 additions & 27 deletions src/Gadget.jl
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
@unit Gadget2Mass "1e10M⊙" Gadget2MassUnit 1e10*u"Msun" false
const uGadget2 = [u"kpc",u"kpc/km*s",u"A",u"K",u"cd",Gadget2Mass,u"mol"]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great design!!!


# Header
mutable struct HeaderGadget2
npart::MVector{6,Int32} # gas, halo, disk, bulge, star, blackhole
Expand Down Expand Up @@ -112,6 +115,27 @@ function read_gadget2_header(f::Union{IOStream,Stream{format"Gadget2"}})
return header
end

function read_gadget2_header(filename::AbstractString)
f = open(filename, "r")

temp = read(f, Int32)
if temp == 256
seekstart(f)
header = read_gadget2_header(f)
elseif temp == 8 # Format 2
name = String(read(f, 4))
temp1 = read(f, Int32)
temp2 = read(f, Int32)
header = read_gadget2_header(f)
else
error("Unsupported Gadget2 Format!")
end

close(f)

return header
end

function read_POS!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray, uLength::Units)
temp1 = read(f, Int32)
Pos = data.Pos
Expand All @@ -131,9 +155,9 @@ function read_VEL!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray
temp1 = read(f, Int32)
Vel = data.Vel
for i in 1:length(data)
vx = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
vy = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
vz = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
vx = read(f, Float32) * 1.0 * uVel
vy = read(f, Float32) * 1.0 * uVel
vz = read(f, Float32) * 1.0 * uVel
Comment on lines +158 to +160
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We are assuming the snapshot velocity unit to be km/s (which is not true in some cases), and there need a unit conversion if uVel is not km/s. So I suggest this to be unchanged:

Suggested change
vx = read(f, Float32) * 1.0 * uVel
vy = read(f, Float32) * 1.0 * uVel
vz = read(f, Float32) * 1.0 * uVel
vx = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
vy = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
vz = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s") #TODO: user defined snapshot unit

Even better, km/s should be assignable (TODO)

Vel[i] = PVector(vx, vy, vz)
end
temp2 = read(f, Int32)
Expand Down Expand Up @@ -173,11 +197,11 @@ function read_MASS!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArra
for type in 1:6
if header.mass[type] == 0.0 # read from file
for i in start:tail
Mass[i] = read(f, Float32) * 1.0e10 * uMass
Mass[i] = read(f, Float32) * uMass
end
else # set using header
for i in start:tail
Mass[i] = header.mass[type] * 1.0e10 * uMass
Mass[i] = header.mass[type] * uMass
Comment on lines +200 to +204
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new version only works for uGadget2. I suggest a more general version:

Suggested change
Mass[i] = read(f, Float32) * uMass
end
else # set using header
for i in start:tail
Mass[i] = header.mass[type] * 1.0e10 * uMass
Mass[i] = header.mass[type] * uMass
Mass[i] = uconvert(uMass, read(f, Float32) * 1.0e10 * u"Msun")
end
else # set using header
for i in start:tail
Mass[i] = uconvert(uMass, header.mass[type] * 1.0e10 * u"Msun") #TODO: user defined snapshot unit

end
end
start += header.npart[type]
Expand Down Expand Up @@ -210,7 +234,7 @@ function read_Density!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructA
temp1 = read(f, Int32)
Density = data.Density
for i in 1:NumGas
Density[i] = read(f, Float32) * 10e10uDensity
Density[i] = read(f, Float32) * 1.0 * uDensity
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Density[i] = read(f, Float32) * 1.0 * uDensity
Density[i] = uconvert(uDensity, read(f, Float32) * 1.0 * u"Msun/kpc^3") #TODO: user defined snapshot unit

end
temp2 = read(f, Int32)
if temp1 != temp2
Expand All @@ -234,7 +258,7 @@ function read_POT!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray
temp1 = read(f, Int32)
Potential = data.Potential
for i in 1:length(data)
Potential[i] = uconvert(uPot, read(f, Float32) * u"km^2/s^2" * data.Mass[i])
Potential[i] = read(f, Float32) * 1.0 * uPot
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Potential[i] = read(f, Float32) * 1.0 * uPot
Potential[i] = uconvert(uPot, read(f, Float32) * u"km^2/s^2" * data.Mass[i]) #TODO: user defined snapshot unit

end
temp2 = read(f, Int32)
if temp1 != temp2
Expand All @@ -246,9 +270,9 @@ function read_ACCE!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArra
temp1 = read(f, Int32)
Acc = data.Acc
for i in 1:length(data)
accx = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
accy = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
accz = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
accx = read(f, Float32) * 1.0 * uAcc
accy = read(f, Float32) * 1.0 * uAcc
accz = read(f, Float32) * 1.0 * uAcc
Comment on lines +273 to +275
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
accx = read(f, Float32) * 1.0 * uAcc
accy = read(f, Float32) * 1.0 * uAcc
accz = read(f, Float32) * 1.0 * uAcc
accx = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
accy = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
accz = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2") #TODO: user defined snapshot unit

Acc[i] = PVector(accx, accy, accz)
end
temp2 = read(f, Int32)
Expand All @@ -257,7 +281,7 @@ function read_ACCE!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArra
end
end

function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro;
function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changing the default unit may cause errors, so I have to add a TODO mark:

Suggested change
function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2;
function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2;
#TODO check uGadget2 unit

acc = false,
pot = false,
)
Expand All @@ -275,18 +299,18 @@ function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, heade
# Read Gas Internal Energy Block
NumGas = header.npart[1]
if NumGas > 0 && header.flag_entropy_instead_u > 0 && !eof(f)
d = data["gases"]
d = data[data.Collection .== GAS]
elehcim marked this conversation as resolved.
Show resolved Hide resolved

if !eof(f)
read_Entropy!(f, data, NumGas, getuEntropy(units))
read_Entropy!(f, d, NumGas, getuEntropy(units))
elehcim marked this conversation as resolved.
Show resolved Hide resolved
end

if !eof(f)
read_Density!(f, data, NumGas, getuDensity(units))
read_Density!(f, d, NumGas, getuDensity(units))
elehcim marked this conversation as resolved.
Show resolved Hide resolved
end

if !eof(f)
read_HSML!(f, data, NumGas, getuLength(units))
read_HSML!(f, d, NumGas, getuLength(units))
elehcim marked this conversation as resolved.
Show resolved Hide resolved
end
end

Expand All @@ -309,7 +333,7 @@ function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, heade
return data
end

function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro)
function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2)
NumGas = header.npart[1]

data = StructArray(Star(units))
Expand All @@ -318,11 +342,18 @@ function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}
append!(data, StructArray([Star(units, collection = GadgetTypes[k]) for i = 1:header.npart[k]]))
end

name = ""
while !eof(f)
try
temp1 = read(f, Int32)
name = String(read(f, 4))
temp2 = read(f, Int32)
skippoint = read(f, Int32)
catch e
if isa(e, EOFError)
return data
end
end

if name == "POS "
read_POS!(f, data, getuLength(units))
Expand All @@ -333,11 +364,11 @@ function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}
elseif name == "MASS"
read_MASS!(f, data, header, getuMass(units))
elseif name == "RHO "
read_Density!(f, data["gases"], NumGas, getuDensity(units))
read_Density!(f, data[data.Collection.==GAS], NumGas, getuDensity(units))
elehcim marked this conversation as resolved.
Show resolved Hide resolved
elseif name == "HSML"
read_HSML!(f, data["gases"], NumGas, getuLength(units))
read_HSML!(f, data[data.Collection.==GAS], NumGas, getuLength(units))
elehcim marked this conversation as resolved.
Show resolved Hide resolved
elseif name == "POT "
read_POT!(f, data, getuEnergy(units))
read_POT!(f, data, getuEnergyUnit(units))
elseif name == "ACCE"
read_ACCE!(f, data, getuAcc(units))
end
Expand All @@ -347,13 +378,13 @@ function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}
end

"""
read_gadget2(filename::AbstractString, units = uAstro; kw...)
read_gadget2(filename::AbstractString, units = uGadget2; kw...)

# Keywords
- acc::Bool = false : read acceleration data if exist
- pot::Bool = false : read potential data if exist
"""
function read_gadget2(filename::AbstractString, units = uAstro; kw...)
function read_gadget2(filename::AbstractString, units = uGadget2; kw...)
f = open(filename, "r")

temp = read(f, Int32)
Expand All @@ -376,7 +407,7 @@ function read_gadget2(filename::AbstractString, units = uAstro; kw...)
return header, data
end

function read_gadget2_pos_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro)
function read_gadget2_pos_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2)
NumTotal = sum(header.npart)
uLength = getuLength(units)
pos = StructArray(PVector(uLength) for i in 1:NumTotal)
Expand All @@ -396,7 +427,7 @@ function read_gadget2_pos_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, hea
return pos
end

function read_gadget2_pos_kernel_format2(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro)
function read_gadget2_pos_kernel_format2(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2)
while !eof(f)
temp1 = read(f, Int32)
name = String(read(f, 4))
Expand All @@ -414,7 +445,7 @@ function read_gadget2_pos_kernel_format2(f::Union{IOStream,Stream{format"Gadget2
return nothing
end

function read_gadget2_pos(filename::AbstractString, units = uAstro)
function read_gadget2_pos(filename::AbstractString, units = uGadget2)
f = open(filename, "r")

temp = read(f, Int32)
Expand Down Expand Up @@ -551,7 +582,7 @@ function write_MASS_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, data::Uni
type = GadgetTypes[i]
for p in data
if p.Collection == type
write(f, Float32(ustrip(u"Msun", p.Mass) / 1.0e10))
write(f, Float32(ustrip(Gadget2Mass, p.Mass)))
end
end
end
Expand Down Expand Up @@ -612,7 +643,7 @@ function write_POT(f::Union{IOStream,Stream{format"Gadget2"}}, data::Union{Array
for type in GadgetTypes
for p in data
if p.Collection == type
write(f, Float32(ustrip(u"km^2/s^2", p.Potential / p.Mass)))
write(f, Float32(ustrip(u"km^2/s^2", p.Potential)))
end
end
end
Expand Down Expand Up @@ -774,7 +805,7 @@ end
# FileIO API
import FileIO: Stream, File

function load(s::Stream{format"Gadget2"}, units = uAstro)
function load(s::Stream{format"Gadget2"}, units = uGadget2)
header = read_gadget2_header(s)
data = read_gadget2_particle(s, header, units)
return header, data
Expand Down
22 changes: 22 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,35 @@ header, data = read_gadget2("gassphere_littleendian.gadget2") # 1472 gas particl

pos = read_gadget2_pos("gadget2.format2")
@test length(pos) == 1472
@test pos[1] == PVector(-0.07133729010820389*u"kpc", -0.35668644309043884*u"kpc", -0.9273847341537476*u"kpc")

pos = read_gadget2_pos("gassphere_littleendian.gadget2")
@test length(pos) == 1472
@test pos[1] == PVector(-0.07133729010820389*u"kpc", -0.35668644309043884*u"kpc", -0.9273847341537476*u"kpc")

uAcc = getuAcc(uGadget2)
uPot = getuEnergyUnit(uGadget2)
uMass = getuMass(uGadget2)

@test 1.0*uMass == 1e10u"Msun"
@test uAcc == u"km^2*kpc^-1*s^-2"
@test uPot == u"km^2*s^-2"

h, d = read_gadget2("pot_acc.format2.gadget2", acc = true, pot = true)
@test d.Acc[20] == PVector(1216.8760986328125*u"km^2*kpc^-1*s^-2",
868.4943237304688*u"km^2*kpc^-1*s^-2",
874.93798828125*u"km^2*kpc^-1*s^-2")
@test d.Potential[20] == -37.11515808105469uPot
@test d.Mass[20] == 9.99999993922529e-9uMass

write_gadget2_format2("pot_acc.format2.test.gadget2", h, d, acc = true, pot = true)

h, d = read_gadget2("pot_acc.format2.test.gadget2", acc = true, pot = true)
@test d.Acc[20] == PVector(1216.8760986328125*u"km^2*kpc^-1*s^-2",
868.4943237304688*u"km^2*kpc^-1*s^-2",
874.93798828125*u"km^2*kpc^-1*s^-2")
@test d.Potential[20] == -37.11515808105469uPot

@test !iszero(norm(average(d, :Acc)))
@test !iszero(norm(average(d, :Potential)))
end
Expand Down