Skip to content
GitLab
Menu
Projects
Groups
Snippets
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
SWIFT
SWIFTsim
Commits
ae7944a5
Commit
ae7944a5
authored
Jan 31, 2019
by
Josh Borrow
Committed by
Folkert Nobels
Feb 01, 2019
Browse files
Add rho t plot script
parent
13118abe
Changes
1
Hide whitespace changes
Inline
Side-by-side
examples/SantaBarbara/rhoTPlot.py
0 → 100644
View file @
ae7944a5
"""
Makes a rho-T plot. Uses the swiftsimio library.
"""
import
matplotlib.pyplot
as
plt
import
numpy
as
np
from
swiftsimio
import
SWIFTDataset
,
SWIFTMetadata
,
SWIFTUnits
from
unyt
import
mh
,
cm
,
Gyr
from
tqdm
import
tqdm
from
matplotlib.colors
import
LogNorm
from
matplotlib.animation
import
FuncAnimation
# Constants; these could be put in the parameter file but are rarely changed.
density_bounds
=
[
1e-8
,
1e4
]
# in nh/cm^3
temperature_bounds
=
[
1e2
,
1e8
]
# in K
bins
=
128
# Plotting controls
cmap
=
"viridis"
def
get_data
(
filename
):
"""
Grabs the data (T in Kelvin and density in mh / cm^3).
"""
data
=
SWIFTDataset
(
filename
)
data
.
gas
.
density
.
convert_to_units
(
mh
/
(
cm
**
3
))
data
.
gas
.
temperature
.
convert_to_cgs
()
return
data
.
gas
.
density
,
data
.
gas
.
temperature
def
make_hist
(
filename
,
density_bounds
,
temperature_bounds
,
bins
):
"""
Makes the histogram for filename with bounds as lower, higher
for the bins and "bins" the number of bins along each dimension.
Also returns the edges for pcolormesh to use.
"""
density_bins
=
np
.
logspace
(
np
.
log10
(
density_bounds
[
0
]),
np
.
log10
(
density_bounds
[
1
]),
bins
)
temperature_bins
=
np
.
logspace
(
np
.
log10
(
temperature_bounds
[
0
]),
np
.
log10
(
temperature_bounds
[
1
]),
bins
)
H
,
density_edges
,
temperature_edges
=
np
.
histogram2d
(
*
get_data
(
filename
),
bins
=
[
density_bins
,
temperature_bins
]
)
return
H
.
T
,
density_edges
,
temperature_edges
def
setup_axes
():
"""
Creates the figure and axis object.
"""
fig
,
ax
=
plt
.
subplots
(
1
,
figsize
=
(
6
,
5
),
dpi
=
300
)
ax
.
set_xlabel
(
"Density [$n_H$ cm$^{-3}$]"
)
ax
.
set_ylabel
(
"Temperature [K]"
)
ax
.
loglog
()
return
fig
,
ax
def
make_single_image
(
filename
,
density_bounds
,
temperature_bounds
,
bins
):
"""
Makes a single image and saves it to rhoTPlot_{filename}.png.
Filename should be given _without_ hdf5 extension.
"""
fig
,
ax
=
setup_axes
()
hist
,
d
,
T
=
make_hist
(
"{:s}.hdf5"
.
format
(
filename
),
density_bounds
,
temperature_bounds
,
bins
)
mappable
=
ax
.
pcolormesh
(
d
,
T
,
hist
,
cmap
=
cmap
,
norm
=
LogNorm
())
fig
.
colorbar
(
mappable
,
label
=
"Number of particles"
,
pad
=
0
)
fig
.
tight_layout
()
fig
.
savefig
(
"rhoTPlot_{:s}.png"
.
format
(
filename
))
return
def
make_movie
(
args
,
density_bounds
,
temperature_bounds
,
bins
):
"""
Makes a movie and saves it to rhoTPlot_{stub}.mp4.
"""
fig
,
ax
=
setup_axes
()
def
grab_metadata
(
n
):
filename
=
"{:s}_{:04d}.hdf5"
.
format
(
args
[
"stub"
],
n
)
data
=
SWIFTMetadata
(
filename
)
return
data
def
grab_data
(
n
):
filename
=
"{:s}_{:04d}.hdf5"
.
format
(
args
[
"stub"
],
n
)
H
,
_
,
_
=
make_hist
(
filename
,
density_bounds
,
temperature_bounds
,
bins
)
# Need to ravel because pcolormesh's set_array takes a 1D array. Might
# as well do it here, beacuse 1d arrays are easier to max() than 2d.
return
H
.
ravel
()
histograms
=
[
grab_data
(
n
)
for
n
in
tqdm
(
range
(
args
[
"initial"
],
args
[
"final"
]
+
1
),
desc
=
"Histogramming data"
)
]
metadata
=
[
grab_metadata
(
n
)
for
n
in
tqdm
(
range
(
args
[
"initial"
],
args
[
"final"
]
+
1
),
desc
=
"Grabbing metadata"
)
]
units
=
SWIFTUnits
(
"{:s}_{:04d}.hdf5"
.
format
(
args
[
"stub"
],
args
[
"initial"
]))
# Need to get a reasonable norm so that we don't overshoot.
max_particles
=
max
([
x
.
max
()
for
x
in
histograms
])
norm
=
LogNorm
(
vmin
=
1
,
vmax
=
max_particles
)
# First, let's make the initial frame (we need this for our d, T values that we
# got rid of in grab_data.
hist
,
d
,
T
=
make_hist
(
"{:s}_{:04d}.hdf5"
.
format
(
args
[
"stub"
],
args
[
"initial"
]),
density_bounds
,
temperature_bounds
,
bins
,
)
mappable
=
ax
.
pcolormesh
(
d
,
T
,
hist
,
cmap
=
cmap
,
norm
=
norm
)
fig
.
colorbar
(
mappable
,
label
=
"Number of particles"
,
pad
=
0
)
fig
.
tight_layout
()
# Once we've rearranged the figure with tight_layout(), we can start laing
# Down the metadata text.
def
format_metadata
(
metadata
:
SWIFTMetadata
):
t
=
metadata
.
t
*
units
.
units
[
"Unit time in cgs (U_t)"
]
t
.
convert_to_units
(
Gyr
)
x
=
"$a$: {:2.2f}
\n
$z$: {:2.2f}
\n
$t$: {:2.2f}"
.
format
(
metadata
.
a
,
metadata
.
z
,
t
)
return
x
text
=
ax
.
text
(
0.025
,
0.975
,
format_metadata
(
metadata
[
0
]),
ha
=
"left"
,
va
=
"top"
,
transform
=
ax
.
transAxes
,
)
ax
.
text
(
0.975
,
0.975
,
metadata
[
0
].
code
[
"Git Revision"
].
decode
(
"utf-8"
),
ha
=
"right"
,
va
=
"top"
,
transform
=
ax
.
transAxes
,
)
def
animate
(
data
):
mappable
.
set_array
(
histograms
[
data
])
text
.
set_text
(
format_metadata
(
metadata
[
data
]))
return
mappable
animation
=
FuncAnimation
(
fig
,
animate
,
range
(
len
(
histograms
)),
fargs
=
[],
interval
=
1000
/
25
)
animation
.
save
(
"rhoTPlot_{:s}.mp4"
.
format
(
args
[
"stub"
]))
return
if
__name__
==
"__main__"
:
import
argparse
as
ap
parser
=
ap
.
ArgumentParser
(
description
=
"""
Plotting script for making a rho-T plot.
Takes the filename handle, start, and (optionally) stop
snapshots. If stop is not given, png plot is produced for
that snapshot. If given, a movie is made.
"""
)
parser
.
add_argument
(
"-i"
,
"--initial"
,
help
=
"""Initial snapshot number. Default: 0."""
,
default
=
0
,
required
=
False
,
type
=
int
,
)
parser
.
add_argument
(
"-f"
,
"--final"
,
help
=
"""Final snapshot number. Default: 0."""
,
default
=
0
,
required
=
False
,
type
=
int
,
)
parser
.
add_argument
(
"-s"
,
"--stub"
,
help
=
"""Stub for the filename (e.g. santabarbara). This is
the first part of the filename for the snapshots,
not including the final underscore. Required."""
,
type
=
str
,
required
=
True
,
)
args
=
vars
(
parser
.
parse_args
())
if
args
[
"final"
]
<=
args
[
"initial"
]:
# Run in single image mode.
filename
=
"{:s}_{:04d}"
.
format
(
args
[
"stub"
],
args
[
"initial"
])
make_single_image
(
filename
,
density_bounds
=
density_bounds
,
temperature_bounds
=
temperature_bounds
,
bins
=
bins
,
)
else
:
# Movie mode!
make_movie
(
args
,
density_bounds
=
density_bounds
,
temperature_bounds
=
temperature_bounds
,
bins
=
bins
,
)
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment