Adding a label or annotation to a pcolormesh plot on mouse click

39 Views Asked by At

I have some code that is generating a pcolormesh plot, and “X” markers at specific coordinates. I would like to be able to click on a location on the plot and move the “X” marker and save the new/updated location. Below is my code so far (working environment: Windows 11, VS Code, .ipynb). I am able to print the coordinates to screen upon click but not store them for later use in the program. It seems like there are a number of ways this could be done, and I am wondering if someone more well versed in the capabilities of matplotlib could suggest a good way to do this? Thanks!

import matplotlib.colors as colors
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl
import os

# DEFINES
ECHOMAP_PIXELS = 1000
N = 734

def mouse_event(event):
    print('x: {} and y: {}'.format(event.xdata, event.ydata))

def echomap_plot_rectangular(d0, d1, x0, x1, save_to, labels, save=False):
    """
    Inputs:
        d0:734 by 170 numpy array of data 
        d1:734 by 170 numpy array of data 
        x0: 170 length 1d numpy array
        x1: 170 length 1d numpy array
        save_to: path to save echomap plots to
        labels: list of tuples of coordinates to label on the plot. x is radial, y is distance (0 to 50m).
    """
    my_dpi = 96
    
    r0 = np.linspace(0,50,N)
    r1 = np.linspace(0,50,N)

   
    fig, (ax0, ax1) = plt.subplots(nrows=1, ncols=2, figsize=(2*(ECHOMAP_PIXELS/my_dpi), ECHOMAP_PIXELS/my_dpi), dpi=my_dpi)
    cid = fig.canvas.mpl_connect('button_press_event', mouse_event)

    # PLOT
    vm = int(np.maximum(np.max(d0), np.max(d1))) + 1

    im = ax0.pcolormesh(x0, r0, d0, vmin=0, vmax=vm)
    fig.colorbar(im, ax=ax0)
    ax0.set_title('plot xyz')
    plt.grid()

    im1 = ax1.pcolormesh(x1, r1, d1, vmin=0, vmax=vm)
    fig.colorbar(im1, ax=ax1)
    ax1.set_title('plot abc')
    plt.grid()
    
    # Unzip the coordinates into separate lists for x and y
    x_values, y_values = zip(*labels)
    y0 = y_values
    y1 = y_values

    # Add target markers to the existing plots
    ax0.scatter(x_values, y0, marker='x', color='black', s=250, linewidth=3)
    ax1.scatter(x_values, y1, marker='x', color='black', s=250, linewidth=3)  
        
    fig.suptitle(f'Echomaps') # adjust spacing between subplots so `ax1` title and `ax0` tick labels don't overlap

    plt.tight_layout()
    if save:
        if not os.path.exists(save_to):
            os.makedirs(save_to)
            
        fig.savefig(save_to + "xyz_echomaps.png")
        print("saved")
        print(save_to)

%matplotlib qt

d0_sample = np.random.rand(N, 170)
d1_sample = np.random.rand(N, 170)
x0_sample = np.linspace(0, 170, 170)
x1_sample = np.linspace(0, 170, 170)
labels_sample = [(45, 15), (140, 22), (85, 40)]

echomap_plot_rectangular(d0_sample, d1_sample, x0_sample, x1_sample, save_to=".", labels=labels_sample)
0

There are 0 best solutions below