src/field_loader.py

"""Load MDN vector fields from metadata JSON + binary grid files."""

import json
import re
from pathlib import Path
from typing import Optional

import numpy as np


def load_field(meta_path: Path, grid_bin_path: Optional[Path] = None) -> dict:
    """
    Load MDN vector grids and metadata. Supports arbitrary K mixture components.
    Returns dict with keys: G, amin, amax, mean, mus, pi, ENT, TRAIN, meta.
    """
    meta = json.loads(Path(meta_path).read_text(encoding="utf-8"))
    G = int(meta["grid"])
    amin = np.array(meta["axis_min"], dtype=np.float32)
    amax = np.array(meta["axis_max"], dtype=np.float32)
    files = dict(meta.get("files") or {})

    def _resolve(p):
        p = Path(p)
        return p if p.exists() else (Path(meta_path).parent / p)

    def _load_xyz3(fname):
        raw = np.fromfile(_resolve(fname), dtype=np.float32)
        return raw.reshape(G, G, G, 3).astype(np.float32)

    def _load_scalar_or_channels(fname):
        raw = np.fromfile(_resolve(fname), dtype=np.float32)
        base = G * G * G
        if raw.size == base:
            return raw.reshape(G, G, G).astype(np.float32)
        if raw.size % base == 0:
            K = raw.size // base
            return raw.reshape(G, G, G, K).astype(np.float32)
        raise ValueError(f"{fname}: size {raw.size} incompatible with G^3={base}")

    # mean field
    if "mean_xyz3" in files:
        V_mean = _load_xyz3(files["mean_xyz3"])
    else:
        if grid_bin_path is None:
            grid_bin_path = Path(meta_path).with_name(
                Path(meta_path).name.replace("_meta.json", ".bin"))
        raw = np.fromfile(grid_bin_path, dtype=np.float32)
        V_mean = raw.reshape(G, G, G, 3).astype(np.float32)

    # collect mu_k components
    mu_items = []
    for k, v in files.items():
        ks = str(k).lower()
        if "xyz3" in ks and "mu" in ks:
            m = re.search(r'(\d+)(?!.*\d)', ks)
            if m:
                mu_items.append((int(m.group(1)), v))
    mu_items.sort(key=lambda kv: kv[0])
    MUS = [_load_xyz3(fname) for _, fname in mu_items]
    K_mu = len(MUS)

    # pi (mixture weights)
    PI = None
    if "pi" in files:
        PI = _load_scalar_or_channels(files["pi"])
        if PI.ndim == 3 and K_mu == 2:
            PI = np.stack([PI, 1.0 - PI], axis=-1).astype(np.float32)

    # entropy
    ENT = None
    if "entropy" in files:
        ENT = _load_scalar_or_channels(files["entropy"])
        if ENT.ndim != 3:
            raise ValueError(f"'entropy' must be (G,G,G); got {ENT.shape}")

    # training points
    train_pts = None
    tp_key = meta.get("training_points_npy")
    if tp_key:
        p = _resolve(tp_key)
        if p.exists():
            train_pts = np.load(p).astype(np.float32)

    return {
        "G": G, "amin": amin, "amax": amax,
        "mean": V_mean,
        "mus": MUS,
        "pi": PI,
        "ENT": ENT,
        "TRAIN": train_pts,
        "meta": meta,
    }


class TriLinearSampler:
    """Trilinear interpolation of 3D vector or scalar fields at arbitrary points."""

    def __init__(self, grid_xyz3, amin, amax):
        self.g = grid_xyz3.astype(np.float32)
        self.G = grid_xyz3.shape[0]
        self.amin = amin.astype(np.float32)
        self.amax = amax.astype(np.float32)
        span = self.amax - self.amin
        self.inv_span = np.where(span > 0, 1.0 / span, 0).astype(np.float32)

    def _idx(self, P):
        f = (P - self.amin) * self.inv_span * (self.G - 1)
        ix = np.floor(f[:, 0]).astype(np.int32)
        tx = (f[:, 0] - ix).astype(np.float32)
        iy = np.floor(f[:, 1]).astype(np.int32)
        ty = (f[:, 1] - iy).astype(np.float32)
        iz = np.floor(f[:, 2]).astype(np.int32)
        tz = (f[:, 2] - iz).astype(np.float32)
        ix = np.clip(ix, 0, self.G - 2)
        iy = np.clip(iy, 0, self.G - 2)
        iz = np.clip(iz, 0, self.G - 2)
        return ix, iy, iz, tx, ty, tz

    def sample_vec(self, P):
        ix, iy, iz, tx, ty, tz = self._idx(P)
        g = self.g
        v000 = g[ix, iy, iz]
        v100 = g[ix + 1, iy, iz]
        v010 = g[ix, iy + 1, iz]
        v110 = g[ix + 1, iy + 1, iz]
        v001 = g[ix, iy, iz + 1]
        v101 = g[ix + 1, iy, iz + 1]
        v011 = g[ix, iy + 1, iz + 1]
        v111 = g[ix + 1, iy + 1, iz + 1]
        vx00 = v000 * (1 - tx)[:, None] + v100 * tx[:, None]
        vx10 = v010 * (1 - tx)[:, None] + v110 * tx[:, None]
        vx01 = v001 * (1 - tx)[:, None] + v101 * tx[:, None]
        vx11 = v011 * (1 - tx)[:, None] + v111 * tx[:, None]
        vxy0 = vx00 * (1 - ty)[:, None] + vx10 * ty[:, None]
        vxy1 = vx01 * (1 - ty)[:, None] + vx11 * ty[:, None]
        out = vxy0 * (1 - tz)[:, None] + vxy1 * tz[:, None]
        return out.astype(np.float32)

    def sample_scalar(self, S, P):
        ix, iy, iz, tx, ty, tz = self._idx(P)
        s000 = S[ix, iy, iz]
        s100 = S[ix + 1, iy, iz]
        s010 = S[ix, iy + 1, iz]
        s110 = S[ix + 1, iy + 1, iz]
        s001 = S[ix, iy, iz + 1]
        s101 = S[ix + 1, iy, iz + 1]
        s011 = S[ix, iy + 1, iz + 1]
        s111 = S[ix + 1, iy + 1, iz + 1]
        sx00 = s000 * (1 - tx) + s100 * tx
        sx10 = s010 * (1 - tx) + s110 * tx
        sx01 = s001 * (1 - tx) + s101 * tx
        sx11 = s011 * (1 - tx) + s111 * tx
        return (sx00 * (1 - ty) + sx10 * ty) * (1 - tz) + (sx01 * (1 - ty) + sx11 * ty) * tz


def load_points_any(path: Path) -> np.ndarray:
    """Load point cloud from .npy, .obj, or .ply file."""
    p = Path(path)
    if p.suffix.lower() == ".npy":
        return np.asarray(np.load(p), np.float32)
    if p.suffix.lower() == ".obj":
        pts = []
        with open(p, "r", encoding="utf-8", errors="ignore") as f:
            for line in f:
                if not line or line[0] != 'v' or (len(line) > 1 and line[1] not in (' ', '\t')):
                    continue
                parts = line.strip().split()
                if len(parts) >= 4:
                    pts.append([float(parts[1]), float(parts[2]), float(parts[3])])
        return np.asarray(pts, np.float32)
    if p.suffix.lower() == ".ply":
        with open(p, "r", encoding="utf-8", errors="ignore") as f:
            header = []
            while True:
                line = f.readline()
                if not line:
                    raise ValueError("Invalid PLY (no end_header)")
                header.append(line.strip())
                if line.strip() == "end_header":
                    break
            nverts = 0
            for h in header:
                if h.startswith("element vertex"):
                    nverts = int(h.split()[-1])
                    break
            pts = []
            for _ in range(nverts):
                parts = f.readline().strip().split()
                if len(parts) >= 3:
                    pts.append([float(parts[0]), float(parts[1]), float(parts[2])])
        return np.asarray(pts, np.float32)
    raise ValueError(f"Unsupported point format: {p.suffix}")