Skip to content

Commit 6efa989

Browse files
islentislent
committed
2 parents c695c52 + c5e784e commit 6efa989

File tree

4 files changed

+90
-28
lines changed

4 files changed

+90
-28
lines changed

‎src/AstroIO.jl‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ export
3838

3939
GadgetKeys,
4040
GadgetTypes,
41+
uGadget2,
4142

4243
# RAMSES
4344
read_ramses,
@@ -83,4 +84,8 @@ include("Houdini.jl")
8384
include("PrettyPrint.jl")
8485
include("Tools.jl")
8586

87+
function __init__()
88+
Unitful.register(AstroIO)
89+
end
90+
8691
end # module

‎src/CSV.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ function write_csv(filename::AbstractString, particles::Union{Array{T,N}, Struct
66
uAcc = getuAcc(units)
77
uMass = getuMass(units)
88
uTime = getuTime(units)
9-
uPotential = getuEnergy(units)
9+
uPotential = getuEnergyUnit(units)
1010
uEntropy = getuEntropy(units)
1111
uDensity2D = getuDensity2D(units)
1212
uPressure = getuPressure(units)
@@ -66,7 +66,7 @@ function write_csv(filename::AbstractString, particles::Union{Array{T,N}, Struct
6666
uAcc = getuAcc(units)
6767
uMass = getuMass(units)
6868
uTime = getuTime(units)
69-
uPotential = getuEnergy(units)
69+
uPotential = getuEnergyUnit(units)
7070
uEntropy = getuEntropy(units)
7171
uDensity = getuDensity(units)
7272
uPressure = getuPressure(units)

‎src/Gadget.jl‎

Lines changed: 56 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1+
@unit Gadget2Mass "1e10M⊙" Gadget2MassUnit 1e10*u"Msun" false
2+
const uGadget2 = [u"kpc",u"kpc/km*s",u"A",u"K",u"cd",Gadget2Mass,u"mol"]
3+
14
# Header
25
mutable struct HeaderGadget2
3-
npart::MVector{6,Int32} # gas, halo, disk, Bulge, star, blackholw
6+
npart::MVector{6,Int32} # gas, halo, disk, bulge, star, blackhole
47
mass::MVector{6,Float64}
58

69
time::Float64
@@ -112,6 +115,27 @@ function read_gadget2_header(f::Union{IOStream,Stream{format"Gadget2"}})
112115
return header
113116
end
114117

118+
function read_gadget2_header(filename::AbstractString)
119+
f = open(filename, "r")
120+
121+
temp = read(f, Int32)
122+
if temp == 256
123+
seekstart(f)
124+
header = read_gadget2_header(f)
125+
elseif temp == 8 # Format 2
126+
name = String(read(f, 4))
127+
temp1 = read(f, Int32)
128+
temp2 = read(f, Int32)
129+
header = read_gadget2_header(f)
130+
else
131+
error("Unsupported Gadget2 Format!")
132+
end
133+
134+
close(f)
135+
136+
return header
137+
end
138+
115139
function read_POS!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray, uLength::Units)
116140
temp1 = read(f, Int32)
117141
Pos = data.Pos
@@ -131,9 +155,9 @@ function read_VEL!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray
131155
temp1 = read(f, Int32)
132156
Vel = data.Vel
133157
for i in 1:length(data)
134-
vx = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
135-
vy = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
136-
vz = uconvert(uVel, read(f, Float32) * 1.0 * u"km/s")
158+
vx = read(f, Float32) * 1.0 * uVel
159+
vy = read(f, Float32) * 1.0 * uVel
160+
vz = read(f, Float32) * 1.0 * uVel
137161
Vel[i] = PVector(vx, vy, vz)
138162
end
139163
temp2 = read(f, Int32)
@@ -173,11 +197,11 @@ function read_MASS!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArra
173197
for type in 1:6
174198
if header.mass[type] == 0.0 # read from file
175199
for i in start:tail
176-
Mass[i] = read(f, Float32) * 1.0e10 * uMass
200+
Mass[i] = read(f, Float32) * uMass
177201
end
178202
else # set using header
179203
for i in start:tail
180-
Mass[i] = header.mass[type] * 1.0e10 * uMass
204+
Mass[i] = header.mass[type] * uMass
181205
end
182206
end
183207
start += header.npart[type]
@@ -210,7 +234,7 @@ function read_Density!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructA
210234
temp1 = read(f, Int32)
211235
Density = data.Density
212236
for i in 1:NumGas
213-
Density[i] = read(f, Float32) * 10e10uDensity
237+
Density[i] = read(f, Float32) * 1.0 * uDensity
214238
end
215239
temp2 = read(f, Int32)
216240
if temp1 != temp2
@@ -220,7 +244,7 @@ end
220244

221245
function read_HSML!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray, NumGas::Int32, uLength::Units)
222246
temp1 = read(f, Int32)
223-
HSML = data.HSML
247+
HSML = data.Hsml
224248
for i in 1:NumGas
225249
HSML[i] = read(f, Float32) * 1.0 * uLength
226250
end
@@ -234,7 +258,7 @@ function read_POT!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArray
234258
temp1 = read(f, Int32)
235259
Potential = data.Potential
236260
for i in 1:length(data)
237-
Potential[i] = uconvert(uPot, read(f, Float32) * u"km^2/s^2" * data.Mass[i])
261+
Potential[i] = read(f, Float32) * 1.0 * uPot
238262
end
239263
temp2 = read(f, Int32)
240264
if temp1 != temp2
@@ -246,9 +270,9 @@ function read_ACCE!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArra
246270
temp1 = read(f, Int32)
247271
Acc = data.Acc
248272
for i in 1:length(data)
249-
accx = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
250-
accy = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
251-
accz = uconvert(uAcc, read(f, Float32) * 1.0 * u"km^2/kpc/s^2")
273+
accx = read(f, Float32) * 1.0 * uAcc
274+
accy = read(f, Float32) * 1.0 * uAcc
275+
accz = read(f, Float32) * 1.0 * uAcc
252276
Acc[i] = PVector(accx, accy, accz)
253277
end
254278
temp2 = read(f, Int32)
@@ -257,7 +281,7 @@ function read_ACCE!(f::Union{IOStream,Stream{format"Gadget2"}}, data::StructArra
257281
end
258282
end
259283

260-
function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro;
284+
function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2;
261285
acc = false,
262286
pot = false,
263287
)
@@ -275,7 +299,6 @@ function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, heade
275299
# Read Gas Internal Energy Block
276300
NumGas = header.npart[1]
277301
if NumGas > 0 && header.flag_entropy_instead_u > 0 && !eof(f)
278-
d = data["gases"]
279302

280303
if !eof(f)
281304
read_Entropy!(f, data, NumGas, getuEntropy(units))
@@ -309,7 +332,7 @@ function read_gadget2_particle(f::Union{IOStream,Stream{format"Gadget2"}}, heade
309332
return data
310333
end
311334

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

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

344+
name = ""
321345
while !eof(f)
346+
try
322347
temp1 = read(f, Int32)
323348
name = String(read(f, 4))
324349
temp2 = read(f, Int32)
325350
skippoint = read(f, Int32)
351+
catch e
352+
if isa(e, EOFError)
353+
return data
354+
end
355+
end
326356

327357
if name == "POS "
328358
read_POS!(f, data, getuLength(units))
@@ -333,11 +363,11 @@ function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}
333363
elseif name == "MASS"
334364
read_MASS!(f, data, header, getuMass(units))
335365
elseif name == "RHO "
336-
read_Density!(f, data["gases"], NumGas, getuDensity(units))
366+
read_Density!(f, data, NumGas, getuDensity(units))
337367
elseif name == "HSML"
338-
read_HSML!(f, data["gases"], NumGas, getuLength(units))
368+
read_HSML!(f, data, NumGas, getuLength(units))
339369
elseif name == "POT "
340-
read_POT!(f, data, getuEnergy(units))
370+
read_POT!(f, data, getuEnergyUnit(units))
341371
elseif name == "ACCE"
342372
read_ACCE!(f, data, getuAcc(units))
343373
end
@@ -347,13 +377,13 @@ function read_gadget2_particle_format2(f::Union{IOStream,Stream{format"Gadget2"}
347377
end
348378

349379
"""
350-
read_gadget2(filename::AbstractString, units = uAstro; kw...)
380+
read_gadget2(filename::AbstractString, units = uGadget2; kw...)
351381
352382
# Keywords
353383
- acc::Bool = false : read acceleration data if exist
354384
- pot::Bool = false : read potential data if exist
355385
"""
356-
function read_gadget2(filename::AbstractString, units = uAstro; kw...)
386+
function read_gadget2(filename::AbstractString, units = uGadget2; kw...)
357387
f = open(filename, "r")
358388

359389
temp = read(f, Int32)
@@ -376,7 +406,7 @@ function read_gadget2(filename::AbstractString, units = uAstro; kw...)
376406
return header, data
377407
end
378408

379-
function read_gadget2_pos_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro)
409+
function read_gadget2_pos_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2)
380410
NumTotal = sum(header.npart)
381411
uLength = getuLength(units)
382412
pos = StructArray(PVector(uLength) for i in 1:NumTotal)
@@ -396,7 +426,7 @@ function read_gadget2_pos_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, hea
396426
return pos
397427
end
398428

399-
function read_gadget2_pos_kernel_format2(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uAstro)
429+
function read_gadget2_pos_kernel_format2(f::Union{IOStream,Stream{format"Gadget2"}}, header::HeaderGadget2, units = uGadget2)
400430
while !eof(f)
401431
temp1 = read(f, Int32)
402432
name = String(read(f, 4))
@@ -414,7 +444,7 @@ function read_gadget2_pos_kernel_format2(f::Union{IOStream,Stream{format"Gadget2
414444
return nothing
415445
end
416446

417-
function read_gadget2_pos(filename::AbstractString, units = uAstro)
447+
function read_gadget2_pos(filename::AbstractString, units = uGadget2)
418448
f = open(filename, "r")
419449

420450
temp = read(f, Int32)
@@ -551,7 +581,7 @@ function write_MASS_kernel(f::Union{IOStream,Stream{format"Gadget2"}}, data::Uni
551581
type = GadgetTypes[i]
552582
for p in data
553583
if p.Collection == type
554-
write(f, Float32(ustrip(u"Msun", p.Mass) / 1.0e10))
584+
write(f, Float32(ustrip(Gadget2Mass, p.Mass)))
555585
end
556586
end
557587
end
@@ -612,7 +642,7 @@ function write_POT(f::Union{IOStream,Stream{format"Gadget2"}}, data::Union{Array
612642
for type in GadgetTypes
613643
for p in data
614644
if p.Collection == type
615-
write(f, Float32(ustrip(u"km^2/s^2", p.Potential / p.Mass)))
645+
write(f, Float32(ustrip(u"km^2/s^2", p.Potential)))
616646
end
617647
end
618648
end
@@ -774,7 +804,7 @@ end
774804
# FileIO API
775805
import FileIO: Stream, File
776806

777-
function load(s::Stream{format"Gadget2"}, units = uAstro)
807+
function load(s::Stream{format"Gadget2"}, units = uGadget2)
778808
header = read_gadget2_header(s)
779809
data = read_gadget2_particle(s, header, units)
780810
return header, data

‎test/runtests.jl‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,15 +23,42 @@ header, data = read_gadget2("gassphere_littleendian.gadget2") # 1472 gas particl
2323
h, d = read_gadget2("gadget2.format2")
2424
@test h.npart[1] == 1472
2525

26+
# getindex
27+
for i in instances(Collection)
28+
@test length(d[i]) == h.npart[Int(i)]
29+
end
30+
2631
pos = read_gadget2_pos("gadget2.format2")
2732
@test length(pos) == 1472
33+
@test pos[1] == PVector(-0.07133729010820389*u"kpc", -0.35668644309043884*u"kpc", -0.9273847341537476*u"kpc")
2834

2935
pos = read_gadget2_pos("gassphere_littleendian.gadget2")
3036
@test length(pos) == 1472
37+
@test pos[1] == PVector(-0.07133729010820389*u"kpc", -0.35668644309043884*u"kpc", -0.9273847341537476*u"kpc")
38+
39+
uAcc = getuAcc(uGadget2)
40+
uPot = getuEnergyUnit(uGadget2)
41+
uMass = getuMass(uGadget2)
42+
43+
@test 1.0*uMass == 1e10u"Msun"
44+
@test uAcc == u"km^2*kpc^-1*s^-2"
45+
@test uPot == u"km^2*s^-2"
3146

3247
h, d = read_gadget2("pot_acc.format2.gadget2", acc = true, pot = true)
48+
@test d.Acc[20] == PVector(1216.8760986328125*u"km^2*kpc^-1*s^-2",
49+
868.4943237304688*u"km^2*kpc^-1*s^-2",
50+
874.93798828125*u"km^2*kpc^-1*s^-2")
51+
@test d.Potential[20] == -37.11515808105469uPot
52+
@test d.Mass[20] == 9.99999993922529e-9uMass
53+
3354
write_gadget2_format2("pot_acc.format2.test.gadget2", h, d, acc = true, pot = true)
55+
3456
h, d = read_gadget2("pot_acc.format2.test.gadget2", acc = true, pot = true)
57+
@test d.Acc[20] == PVector(1216.8760986328125*u"km^2*kpc^-1*s^-2",
58+
868.4943237304688*u"km^2*kpc^-1*s^-2",
59+
874.93798828125*u"km^2*kpc^-1*s^-2")
60+
@test d.Potential[20] == -37.11515808105469uPot
61+
3562
@test !iszero(norm(average(d, :Acc)))
3663
@test !iszero(norm(average(d, :Potential)))
3764
end

0 commit comments

Comments
 (0)