Pythonで二次元データの重心とピーク位置を計算して可視化する例

# はじめに

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

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

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


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

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

https://colab.research.google.com/drive/1r5me6UkMwNlTYmmzRqXPOqquC1wmK2VJ?usp=sharing

で再現できます。

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

https://github.com/yamadasuzaku/rksysoft/blob/main/xtend/xtend_pileup_gen_plfraction.py


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

## **概要**

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

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

## コード

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

``` python:

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

```

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

``` python:
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}")
``` 

## **重要な処理**

### 探索範囲の設定

```python
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'` により、行優先（数学的な行列表現）でグリッドを作成。

### 重心の計算

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

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

### 簡易ピークの特定

```python
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` 関数

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

```python:
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()
```

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

## 動作確認と結果の可視化

### **コード全体の実行例**

```python
# ダミーデータの作成
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)
```
### **出力例**

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

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

![スクリーンショット 2024-11-23 0.48.21.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/179652/cbfc2402-d2cb-8998-6185-eb05fbf6e39e.png)


# まとめ

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