FEniCSx用メッシュファイルを取り出す

"""
Export-only program for the six mesh families from Section 9.

This script
  1. generates one of the six mesh families (I)-(VI),
  2. builds a FEniCSx mesh object,
  3. tags exterior boundary facets by side of the unit square:
       bottom = 1, right = 2, top = 3, left = 4,
  4. writes the mesh and boundary tags to XDMF,
  5. optionally writes one indicator CSV for the chosen (N, DELTA, EPSILON).

Poisson solves are intentionally NOT included here.
They are meant for a separate blog / script.
"""

from __future__ import annotations

from pathlib import Path
import numpy as np
import pandas as pd


# -------------------------
# Global configuration
# -------------------------
# Export either one family or all six.
EXPORT_ALL = True
KIND_DEFAULT = "IV"  # used when EXPORT_ALL = False

# One case = one choice of (N, DELTA, EPSILON)
N_DEFAULT = 16
DELTA = 1 / 128
EPSILON = 2.0

# Output settings
OUTDIR = "fenicsx_mesh_data"
WRITE_INDICATOR_CSV = True
WRITE_MANIFEST_CSV = True
BOUNDARY_TAG_NAME = "boundary_facets"

# Boundary labels on the unit square
BOTTOM_TAG_VALUE = 1
RIGHT_TAG_VALUE = 2
TOP_TAG_VALUE = 3
LEFT_TAG_VALUE = 4


# -------------------------
# Utility: parameter handling
# -------------------------
def resolve_parameters(N=None, delta=None, epsilon=None):
    if N is None:
        N = N_DEFAULT
    if delta is None:
        delta = DELTA
    if epsilon is None:
        epsilon = EPSILON
    return int(N), float(delta), float(epsilon)


def make_case(N=None, delta=None, epsilon=None):
    N, delta, epsilon = resolve_parameters(N=N, delta=delta, epsilon=epsilon)
    return {"N": N, "DELTA": delta, "EPSILON": epsilon}


def case_to_stem(case):
    return (
        f"N{case['N']}_"
        f"delta{case['DELTA']:.6g}_"
        f"eps{case['EPSILON']:.6g}"
    ).replace("/", "_")


# -------------------------
# Section 9.1: good families (I)-(IV)
# -------------------------
def x_standard(N=None):
    N, _, _ = resolve_parameters(N=N)
    return np.linspace(0.0, 1.0, N + 1)


def x_shishkin_y(N=None, delta=None):
    N, delta, _ = resolve_parameters(N=N, delta=delta)
    tau = 2 * delta * abs(np.log(N))
    tau = min(max(tau, 0.0), 1.0)

    y = np.zeros(N + 1)
    half = N // 2
    for i in range(N + 1):
        if i <= half:
            y[i] = tau * i / half if half > 0 else 0.0
        else:
            y[i] = tau + (1 - tau) * (i - half) / (N - half)
    return y


def x_cosine_cluster(N=None):
    N, _, _ = resolve_parameters(N=N)
    i = np.arange(N + 1, dtype=float)
    return 0.5 * (1.0 - np.cos(i * np.pi / N))


def x_graded_power(N=None, epsilon=None):
    N, _, epsilon = resolve_parameters(N=N, epsilon=epsilon)
    i = np.arange(N + 1, dtype=float)
    return (i / N) ** epsilon


def section9_lines(kind, N=None, delta=None, epsilon=None):
    N, delta, epsilon = resolve_parameters(N=N, delta=delta, epsilon=epsilon)

    if kind == "I":
        x1 = x_standard(N)
        x2 = x_standard(N)
    elif kind == "II":
        x1 = x_standard(N)
        x2 = x_shishkin_y(N, delta=delta)
    elif kind == "III":
        x1 = x_cosine_cluster(N)
        x2 = x_cosine_cluster(N)
    elif kind == "IV":
        x1 = x_standard(N)
        x2 = x_graded_power(N, epsilon=epsilon)
    else:
        raise ValueError("kind must be one of 'I', 'II', 'III', 'IV'.")
    return x1, x2


def build_tri_mesh_from_tensor_grid(x1, x2, diagonal="right", flip_corner_cells=False):
    x1 = np.asarray(x1, dtype=float)
    x2 = np.asarray(x2, dtype=float)
    nx = len(x1) - 1
    ny = len(x2) - 1

    nodes = []
    for j in range(ny + 1):
        for i in range(nx + 1):
            nodes.append([x1[i], x2[j]])
    nodes = np.array(nodes, dtype=float)

    def nid(i, j):
        return j * (nx + 1) + i

    elements = []
    for j in range(ny):
        for i in range(nx):
            n00 = nid(i, j)
            n10 = nid(i + 1, j)
            n01 = nid(i, j + 1)
            n11 = nid(i + 1, j + 1)

            use_diagonal = diagonal
            if flip_corner_cells:
                is_top_left = (i == 0 and j == ny - 1)
                is_bottom_right = (i == nx - 1 and j == 0)
                if is_top_left or is_bottom_right:
                    use_diagonal = "left" if diagonal == "right" else "right"

            if use_diagonal == "right":
                elements.append([n00, n10, n11])
                elements.append([n00, n11, n01])
            elif use_diagonal == "left":
                elements.append([n00, n10, n01])
                elements.append([n10, n11, n01])
            else:
                raise ValueError("diagonal must be 'right' or 'left'.")

    return nodes, np.array(elements, dtype=np.int64)


# -------------------------
# Indicators (optional metadata)
# -------------------------
def tri_edge_lengths(pts):
    L = [
        np.linalg.norm(pts[1] - pts[0]),
        np.linalg.norm(pts[2] - pts[1]),
        np.linalg.norm(pts[0] - pts[2]),
    ]
    return np.sort(np.array(L, dtype=float))  # L1 <= L2 <= L3


def tri_area(pts):
    a = pts[1] - pts[0]
    b = pts[2] - pts[0]
    return 0.5 * abs(a[0] * b[1] - a[1] * b[0])


def section9_indicators(nodes, elements):
    min_like = 0.0
    max_like = 0.0
    for tri in elements:
        pts = nodes[np.asarray(tri, dtype=int)]
        L1, L2, L3 = tri_edge_lengths(pts)
        area = tri_area(pts)
        min_like = max(min_like, (L3 ** 2) / area)
        max_like = max(max_like, (L1 * L2) / area)
    return float(min_like), float(max_like)


# -------------------------
# Section 9.3: bad families (V)-(VI), x-direction from mesh1.h style
# -------------------------
def section93_lines(kind, N=None, delta=None, epsilon=None):
    N, delta, epsilon = resolve_parameters(N=N, delta=delta, epsilon=epsilon)
    if kind == "V":
        return x_shishkin_y(N, delta=delta)
    elif kind == "VI":
        return x_graded_power(N, epsilon=epsilon)
    else:
        raise ValueError("kind must be 'V' or 'VI'.")


def mesh1_style_rows(N=None):
    N, _, _ = resolve_parameters(N=N)
    x_base = np.linspace(0.0, 1.0, N + 1)
    mids = 0.5 * (x_base[:-1] + x_base[1:])
    rows = []
    for j in range(N + 1):
        if j % 2 == 0:
            rows.append(x_base.copy())
        else:
            rows.append(np.concatenate(([x_base[0]], mids, [x_base[-1]])))
    return rows


def build_mesh1_style_mesh(N=None, y_lines=None):
    N, _, _ = resolve_parameters(N=N)
    y_lines = np.asarray(y_lines, dtype=float)
    if len(y_lines) != N + 1:
        raise ValueError("y_lines must have length N+1")

    rows_x = mesh1_style_rows(N)

    nodes = []
    row_starts = []
    count = 0
    for j in range(N + 1):
        row_starts.append(count)
        for x in rows_x[j]:
            nodes.append([x, y_lines[j]])
            count += 1
    nodes = np.array(nodes, dtype=float)

    def nid(row, k):
        return row_starts[row] + k

    elements = []
    for j in range(N):
        if j % 2 == 0:
            # odd strip in original mesh1.h (0-based => j even)
            for i in range(N):
                p1 = nid(j, i)
                p2 = nid(j, i + 1)
                p3 = nid(j + 1, i + 1)
                elements.append([p1, p2, p3])
                p1 = nid(j, i)
                p2 = nid(j + 1, i + 1)
                p3 = p2 - 1
                elements.append([p1, p2, p3])
            p1 = nid(j, N)
            p2 = nid(j + 1, N + 1)
            p3 = p2 - 1
            elements.append([p1, p2, p3])
        else:
            # even strip in original mesh1.h (0-based => j odd)
            p1 = nid(j, 0)
            p2 = p1 + 1
            p3 = nid(j + 1, 0)
            elements.append([p1, p2, p3])
            for i in range(N):
                p1 = nid(j, i + 1)
                p2 = p1 + 1
                p3 = nid(j + 1, i + 1)
                elements.append([p1, p2, p3])
                p1 = nid(j, i + 1)
                p2 = nid(j + 1, i + 1)
                p3 = p2 - 1
                elements.append([p1, p2, p3])

    return nodes, np.array(elements, dtype=np.int64)


def section93_reconstruction(kind, N=None, delta=None, epsilon=None):
    N, delta, epsilon = resolve_parameters(N=N, delta=delta, epsilon=epsilon)
    y_lines = section93_lines(kind, N=N, delta=delta, epsilon=epsilon)
    return build_mesh1_style_mesh(N=N, y_lines=y_lines)


# -------------------------
# Unified family interface
# -------------------------
def build_family_mesh(kind, N=None, delta=None, epsilon=None):
    N, delta, epsilon = resolve_parameters(N=N, delta=delta, epsilon=epsilon)

    if kind in ["I", "II", "III", "IV"]:
        x1, x2 = section9_lines(kind, N=N, delta=delta, epsilon=epsilon)
        nodes, elements = build_tri_mesh_from_tensor_grid(
            x1, x2, diagonal="right", flip_corner_cells=True
        )
    elif kind in ["V", "VI"]:
        nodes, elements = section93_reconstruction(kind, N=N, delta=delta, epsilon=epsilon)
    else:
        raise ValueError("kind must be one of 'I', 'II', 'III', 'IV', 'V', 'VI'.")

    return nodes, elements


# -------------------------
# FEniCSx helpers for export
# -------------------------
def _require_fenicsx():
    try:
        from mpi4py import MPI
        import basix.ufl
        import ufl
        from dolfinx import mesh as dmesh, io
    except Exception as e:
        raise ImportError(
            "This export script requires FEniCSx, basix, UFL and mpi4py. "
            "Run it in your FEniCSx / Colab environment."
        ) from e
    return MPI, basix, ufl, dmesh, io


def build_fenicsx_mesh_object(nodes, elements):
    MPI, basix, ufl, dmesh, io = _require_fenicsx()

    x = np.asarray(nodes, dtype=np.float64)
    cells = np.asarray(elements, dtype=np.int64)
    domain = ufl.Mesh(
        basix.ufl.element("Lagrange", "triangle", 1, shape=(2,), dtype=np.float64)
    )
    msh = dmesh.create_mesh(MPI.COMM_WORLD, cells, domain, x)
    msh.name = "mesh"
    return msh


def build_boundary_facets_meshtags_square(msh, name=BOUNDARY_TAG_NAME, tol=1.0e-12):
    """
    Tag exterior facets of the unit square by midpoint:
      bottom = 1, right = 2, top = 3, left = 4
    """
    MPI, basix, ufl, dmesh, io = _require_fenicsx()

    tdim = msh.topology.dim
    fdim = tdim - 1

    msh.topology.create_connectivity(fdim, tdim)
    exterior_facets = dmesh.exterior_facet_indices(msh.topology)
    midpoints = dmesh.compute_midpoints(msh, fdim, exterior_facets)

    values = np.zeros(len(exterior_facets), dtype=np.int32)

    for k, midpoint in enumerate(midpoints):
        x = float(midpoint[0])
        y = float(midpoint[1])

        if abs(y - 0.0) <= tol:
            values[k] = BOTTOM_TAG_VALUE
        elif abs(x - 1.0) <= tol:
            values[k] = RIGHT_TAG_VALUE
        elif abs(y - 1.0) <= tol:
            values[k] = TOP_TAG_VALUE
        elif abs(x - 0.0) <= tol:
            values[k] = LEFT_TAG_VALUE
        else:
            raise ValueError(
                f"Exterior facet midpoint {(x, y)} is not on the expected unit-square boundary."
            )

    order = np.argsort(exterior_facets)
    tags = dmesh.meshtags(msh, fdim, exterior_facets[order], values[order])
    tags.name = name
    return tags


def build_fenicsx_family_mesh(kind, case=None):
    if case is None:
        case = make_case()

    nodes, elements = build_family_mesh(
        kind,
        N=case["N"],
        delta=case["DELTA"],
        epsilon=case["EPSILON"],
    )
    msh = build_fenicsx_mesh_object(nodes, elements)
    facet_tags = build_boundary_facets_meshtags_square(msh)
    return msh, facet_tags, nodes, elements


# -------------------------
# Export routines
# -------------------------
def export_family_xdmf(kind, case=None, outdir=None):
    MPI, basix, ufl, dmesh, io = _require_fenicsx()

    if case is None:
        case = make_case()
    if outdir is None:
        outdir = OUTDIR

    outdir = Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    msh, facet_tags, nodes, elements = build_fenicsx_family_mesh(kind, case=case)
    xdmf_path = outdir / f"family_{kind}_{case_to_stem(case)}.xdmf"

    with io.XDMFFile(msh.comm, str(xdmf_path), "w") as xdmf:
        xdmf.write_mesh(msh)
        xdmf.write_meshtags(facet_tags, msh.geometry)

    return xdmf_path


def export_selected_or_all(case=None, outdir=None):
    if case is None:
        case = make_case()
    if outdir is None:
        outdir = OUTDIR

    if EXPORT_ALL:
        kinds = ["I", "II", "III", "IV", "V", "VI"]
    else:
        kinds = [KIND_DEFAULT]

    paths = {}
    for kind in kinds:
        paths[kind] = export_family_xdmf(kind, case=case, outdir=outdir)
    return paths


# -------------------------
# Optional metadata CSVs
# -------------------------
def build_indicator_table(case=None):
    if case is None:
        case = make_case()

    rows = []
    for kind in ["I", "II", "III", "IV", "V", "VI"]:
        nodes, elements = build_family_mesh(
            kind,
            N=case["N"],
            delta=case["DELTA"],
            epsilon=case["EPSILON"],
        )
        min_like, max_like = section9_indicators(nodes, elements)
        rows.append(
            {
                "Mesh": kind,
                "N": case["N"],
                "DELTA": case["DELTA"],
                "EPSILON": case["EPSILON"],
                "MinAngle": min_like,
                "MaxAngle": max_like,
            }
        )
    return pd.DataFrame(rows)


def write_indicator_csv(case=None, outdir=None):
    if case is None:
        case = make_case()
    if outdir is None:
        outdir = OUTDIR

    outdir = Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)
    df = build_indicator_table(case)
    csv_path = outdir / f"indicator_table_{case_to_stem(case)}.csv"
    df.to_csv(csv_path, index=False)
    return csv_path


def write_manifest_csv(exported_paths, case=None, outdir=None):
    if case is None:
        case = make_case()
    if outdir is None:
        outdir = OUTDIR

    outdir = Path(outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    rows = []
    for kind, path in exported_paths.items():
        rows.append(
            {
                "Mesh": kind,
                "N": case["N"],
                "DELTA": case["DELTA"],
                "EPSILON": case["EPSILON"],
                "XDMF": str(path),
            }
        )
    manifest = pd.DataFrame(rows)
    csv_path = outdir / f"manifest_{case_to_stem(case)}.csv"
    manifest.to_csv(csv_path, index=False)
    return csv_path


# -------------------------
# Main
# -------------------------
def main():
    case = make_case(N=N_DEFAULT, delta=DELTA, epsilon=EPSILON)
    print("Using case:", case)
    print("EXPORT_ALL =", EXPORT_ALL)
    print("OUTDIR     =", OUTDIR)
    print("Boundary tags: bottom=1, right=2, top=3, left=4")

    exported = export_selected_or_all(case=case, outdir=OUTDIR)
    print()
    print("Exported XDMF files:")
    for kind, path in exported.items():
        print(f"  {kind}: {path}")

    if WRITE_INDICATOR_CSV:
        indicator_csv = write_indicator_csv(case=case, outdir=OUTDIR)
        print()
        print(f"Indicator CSV: {indicator_csv}")

    if WRITE_MANIFEST_CSV:
        manifest_csv = write_manifest_csv(exported, case=case, outdir=OUTDIR)
        print(f"Manifest CSV : {manifest_csv}")

    print()
    print("Done. This script only exports mesh files and metadata.")
    print("Poisson solves are intentionally left for a separate script / blog post.")


if __name__ == "__main__":
    main()