def draw_samples(ax: plt.axis,
points,
p_val=None,
detection_flag: bool = False,
detection_threshold: float = -140,
*args,
**kwargs):
scat = None
if detection_flag is True:
p_val = p_val[:, 0]
p_val[p_val < detection_threshold] = detection_threshold
if p_val is None:
scat = ax.scatter(points[:, 0],
points[:, 1],
c='r',
marker='x',
alpha=0.4,
*args,
**kwargs)
else:
scat = ax.scatter(points[:, 0],
points[:, 1],
c=p_val,
norm=colors.PowerNorm(gamma=2),
marker='x',
alpha=0.4,
*args,
**kwargs)
return scat
def decision_boundary(self, x, y, name: str = '', ax: plt.axis = None):
if ax is None:
fig, ax = plt.subplots(figsize=(10, 10))
h = 0.02
x_min, x_max = x[:, 0].min() - 1, x[:, 0].max() + 1
y_min, y_max = x[:, 1].min() - 1, x[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
np.arange(y_min, y_max, h))
print(np.c_[xx.ravel(), yy.ravel()].shape)
z = self.predict(np.c_[xx.ravel(), yy.ravel()])
z = z.reshape(xx.shape)
ax.contourf(xx, yy, z, cmap=plt.cm.coolwarm, alpha=0.8)
ax.scatter(x[:, 0], x[:, 1], c=y, cmap=plt.cm.coolwarm)
ax.set_xlim(xx.min(), xx.max())
ax.set_ylim(yy.min(), yy.max())
ax.set_xticks(())
ax.set_yticks(())
if name:
plt.savefig(os.path.join(os.getcwd(), name + '.png'))
plt.clf()
else:
plt.show()
plt.clf()
def draw(self, ax: plt.axis):
ax.scatter(
*self.p0,
s=200,
lw=1,
ec=[0.3, 0.3, 0.3],
color=self.color,
zorder=100,
)
if self.last:
ax.scatter(
*self.p1,
s=200,
lw=1,
ec=[0.3, 0.3, 0.3],
color=self.color,
zorder=100,
)
ax.plot(
[self.p0[0], self.p1[0]],
[self.p0[1], self.p1[1]],
lw=6,
color=[0.3, 0.3, 0.3],
zorder=98,
)
ax.plot(
[self.p0[0], self.p1[0]],
[self.p0[1], self.p1[1]],
lw=5,
color=self.color,
zorder=99,
)
def decision_boundary(self, x, y, ax: plt.axis = None):
"""
Plot decision boundary and labeled data
:param x: data
:param y: true labels
:param ax: matplotlib axes (optional)
:return: None
"""
if ax is None:
fig, ax = plt.subplots(figsize=(10, 10))
h = 0.02
x_min, x_max = x[:, 0].min() - 1, x[:, 0].max() + 1
y_min, y_max = x[:, 1].min() - 1, x[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
np.arange(y_min, y_max, h))
z = self.predict(np.c_[xx.ravel(), yy.ravel()])
z = z.reshape(xx.shape)
ax.contourf(xx, yy, z, cmap=plt.cm.coolwarm, alpha=0.8)
ax.scatter(x[:, 0], x[:, 1], c=y, cmap=plt.cm.coolwarm)
ax.set_xlim(xx.min(), xx.max())
ax.set_ylim(yy.min(), yy.max())
ax.set_xticks(())
ax.set_yticks(())
plt.show()
def draw(
self,
ax: plt.axis = None,
tracking: dict = None,
colorby: str = "segment",
):
ax = ax or plt.subplots(figsize=(9, 9))[1]
# draw each segment
for segment in self.segments:
segment.draw(ax)
ax.scatter(
self.trace[:, 0],
self.trace[:, 1],
s=20,
zorder=101,
color=[0.4, 0.4, 0.4],
)
# draw tracking data
if tracking is not None:
if colorby == "segment":
c = self.colors_from_segment(tracking["segment"])
else:
c = self.colors_from_global_coordinates(
tracking["global_coord"])
ax.scatter(
tracking["x"][::5],
tracking["y"][::5],
c=c[::5],
s=50,
alpha=0.8,
)
def regplot(
data: Union[pd.DataFrame, pd.Series, dict],
x: str,
y: str,
ax: plt.axis,
scatter_sample: int = 10,
**kwargs,
):
ax.scatter(data[x][::scatter_sample], data[y][::scatter_sample], **kwargs)
def plot_bouts_1d(
tracking: Union[dict, pd.DataFrame],
bouts: pd.DataFrame,
ax: plt.axis,
direction: bool = None,
zorder: int = 100,
lw: float = 2,
alpha: float = 1,
**kwargs,
):
# select bouts by direction
if direction is not None:
bouts = bouts.loc[bouts.direction == direction]
# get coords
x = tracking["global_coord"]
y = np.linspace(1, 0, len(x))
# plot
for i, bout in bouts.iterrows():
_x = x[bout.start_frame:bout.end_frame]
_y = y[bout.start_frame:bout.end_frame]
ax.plot(
_x,
_y,
color=colors.bout_direction_colors[bout.direction],
zorder=zorder,
lw=lw,
alpha=alpha,
**kwargs,
)
ax.scatter(
_x[0],
_y[0],
color="white",
lw=1,
ec=colors.bout_direction_colors[bout.direction],
s=30,
zorder=101,
alpha=0.85,
**kwargs,
)
ax.scatter(
_x[-1],
_y[-1],
color=[0.2, 0.2, 0.2],
lw=1,
ec=colors.bout_direction_colors[bout.direction],
s=30,
zorder=101,
alpha=0.85,
**kwargs,
)
def plot_tracking_linearized(
tracking: Union[dict, pd.DataFrame],
ax: plt.axis = None,
plot: bool = True,
**kwargs,
):
ax = ax or plt.subplots(figsize=(9, 9))[1]
x = tracking["global_coord"]
y = np.linspace(1, 0, len(x))
if not plot:
ax.scatter(x, y, **kwargs)
else:
ax.plot(x, y, **kwargs)
def plot_probe_electrodes(
rsites: pd.DataFrame,
ax: plt.axis,
TARGETS: list = [],
annotate_every: Union[int, bool] = 5,
x_shift: bool = True,
s: int = 25,
lw: float = 0.25,
x_pos_delta: float = 0,
):
x = np.ones(len(rsites)) * 1.025 + x_pos_delta
if x_shift:
x[::2] = 1.025 + x_pos_delta - 0.05
x[2::4] = 1.025 + x_pos_delta - 0.025
x[1::4] = 1.025 + x_pos_delta + 0.025
else:
x = (x_pos_delta +
np.tile([0.975, 1.025, 1.0, 1.05], np.int(np.ceil(
len(rsites) / 4)))[:len(rsites)])
if TARGETS is not None:
colors = [
rs.color if rs.brain_region in TARGETS else
([0.3, 0.3, 0.3] if rs.brain_region in ("unknown",
"OUT") else blue_grey)
for i, rs in rsites.iterrows()
]
else:
colors = [rs.color for i, rs in rsites.iterrows()]
ax.scatter(
x,
rsites.probe_coordinates,
s=s,
lw=lw,
ec=[0.3, 0.3, 0.3],
marker="s",
c=colors,
)
if annotate_every:
for i in np.arange(len(x))[::annotate_every]:
ax.annotate(
f"{rsites.site_id.iloc[i]} - {rsites.brain_region.iloc[i]}",
(0.6, rsites.probe_coordinates.iloc[i]),
color=colors[i],
)
ax.set(xlim=[0.5, 1.25], ylabel="probe coordinates (um)")
def plot_raster(
spikes: np.ndarray,
events: Union[np.ndarray, list],
ax: plt.axis,
window: int = 120,
s=5,
color=grey_darker,
kde: bool = True,
bw: int = 6,
**kwargs,
):
"""
Plots a raster plot of spikes aligned to timestamps events
It assumes that all event and spike times are in frames and framerate is 60
"""
half_window = window / 2
yticks_step = int(np.ceil(len(events) / 8)) if len(events) > 8 else 2
X = []
for n, event in enumerate(events):
event_spikes = (spikes[(spikes >= event - half_window)
& (spikes <= event + half_window)] - event)
X.extend(list(event_spikes))
y = np.ones_like(event_spikes) * n
ax.scatter(event_spikes, y, s=5, color=color, **kwargs)
ax.axvline(0, ls=":", color="k", lw=0.75)
# plot KDE
if kde:
raise NotImplementedError("KDE env setup incorrect")
# plot_kde(
# ax=ax,
# z=-len(events) / 4,
# data=X,
# normto=len(events) / 5,
# color=blue_grey_dark,
# kde_kwargs=dict(bw=bw, cut=0),
# alpha=0.6,
# invert=False,
# )
# set x axis properties
ax.set(
yticks=np.arange(0, len(events), yticks_step),
ylabel="event number",
**get_window_ticks(window),
)
def mpl_plot(self, ax: plt.axis = None, temp_unit: str = "GK", **kwargs):
ax = ax or plt.gca()
if temp_unit is "GK":
ax.errorbar(
Temperature(np.array(self.temperature),
unit=self.temp_unit).gk,
self.rr,
yerr=self.err,
color=self.colour,
label="{0} {1}".format(self.label, self.__str__()),
**kwargs,
)
ax.scatter(
Temperature(np.array(self.temperature),
unit=self.temp_unit).gk,
self.rr,
color=self.colour,
)
elif temp_unit is "KeV":
ax.errorbar(
Temperature(np.array(self.temperature),
unit=self.temp_unit).kev,
self.rr,
yerr=self.err,
color=self.colour,
label="{0} {1}".format(self.label, self.__str__()),
**kwargs,
)
ax.scatter(
Temperature(np.array(self.temperature),
unit=self.temp_unit).kev,
self.rr,
color=self.colour,
)
ax.set_yscale("log")
ax.legend()
ax = super().mpl_plot(ax, temp_unit=temp_unit)
return ax
def plot_balls_errors(
x: np.ndarray,
y: np.ndarray,
yerr: np.ndarray,
ax: plt.axis,
s: int = 150,
colors: Union[list, str] = None,
):
"""
Given a serires of XY values and Y errors it plots a scatter for each XY point and a line
to mark each Y error
"""
if colors is None:
colors = [blue_grey] * len(x)
elif isinstance(colors, str):
colors = [colors] * len(x)
ax.scatter(x, y, s=s, c=colors, zorder=100, lw=1, ec=[0.3, 0.3, 0.3])
ax.plot(x, y, lw=3, color=colors[0], zorder=-1)
if yerr is not None:
for n in range(len(x)):
ax.plot(
[x[n], x[n]],
[y[n] - yerr[n], y[n] + yerr[n]],
lw=4,
color=[0.3, 0.3, 0.3],
zorder=96,
solid_capstyle="round",
)
ax.plot(
[x[n], x[n]],
[y[n] - yerr[n], y[n] + yerr[n]],
lw=2,
color=colors[n],
zorder=98,
solid_capstyle="round",
)
def plot_tracking_xy(
tracking: Union[dict, pd.DataFrame],
key: str = None,
skip_frames: int = 1,
ax: plt.axis = None,
plot: bool = False,
**kwargs,
):
ax = ax or plt.subplots(figsize=(9, 9))[1]
if key is None:
if not plot:
ax.scatter(
tracking["x"][::skip_frames],
tracking["y"][::skip_frames],
color=[0.3, 0.3, 0.3],
**kwargs,
)
else:
ax.plot(
tracking["x"][::skip_frames],
tracking["y"][::skip_frames],
**kwargs,
)
else:
ax.scatter(
tracking["x"][::skip_frames],
tracking["y"][::skip_frames],
c=tracking[key][::skip_frames],
**kwargs,
)
if "orientation" in key or "angle" in key:
# draw arrows to mark the angles/colors mapping
angles = np.linspace(0, 2 * np.pi, 16)
x = 2 * np.cos(angles[::-1] + np.pi / 2) + 25
y = 2 * np.sin(angles + np.pi / 2) + 2
ax.scatter(x,
y,
s=80,
zorder=50,
c=np.degrees(angles),
alpha=1,
**kwargs)
def word_group_visualization(
transformed_word_embeddings: np.ndarray,
words: np.ndarray,
word_groups: dict,
xlabel: str,
ylabel: str,
emphasis_words: list = None,
alpha: float = 1,
non_group_words_color: str = "#ccc",
scatter_set_rasterized: bool = False,
rasterization_threshold: int = 1000,
ax: plt.axis = None,
show_plot: bool = True,
) -> None:
"""
Visualizes one or more word groups by plotting its word embeddings in 2D.
Parameters
----------
transformed_word_embeddings : np.ndarray
Transformed word embeddings.
words : np.ndarray
Numpy array containing all words from vocabulary.
word_groups : dict
Dictionary containing word groups to visualize.
xlabel : str
X-axis label.
ylabel : str
Y-axis label.
emphasis_words : list, optional
List representing words to emphasize in the visualization (defaults to None).
Entries can be either be strings (words) or tuples, consisting of the word, x-offset
and y-offset.
alpha : float
Scatter plot alpha value (defaults to 1).
non_group_words_color : str
Color for words outside groups (defaults to #ccc).
scatter_set_rasterized : bool
Whether or not to enable rasterization on scatter plotting (defaults to False).
rasterization_threshold : int
The least number of data points to enable rasterization, given that
`scatter_set_rasterized` is set to True (defaults to 1000).
ax : plt.axis
Axis (defaults to None).
show_plot : bool
Whether or not to call plt.show() (defaults to True).
"""
# Filter and restrict words in word groups
word_group_words_restricted = {}
for group_key, group_data in word_groups.items():
group_words = group_data["words"]
group_words = np.array([word for word in group_words if word in words])
group_words_indices = np.array(
[np.where(words == word)[0][0] for word in group_words])
group_word_embeddings = transformed_word_embeddings[
group_words_indices]
boundaries = group_data.get("boundaries", {})
if boundaries.get("xmin") is None:
boundaries["xmin"] = group_word_embeddings[:, 0].min()
if boundaries.get("xmax") is None:
boundaries["xmax"] = group_word_embeddings[:, 0].max()
if boundaries.get("ymin") is None:
boundaries["ymin"] = group_word_embeddings[:, 1].min()
if boundaries.get("ymax") is None:
boundaries["ymax"] = group_word_embeddings[:, 1].max()
group_word_embeddings_boundaries_mask = [
(boundaries["xmin"] <= word_vec[0] <= boundaries["xmax"])
and (boundaries["ymin"] <= word_vec[1] <= boundaries["ymax"])
for i, word_vec in enumerate(group_word_embeddings)
]
word_group_words_restricted[group_key] = group_words[
group_word_embeddings_boundaries_mask]
# Find words not in groups
words_not_in_groups_mask = [
i for i, word in enumerate(words)
for group_words in word_group_words_restricted.values()
if word not in group_words
]
if ax is None:
_, ax = plt.subplots(figsize=(12, 7))
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
# Plot non-group words
non_grp_scatter_handle = ax.scatter(
x=transformed_word_embeddings[words_not_in_groups_mask][:, 0],
y=transformed_word_embeddings[words_not_in_groups_mask][:, 1],
s=10,
alpha=alpha,
c=non_group_words_color,
)
if (scatter_set_rasterized
and len(words_not_in_groups_mask) >= rasterization_threshold):
non_grp_scatter_handle.set_rasterized(True)
# Plot group words
for group_key, group_words in word_group_words_restricted.items():
group_words_indices = np.array(
[np.where(words == word)[0][0] for word in group_words])
group_word_embeddings = transformed_word_embeddings[
group_words_indices]
grp_scatter_handle = ax.scatter(
x=group_word_embeddings[:, 0],
y=group_word_embeddings[:, 1],
s=15,
alpha=alpha,
c=word_groups[group_key]["color"],
label=word_groups[group_key]["label"],
)
if (scatter_set_rasterized
and len(group_word_embeddings) >= rasterization_threshold):
grp_scatter_handle.set_rasterized(True)
# Visualize emphasized words
if emphasis_words is not None:
emphasis_words = [(entry, 0, 0) if type(entry) == str else entry
for entry in emphasis_words]
for emphasis_word, x_offset, y_offset in emphasis_words:
word_group_key = None
for group_key, group_data in word_groups.items():
if emphasis_word in group_data["words"]:
word_group_key = group_key
break
if word_group_key is None:
word_color = non_group_words_color
else:
word_color = word_groups[group_key]["color"]
word_idx = [
i for i, word in enumerate(words) if word == emphasis_word
][0]
emphasis_scatter_handle = ax.scatter(
x=transformed_word_embeddings[word_idx, 0],
y=transformed_word_embeddings[word_idx, 1],
s=40,
alpha=alpha,
c=word_color,
)
if (scatter_set_rasterized
and len(emphasis_words) >= rasterization_threshold):
emphasis_scatter_handle.set_rasterized(True)
# Annotate emphasis word with a text box
offsetbox = TextArea(emphasis_word)
ab = AnnotationBbox(
offsetbox,
tuple(transformed_word_embeddings[word_idx]),
xybox=(x_offset, 40 + y_offset),
xycoords="data",
boxcoords="offset points",
arrowprops=dict(arrowstyle="->", color="black", linewidth=2),
)
ax.add_artist(ab)
ax.legend()
if show_plot:
plt.show()
def plot(self, ax: plt.axis):
ax.scatter(self.points[:, 0],
self.points[:, 1],
facecolor=np.array(self.color) / 255.0,
label=self.name)
def plot_word_vectors(
transformed_word_embeddings: np.ndarray,
words: list,
title: str,
x_label: str,
y_label: str,
word_colors: np.ndarray = None,
ax: plt.axis = None,
show_plot: bool = True,
interactive: bool = False,
continuous_word_colors: bool = False,
) -> None:
"""
Plots word vectors in transformed 2D space.
Parameters
----------
transformed_word_embeddings : np.ndarray
Word embeddings transformed into 2D space.
words : list
List of words to plot
title : str,
Title to use for the plot.
x_label : str,
Label to use for the x-axis.
y_label : str
Label to use for the y-axis
word_colors : np.ndarray, optional
Numpy array consisting of unique labels for each word (i.e. cluster labels),
(defaults to None).
ax : plt.axis
Matplotlib axis (defaults to None)
show_plot : bool
Whether or not to call plt.show() (defaults to True)
interactive : bool
Whether or not to make the visualization interactive
using Plotly (defaults to False).
continuous_word_colors : bool
Whether or not to make the word color continuous (defaults to False).
"""
if interactive:
# Plot interactive plot
fig = px.scatter(
x=transformed_word_embeddings[:, 0],
y=transformed_word_embeddings[:, 1],
title=title,
labels={
"x": x_label,
"y": y_label
},
color=[
cluster_label if continuous_word_colors else str(cluster_label)
for cluster_label in word_colors
] if word_colors is not None else None,
hover_data={"word": words},
)
fig.show()
else:
if ax is None:
_, ax = plt.subplots()
ax.scatter(
transformed_word_embeddings[:, 0],
transformed_word_embeddings[:, 1],
c=word_colors,
)
ax.set_title(title)
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
if show_plot:
plt.show()
def plot_cluster_metric_scores(
metric_scores: list,
hyperparameters: list,
best_score_idx: int,
metric_name: str,
scatter: bool = True,
set_xticks: bool = True,
set_xtickslabels: bool = True,
xtickslabels_rotation: int = 90,
ax: plt.axis = None,
xlabel: str = "Hyperparameters",
xrange: range = None,
show_plot: bool = True,
) -> None:
"""
Plots internal cluster validation metric scores
Parameters
----------
metric_scores : list
List of scores computed using metric
hyperparameters : list
List of hyperparameters used to compute the scores
best_score_idx : int
Best score index
metric_name : str
Name of the internal cluster validation metric
scatter : bool
Whether or not to scatter points (defaults to True)
set_xticks : bool
Whether or not to set the ticks on the x-axis
set_xtickslabels : bool
Whether or not to set the labels on the x-axis
xtickslabels_rotation : int
Sets the xticks labels rotation (defaults to 90), set_xtickslabels
must be set to True to have an effect.
ax : plt.axis
Matplotlib axis (defaults to None)
xlabel : str
X-axis label (defaults to "Hyperparameters")
xrange : range
Range to use for the x-axis (default starts from 0 to)
show_plot : bool
Whether or not to call plt.show() (defaults to True)
"""
if ax is None:
_, ax = plt.subplots()
if xrange is None:
xrange = range(len(hyperparameters))
ax.plot(xrange, metric_scores)
if scatter:
ax.scatter(xrange, metric_scores)
ax.scatter(xrange[best_score_idx],
metric_scores[best_score_idx],
c="r",
s=72,
zorder=10)
if set_xticks:
ax.set_xticks(xrange)
if set_xtickslabels:
ax.set_xticklabels(hyperparameters,
rotation=xtickslabels_rotation,
ha="center")
ax.set_xlabel(xlabel)
ax.set_ylabel(f"{metric_name} score")
if show_plot:
plt.tight_layout()
plt.show()
def draw_gw(ax: plt.axis, gw_pos, *args, **kwargs):
return ax.scatter(gw_pos[0], gw_pos[1], c='r', marker='D')
def plot_bouts_2d(
tracking: Union[dict, pd.DataFrame],
bouts: pd.DataFrame,
ax: plt.axis,
direction: bool = None,
zorder: int = 100,
lw: float = 2,
c: str = None,
unit: pd.Series = None,
**kwargs,
):
# select bouts by direction
if direction is not None:
bouts = bouts.loc[bouts.direction == direction]
# plot
for i, bout in bouts.iterrows():
# prepare data
x = tracking["x"][bout.start_frame:bout.end_frame]
y = tracking["y"][bout.start_frame:bout.end_frame]
if c is None:
color = colors.bout_direction_colors[bout.direction]
else:
color = c
# plot tracking
ax.plot(x, y, color=color, zorder=zorder, lw=lw, **kwargs)
if unit is None:
# mark start and end
ax.scatter(
x[0],
y[0],
color="white",
lw=1,
ec=color,
s=25,
zorder=101,
**kwargs,
)
ax.scatter(
x[-1],
y[-1],
color=[0.2, 0.2, 0.2],
lw=1,
ec=color,
s=25,
zorder=101,
**kwargs,
)
else:
# mark unit spikes
spikes = unit.spikes[(unit.spikes > bout.start_frame)
& (unit.spikes < bout.end_frame)]
ax.scatter(
tracking["x"][spikes],
tracking["y"][spikes],
s=15,
zorder=101,
color=unit.color,
)
def visualize_word_cluster_groups(
transformed_word_embeddings: np.ndarray,
words: np.ndarray,
word_groups: dict,
visualize_non_group_words: bool,
xlabel: str,
ylabel: str,
non_group_words_color: str = "#ccc",
ax: plt.axis = None,
show_plot: bool = True,
alpha: float = 1,
interactive: bool = False,
) -> None:
"""
Visualizes word cluster groups.
Parameters
----------
transformed_word_embeddings : np.ndarray
Transformed word embeddings.
words : np.ndarray
Numpy array containing all words from vocabulary.
word_groups : dict
Dictionary containing word groups to visualize.
visualize_non_group_words : bool
Whether or not to visualize words outside word groups
xlabel : str
X-axis label
ylabel : str
Y-axis label
non_group_words_color : str
Color for words outside groups (defaults to #ccc)
ax : plt.axis
Matplotlib axis (defaults to None)
show_plot : bool
Whether or not to call plt.show() (defaults to True)
alpha : float
Scatter plot alpha value (defaults to 1)
interactive : bool
Whether or not to make the visualization interactive
using Plotly (defaults to False).
"""
if ax is None and not interactive:
_, ax = plt.subplots(figsize=(12, 7))
if interactive:
fig = go.Figure(
layout=dict(xaxis=dict(title=xlabel), yaxis=dict(title=ylabel)))
if visualize_non_group_words:
# Create boolean mask for words outside groups
words_in_groups = []
for group_name in word_groups.keys():
words_in_groups.extend(word_groups[group_name]["words"])
words_not_in_groups_mask = [
word not in words_in_groups for word in words
]
words_not_in_groups_sorted = [
word for word in words if word not in words_in_groups
]
# Plot words outside word group
if interactive:
fig.add_trace(
go.Scatter(
x=transformed_word_embeddings[words_not_in_groups_mask][:,
0],
y=transformed_word_embeddings[words_not_in_groups_mask][:,
1],
mode="markers",
marker=dict(color=non_group_words_color),
hovertext=words_not_in_groups_sorted,
hoverinfo="x+y+text",
name="Non group words",
opacity=alpha,
))
else:
ax.scatter(
x=transformed_word_embeddings[words_not_in_groups_mask][:, 0],
y=transformed_word_embeddings[words_not_in_groups_mask][:, 1],
c=non_group_words_color,
alpha=alpha,
)
# Visualize words in groups
for group_name, word_group in word_groups.items():
words_in_group = word_group["words"]
words_in_group_mask = [word in words_in_group for word in words]
words_in_group_sorted = [
word for word in words if word in words_in_group
]
word_group_color = word_group["color"]
# Plot words inside word group
if interactive:
fig.add_trace(
go.Scatter(
x=transformed_word_embeddings[words_in_group_mask][:, 0],
y=transformed_word_embeddings[words_in_group_mask][:, 1],
mode="markers",
marker=dict(color=word_group_color),
hovertext=words_in_group_sorted,
hoverinfo="x+y+text",
name=f"Words in {group_name}",
opacity=alpha,
))
else:
ax.scatter(
x=transformed_word_embeddings[words_in_group_mask][:, 0],
y=transformed_word_embeddings[words_in_group_mask][:, 1],
c=word_group_color,
alpha=alpha,
)
if interactive:
fig.show()
else:
ax_legends = ["Non group words"]
ax_legends.extend([
f"Words which are {group_name}"
for group_name in word_groups.keys()
])
ax.legend(ax_legends)
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
if show_plot:
plt.show()