Updated script accordingly

fontsize_medium = 16

fontsize_large = 18

def plane_3d(xy, ra_grad, dec_grad, offset):

    """3D plane function matching your fitting routine."""

    return ra_grad * xy[0] - dec_grad * xy[1] + offset

    velocities = valid_data[:, 0]

    x_pix = valid_data[:, 1]

    y_pix = valid_data[:, 2]

    x_coords = valid_data[:, 3]

    x_coords = -valid_data[:, 3]

    y_coords = valid_data[:, 4]

    # 2. Re-run the Bayesian Gaussian Mixture clustering

    # 4. Fit the plane to the main component to extract the gradient vector

    pixel_size_pc = source_properties[source]['distance'] * source_properties[source]['pixel_size'] / (

            360 * 3600 / (2 * np.pi))

    x_fit = (x_pix[main_mask] - 255) * pixel_size_pc

    x_fit = -(x_pix[main_mask] - 255) * pixel_size_pc

    y_fit = (y_pix[main_mask] - 255) * pixel_size_pc

    v_fit = velocities[main_mask]

    ax.tick_params(labelsize=fontsize_small)

    # Plot the main component

    # Center the colormap exactly on the median velocity of the main cluster

    v_mid = np.median(v_fit)

    v_range = max(np.abs(np.percentile(v_fit, 95) - v_mid),

                  np.abs(np.percentile(v_fit, 5) - v_mid))

    # Use square markers and remove black edges to tile it cleanly like an image

    scatter = ax.scatter(x_coords[main_mask], y_coords[main_mask],

                         c=v_fit, cmap='viridis', s=20, marker='s', zorder=2,

                         vmin=v_mid - v_range, vmax=v_mid + v_range,

    cbar.ax.tick_params(labelsize=fontsize_small)

    # 6. Add the Gradient Direction Arrow

    # Plane eq: V = ra_grad * x - dec_grad * y + offset

    # Vector of steepest velocity increase: (ra_grad, -dec_grad)

    arrow_length = np.max(np.abs(x_coords)) * 0.5  # Scale arrow to the plot size

    arrow_length = np.max(np.abs(x_coords)) * 0.5

    grad_norm = np.sqrt(ra_grad ** 2 + (-dec_grad) ** 2)

    dx = (ra_grad / grad_norm) * arrow_length

    dy = (-dec_grad / grad_norm) * arrow_length

    # Start the arrow near the center of the main cluster

    center_x, center_y = np.median(x_fit), np.median(y_fit)

    # We subtract dx/2 to center the arrow visually on the core

    ax.arrow(center_x - dx / 2, center_y - dy / 2, dx, dy,

             color='black', width=0.005, head_width=0.03,

             length_includes_head=True, zorder=5)

    # Add a dummy line for the legend to explain the arrow

    ax.plot([], [], color='black', marker=r'$\rightarrow$', markersize=15,

            linestyle='None', label='Gradient Direction')

    ax.set_title(f'{source} ({species_map[species]})', fontsize=fontsize_large)

    ax.set_aspect('equal')

    # FORCE identical plot box sizes by setting a fixed physical window for all sources

    plot_radius_pc = 0.8  # Adjust this value so it covers your largest clump

    ax.set_xlim(plot_radius_pc, -plot_radius_pc)  # Standard astronomical convention (RA increases left)

    plot_radius_pc = 0.8

    ax.set_xlim(plot_radius_pc, -plot_radius_pc)

    ax.set_ylim(-plot_radius_pc, plot_radius_pc)

    # ---> FIX: Only plot the legend on the first source (G332.96), and put it in the lower right <---

    if source == 'G332.96':

        ax.legend(loc='upper right', fontsize=fontsize_small, framealpha=0.9)

    ax.grid(alpha=0.3)

    plt.tight_layout()

    output_filename = f'projected_2d_kinematics_{source}_{species}.png'

    plt.savefig(output_filename, dpi=300, bbox_inches='tight')

    print(f"Saved plot to {output_filename}\n")

    # plt.show() # Uncomment if you want to view them interactively

if __name__ == '__main__':

    # You can change these to loop over your specific sources/species

    sources = ['G28', 'G332.96', 'G14']

    species_list = ['CH3OH', 'N2H+']