はじめに

二次元データの重心とピーク位置を計算して可視化する方法を紹介します。

本記事で扱うコードは、以下の2つの機能を持ちます。

二次元データの重心（バリセンター）と簡易ピーク（最大値）を計算
計算結果を可視化して正しい動作を検証

これらを通じて、二次元データの重心計算の方法を確認をしましょう。

この記事のコードは、google colab

で再現できます。

また、本記事の calc_center 関数は、xtend_pileup_gen_plfraction.py というコードの中で使われているものを解説したものです。

コアとなる重心計算のコード

概要

calc_center 関数は、以下を実現します：

指定した初期位置を中心とした探索範囲内で簡易ピーク（最大値）を特定
同じ範囲内でデータ値の重みに基づいた重心を計算
探索範囲内の総カウント値を計算

コード

重心計算のコードは短いので、まずはコードをみてみましょう。


import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as cr

params = {'xtick.labelsize': 10, 'ytick.labelsize': 10, 'legend.fontsize': 8}
plt.rcParams['font.family'] = 'serif'
plt.rcParams.update(params)

def calc_center(data, x_center, y_center, search_radius=10):
    """
    Calculate the simple peak (just for consistency check) and barycentric peak.
    
    Parameters:
    data (2D array): The input data array.
    x_center (int): The x-coordinate of the initial guess for the peak.
    y_center (int): The y-coordinate of the initial guess for the peak.
    search_radius (int): The radius around the initial guess to search for the peak. Default is 10.
    
    Returns:
    tuple: The coordinates of the barycentric peak (barycenter_x, barycenter_y) and the total counts (total_counts).
    """
    transposed_data = data.T
    x_indices, y_indices = np.meshgrid(np.arange(-search_radius, search_radius), 
                                       np.arange(-search_radius, search_radius), 
                                       indexing='ij')
    x_indices += x_center
    y_indices += y_center
    
    values = transposed_data[x_indices, y_indices]
    total_counts = np.sum(values)
    
    if total_counts > 0:
        barycenter_x = np.sum(values * x_indices) / total_counts
        barycenter_y = np.sum(values * y_indices) / total_counts
    else:
        print("[ERROR] in calc_center: Total counts are non-positive. Something went wrong.")
        sys.exit(1)

    peak_idx = np.unravel_index(np.argmax(values), values.shape)
    peak_x, peak_y = x_indices[peak_idx], y_indices[peak_idx]
    peak_value = values[peak_idx]

    # Print simple peak and barycenter information
    color_print(f"..... in calc_center: [simple peak] x_max, y_max, z_max = {peak_x}, {peak_y}, {peak_value}", ConsoleColors.OKCYAN)
    color_print(f"..... in calc_center: [barycenter] barycenter_x = {barycenter_x:.2f}, barycenter_y = {barycenter_y:.2f}", ConsoleColors.OKCYAN)
    color_print(f"..... in calc_center: [total counts in search radius] total_counts = {total_counts}", ConsoleColors.OKCYAN)

    if peak_value < 0:
        print("[ERROR] in calc_center: Peak value not found. Something went wrong.")
        sys.exit(1)
        
    return barycenter_x, barycenter_y, total_counts

おまけで、標準出力の色をつけるコードは下記を利用しています。

class ConsoleColors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def color_print(msg, color):
    print(f"{color}{msg}{ConsoleColors.ENDC}")

重要な処理

探索範囲の設定

x_indices, y_indices = np.meshgrid(
    np.arange(-search_radius, search_radius),
    np.arange(-search_radius, search_radius),
    indexing='ij'
)
x_indices += x_center
y_indices += y_center

np.meshgrid を使用して探索範囲の座標を生成。
indexing='ij' により、行優先（数学的な行列表現）でグリッドを作成。

重心の計算

if total_counts > 0:
    barycenter_x = np.sum(values * x_indices) / total_counts
    barycenter_y = np.sum(values * y_indices) / total_counts

探索範囲内のデータ値を重みとして重心を計算。

簡易ピークの特定

peak_idx = np.unravel_index(np.argmax(values), values.shape)
peak_x, peak_y = x_indices[peak_idx], y_indices[peak_idx]

np.argmax でフラット化した配列の最大値のインデックスを取得。
np.unravel_index で元の多次元インデックスに変換。

可視化用の `visualize_result` 関数

計算結果を可視化するための補助関数を用意します。

import matplotlib.pyplot as plt

# 二次元データの可視化
def visualize_result(data, x_center, y_center, search_radius, barycenter_x, barycenter_y, peak_x, peak_y):
    """
    Visualize the data and the results of calc_center.
    
    Parameters:
    data (2D array): The input data array.
    x_center (int): The x-coordinate of the initial guess for the peak.
    y_center (int): The y-coordinate of the initial guess for the peak.
    search_radius (int): The radius around the initial guess to search for the peak.
    barycenter_x (float): The x-coordinate of the calculated barycenter.
    barycenter_y (float): The y-coordinate of the calculated barycenter.
    peak_x (int): The x-coordinate of the simple peak.
    peak_y (int): The y-coordinate of the simple peak.
    """
    fig, ax = plt.subplots(figsize=(8, 6))
    ax.imshow(data, origin='lower', cmap='viridis', extent=(0, data.shape[1], 0, data.shape[0]))
    ax.set_title("Data Visualization with Peak and Barycenter")
    ax.set_xlabel("X")
    ax.set_ylabel("Y")

    # 探索範囲の描画
    circle = plt.Circle((x_center, y_center), search_radius, color='cyan', fill=False, linestyle='--', label="Search Radius")
    ax.add_artist(circle)

    # 簡易ピークのプロット
    ax.scatter(peak_x, peak_y, color='red', label='Simple Peak', zorder=5)

    # 重心のプロット
    ax.scatter(barycenter_x, barycenter_y, color='yellow', label='Barycenter', zorder=5)

    # 凡例の表示
    ax.legend()
    plt.colorbar(ax.imshow(data, origin='lower', cmap='viridis', extent=(0, data.shape[1], 0, data.shape[0])))
    plt.show()

可視化の特徴

入力データをカラーマップで描画
探索範囲を円で表示
簡易ピークを赤点、重心を黄色点でプロット
凡例とカラーバーを追加して視覚的にわかりやすく

動作確認と結果の可視化

コード全体の実行例

# ダミーデータの作成
data = np.random.random((100, 100))
x_center, y_center = 50, 50
search_radius = 10

# calc_center を実行
barycenter_x, barycenter_y, total_counts = calc_center(data, x_center, y_center, search_radius)

# 簡易ピークの計算
peak_idx = np.unravel_index(
    np.argmax(data[x_center-search_radius:x_center+search_radius,
                   y_center-search_radius:y_center+search_radius]),
    (2 * search_radius, 2 * search_radius)
)
peak_x = x_center - search_radius + peak_idx[0]
peak_y = y_center - search_radius + peak_idx[1]

# 可視化の実行
visualize_result(data, x_center, y_center, search_radius, barycenter_x, barycenter_y, peak_x, peak_y)

出力例

以下の図が表示されます：

背景: データ値のカラーマップ。
青い円: 探索範囲。
赤い点: 簡易ピーク（最大値）。
黄色い点: 重心。

まとめ

探索範囲内での重心を計算する方法を紹介しました。