# plot_phi_time_averaged.py¶

plot_phi_time_averaged plots the azimuthal and time average of a 2D grid function. The scripts interpolates the quantity over concentric rings around a given center, and average all the values at fixed radius. Then, these are averaged over a time window defined by passing --tmin, --tmax, and --time-every. This last parameter specifies the frequency of snapshots to use for the time average (i.e., for the time average, use one snapshot every N available).

#!/usr/bin/env python3
# PYTHON_ARGCOMPLETE_OK

# Copyright (C) 2021-2022 Gabriele Bozzola
#
# This program is free software; you can redistribute it and/or modify it under
# 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
import numpy as np

from kuibit import argparse_helper as kah
from kuibit.simdir import SimDir
from kuibit.visualize_matplotlib import (
save_from_dir_filename_ext,
set_axis_limits,
setup_matplotlib,
)

if __name__ == "__main__":

desc = """\

{kah.get_program_name()} plots the azimuthal and time average of a given grid
function around a given center. The script interpolates the data onto concentric
rings and takes the average at a fixed radius. Then, this is averaged over a
window of time defined by tmin and tmax."""

parser = kah.init_argparse(desc)

"--variable", type=str, required=True, help="Variable to plot."
)
"--plane",
type=str,
choices=["xy", "xz", "yz"],
default="xy",
help="Plane to consider (default: %(default)s).",
)
"--tmin",
type=float,
required=True,
help="Minimum time in the time window.",
)
"--tmax",
type=float,
required=True,
help="Minimum time in the time window.",
)
"--resolution",
type=int,
default=300,
help="Resolution to use for the x axis.",
)
"--phi-resolution",
type=int,
default=100,
help="Resolution to use for the phi averaging.",
)
"--time-every",
type=int,
default=1,
help="Use one snapshot every this number of available ones (default: %(default)s).",
)
"--logscale", help="Use a logarithmic y scale.", action="store_true"
)
"--absolute",
action="store_true",
help="Whether to take the absolute value.",
)
"--center",
type=float,
nargs=2,
default=(0, 0),
help="Center around which to perform the average (default: %(default)s)",
)
type=float,
default=1,
help="Maximum radius at which perform the average (default: %(default)s)",
)
args = kah.get_args(parser)
setup_matplotlib(rc_par_file=args.mpl_rc_file)

tmin, tmax, var_name, = (
args.tmin,
args.tmax,
args.variable,
)
X0, Y0 = args.center

logger = logging.getLogger(__name__)

if args.verbose:
logging.basicConfig(format="%(asctime)s - %(message)s")
logger.setLevel(logging.DEBUG)

if args.figname is None:
figname = f"{var_name}_{args.plane}_phi_time_averaged"
else:
figname = args.figname

with SimDir(
pickle_file=args.pickle_file,
) as sim:
logger.debug("Prepared SimDir")

times = var.available_times
logger.debug(f"Available time {times}")

selected_times = [time for time in times if tmin <= time <= tmax]
selected_times = selected_times[:: args.time_every]
logger.debug(f"Selected times {selected_times}")

# The equations are
# X = X0 + R * cos(phi)
# Y = Y0 + R * sin(phi)

# We prepare the polar grid that we want evaluate
angles = np.linspace(0, 2 * np.pi, args.phi_resolution)
# R_min is not set to zero because it would be pointless

# These are the points where we are going to evaluate the variable. They
# are a series of concentric rings.
points = np.asarray(
[
[
(X0 + R * np.cos(phi), Y0 + R * np.sin(phi))
for phi in angles
]
]
)

# We are going to collect all the results in values, and the end we do
# the time averaging. So first, we do the phi averaging, then the time
# one. The variable ret_values contains the phi-averaged quantity.
ret_values = np.zeros(args.resolution)

for time in selected_times:
logger.debug(f"Working on time {time}")

var_at_time = var.get_time(time)

values_at_time = var_at_time(points)

if args.logscale:
val_time = np.log10(values_at_time)

# Phi-averaging
phi_averaged_at_time = np.mean(values_at_time, axis=1)

ret_values += phi_averaged_at_time

# Time-averaging
ret_values /= len(selected_times)

if args.logscale:
label = f"log10({var_name})"
else:
label = f"{var_name}"

logger.debug(f"Using label {label}")

logger.debug(f"Plotting variable {var_name}")

rf"$t \in ({selected_times[0]:.3f}, {selected_times[-1]:.3f})$"
)