Skip to content

Commit

Permalink
Merge pull request #23 from Ambreghisalberti/main
Browse files Browse the repository at this point in the history 
Add the possibility of giving as input a figure and axes of dimension 0 to 2 in plot functions.
  • Loading branch information
@nicolasaunai
nicolasaunai authored Mar 25, 2024
2 parents 62a1971 + fdf0e9b commit 05905a2
Showing 1 changed file with 117 additions and 116 deletions.
233 changes: 117 additions & 116 deletions space/plot/planet_env.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,94 +3,71 @@
from scipy.optimize import root_scalar, root



def _make_figure(**kwargs):
ncols = np.sum([1 for arg in ['x_slice', 'y_slice', 'z_slice'] if arg in kwargs])
if ncols == 0:
ncols = 1
figsize = kwargs.get('figsize', (5 * ncols, 4.5))
fig, ax = plt.subplots(nrows=1, ncols=ncols, figsize=figsize, constrained_layout=True)
if 'color_background' in kwargs:
if isinstance(ax,np.ndarray) :
for i in range(len(ax)) :
ax[i].set_facecolor('xkcd:{}'.format(kwargs['color_background']))
else :
ax.set_facecolor('xkcd:{}'.format(kwargs['color_background']))

return fig, ax

if ('axes' in kwargs) and ('figure' in kwargs):
figure = kwargs.pop("figure")
axes = kwargs.pop("axes")
if isinstance(axes, np.ndarray) is False:
axes = np.array([[axes]])
elif len(axes.shape) == 1:
axes = np.array([axes])
assert axes.shape[1] == ncols, "The axes given as input must have the same length as the asked number of cuts."
print('Figure given as parameter.')
else:
figsize = kwargs.get('figsize', (5 * ncols, 4.5))
figure, axes = plt.subplots(nrows=1, ncols=ncols, figsize=figsize, constrained_layout=True)

def _set_infos_earth_env(fig, ax, **kwargs):
if isinstance(ax, np.ndarray) is False:
ax = [ax]
c = 0
if 'z_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('x_label', 'X (Re)'))
ax[c].set_ylabel(kwargs.get('y_label', 'Y (Re)'))
ax[c].set_xlim(kwargs.get('x_lim', (-20, 15)))
ax[c].set_ylim(kwargs.get('y_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
c += 1
if 'y_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('x_label', 'X (Re)'))
ax[c].set_ylabel(kwargs.get('z_label', 'Z (Re)'))
ax[c].set_xlim(kwargs.get('x_lim', (-20, 15)))
ax[c].set_ylim(kwargs.get('z_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
c += 1
if 'x_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('y_label', 'Y (Re)'))
ax[c].set_ylabel(kwargs.get('z_label', 'Z (Re)'))
ax[c].set_xlim(kwargs.get('y_lim', (-30, 30)))
ax[c].set_ylim(kwargs.get('z_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
if 'title' in kwargs:
fig.suptitle(kwargs.get('title'))
return fig, ax


def _make_figure(**kwargs):
ncols = np.sum([1 for arg in ['x_slice', 'y_slice', 'z_slice'] if arg in kwargs])
if ncols == 0:
ncols = 1
figsize = kwargs.get('figsize', (5 * ncols, 4.5))
fig, ax = plt.subplots(nrows=1, ncols=ncols, figsize=figsize, constrained_layout=True)
return fig, ax

if 'color_background' in kwargs:
if isinstance(axes, np.ndarray):
for i in range(len(axes)):
axes[i].set_facecolor('xkcd:{}'.format(kwargs['color_background']))
else:
axes.set_facecolor('xkcd:{}'.format(kwargs['color_background']))

return figure, axes


def _set_infos_earth_env(fig, axis, **kwargs):
if isinstance(axis, np.ndarray) is False:
axis = np.array([[axis]])
elif len(axis.shape) == 1:
axis = np.array([axis])

for ax in axis:
c = 0
if 'z_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('x_label', 'X (Re)'))
ax[c].set_ylabel(kwargs.get('y_label', 'Y (Re)'))
ax[c].set_xlim(kwargs.get('x_lim', (-20, 15)))
ax[c].set_ylim(kwargs.get('y_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
c += 1

if 'y_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('x_label', 'X (Re)'))
ax[c].set_ylabel(kwargs.get('z_label', 'Z (Re)'))
ax[c].set_xlim(kwargs.get('x_lim', (-20, 15)))
ax[c].set_ylim(kwargs.get('z_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
c += 1

if 'x_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('y_label', 'Y (Re)'))
ax[c].set_ylabel(kwargs.get('z_label', 'Z (Re)'))
ax[c].set_xlim(kwargs.get('y_lim', (-30, 30)))
ax[c].set_ylim(kwargs.get('z_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)

def _set_infos_earth_env(fig, ax, **kwargs):
if isinstance(ax, np.ndarray) is False:
ax = [ax]
c = 0
if 'z_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('x_label', 'X (Re)'))
ax[c].set_ylabel(kwargs.get('y_label', 'Y (Re)'))
ax[c].set_xlim(kwargs.get('x_lim', (-20, 15)))
ax[c].set_ylim(kwargs.get('y_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
c += 1
if 'y_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('x_label', 'X (Re)'))
ax[c].set_ylabel(kwargs.get('z_label', 'Z (Re)'))
ax[c].set_xlim(kwargs.get('x_lim', (-20, 15)))
ax[c].set_ylim(kwargs.get('z_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
c += 1
if 'x_slice' in kwargs:
ax[c].set_xlabel(kwargs.get('y_label', 'Y (Re)'))
ax[c].set_ylabel(kwargs.get('z_label', 'Z (Re)'))
ax[c].set_xlim(kwargs.get('y_lim', (-30, 30)))
ax[c].set_ylim(kwargs.get('z_lim', (-30, 30)))
ax[c].axhline(0, color='k', ls='dotted', alpha=0.4)
ax[c].axvline(0, color='k', ls='dotted', alpha=0.4)
if 'title' in kwargs:
fig.suptitle(kwargs.get('title'))
return fig, ax
return fig, axis


def _find_theta_for_x_slice(boundary_model, phi, **kwargs):
Expand All @@ -100,50 +77,63 @@ def eq(t, x):
n_pts = kwargs.get('n_pts', 300)
xs = np.ones(n_pts) * kwargs.get('x_slice', 0)
x0 = np.ones(
n_pts) * np.pi / 3 # initial guess of theta found empiricly, will enable to find a positive value of theta to make the x_slice with phi between [0;2pi]
n_pts) * np.pi / 3 # initial guess of theta found empirically, will enable to
# find a positive value of theta to make the x_slice with phi between [0;2pi]
return root(eq, args=[phi, xs], x0=x0, jac=False, method='lm').x


def _find_phi_for_z_slice(boundary_model, theta, **kwargs):
def eq(p, x):
if (p <= np.pi) & (
p >= 0): # cos(phi) must be positive because y=r*sin(theta)*cos(phi) and it's easier to get z positive or negative with theta than with phi (multiple roots with cos(roots) positive or negative)
if (p <= np.pi) & (p >= 0):
# cos(phi) must be positive because y=r*sin(theta)*cos(phi) and it's
# easier to get z positive or negative with theta than with phi (multiple roots
# with cos(roots) positive or negative)
return boundary_model(x[0], p, coord_sys='cartesian', **kwargs)[2] - x[1]
else:
return 1000

xs = kwargs.get('z_slice', 0)
phi = np.array([root_scalar(eq, args=[t, xs], x0=np.pi / 3, x1=np.pi / 4, method='secant').root for t in
theta]) # x0, x1 initial guesses of phi that will enable to find a cos(phi) positive. -sign(theta) will allow to have an initial guess closer to the wanted root. Found empirically
phi = np.array([root_scalar(eq, args=[t, xs], x0=np.pi / 3,
x1=np.pi / 4, method='secant').root for t in
theta]) # x0, x1 initial guesses of phi that will enable
# to find a cos(phi) positive. -sign(theta) will allow to have an initial
# guess closer to the wanted root. Found empirically
return phi


def _find_phi_for_y_slice(boundary_model, theta, **kwargs):
def eq(p, x):
if (p <= np.pi / 2) & (
p >= -np.pi / 2): # sin(phi) must be positive because z=r*sin(theta)*sin(phi) and it's simplier to get z positive or negative with theta than with phi (multiple roots with sin(roots) positive or negative)
if (p <= np.pi / 2) & (p >= -np.pi / 2):
# sin(phi) must be positive because z=r*sin(theta)*sin(phi)
# and it's simpler to get z positive or negative with theta than with phi
# (multiple roots with sin(roots) positive or negative)
return boundary_model(x[0], p, coord_sys='cartesian', **kwargs)[1] - x[1]
else:
return 1000

xs = kwargs.get('y_slice', 0)
phi = np.array(
[root_scalar(eq, args=[t, xs], x0=-np.sign(t) * np.pi / 3, x1=-np.sign(t) * np.pi / 4, method='secant').root for
t in theta]) # x0, x1 initial guesses of phi that will enable to find a sin(phi) positive, found empirically
[root_scalar(eq, args=[t, xs], x0=-np.sign(t) * np.pi / 3,
x1=-np.sign(t) * np.pi / 4, method='secant').root for
t in theta]) # x0, x1 initial guesses of phi that will enable to
# find a sin(phi) positive, found empirically
return phi


def _find_theta_lim_y_slice(boundary_model,
**kwargs): # find the smallest value of theta allowed in the considerate y_slice
**kwargs): # find the smallest value of theta allowed in the
# considerate y_slice
def eq(t, a):
return boundary_model(t, a[0], coord_sys='cartesian', **kwargs)[1] - a[1]

phi_p = np.pi / 2 # phi for y positive : r*cos(theta) - y_slice = 0
phi_n = -np.pi / 2 # phi for y negative : -r*cos(theta) - y_slice = 0
t_lim1 = root_scalar(eq, args=[phi_p, kwargs['y_slice']], x0=0.01, x1=np.pi / 3,
method='secant').root # x0, x1 initial guesses of theta that will enable to find the smallest theta giving an y positive
method='secant').root # x0, x1 initial guesses of theta that
# will enable to find the smallest theta giving a y positive
t_lim2 = root_scalar(eq, args=[phi_n, kwargs['y_slice']], x0=0.01, x1=np.pi / 3,
method='secant').root # x0, x1 initial guesses of theta that will enable to find the smallest theta giving an y negative
method='secant').root # x0, x1 initial guesses of theta that
# will enable to find the smallest theta giving a y negative
return t_lim1, t_lim2


Expand All @@ -155,9 +145,11 @@ def eq(t, a):
phi_p = 0 # phi for z positive : r*cos(theta) - z_slice = 0
phi_n = np.pi # phi for z negative : -r*cos(theta) - z_slice = 0
t_lim1 = root_scalar(eq, args=[phi_p, kwargs['z_slice']], x0=0.01, x1=np.pi / 3,
method='secant').root # x0, x1 initial guesses of theta that will enable to find the smallest theta giving a z positive
method='secant').root # x0, x1 initial guesses of theta that
# will enable to find the smallest theta giving a z positive
t_lim2 = root_scalar(eq, args=[phi_n, kwargs['z_slice']], x0=0.01, x1=np.pi / 3,
method='secant').root # x0, x1 initial guesses of theta that will enable to find the smallest theta giving a z negative
method='secant').root # x0, x1 initial guesses of theta that
# will enable to find the smallest theta giving a z negative
return t_lim1, t_lim2


Expand All @@ -179,42 +171,51 @@ def make_theta_and_phi(boundary_model, fct_theta, fct_phi, **kwargs):
def check_validity_of_asked_slice(boundary_model, **kwargs):
theta0 = np.pi / 2 # theta to have x=0 (terminator)
if 'z_slice' in kwargs:
range_phi = np.linspace(-np.pi / 4, np.pi / 4, 100) # range of phi containing the greatest value of z for x=0
range_phi = np.linspace(-np.pi / 4, np.pi / 4, 100)
# range of phi containing the greatest value of z for x=0
z_max = np.max(boundary_model(theta0, range_phi, **kwargs)[2])
z_min = np.min(boundary_model(-theta0, range_phi, **kwargs)[2])
if (kwargs['z_slice'] > z_max) | (kwargs['z_slice'] < z_min):
raise ValueError(f"z_slice value must be between [{round(z_min, 2)},{round(z_max, 2)}] to be able to plot the boudary")
raise ValueError(f"z_slice value must be between [{round(z_min, 2)},{round(z_max, 2)}] to be able to plot "
f"the boundary")

if 'y_slice' in kwargs:
range_phi = np.linspace(np.pi / 4, 3 * np.pi / 4, 100) # range of phi containing the greatest value of y for x=0
range_phi = np.linspace(np.pi / 4, 3 * np.pi / 4, 100)
# range of phi containing the greatest value of y for x=0
y_max = np.max(boundary_model(theta0, range_phi, **kwargs)[1])
y_min = np.min(boundary_model(-theta0, range_phi, **kwargs)[1])
if (kwargs['y_slice'] > y_max) | (kwargs['y_slice'] < y_min):
raise ValueError(f" y_slice value must be between [{round(y_min, 2)},{round(y_max, 2)}] to be able to plot the boundary")
raise ValueError(f" y_slice value must be between [{round(y_min, 2)},{round(y_max, 2)}] to be able to plot "
f"the boundary")


def plot_boundary(boundary_model, fig, axis, **kwargs):
if isinstance(axis, np.ndarray) is False:
axis = np.array([[axis]])
elif len(axis.shape) == 1:
axis = np.array([axis])

def plot_boundary(boundary_model, fig, ax, **kwargs):
c = 0
check_validity_of_asked_slice(boundary_model, **kwargs)
if 'z_slice' in kwargs: # cut z axis to plot xy plane
theta, phi = make_theta_and_phi(boundary_model, _find_theta_lim_z_slice, _find_phi_for_z_slice, **kwargs)
x, y = boundary_model(theta, phi, **kwargs)[:2]
ax[c].plot(x, y, kwargs['style'], alpha= kwargs['alpha'])
c += 1
for ax in axis:
c = 0
check_validity_of_asked_slice(boundary_model, **kwargs)
if 'z_slice' in kwargs: # cut z axis to plot xy plane
theta, phi = make_theta_and_phi(boundary_model, _find_theta_lim_z_slice, _find_phi_for_z_slice, **kwargs)
x, y = boundary_model(theta, phi, **kwargs)[:2]
ax[c].plot(x, y, kwargs['style'], alpha=kwargs['alpha'])
c += 1

if 'y_slice' in kwargs: # cut y axis to plot xz plane
theta, phi = make_theta_and_phi(boundary_model, _find_theta_lim_y_slice, _find_phi_for_y_slice, **kwargs)
x, z = boundary_model(theta, phi, **kwargs)[::2]
ax[c].plot(x, z, kwargs['style'], alpha= kwargs['alpha'])
c += 1
if 'y_slice' in kwargs: # cut y axis to plot xz plane
theta, phi = make_theta_and_phi(boundary_model, _find_theta_lim_y_slice, _find_phi_for_y_slice, **kwargs)
x, z = boundary_model(theta, phi, **kwargs)[::2]
ax[c].plot(x, z, kwargs['style'], alpha=kwargs['alpha'])
c += 1

if 'x_slice' in kwargs: # cut x axis to plot yz plane
theta, phi = make_theta_and_phi(boundary_model, _find_theta_for_x_slice, None, **kwargs)
y, z = boundary_model(theta, phi, **kwargs)[1:]
ax[c].plot(y, z, kwargs['style'], alpha= kwargs['alpha'])
if 'x_slice' in kwargs: # cut x axis to plot yz plane
theta, phi = make_theta_and_phi(boundary_model, _find_theta_for_x_slice, None, **kwargs)
y, z = boundary_model(theta, phi, **kwargs)[1:]
ax[c].plot(y, z, kwargs['style'], alpha=kwargs['alpha'])

return fig, ax
return fig, axis


def layout_earth_env(magnetosheath, **kwargs):
Expand All @@ -228,5 +229,5 @@ def layout_earth_env(magnetosheath, **kwargs):
kwargs['style'] = kwargs.get('style_bs', '-k')
kwargs['alpha'] = kwargs.get('alpha_bs', 1)
fig, ax = plot_boundary(magnetosheath.bow_shock, fig, ax, **kwargs)
plt.tight_layout()
return fig, ax

0 comments on commit 05905a2

Please sign in to comment.