space_lib/simulate.py

from __future__ import annotations

import importlib.util
import os
import sys
from pathlib import Path
from typing import Dict, List

import numpy as np


def _default_data_generation_dir() -> Path:
    env = os.environ.get("MG_QUARTER_DATA_GEN_DIR")
    if env:
        return Path(env).resolve()
    return (Path(__file__).resolve().parent.parent / "assets" / "quarter_sim").resolve()


def _import_module_from_path(module_name: str, path: Path):
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    assert spec and spec.loader
    spec.loader.exec_module(module)
    return module


def format_vol_fraction(vf: float) -> str:
    return f"{float(vf):.4f}"


def simulate_instances(
    *,
    curve_dir: str,
    mat_type: str,
    vf: float,
    quarter_angles: List[float],
    instances: List[int],
    num_output_points: int = 20,
) -> Dict[int, Dict[str, Dict[str, np.ndarray]]]:
    data_gen_dir = _default_data_generation_dir()
    if not data_gen_dir.exists():
        raise FileNotFoundError(f"Quarter data-generation directory not found: {data_gen_dir}")

    if str(data_gen_dir) not in sys.path:
        sys.path.insert(0, str(data_gen_dir))

    gen_mod = _import_module_from_path("quarter_gen_v2", data_gen_dir / "generate_data_mp_version_2.py")
    lam2 = _import_module_from_path("quarter_lam2", data_gen_dir / "lam2.py")

    curve_dir_path = Path(curve_dir).resolve()
    if not curve_dir_path.exists():
        raise FileNotFoundError(f"curve_dir does not exist: {curve_dir_path}")

    vf_str = format_vol_fraction(vf)
    prefix = f"{mat_type}_{vf_str}"

    original_curve_dir = lam2.CURVE_DIR
    lam2.CURVE_DIR = curve_dir_path
    try:
        mat = lam2.load_ud_material_from_files(prefix)
    finally:
        lam2.CURVE_DIR = original_curve_dir

    full_angles = list(gen_mod.build_full_from_quarter_angles([float(a) for a in quarter_angles]))

    out: Dict[int, Dict[str, Dict[str, np.ndarray]]] = {}
    for inst in instances:
        sim_modes: Dict[str, Dict[str, np.ndarray]] = {}
        for mode in ("11", "22", "12"):
            try:
                ex, sx, ey, _gxy, ezz, _g23, _g13, e11 = gen_mod.run_simulation_with_mat(prefix, full_angles, mode, mat)
                if len(ex) < 2:
                    continue
                x_out = np.linspace(float(ex[0]), float(ex[-1]), int(num_output_points), dtype=np.float32)
                sx_out = np.interp(x_out, ex, sx / 1e6).astype(np.float32)
                if mode == "11":
                    lateral = np.interp(x_out, ex, ey).astype(np.float32)
                    eps33 = np.interp(x_out, ex, ezz).astype(np.float32)
                elif mode == "22":
                    lateral = np.interp(x_out, ex, e11).astype(np.float32)
                    eps33 = np.interp(x_out, ex, ezz).astype(np.float32)
                else:
                    lateral = None
                    eps33 = None

                sim_modes[mode] = {
                    "strain": x_out,
                    "stress": sx_out,
                    "lateral": lateral,
                    "eps33": eps33,
                }
            except Exception:
                continue
        if sim_modes:
            out[int(inst)] = sim_modes

    if not out:
        raise RuntimeError(f"No simulations succeeded for {mat_type} vf={vf_str} angles={quarter_angles}")
    return out