Skip to content

Add ability to add taste names to plots #38

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
abuzarmahmood wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add ability to add taste names to plots #38

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 32 additions & 9 deletions src/visualize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,6 +118,7 @@ def plot_raster(
def generate_raster_plot(
segments_frame: pd.DataFrame,
session_name: str = None,
taste_names: list = None,
):
"""
Generate a raster plot of the segments frame
Expand All @@ -131,6 +132,8 @@ def generate_raster_plot(
DataFrame containing the segments with classification results
session_name : str
Name of the session to display in the plot title
taste_names : list, optional
List of taste names corresponding to taste indices. If None, uses "Taste N" labels.

Outputs:
fig, ax : matplotlib figure and axis objects
Expand All @@ -154,8 +157,12 @@ def generate_raster_plot(
ax = [ax]

for taste in range(len(taste_pred_array_list)):
ax[taste], im = plot_raster(taste_pred_array_list[taste], ax = ax[taste])
ax[taste].set_ylabel(f'Taste {taste}\nTrial #')
ax[taste], im = plot_raster(taste_pred_array_list[taste], ax = ax[taste])
if taste_names is not None and taste < len(taste_names):
taste_label = taste_names[taste]
else:
taste_label = f'Taste {taste}'
ax[taste].set_ylabel(f'{taste_label}\nTrial #')

ax[0].set_title('Movement Classification')
ax[-1].set_xlabel('Time (ms)')
Expand All @@ -169,7 +176,7 @@ def generate_raster_plot(
fig.suptitle(session_name)
return fig, ax

def generate_detailed_plot(segments_frame, raw_emg=None, trial_idx=0, taste_idx=0):
def generate_detailed_plot(segments_frame, raw_emg=None, trial_idx=0, taste_idx=0, taste_names=None):
"""
Generate a detailed plot showing raw EMG signal and movement classifications.

Expand All @@ -182,6 +189,7 @@ def generate_detailed_plot(segments_frame, raw_emg=None, trial_idx=0, taste_idx=
raw_emg (np.ndarray, optional): Raw EMG signal data. If None, only shows classifications.
trial_idx (int, optional): Trial index to visualize. Defaults to 0.
taste_idx (int, optional): Taste index to visualize. Defaults to 0.
taste_names (list, optional): List of taste names corresponding to taste indices. If None, uses "Taste N" labels.

Returns:
tuple: (fig, ax) matplotlib figure and axes objects
Expand Down Expand Up @@ -240,7 +248,11 @@ def generate_detailed_plot(segments_frame, raw_emg=None, trial_idx=0, taste_idx=

# Set labels and title
class_ax.set_ylabel('Movement Type')
class_ax.set_title(f'Taste {taste_idx}, Trial {trial_idx} - Movement Classifications')
if taste_names is not None and taste_idx < len(taste_names):
taste_label = taste_names[taste_idx]
else:
taste_label = f'Taste {taste_idx}'
class_ax.set_title(f'{taste_label}, Trial {trial_idx} - Movement Classifications')

if raw_emg is None:
class_ax.set_xlabel('Time (ms)')
Expand All @@ -253,6 +265,7 @@ def plot_env_pred_overlay(
raw_emg,
cmap = None,
mad_scale = 5,
taste_names = None,
):
"""
Create a comprehensive grid plot showing raw EMG signals with overlaid predictions.
Expand All @@ -278,6 +291,7 @@ def plot_env_pred_overlay(
- Blue (#3B75AF): MTMs
Defaults to None.
mad_scale (float, optional): Scaling factor for median absolute deviation
taste_names (list, optional): List of taste names corresponding to taste indices. If None, uses "Taste N" labels.

Returns:
tuple: Contains two elements:
Expand Down Expand Up @@ -326,7 +340,11 @@ def plot_env_pred_overlay(
if trial == raw_emg.shape[1] - 1:
this_ax.set_xlabel('Time (ms)')
if trial == 0:
this_ax.set_title(f'Taste {taste}')
if taste_names is not None and taste < len(taste_names):
taste_label = taste_names[taste]
else:
taste_label = f'Taste {taste}'
this_ax.set_title(taste_label)
# Overlay predictions
for _, segment in segments_frame.iterrows():
if segment.taste == taste and segment.trial == trial:
Expand Down Expand Up @@ -383,17 +401,22 @@ def plot_env_pred_overlay(
}
cmap = ListedColormap(list(event_color_map.values()), name = 'NBT_cmap')

# Generate the plot
# Example taste names
taste_names = ['Sucrose', 'NaCl']

# Generate the plot with taste names
fig, ax = generate_raster_plot(
segments_frame = sample_frame,
session_name = "Example Visualization"
session_name = "Example Visualization",
taste_names = taste_names
)
plt.show()

# Example of detailed plot
# Example of detailed plot with taste names
fig, ax = generate_detailed_plot(
segments_frame = sample_frame,
trial_idx = 0,
taste_idx = 0
taste_idx = 0,
taste_names = taste_names
)
plt.show()