XIAN-FEM-2026June/3D opticsfem-master/tools/export_ab_coo.py

45 lines
1.3 KiB
Python

#!/usr/bin/env python3
"""Export sparse A/b in OpticsFEM COO text format (matches C++ Test_OutputMatrix)."""
from __future__ import annotations
from pathlib import Path
import numpy as np
from scipy import sparse
def export_ab_coo(a: sparse.spmatrix, b: np.ndarray, out_dir: Path) -> None:
out_dir.mkdir(parents=True, exist_ok=True)
a = a.tocoo()
order = np.lexsort((a.col, a.row))
ai = a.row[order].astype(np.int64)
aj = a.col[order].astype(np.int64)
av = a.data[order]
np.savetxt(out_dir / "Ai.txt", ai, fmt="%d")
np.savetxt(out_dir / "Aj.txt", aj, fmt="%d")
with open(out_dir / "Av.txt", "w", encoding="utf-8") as f:
for v in av:
f.write(f"({v.real:.12g},{v.imag:.12g})\n")
b = np.asarray(b).reshape(-1)
np.savetxt(out_dir / "Bv_real.txt", b.real, fmt="%.12g")
np.savetxt(out_dir / "Bv_imag.txt", b.imag, fmt="%.12g")
meta = out_dir / "meta.txt"
meta.write_text(
"\n".join(
[
"source=export_ab_coo.py",
f"n={a.shape[0]}",
f"nnz={ai.size}",
f"|b|={np.linalg.norm(b):.15g}",
]
)
+ "\n",
encoding="utf-8",
)
print(f"Exported A/b to {out_dir}")
print(f" n={a.shape[0]}, nnz={ai.size}, |b|={np.linalg.norm(b):.6g}")