はじめに
二次元データの重心とピーク位置を計算して可視化する方法を紹介します。
本記事で扱うコードは、以下の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)
出力例
以下の図が表示されます：
- 背景: データ値のカラーマップ。
- 青い円: 探索範囲。
- 赤い点: 簡易ピーク（最大値）。
- 黄色い点: 重心。
まとめ
探索範囲内での重心を計算する方法を紹介しました。