process_sdfs.py

Go to the documentation of this file.

#!/usr/bin/env python3
"""
Small helper script for working with SDF functions.
 
It:
- parses two arguments: -sdf_filename and -log_filename
- checks that both paths exist
- dynamically loads the given SDF Python file
- verifies that it defines a function:
 
    def sdf(delta_t, t, x, y, z, tx, ty, tz, block_id)
"""
 
import argparse
import inspect
import importlib.util
import sys
import re
from pathlib import Path
from multiprocessing import Pool, cpu_count
 
import numpy as np
import pyvista as pv
 
ISO_VALUE = 0.0
EXPECTED_SDF_ARGS = ["delta_t", "t", "x", "y", "z", "tx", "ty", "tz", "block_id"]
 
 
def create_uniform_grid():
    try:
        return pv.UniformGrid()
    except AttributeError:
        # Fallback for older / minimal PyVista builds
        from vtkmodules.vtkCommonDataModel import vtkImageData
        return pv.wrap(vtkImageData())
    
def make_grid(bounds, n):
    x_min, x_max, y_min, y_max, z_min, z_max = bounds
 
    x = np.linspace(x_min, x_max, n, dtype=float)
    y = np.linspace(y_min, y_max, n, dtype=float)
    z = np.linspace(z_min, z_max, n, dtype=float)
 
    X, Y, Z = np.meshgrid(x, y, z, indexing="ij")
 
    grid = create_uniform_grid()
    grid.dimensions = (n, n, n)
    grid.origin = (x[0], y[0], z[0])
    grid.spacing = (x[1] - x[0], y[1] - y[0], z[1] - z[0])
 
    return grid, X, Y, Z
 
def sdf_block_to_surface(
    sdf_func,
    delta_t,
    t,
    block_id,
    grid_template,
    X,
    Y,
    Z,
    tx=0.0,
    ty=0.0,
    tz=0.0,
):
    grid = grid_template.copy()
 
    sdf_vals = sdf_func(delta_t, t, X, Y, Z, tx, ty, tz, block_id)
    sdf_vals = np.asarray(sdf_vals, dtype=float)
 
    # Accept either 3D or flat array
    if sdf_vals.shape != X.shape:
        if sdf_vals.size == X.size:
            # reshape to 3D; change order="C"/"F" if your sdf uses different flattening
            sdf_vals = sdf_vals.reshape(X.shape)
        else:
            raise RuntimeError(
                f"sdf() returned array with shape {sdf_vals.shape}, "
                f"but expected {X.shape} or flat of size {X.size}"
            )
 
    grid["SDF"] = sdf_vals.ravel(order="F")
 
    try:
        surface = grid.contour([ISO_VALUE], scalars="SDF")
    except Exception:
        return None
 
    if surface.n_points == 0:
        return None
 
    return surface
 
def chunk_indices(n, n_chunks):
    """Yield index lists splitting range(n) into n_chunks chunks."""
    if n_chunks <= 0:
        n_chunks = 1
    chunk_size = (n + n_chunks - 1) // n_chunks
    for i in range(n_chunks):
        start = i * chunk_size
        end = min(start + chunk_size, n)
        if start < end:
            yield list(range(start, end))
 
def generate_sdf_meshes_parallel(
    sdf_path,
    block_ids,
    load_steps,
    total_times,
    delta_times,
    output_prefix,
    surface_only,
    grid_size,
    resolution,
    n_procs,
):
    n_steps = len(load_steps)
    if n_steps == 0:
        print("No time steps to process.")
        return
 
    # n_procs = max(1, min(n_procs, cpu_count()))
    print(f"Using {n_procs} worker process(es) for {n_steps} time steps.")
 
    bounds = (-grid_size, grid_size, -grid_size, grid_size, -grid_size, grid_size)
    sdf_path_str = str(sdf_path)
 
    tasks = []
    for indices in chunk_indices(n_steps, n_procs):
        task_args = (
            sdf_path_str,
            bounds,
            resolution,
            surface_only,
            output_prefix,
            block_ids,
            load_steps,
            total_times,
            delta_times,
            indices,
        )
        tasks.append(task_args)
 
    if n_procs == 1:
        # No Pool overhead, just run tasks in this process
        for targs in tasks:
            _process_time_steps_chunk(targs)
    else:
        with Pool(processes=n_procs) as pool:
            pool.map(_process_time_steps_chunk, tasks)
 
def _process_time_steps_chunk(args):
    """
    Worker function for multiprocessing.
 
    args = (
        sdf_path_str,
        bounds,
        resolution,
        surface_only,
        output_prefix,
        block_ids,
        load_steps,
        total_times,
        delta_times,
        indices,
    )
    """
    (
        sdf_path_str,
        bounds,
        resolution,
        surface_only,
        output_prefix,
        block_ids,
        load_steps,
        total_times,
        delta_times,
        indices,
    ) = args
 
    sdf_func = load_sdf_function(Path(sdf_path_str))
 
    # Each worker builds its own grid and coordinate arrays
    grid, X, Y, Z = make_grid(bounds=bounds, n=resolution)
 
    combined_sdf = None
 
    for idx in indices:
        ls = load_steps[idx]
        t = total_times[idx]
        dt = delta_times[idx]
 
        first_block = True
 
        for block_id in block_ids:
            vals = sdf_func(dt, t, X, Y, Z, 0.0, 0.0, 0.0, block_id)
            vals = np.asarray(vals)  # keep native dtype
 
            # Accept flat or 3D
            if vals.shape != X.shape:
                if vals.size == X.size:
                    vals = vals.reshape(X.shape)
                else:
                    raise RuntimeError(
                        f"sdf() returned shape {vals.shape}, "
                        f"expected {X.shape} or flat size {X.size}"
                    )
 
            if combined_sdf is None:
                combined_sdf = np.empty_like(vals)
 
            if first_block:
                combined_sdf[...] = vals
                first_block = False
            else:
                np.minimum(combined_sdf, vals, out=combined_sdf)
 
        if combined_sdf is None or first_block:
            print(f"[step {ls}] no SDF values, skipping")
            continue
 
        # SURFACE-ONLY MODE (STL)
        if surface_only:
            grid["SDF"] = combined_sdf.ravel(order="F")
            try:
                surface = grid.contour([ISO_VALUE], scalars="SDF")
            except Exception:
                print(f"[step {ls}] contour generation failed, skipping")
                continue
 
            if surface.n_points == 0:
                print(f"[step {ls}] empty surface, skipping")
                continue
 
            filename = f"{output_prefix}{ls}.stl"
            surface.save(filename)
            print(f"[step {ls}] saved STL surface to {filename}")
 
        # FULL VOLUME MODE (VTI)
        else:
            grid["SDF"] = combined_sdf.ravel(order="F")
            filename = f"{output_prefix}{ls}.vti"
            grid.save(filename)
            print(f"[step {ls}] saved VTI volume to {filename}")
 
def str2bool(v):
    if isinstance(v, bool):
        return v
    if v.lower() in ("yes", "true", "t", "1"):
        return True
    if v.lower() in ("no", "false", "f", "0"):
        return False
    raise argparse.ArgumentTypeError("Boolean value expected.")
 
def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Load an SDF definition and a log file."
    )
    parser.add_argument(
        "-sdf_filename",
        required=True,
        type=Path,
        help="Path to Python file defining sdf(delta_t, t, x, y, z, tx, ty, tz, block_id).",
    )
    parser.add_argument(
        "-log_filename",
        required=True,
        type=Path,
        help="Path to log text file.",
    )
    parser.add_argument(
        "-surface_only",
        type=str2bool,
        default=True,
        help="Generate only STL surfaces (fast). "
             "If not set, full VTI volumes are generated.",
    )
    parser.add_argument(
        "-grid_size",
        type=float,
        default=50.0,
        help="Half-size of the domain in each direction. "
             "Domain will be (-grid_size, grid_size)³. Default = 50.",
    )
    parser.add_argument(
        "-resolution",
        type=int,
        default=100,
        help="Grid resolution in each axis. Default = 100.",
    )
    parser.add_argument(
        "-np",
        type=int,
        default=4,
        help="Number of worker processes (cores) to use. Default = 4 (no multiprocessing).",
    )
 
    return parser.parse_args()
 
 
def check_paths(sdf_path: Path, log_path: Path) -> None:
    if not sdf_path.is_file():
        raise FileNotFoundError(f"SDF file does not exist or is not a file: {sdf_path}")
 
    if not log_path.is_file():
        raise FileNotFoundError(f"Log file does not exist or is not a file: {log_path}")
 
 
def extract_time_data(log_path: Path):
    # Pattern A (ep)
    pat_loadstep = re.compile(
        r"Load step\s+(\d+)\s+Time\s+([0-9.+-eE]+)\s+delta time\s+([0-9.+-eE]+)"
    )
 
    load_steps = []
    total_times = []
    delta_times = []
 
    with log_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = re.sub(r"\x1b\[[0-9;]*m", "", line)  # <-- ignore coloring artifacts
            m = pat_loadstep.search(line)
            if not m:
                continue
            load_steps.append(int(m.group(1)))
            total_times.append(float(m.group(2)))
            delta_times.append(float(m.group(3)))
 
    if load_steps:
        print(f"Extracted {len(load_steps)} time entries from 'Load step' lines.")
        return load_steps, total_times, delta_times
 
    # Fallback Pattern B (PETSc TS)
    pat_ts = re.compile(
        r"\[petsc\]\s+(\d+)\s+TS\s+dt\s+([0-9.+-eE]+)\s+time\s+([0-9.+-eE]+)",
        re.IGNORECASE,
    )
 
    load_steps = []
    total_times = []
    delta_times = []
 
    with log_path.open("r", encoding="utf-8") as f:
        for line in f:
            line = re.sub(r"\x1b\[[0-9;]*m", "", line)  # <-- ignore coloring artifacts
            m = pat_ts.search(line)
            if not m:
                continue
            load_steps.append(int(m.group(1)))     # TS step index
            delta_times.append(float(m.group(2)))  # dt
            total_times.append(float(m.group(3)))  # time
 
    if not load_steps:
        raise RuntimeError(
            "No time-step data found. Tried patterns:\n"
            "  - 'Load step <n> Time <t> delta time <dt>'\n"
            "  - '[petsc] <n> TS dt <dt> time <t>'"
        )
 
    print(f"Extracted {len(load_steps)} time entries from PETSc TS lines.")
    return load_steps, total_times, delta_times
 
def extract_sdf_block_ids(log_path: Path):
    """
    Extract msId values for blocks whose name contains 'SDF',
    but stop scanning once the log reaches the line:
        "Initialise MoFEM database <- done"
    """
 
    stop_phrase = "Initialise MoFEM database <- done"
 
    pattern = re.compile(
        r"type BLOCKSET.*?msId\s+(\d+).*?name\s+([A-Za-z0-9_]+)",
        re.IGNORECASE,
    )
 
    ids = []
 
    with log_path.open("r", encoding="utf-8") as f:
        for line in f:
            if stop_phrase in line:
                break  # stop scanning entirely
 
            m = pattern.search(line)
            if not m:
                continue
 
            ms_id = int(m.group(1))
            name = m.group(2)
 
            if "SDF" in name.upper():
                ids.append(ms_id)
                continue
            if "CONTACT" in name.upper():
                ids.append(ms_id)
 
    if not ids:
        raise RuntimeError("No SDF block IDs found in the log file.")
    print(f"Extracted SDF block IDs:  {ids}")
 
    return ids
 
def load_sdf_function(sdf_path: Path):
 
    module_name = "sdf_module_from_file"
 
    spec = importlib.util.spec_from_file_location(module_name, sdf_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot create import spec for {sdf_path}")
 
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # optional, but can help with imports inside the SDF file
    spec.loader.exec_module(module)  # type: ignore[union-attr]
 
    if not hasattr(module, "sdf"):
        raise RuntimeError(f"The SDF file {sdf_path} does not define a function named 'sdf'.")
 
    sdf_func = getattr(module, "sdf")
 
    if not callable(sdf_func):
        raise RuntimeError(f"'sdf' in {sdf_path} is not callable.")
 
    # Check the function signature
    sig = inspect.signature(sdf_func)
    param_names = list(sig.parameters.keys())
 
    if param_names != EXPECTED_SDF_ARGS:
        raise RuntimeError(
            f"sdf() in {sdf_path} has wrong parameters.\n"
            f"Expected: ({', '.join(EXPECTED_SDF_ARGS)})\n"
            f"Got:      ({', '.join(param_names)})"
        )
 
    return sdf_func
 
 
def main():
    args = parse_args()
 
    try:
        check_paths(args.sdf_filename, args.log_filename)
    except FileNotFoundError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)
 
    try:
        sdf = load_sdf_function(args.sdf_filename)
    except RuntimeError as exc:
        print(f"Error: {exc}", file=sys.stderr)
        sys.exit(1)
 
    print(f"SDF function loaded from: {args.sdf_filename}")
    print(f"Log file path:            {args.log_filename}")
    print(f"SDF signature:            {inspect.signature(sdf)}")
    print(f"Grid bounds:              {(-args.grid_size, args.grid_size)*3}")
    print(f"Grid resolution:          {args.resolution} points per axis")
    print(f"Surface only mode:        {args.surface_only}")
 
    load_steps, total_times, delta_times = extract_time_data(args.log_filename)
    block_ids = extract_sdf_block_ids(args.log_filename)
 
    print("All checks passed.")
    print("Generating SDF meshes for all time steps...")
 
    generate_sdf_meshes_parallel(
        args.sdf_filename,
        block_ids,
        load_steps,
        total_times,
        delta_times,
        "out_sdf_",
        args.surface_only,
        args.grid_size,
        args.resolution,
        args.np,
    )
    print("Done.")
 
if __name__ == "__main__":
    main()