2022-07-14 15:14:26 +02:00
|
|
|
from pathlib import Path
|
|
|
|
from sys import argv
|
2022-05-25 17:17:24 +02:00
|
|
|
from typing import List
|
|
|
|
|
|
|
|
import h5py
|
|
|
|
import numpy as np
|
|
|
|
from matplotlib import pyplot as plt
|
2022-06-01 12:12:45 +02:00
|
|
|
from matplotlib.axes import Axes
|
2022-05-25 17:17:24 +02:00
|
|
|
from matplotlib.colors import LogNorm
|
|
|
|
from matplotlib.figure import Figure
|
2022-06-01 11:31:42 +02:00
|
|
|
from matplotlib.patches import Circle
|
2022-07-14 15:14:26 +02:00
|
|
|
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
|
2022-05-25 17:17:24 +02:00
|
|
|
|
2022-07-05 15:40:17 +02:00
|
|
|
from halo_vis import Coords
|
2022-07-14 15:14:26 +02:00
|
|
|
from paths import base_dir, vis_datafile, has_1024_simulations
|
2022-06-01 11:31:42 +02:00
|
|
|
from read_vr_files import read_velo_halos
|
2022-07-18 19:27:56 +02:00
|
|
|
from utils import figsize_from_page_fraction, rowcolumn_labels
|
2022-06-01 11:31:42 +02:00
|
|
|
|
|
|
|
|
2022-07-05 15:40:17 +02:00
|
|
|
def coord_to_2d_extent(coords: Coords):
|
|
|
|
radius, X, Y, Z = coords
|
|
|
|
return X - radius, X + radius, Y - radius, Y + radius
|
2022-06-01 11:31:42 +02:00
|
|
|
|
|
|
|
|
2022-07-05 15:40:17 +02:00
|
|
|
def in_area(coords: Coords, xobj, yobj, zobj, factor=1.3) -> bool:
|
|
|
|
radius, xcenter, ycenter, zcenter = coords
|
|
|
|
radius *= factor
|
|
|
|
return (
|
|
|
|
(xcenter - radius < xobj < xcenter + radius)
|
|
|
|
and
|
|
|
|
(ycenter - radius < yobj < ycenter + radius)
|
|
|
|
and
|
|
|
|
(zcenter - radius < zobj < zcenter + radius)
|
|
|
|
)
|
2022-06-01 11:31:42 +02:00
|
|
|
|
2022-05-25 17:17:24 +02:00
|
|
|
|
|
|
|
def main():
|
2022-05-30 19:12:59 +02:00
|
|
|
rows = ["shannon", "DB8", "DB4", "DB2"]
|
2022-05-31 17:22:57 +02:00
|
|
|
offset = 2
|
2022-05-25 17:17:24 +02:00
|
|
|
columns = [128, 256, 512]
|
2022-07-14 15:14:26 +02:00
|
|
|
initial_halo_id = int(argv[1])
|
|
|
|
if has_1024_simulations:
|
|
|
|
columns.append(1024)
|
|
|
|
fig: Figure = plt.figure(figsize=figsize_from_page_fraction(columns=2, height_to_width=1))
|
2022-07-05 15:40:17 +02:00
|
|
|
axes: List[List[Axes]] = fig.subplots(len(rows), len(columns), sharex="row", sharey="row")
|
2022-06-01 11:31:42 +02:00
|
|
|
with h5py.File(vis_datafile) as vis_out:
|
2022-07-14 15:14:26 +02:00
|
|
|
halo_group = vis_out[str(initial_halo_id)]
|
|
|
|
|
|
|
|
vmin, vmax = halo_group["vmin_vmax"]
|
2022-05-25 17:17:24 +02:00
|
|
|
print(vmin, vmax)
|
|
|
|
for i, waveform in enumerate(rows):
|
|
|
|
for j, resolution in enumerate(columns):
|
2022-06-01 11:31:42 +02:00
|
|
|
dir = base_dir / f"{waveform}_{resolution}_100"
|
|
|
|
halos = read_velo_halos(dir)
|
2022-05-25 17:17:24 +02:00
|
|
|
ax = axes[i][j]
|
2022-07-14 15:14:26 +02:00
|
|
|
dataset_group = halo_group[f"{waveform}_{resolution}"]
|
|
|
|
rho = np.asarray(dataset_group["rho"])
|
2022-07-05 15:40:17 +02:00
|
|
|
# radius, X, Y, Z
|
2022-07-14 15:14:26 +02:00
|
|
|
coords: Coords = tuple(dataset_group["coords"])
|
|
|
|
mass = dataset_group["mass"][()] # get scalar value from Dataset
|
|
|
|
main_halo_id = dataset_group["halo_id"][()]
|
2022-05-25 17:17:24 +02:00
|
|
|
vmin_scaled = (vmin + offset) * mass
|
|
|
|
vmax_scaled = (vmax + offset) * mass
|
|
|
|
rho = (rho + offset) * mass
|
2022-07-05 15:40:17 +02:00
|
|
|
extent = coord_to_2d_extent(coords)
|
2022-06-01 11:31:42 +02:00
|
|
|
img = ax.imshow(rho.T, norm=LogNorm(vmin=vmin_scaled, vmax=vmax_scaled), extent=extent,
|
|
|
|
origin="lower")
|
2022-07-05 15:40:17 +02:00
|
|
|
found_main_halo = False
|
2022-06-01 11:31:42 +02:00
|
|
|
for halo_id, halo in halos.iterrows():
|
2022-07-14 15:48:34 +02:00
|
|
|
if halo["Mvir"] > 20:
|
2022-07-05 15:40:17 +02:00
|
|
|
if in_area(coords, halo.X, halo.Y, halo.Z):
|
2022-07-14 15:14:26 +02:00
|
|
|
if halo_id == main_halo_id:
|
|
|
|
color = "red"
|
|
|
|
elif halo["Structuretype"] > 10:
|
2022-07-14 15:48:34 +02:00
|
|
|
color = "lightgreen"
|
2022-07-14 15:14:26 +02:00
|
|
|
else:
|
2022-07-14 15:48:34 +02:00
|
|
|
color = "lightblue"
|
2022-06-01 11:31:42 +02:00
|
|
|
if halo_id == main_halo_id:
|
2022-07-05 15:40:17 +02:00
|
|
|
found_main_halo = True
|
2022-06-01 11:31:42 +02:00
|
|
|
print("plotting main halo")
|
|
|
|
circle = Circle(
|
|
|
|
(halo.X, halo.Y),
|
|
|
|
halo["Rvir"], zorder=10,
|
|
|
|
linewidth=1, edgecolor=color, fill=None, alpha=.2
|
|
|
|
)
|
|
|
|
ax.add_artist(circle)
|
2022-07-14 15:14:26 +02:00
|
|
|
# assert found_main_halo
|
2022-05-25 17:17:24 +02:00
|
|
|
print(img)
|
2022-07-14 15:14:26 +02:00
|
|
|
|
|
|
|
# ax.set_axis_off()
|
|
|
|
ax.set_xticks([])
|
|
|
|
ax.set_yticks([])
|
|
|
|
if j == 0:
|
|
|
|
scalebar = AnchoredSizeBar(
|
|
|
|
ax.transData,
|
|
|
|
1, '1 Mpc', 'lower left',
|
|
|
|
# pad=0.1,
|
|
|
|
color='white',
|
|
|
|
frameon=False,
|
|
|
|
# size_vertical=1
|
|
|
|
)
|
|
|
|
ax.add_artist(scalebar)
|
2022-06-01 11:31:42 +02:00
|
|
|
# break
|
|
|
|
# break
|
2022-07-18 19:27:56 +02:00
|
|
|
rowcolumn_labels(axes, rows, isrow=True)
|
|
|
|
rowcolumn_labels(axes, columns, isrow=False)
|
2022-07-14 15:14:26 +02:00
|
|
|
|
2022-05-25 17:17:24 +02:00
|
|
|
fig.tight_layout()
|
|
|
|
fig.subplots_adjust(right=0.825)
|
2022-07-14 15:14:26 +02:00
|
|
|
cbar_ax = fig.add_axes([0.85, 0.05, 0.05, 0.9])
|
2022-05-25 17:17:24 +02:00
|
|
|
fig.colorbar(img, cax=cbar_ax)
|
|
|
|
|
2022-07-14 15:14:26 +02:00
|
|
|
fig.savefig(Path(f"~/tmp/halo_plot_{initial_halo_id}.pdf").expanduser())
|
|
|
|
fig.savefig(Path(f"~/tmp/halo_plot_{initial_halo_id}.png").expanduser())
|
|
|
|
|
2022-05-25 17:17:24 +02:00
|
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
main()
|