plot_ah_trajectories.pyΒΆ
plot_ah_trajectories plots the coordinate trajectories on a plane. It can
also draw the outline of the horizons and force the center of mass to be in the
origin of the coordinate system.
#!/usr/bin/env python3
# Copyright (C) 2020-2024 Gabriele Bozzola
#
# This program is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free Software
# Foundation; either version 3 of the License, or (at your option) any later
# version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, see <https://www.gnu.org/licenses/>.
import logging
import matplotlib.pyplot as plt
from kuibit import argparse_helper as kah
from kuibit.series import sample_common
from kuibit.simdir import SimDir
from kuibit.visualize_matplotlib import (
add_text_to_corner,
get_figname,
plot_horizon_on_plane_at_time,
save_from_dir_filename_ext,
setup_matplotlib,
)
if __name__ == "__main__":
desc = f"""{kah.get_program_name()} plots the trajectories of given apparent horizons on
a plane. Optionally, it also plots the outline of the horizons."""
parser = kah.init_argparse(desc)
kah.add_figure_to_parser(parser)
parser.add_argument(
"--plane",
type=str,
choices=["xy", "xz", "yz"],
default="xy",
help="Plane to plot (default: %(default)s)",
)
parser.add_argument(
"-a",
"--horizons",
type=int,
required=True,
help="Apparent horizons to consider.",
nargs="+",
)
parser.add_argument(
"--draw-ah",
help="Draw the outlines of the apparent horizons.",
action="store_true",
)
parser.add_argument(
"--force-com-at-origin",
help="Force the Newtonian center of mass to be in the origin "
"(this uses the irreducible masses). In this case, everything "
"will be resampled to times that are common across the horizons.",
action="store_true",
)
args = kah.get_args(parser)
setup_matplotlib(rc_par_file=args.mpl_rc_file)
# Parse arguments
logger = logging.getLogger(__name__)
if args.verbose:
logging.basicConfig(format="%(asctime)s - %(message)s")
logger.setLevel(logging.DEBUG)
# For the figure name, we chain all the horizon numbers
horizon_name = "_".join([str(h) for h in args.horizons])
figname = get_figname(
args, default=f"ah_{horizon_name}_trajectory_{args.plane}"
)
logger.debug(f"Figname: {figname}")
with SimDir(
args.datadir,
ignore_symlinks=args.ignore_symlinks,
pickle_file=args.pickle_file,
) as sim:
logger.debug("Prepared SimDir")
sim_hor = sim.horizons
logger.debug(
f"Apparent horizons available: {sim_hor.available_apparent_horizons}"
)
# Check that the horizons are available
for ah in args.horizons:
if ah not in sim_hor.available_apparent_horizons:
raise ValueError(f"Apparent horizons {ah} is not available")
# Now, we prepare the list with the centroids. The keys are the
# horizon numbers and the values are TimeSeries. We prepare this so that
# we can compute the center of mass. For that, we also need the masses
# (x_cm = sum m_i / M x_i). We use the irreducible mass for this.
ah_coords = {"x": [], "y": [], "z": []}
masses = []
# If force_com_at_origin, these objects will be rewritten so that we can
# assume that they house the quantities we want to plot.
for ah in args.horizons:
logger.debug(f"Reading horizon {ah}")
hor = sim_hor.get_apparent_horizon(ah)
for coord, ah_coord in ah_coords.items():
ah_coord.append(hor.ah[f"centroid_{coord}"])
if args.force_com_at_origin:
logger.debug("Reading mass")
masses.append(hor.ah.m_irreducible)
if args.force_com_at_origin:
# x_cm = sum m_i / M x_i
logger.debug("Computing center of mass")
logger.debug("Resampling to common times")
# We have to resample everything to a common time interval. This is
# because we are going to combine the various objects with
# mathematical operations.
# Loop over the three coordinates and overwrite the list of
# TimeSeries (note that ahs here are lists of TimeSeries and
# sample_common(ahs) returns a new list of TimeSeries)
ah_coords = {
coord: sample_common(ahs) for coord, ahs in ah_coords.items()
}
masses = sample_common(masses)
# First, we compute the total mass (as TimeSeries)
total_mass = sum(mass for mass in masses)
# Loop over the three coordinates
for coord, ah_coord in ah_coords.items():
# For each coordinate, compute the center of mass along that
# coordinate
com = sum(
mass * ah / total_mass
for mass, ah in zip(masses, ah_coords[coord])
)
# Now, we update ah_coords over that coordinate by subtracting
# the center of mass from each apparent horizon
ah_coord = [ah - com for ah in ah_coords[coord]]
to_plot_x, to_plot_y = args.plane
logger.debug(f"Plotting on the x axis {to_plot_x}")
logger.debug(f"Plotting on the y axis {to_plot_y}")
# Now we loop over all the horizons
for ind, ah in enumerate(args.horizons):
plt.plot(
ah_coords[to_plot_x][ind].y,
ah_coords[to_plot_y][ind].y,
label=f"Horizon {ah}",
)
# We save the time to plot the horizon outline
time = ah_coords[to_plot_x][ind].tmax
# Try to draw the shape of the horizon
if args.draw_ah:
logger.debug(f"Drawing shape at time {time} for ah {ah}")
plot_horizon_on_plane_at_time(
sim_hor.get_apparent_horizon(ah), time, args.plane
)
if args.force_com_at_origin:
xlabel = f"{to_plot_x} - {to_plot_x}_CM"
ylabel = f"{to_plot_y} - {to_plot_y}_CM"
else:
xlabel, ylabel = to_plot_x, to_plot_y
plt.xlabel(xlabel)
plt.ylabel(ylabel)
plt.gca().set_aspect("equal")
plt.legend()
add_text_to_corner(rf"$t = {time:.3f}$")
logger.debug("Saving")
save_from_dir_filename_ext(
args.outdir,
figname,
args.fig_extension,
tikz_clean_figure=args.tikz_clean_figure,
)
logger.debug("DONE")
