def cont_scatterplot(data: pd.DataFrame, x: str, y: str, z: str or None,
label: str, cmap: str, size: int or str or None,
fig: plt.Figure, cbar_kwargs: dict, **kwargs):
"""
Scatterplot with continuous label
Parameters
----------
data: Pandas.DataFrame
x: str
y: str
z: str, optional
label: str
cmap: str
size: int or str, optional
fig: Matplotlib.Figure
cbar_kwargs: dict
Keyword arguments passed to colorbar
kwargs:
Additional keyword arguments passed to Matplotlib.Axes.scatter call
Returns
-------
Matplotlib.Axes
"""
if isinstance(size, str):
size = data[size].values
if z is not None:
ax = fig.add_subplot(111, projection="3d")
im = ax.scatter(data[x].values,
data[y].values,
data[z].values,
c=data[label].values,
s=size,
cmap=cmap,
**kwargs)
else:
ax = fig.add_subplot(111)
im = ax.scatter(data[x].values,
data[y].values,
c=data[label].values,
s=size,
cmap=cmap,
**kwargs)
fig.colorbar(im, ax=ax, **cbar_kwargs)
return ax
def draw_graph_edges(figure: plt.Figure, axis: plt.Axes, graph: nx.Graph):
"""draws graph edges from node to node, colored by weight
Parameters
----------
figure: plt.Figure
a matplotlib Figure
axis: plt.Axes
a matplotlib Axes, part of Figure
graphs: nx.Graph
a networkx graph, assumed to have edges formed like
graph.add_edge((0, 1), (0, 2), weight=1.234)
Notes
-----
modifes figure and axis in-place
"""
weights = np.array([i["weight"] for i in graph.edges.values()])
# graph is (row, col), transpose to get (x, y)
segments = []
for edge in graph.edges:
segments.append([edge[0][::-1], edge[1][::-1]])
line_coll = LineCollection(segments,
linestyle='solid',
cmap="plasma",
linewidths=0.3)
line_coll.set_array(weights)
vals = np.concatenate(line_coll.get_segments())
mnvals = vals.min(axis=0)
mxvals = vals.max(axis=0)
ppvals = vals.ptp(axis=0)
buffx = 0.02 * ppvals[0]
buffy = 0.02 * ppvals[1]
line_coll.set_linewidth(0.3 * 512 / ppvals[0])
axis.add_collection(line_coll)
axis.set_xlim(mnvals[0] - buffx, mxvals[0] + buffx)
axis.set_ylim(mnvals[1] - buffy, mxvals[1] + buffy)
# invert yaxis for image-like orientation
axis.invert_yaxis()
axis.set_aspect("equal")
divider = make_axes_locatable(axis)
cax = divider.append_axes("right", size="5%", pad=0.05)
figure.colorbar(line_coll, ax=axis, cax=cax)
axis.set_title("PCC neighbor graph")
def stats_plot(stats: Stats,
fig: plt.Figure,
ax: plt.Axes,
colorbar: bool = False,
label=''):
# norm = None
norm = LogNorm()
im = ax.imshow(stats.hist.T,
norm=norm,
origin='lower',
extent=stats.hist_extent)
if colorbar:
fig.colorbar(im, ax=ax)
ax.plot(stats.x, stats.thresh_pts_upper, 'k', scaley=False)
ax.plot(stats.x, stats.thresh_pts_lower, 'k', scaley=False)
ax.plot(stats.x, stats.thresh_upper(stats.x), 'r', scaley=False)
ax.plot(stats.x, stats.thresh_lower(stats.x), 'r', scaley=False)
ax.set_title(label)
def init_frame(img: ndarray, ratio: float, fig: Figure, ax: Axes,
title) -> Tuple[AxesImage, Colorbar, Text]:
image_slice = get_slice(img, ratio=ratio)
# the bigger alpha, the image would become more black
true_args = dict(vmin=0, vmax=255, cmap="bone", alpha=0.8)
im = ax.imshow(image_slice, animated=True, **true_args)
# im = ax.imshow(image_slice, animated=True)
ax.set_xticks([])
ax.set_yticks([])
title = ax.set_title(title)
cb = fig.colorbar(im, ax=ax)
return im, cb, title
def visualize(self,
initial_point: np.ndarray,
final_point: np.ndarray,
fig: plt.Figure = None,
ax: plt.Axes = None,
show: bool = False,
**kwargs) -> Optional[Tuple[plt.Figure, plt.Axes]]:
"""
Given an initial and final point (and some formatting parameters), visualize the
atmosphere on a matplotlib graph
:param initial_point : np.ndarray, shape (3, )
3-vec corresponding to the starting coordinate of the path in
cartesian coordinates. Used to determine the interval in which
to calculate the atmosphere
:param final_point : np.ndarray, shape (3, )
3-vec corresponding to the starting coordinate of the path in
cartesian coordinates. Used to determine the interval in which
to calculate the atmosphere
:param fig : Figure, optional
If fig and ax are both provided,
display the atmosphere colormap on top of the old axes
otherwise, create a new figure and axes to work with
:param ax : Axes, optional
If fig and ax are both provided,
display the atmosphere colormap on top of the old axes
otherwise, create a new figure and axes to work with
:param show : boolean, optional
If show is true, display the plotted atmosphere immediately and return nothing.
Otherwise, don't display and instead return the computed figure and axes
:param kwargs : dict, optional
Any additional kwargs are passed to the imshow function
:returns : (Figure, Axes), optional
If show is False, return the computed figure and axes, otherwise return nothing.
"""
total_angle = Vector.angle_between(initial_point, final_point)
normal_vec = Vector.unit_vector(np.cross(initial_point, final_point))
point_number = 500
if ax is None or fig is None:
fig, ax = plt.subplots(1, 1, figsize=(6, 4.5))
radii = np.linspace(Constants.EARTH_RADIUS,
Constants.EARTH_RADIUS + 400E3, point_number)
alpha = np.linspace(0, total_angle, point_number)
r_1 = Rotation.from_rotvec(normal_vec * alpha.reshape(-1, 1))
v_1 = r_1.apply(initial_point)
v_1 = Vector.cartesian_to_spherical(v_1)
frequency_grid = np.zeros((point_number, point_number))
for i in range(point_number):
plotted_vecs = np.repeat(v_1[i].reshape(-1, 1),
point_number,
axis=1).T
plotted_vecs[:, 0] = radii
frequency_grid[:, i] = self.plasma_frequency(
Vector.spherical_to_cartesian(plotted_vecs)) / 1E6
image = ax.imshow(
frequency_grid,
cmap='gist_rainbow',
interpolation='bilinear',
origin='lower',
alpha=1,
aspect='auto',
extent=[0, total_angle * Constants.EARTH_RADIUS / 1000, 0, 400],
**kwargs)
ax.yaxis.set_ticks_position('both')
color_bar = fig.colorbar(image, ax=ax)
color_bar.set_label("Plasma Frequency (MHz)")
if show:
ax.set_title("Chapman Layers Atmosphere")
plt.show()
else:
return fig, ax
