如何在JupyterLab中用Python实现Mathematica的散点图多边形内点索引提取功能？

Hey there! Let's translate that Mathematica functionality into Python step by step. We'll use numpy for data handling, matplotlib for plotting/interactivity, and ipywidgets for buttons. First, make sure you have these installed:

pip install numpy matplotlib ipywidgets seaborn

First, let's recap what the original code does:

• Generates random scatter points
• Shows a density heatmap of the points
• Lets you draw a polygon interactively
• Highlights points inside the polygon and shows their indices
• Has buttons to copy indices and reset the polygon

Step 1: Generate Sample Data #

First, let's replicate the random point generation from Mathematica:

import numpy as np

# Generate x-coordinates: 20 points between 0-5, 10 between 4-4.5
x = np.concatenate([np.random.uniform(0, 5, 20), np.random.uniform(4, 4.5, 10)])
# Generate y-coordinates: 20 points between 0-1, 10 between 1.5-2
y = np.concatenate([np.random.uniform(0, 1, 20), np.random.uniform(1.5, 2, 10)])
# Combine and shuffle the points
points = np.column_stack([x, y])
np.random.shuffle(points)

Step 2: Implement the Winding Number Algorithm #

This is the core logic to check if a point is inside a polygon. Here's the Python equivalent of your Mathematica winding function:

def winding_number(poly, pt):
    """Calculate winding number to determine if point is inside polygon."""
    # Get vectors from polygon vertices to the point
    vectors = pt - poly
    # Calculate angles of each vector
    angles = np.arctan2(vectors[:, 1], vectors[:, 0])
    # Compute differences between consecutive angles (wrap around for last-first)
    angle_diffs = np.diff(angles, append=angles[0])
    # Adjust differences to be in [-π, π] range
    angle_diffs = np.mod(angle_diffs + np.pi, 2*np.pi) - np.pi
    # Sum the differences and divide by 2π, round to get winding number
    return round(np.sum(angle_diffs) / (2 * np.pi))

Step 3: Build the Interactive Interface #

We'll use matplotlib's interactive backend in JupyterLab, plus PolygonSelector to draw the polygon, and ipywidgets buttons for copy/reset.

First, enable the matplotlib widget backend in Jupyter:

%matplotlib widget

Now the full interactive code:

import matplotlib.pyplot as plt
from matplotlib.widgets import PolygonSelector
import ipywidgets as widgets
from IPython.display import display
import seaborn as sns

# Initialize figure and axes
fig, ax = plt.subplots(figsize=(8, 6))
plt.subplots_adjust(bottom=0.2)  # Make space for buttons

# Plot the density heatmap (equivalent to SmoothDensityHistogram)
sns.kdeplot(x=points[:,0], y=points[:,1], cmap='coolwarm', fill=True, ax=ax, alpha=0.5)
# Plot all points initially in black
scatter_all = ax.scatter(points[:,0], points[:,1], color='black', s=30)
# Scatter object for highlighted points (initially empty)
scatter_highlight = ax.scatter([], [], color='darkgreen', s=50)

selected_indices = []
polygon = []

def update_highlight():
    """Update highlighted points based on the current polygon."""
    global selected_indices
    if not polygon:
        # Reset if no polygon
        scatter_highlight.set_offsets([])
        selected_indices = []
        return
    
    # Calculate winding number for each point
    winding_nums = [winding_number(np.array(polygon), pt) for pt in points]
    # Select points where winding number is non-zero (inside polygon)
    selected_indices = np.where(np.array(winding_nums) != 0)[0].tolist()
    # Update highlighted points
    if selected_indices:
        scatter_highlight.set_offsets(points[selected_indices])
    else:
        scatter_highlight.set_offsets([])
    fig.canvas.draw_idle()

def on_polygon_select(verts):
    """Callback when polygon is drawn/updated."""
    global polygon
    polygon = verts
    update_highlight()

# Create PolygonSelector tool
selector = PolygonSelector(ax, on_polygon_select, lineprops={'color': 'red', 'linestyle': 'dashed'})

# Create buttons
copy_btn = widgets.Button(description='Copy Points')
reset_btn = widgets.Button(description='Reset')

def on_copy_click(b):
    """Copy selected indices to output."""
    print("Selected point indices:", selected_indices)

def on_reset_click(b):
    """Reset polygon and highlighted points."""
    global polygon, selected_indices
    polygon = []
    selected_indices = []
    selector.clear()
    update_highlight()

copy_btn.on_click(on_copy_click)
reset_btn.on_click(on_reset_click)

# Display buttons below the plot
button_box = widgets.HBox([copy_btn, reset_btn])
display(button_box)

plt.show()

How to Use This: #

1. Run all cells in JupyterLab (make sure you enabled %matplotlib widget)
2. Click on the plot to add polygon vertices (click the first vertex again to close the polygon)
3. Points inside the polygon will turn dark green
4. Click "Copy Points" to print the indices of selected points in the output
5. Click "Reset" to clear the polygon and highlights

Key Notes: #

• The winding_number function works exactly like your Mathematica version to handle complex polygons (even self-intersecting ones)
• The density heatmap uses seaborn's kdeplot which is similar to Mathematica's SmoothDensityHistogram
• The PolygonSelector tool lets you draw polygons interactively, just like the Locator with LocatorAutoCreate in Mathematica

内容的提问来源于stack exchange，提问作者nitish pal