""" Makes an image of the Santa Barbara cluster. Requires py-sphviewer. Invoke as follows: python3 makeImage.py \ \ \ \ """ import numpy as np import matplotlib.pyplot as plt import h5py import matplotlib from sphviewer.tools import QuickView from matplotlib.patches import Rectangle from typing import Tuple from collections import namedtuple # Set up our simulation data collection to keep stuff together SimulationData = namedtuple( "SimulationData", ["coordinates", "masses", "sph_name", "dark_matter_mass", "swift_name", "boxsize"], ) def latex_float(f): """ Taken from: https://stackoverflow.com/questions/13490292/format-number-using-latex-notation-in-python. Formats a float to LaTeX style. """ float_str = "{0:.2g}".format(f) if "e" in float_str: base, exponent = float_str.split("e") return r"{0} \times 10^{{{1}}}".format(base, int(exponent)) else: return float_str def read_data_from_file(filename: str, part_type=0) -> SimulationData: """ Reads the relevant data from the HDF5 file. """ part_type_name = f"PartType{part_type}" with h5py.File(filename, "r") as file: coordinates, boxsize = move_box(file[f"{part_type_name}/Coordinates"][...].T) masses = file[f"{part_type_name}/Masses"][...] sph_name = file["HydroScheme"].attrs["Scheme"].decode("utf-8") unit_mass = ( float(file["Units"].attrs["Unit mass in cgs (U_M)"]) / 2e33 ) # in M_sun dark_matter_mass = float(file["PartType1/Masses"][0]) * unit_mass code_revision = file["Code"].attrs["Git Revision"].decode("utf-8") swift_name = f"SWIFT {code_revision}" data = SimulationData( coordinates=coordinates, masses=masses, sph_name=sph_name, dark_matter_mass=dark_matter_mass, swift_name=swift_name, boxsize=boxsize, ) return data def move_box(coordinates: np.ndarray) -> np.ndarray: """ Takes the coordinates and moves them in the x-y plane. This moves them 20 code units to the left/right to ensure that the zoomed-out version of the cluster image is nicely shown """ boxsize = np.max(coordinates[0]) coordinates[0] -= 20 coordinates[1] -= 20 coordinates[0] %= boxsize coordinates[1] %= boxsize return coordinates, boxsize def generate_views(data: SimulationData, res=2048) -> Tuple[np.ndarray]: """ Generates the views on the data from py-sphviewer. Returns the overall image for the whole box and then a zoomed region. """ qv_all = QuickView( data.coordinates, data.masses, r="infinity", plot=False, xsize=res, ysize=res, logscale=False, p=0, np=48, ) zoomed_res = (res * 6) // 10 mask = np.logical_and( np.logical_and( data.coordinates[0] > (data.boxsize / 2 - 4 - 20), data.coordinates[0] < (data.boxsize / 2 + 6 - 20), ), np.logical_and( data.coordinates[1] > (data.boxsize / 2 - 3.5 - 20), data.coordinates[1] < (data.boxsize / 2 + 6.5 - 20), ), ) qv_zoomed = QuickView( data.coordinates.T[mask].T, data.masses[mask], r="infinity", plot=False, xsize=zoomed_res, ysize=zoomed_res, logscale=False, np=48, ) return qv_all.get_image(), qv_zoomed.get_image() def create_plot(data: SimulationData, res=2048, cmap="viridis", text_color="white"): """ Creates a figure and axes object and returns them for you to do with what you wish. """ img_all, img_zoomed = generate_views(data, res) fig, ax = plt.subplots(figsize=(8, 8)) # Set up in "image" mode ax.axis("off") fig.subplots_adjust(0, 0, 1, 1) ax.imshow( np.log10(img_all + np.min(img_all[img_all != 0])), origin="lower", extent=[-1, 1, -1, 1], cmap=cmap, ) lower_left = [(-24 / (0.5 * data.boxsize)), (-23.5 / (0.5 * data.boxsize))] zoom_rect = Rectangle( lower_left, 10 / (0.5 * data.boxsize), 10 / (0.5 * data.boxsize), linewidth=2, edgecolor=text_color, facecolor="none", ) ax.add_patch(zoom_rect) # Remove ticks as we want "image mode" ax2 = fig.add_axes([0.35, 0.35, 0.6, 0.6], frame_on=True, xticks=[], yticks=[]) ax2.imshow( np.log10(img_zoomed + np.min(img_zoomed[img_zoomed != 0])), origin="lower", extent=[-1, 1, -1, 1], cmap=cmap, ) # This ugly hack sets the box around the subfigure to be white for child in ax2.get_children(): if isinstance(child, matplotlib.spines.Spine): child.set_color(text_color) child.set_linewidth(2) # Draw lines between boxes # Bottom Right ax.plot( [(-14 / (0.5 * data.boxsize)), 0.9], [(-23.5 / (0.5 * data.boxsize)), -0.3], lw=2, color=text_color, ) # Top Left ax.plot( [(-24 / (0.5 * data.boxsize)), -0.3], [(-13.5 / (0.5 * data.boxsize)), 0.9], lw=2, color=text_color, ) ax.text(0.95, -0.95, data.swift_name, color=text_color, ha="right") formatted_dark_matter_mass = latex_float(data.dark_matter_mass) ax.text( -0.95, 0.95, rf"M$_{{\rm DM}} = {formatted_dark_matter_mass}$ M$_\odot$", color=text_color, va="top", ) ax.text( -0.95, -0.95, data.sph_name + "\n" + r"Santa Barbara Cluster (re-ran from Frenk+ 1999)", color=text_color, ) return fig, ax if __name__ == "__main__": import sys try: filename = sys.argv[1] except IndexError: filename = "santabarbara_0153.hdf5" try: part_type = int(sys.argv[2]) except IndexError: part_type = 0 try: cmap = sys.argv[3] except IndexError: cmap = "viridis" try: text_color = sys.argv[4] except IndexError: text_color = "white" try: res = int(sys.argv[5]) except IndexError: res = 2048 # Read in the data from file try: data = read_data_from_file(filename, part_type) except IndexError: # Must be a dark matter only run part_type = 1 data = read_data_from_file(filename, part_type) # Make the plot fig, ax = create_plot(data, res, cmap, text_color) fig.savefig( f"SantaBarbara_{data.sph_name[:8]}_{cmap}_PartType{part_type}_res{res}.png", dpi=res // 8, )