from nonos.api import GasDataSet
import matplotlib.pyplot as plt

plt.close("all")
# We use GasDataSet which takes as argument the output number of the output file given by idefix/pluto/fargo
# contains in particular a dictionary with the different fields.
ds = GasDataSet(43, geometry="polar", directory="nonos/tests/data/idefix_planet3d")
# We select the GasField "RHO", then
# we perform a vertical slice in the midplane,
# and make the result plotable in the xy plane,
# rotating the grid given the planet number 0
# (which orbit is described in the planet0.dat file).
dsop = ds["RHO"].vertical_at_midplane().map("x", "y", planet_corotation=0)
fig, ax = plt.subplots()
# dsop is now a Plotable object.
# We represent its log10, with a given colormap,
# and we display the colorbar by adding the argument title.
dsop.plot(fig, ax, log=True, cmap="inferno", title=r"$\rho_{\rm mid}$")
ax.set_aspect("equal")

# This time, we perform a latitudinal projection,
# i.e. the integral of "RHO" between -theta and theta,
# and then an azimuthal average,
# before mapping it in the radial ("R") direction.
dsop = ds["RHO"].latitudinal_projection(theta=3*0.05).azimuthal_average().map("R")
fig, ax = plt.subplots()
# We display the y-axis by adding the argument title.
dsop.plot(fig, ax, c="k", title=r"$\Sigma$")
plt.show()

import matplotlib.pyplot as plt
plt.style.use("nonos.default")