goodanalysis2.py

"""LAMMPS トラジェクトリから「図の状態」の N 個数を数える。

ユーザー要件に合わせて vacancy は考慮せず、N 原子の空間配置のみで判定する。

ここでの「図の状態」は次の条件を満たす N を指す:

1) N 原子（このデータでは dump の type=3）
2) 表面からの深さ depth = surface_z - z が一定以上（デフォルト 5 nm = 50 Å）
3) その N が、別の N と 5–15 nm の距離にある（磁気双極子相互作用を仮定）
4) 上の(3)を満たす 2つの N が、それ以外の N から 15 nm 以上離れている

入力:
- dump_run_1_min0K.lammpstrj (LAMMPS dump 形式)

例:

  # 「図の状態」のN個数をカウント
  python goodanalysis2.py \
	--dump "10Ncluster_implantation to C_0.5keV(final)\\dump_run_1_min0K.lammpstrj" \
	--verbose

  # N-N 原子間距離のヒストグラム（横軸:距離[nm], 縦軸:個数）を保存
  python goodanalysis2.py \
	--dump "10Ncluster_implantation to C_0.5keV(final)\\dump_run_1_min0K.lammpstrj" \
	--plot-dist-hist --hist-output "dist_hist.png"
"""

from __future__ import annotations

import argparse
import math
import sys
from pathlib import Path


# matplotlib はプロット時のみ必要（未インストールでもカウントは動くようにする）


ANGSTROM_PER_NM = 10.0


def _ensure_utf8_stdio() -> None:
	"""Windows端末の既定エンコーディング差で help 表示が落ちるのを防ぐ。"""
	for stream in (sys.stdout, sys.stderr):
		try:
			stream.reconfigure(encoding="utf-8", errors="replace")  # type: ignore[attr-defined]
		except Exception:
			pass


def _to_float(token: str) -> float | None:
	try:
		return float(token)
	except ValueError:
		return None


def _pick_idx(cols_lower: list[str], candidates: list[str]) -> int | None:
	for name in candidates:
		if name in cols_lower:
			return cols_lower.index(name)
	return None


def read_last_dump_n_positions(
	dump_path: Path,
	nitrogen_type: int,
) -> list[tuple[float, float, float]]:
	"""LAMMPS dump から N(type) の (x,y,z) を読む（最後のスナップショット）。

	このリポジトリの dump は通常 `ITEM: ATOMS id type x y z` 形式なので、
	PBC/minimum image や scaled 座標の変換は行わない（ユーザー要望で簡略化）。
	"""

	positions: list[tuple[float, float, float]] = []
	current: list[tuple[float, float, float]] = []
	reading_atoms = False
	x_idx = y_idx = z_idx = None
	type_idx: int | None = None

	with dump_path.open("r", encoding="utf-8", errors="ignore") as handle:
		for raw in handle:
			line = raw.strip()
			if not line:
				continue

			if line.startswith("ITEM:"):
				if line.startswith("ITEM: ATOMS"):
					cols = line.split()[2:]
					cols_l = [c.lower() for c in cols]
					type_idx = cols_l.index("type") if "type" in cols_l else None
					x_idx = _pick_idx(cols_l, ["x", "xu"])
					y_idx = _pick_idx(cols_l, ["y", "yu"])
					z_idx = _pick_idx(cols_l, ["z", "zu"])
					current = []
					reading_atoms = True
					continue
				# 別の ITEM に入ったら、その時点の ATOMS を最後として採用
				if reading_atoms:
					positions = current
				reading_atoms = False
				continue

			if not reading_atoms or type_idx is None or x_idx is None or y_idx is None or z_idx is None:
				continue

			parts = line.split()
			if len(parts) <= max(type_idx, x_idx, y_idx, z_idx):
				continue

			t_val = _to_float(parts[type_idx])
			if t_val is None or int(round(t_val)) != nitrogen_type:
				continue
			x = _to_float(parts[x_idx])
			y = _to_float(parts[y_idx])
			z = _to_float(parts[z_idx])
			if x is None or y is None or z is None:
				continue
			current.append((x, y, z))

	# EOF: ATOMS の最中なら最後を採用
	if reading_atoms:
		positions = current

	return positions


def distance(a: tuple[float, float, float], b: tuple[float, float, float]) -> float:
	dx = a[0] - b[0]
	dy = a[1] - b[1]
	dz = a[2] - b[2]
	return math.sqrt(dx * dx + dy * dy + dz * dz)


def pairwise_distances(
	positions: list[tuple[float, float, float]],
	indices: list[int] | None = None,
) -> list[float]:
	"""指定した点集合の全ペア距離を返す（Å）。

	N注入のように点数が小さいケースを想定し、単純な二重ループで計算する。
	"""
	idxs = indices if indices is not None else list(range(len(positions)))
	dists: list[float] = []
	for a_i in range(len(idxs)):
		i = idxs[a_i]
		for a_j in range(a_i + 1, len(idxs)):
			j = idxs[a_j]
			dists.append(distance(positions[i], positions[j]))
	return dists


def plot_distance_histogram(
	distances_a: list[float],
	*,
	bins: int,
	output: Path | None,
	range_nm: tuple[float, float] | None,
	density: bool,
) -> None:
	if len(distances_a) == 0:
		print("距離データが0件のため、ヒストグラムは作成しませんでした。")
		return

	try:
		import matplotlib.pyplot as plt
	except ModuleNotFoundError as exc:
		raise SystemExit(
			"matplotlib が見つかりません。プロットするには `pip install matplotlib` を実行してください。"
		) from exc

	distances_nm = [d / ANGSTROM_PER_NM for d in distances_a]

	plt.figure()
	plt.hist(
		distances_nm,
		bins=int(bins),
		range=range_nm,
		density=bool(density),
	)
	# 日本語フォントが無い環境でも警告が出ないように英語ラベルにする
	plt.xlabel("Interatomic distance (nm)")
	plt.ylabel("Density" if density else "Count")
	plt.tight_layout()

	if output is not None:
		output.parent.mkdir(parents=True, exist_ok=True)
		plt.savefig(output, dpi=200)
		plt.close()
	else:
		plt.show()


def count_n_in_image_state(
	dump_path: Path,
	*,
	nitrogen_type: int,
	surface_z: float,
	min_depth_a: float,
	min_pair_dist_a: float,
	max_pair_dist_a: float,
	isolation_dist_a: float,
) -> tuple[int, dict[str, int]]:
	n_positions = read_last_dump_n_positions(dump_path, nitrogen_type)

	# 深さフィルタ（ペア探索対象）: depth >= min_depth
	deep_indices: list[int] = []
	for idx, (_, _, z) in enumerate(n_positions):
		depth = surface_z - z
		if depth >= min_depth_a:
			deep_indices.append(idx)

	# (3) 5–15 nm の距離にある N ペアを列挙
	candidate_pairs: list[tuple[int, int]] = []
	for a_i in range(len(deep_indices)):
		i = deep_indices[a_i]
		for a_j in range(a_i + 1, len(deep_indices)):
			j = deep_indices[a_j]
			d = distance(n_positions[i], n_positions[j])
			if min_pair_dist_a <= d <= max_pair_dist_a:
				candidate_pairs.append((i, j))

	# (4) その2つが「それ以外のN」から 15nm 以上離れているか
	qualifying: set[int] = set()
	for i, j in candidate_pairs:
		isolated = True
		for k in range(len(n_positions)):
			if k == i or k == j:
				continue
			dik = distance(n_positions[i], n_positions[k])
			if dik < isolation_dist_a:
				isolated = False
				break
			djk = distance(n_positions[j], n_positions[k])
			if djk < isolation_dist_a:
				isolated = False
				break
		if isolated:
			qualifying.add(i)
			qualifying.add(j)

	stats = {
		"N_total_in_dump": len(n_positions),
		"N_deep": len(deep_indices),
		"pair_candidates_5to15nm": len(candidate_pairs),
		"N_in_image_state": len(qualifying),
	}
	return len(qualifying), stats


def main() -> None:
	_ensure_utf8_stdio()
	parser = argparse.ArgumentParser(description="Count nitrogen atoms that match the figure-like conditions.")
	parser.add_argument(
		"--dump",
		type=Path,
		default=Path(
			"10Ncluster_implantation to C_0.5keV(final)/dump_run_1_min0K.lammpstrj"
		),
		help="LAMMPS dump (.lammpstrj) path",
	)
	parser.add_argument("--nitrogen-type", type=int, default=3, help="Atom type ID for nitrogen in dump (default: 3)")
	parser.add_argument(
		"--surface-z",
		type=float,
		default=125.0,
		help="Surface z in Å used as depth = surface_z - z (default: 125.0)",
	)
	parser.add_argument(
		"--min-depth-nm",
		type=float,
		default=5.0,
		help="Minimum depth from surface in nm (default: 5.0)",
	)
	parser.add_argument(
		"--pair-min-nm",
		type=float,
		default=5.0,
		help="N-N 相互作用距離の下限 (nm) (default: 5.0)",
	)
	parser.add_argument(
		"--pair-max-nm",
		type=float,
		default=15.0,
		help="N-N 相互作用距離の上限 (nm) (default: 15.0)",
	)
	parser.add_argument(
		"--isolation-nm",
		type=float,
		default=15.0,
		help="(4) 2つの相互作用Nが他のNから離れている距離の下限 (nm) (default: 15.0)",
	)
	parser.add_argument(
		"--verbose",
		action="store_true",
		help="詳細な内訳も表示",
	)
	parser.add_argument(
		"--plot-dist-hist",
		action="store_true",
		help="N-N の原子間距離ヒストグラム（横軸:距離, 縦軸:個数）をプロット",
	)
	parser.add_argument(
		"--hist-scope",
		choices=["all", "deep"],
		default="all",
		help="ヒストグラムに使うNの範囲: all=全N, deep=深さ条件を満たすN (default: all)",
	)
	parser.add_argument(
		"--hist-bins",
		type=int,
		default=30,
		help="ヒストグラムのビン数 (default: 30)",
	)
	parser.add_argument(
		"--hist-density",
		action="store_true",
		help="個数ではなく確率密度（正規化）で表示する。N数が揃わない比較に有効",
	)
	parser.add_argument(
		"--hist-range-nm",
		nargs=2,
		type=float,
		default=None,
		metavar=("MIN", "MAX"),
		help="ヒストグラム表示範囲 (nm)。例: --hist-range-nm 0 20",
	)
	parser.add_argument(
		"--hist-output",
		type=Path,
		default=None,
		help="保存先（例: dist_hist.png）。未指定ならウィンドウ表示",
	)

	args = parser.parse_args()
	dump_path: Path = args.dump
	if not dump_path.exists():
		raise SystemExit(f"dump が見つかりません: {dump_path}")

	min_depth_a = float(args.min_depth_nm) * ANGSTROM_PER_NM
	min_pair_dist_a = float(args.pair_min_nm) * ANGSTROM_PER_NM
	max_pair_dist_a = float(args.pair_max_nm) * ANGSTROM_PER_NM
	isolation_dist_a = float(args.isolation_nm) * ANGSTROM_PER_NM

	n_count, stats = count_n_in_image_state(
		dump_path,
		nitrogen_type=int(args.nitrogen_type),
		surface_z=float(args.surface_z),
		min_depth_a=min_depth_a,
		min_pair_dist_a=min_pair_dist_a,
		max_pair_dist_a=max_pair_dist_a,
		isolation_dist_a=isolation_dist_a,
	)

	print(n_count)
	if args.verbose:
		for k, v in stats.items():
			print(f"{k}: {v}")

	if args.plot_dist_hist:
		n_positions = read_last_dump_n_positions(dump_path, int(args.nitrogen_type))
		if args.hist_scope == "deep":
			deep_indices: list[int] = []
			for idx, (_, _, z) in enumerate(n_positions):
				depth = float(args.surface_z) - z
				if depth >= min_depth_a:
					deep_indices.append(idx)
			dists_a = pairwise_distances(n_positions, deep_indices)
			n_used = len(deep_indices)
		else:
			dists_a = pairwise_distances(n_positions, None)
			n_used = len(n_positions)

		# 比較・デバッグ用の要約（N数が揃っていないときの混乱を避ける）
		if len(dists_a) > 0:
			d_nm = [d / ANGSTROM_PER_NM for d in dists_a]
			d_nm_sorted = sorted(d_nm)
			median_nm = d_nm_sorted[len(d_nm_sorted) // 2]
			print(
				f"hist_scope={args.hist_scope}, N_used={n_used}, pair_count={len(dists_a)}, "
				f"mean_nm={sum(d_nm)/len(d_nm):.3f}, median_nm={median_nm:.3f}, "
				f"min_nm={d_nm_sorted[0]:.3f}, max_nm={d_nm_sorted[-1]:.3f}"
			)

		range_nm = tuple(args.hist_range_nm) if args.hist_range_nm is not None else None
		plot_distance_histogram(
			dists_a,
			bins=int(args.hist_bins),
			output=args.hist_output,
			range_nm=range_nm,
			density=bool(args.hist_density),
		)


if __name__ == "__main__":
	main()