diff --git a/.github/workflows/publish-witwin.yml b/.github/workflows/publish-witwin.yml index 8186784..c9215e7 100644 --- a/.github/workflows/publish-witwin.yml +++ b/.github/workflows/publish-witwin.yml @@ -83,14 +83,13 @@ jobs: PY build_cuda_wheels: - name: Build CUDA wheel / ${{ matrix.os }} / py${{ matrix.python_version }} + name: Build stable CUDA wheel / ${{ matrix.os }} needs: validate_release runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: os: [ubuntu-22.04, windows-2022] - python_version: ["3.10", "3.11", "3.12"] steps: - name: Check out repository @@ -99,7 +98,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: ${{ matrix.python_version }} + python-version: "3.10" - name: Free disk space (Linux) if: runner.os == 'Linux' @@ -113,11 +112,11 @@ jobs: echo "Disk after cleanup:"; df -h / - name: Install CUDA Toolkit - uses: Jimver/cuda-toolkit@v0.2.29 + uses: Jimver/cuda-toolkit@v0.2.35 with: - cuda: "12.5.0" + cuda: "12.8.1" method: local - log-file-suffix: "${{ runner.os }}-py${{ matrix.python_version }}" + log-file-suffix: "${{ runner.os }}-stable-abi" - name: Set up MSVC if: runner.os == 'Windows' @@ -149,19 +148,28 @@ jobs: - name: Install CUDA extension build dependencies run: | python -m pip install --upgrade pip setuptools wheel ninja numpy - python -m pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124 + python -m pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cu128 python -m pip install build - - name: Build packaged mesh SDF CUDA extension + - name: Build packaged stable mesh SDF CUDA extension shell: bash env: MAX_JOBS: "2" - TORCH_CUDA_ARCH_LIST: "7.0;8.0;8.6;8.9;9.0+PTX" - WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/witwin_core_mesh_sdf_cuda + WITWIN_CUDA_GENCODE_ARCHES: "7.0;7.5;8.0;8.6;8.9;9.0;10.0;10.1;12.0+PTX" PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/witwin_core_mesh_sdf_cuda/stable_abi_v1 run: | python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose + - name: Verify compiled CUDA architectures + shell: bash + env: + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + run: | + python "$PACKAGE_DIR"/scripts/verify_cuda_binary_arches.py \ + --stem witwin_core_mesh_sdf_cuda \ + "$PACKAGE_DIR"/witwin/core/geometry/cuda/prebuilt + - name: Build platform wheel shell: bash env: @@ -169,21 +177,25 @@ jobs: run: | python -m build --wheel "$PACKAGE_DIR" - - name: Normalize Linux wheel platform tag + - name: Repair Linux wheel platform tag if: runner.os == 'Linux' shell: bash env: PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} run: | - # PyPI rejects the raw linux_x86_64 tag; the extension is built on the - # ubuntu-22.04 runner (glibc 2.35), so retag to the matching manylinux - # tag. torch/CUDA runtime libraries are runtime dependencies and stay - # unbundled, so no auditwheel repair is required. - shopt -s nullglob + python -m pip install auditwheel patchelf + mkdir -p "$PACKAGE_DIR"/wheelhouse for whl in "$PACKAGE_DIR"/dist/*-linux_x86_64.whl; do - echo "Retagging $whl" - python -m wheel tags --platform-tag manylinux_2_35_x86_64 --remove "$whl" + python -m auditwheel repair "$whl" \ + --plat manylinux_2_35_x86_64 \ + --wheel-dir "$PACKAGE_DIR"/wheelhouse \ + --exclude libc10.so --exclude libc10_cuda.so \ + --exclude libtorch.so --exclude libtorch_cpu.so \ + --exclude libtorch_cuda.so --exclude libtorch_python.so \ + --exclude libcudart.so.12 --exclude libcuda.so.1 done + rm -f "$PACKAGE_DIR"/dist/*-linux_x86_64.whl + mv "$PACKAGE_DIR"/wheelhouse/*.whl "$PACKAGE_DIR"/dist/ ls -l "$PACKAGE_DIR"/dist - name: Smoke install prebuilt wheel @@ -197,9 +209,18 @@ jobs: python -m pip install --no-deps "$PACKAGE_DIR"/dist/*.whl python - <<'PY' import importlib.util + import os import site + import zipfile from pathlib import Path + wheel = next((Path(os.environ["PACKAGE_DIR"]) / "dist").glob("*.whl")) + names = zipfile.ZipFile(wheel).namelist() + native = [name for name in names if "witwin_core_mesh_sdf_cuda" in name and name.endswith((".so", ".pyd"))] + assert len(native) == 1, native + assert "/prebuilt/" in f"/{native[0]}" + assert "torch_" not in native[0] + candidates = [] for root in site.getsitepackages() + [site.getusersitepackages()]: candidates.append(Path(root) / "witwin" / "core" / "geometry" / "cuda" / "build.py") @@ -212,6 +233,7 @@ jobs: path = module.prebuilt_extension_path() assert path.exists(), f"missing packaged prebuilt extension: {path}" + assert path.parent.name == "prebuilt" extension = module.build_extension() print(f"Loaded packaged mesh SDF CUDA extension from {extension.__file__}") PY @@ -219,10 +241,83 @@ jobs: - name: Upload wheel artifact uses: actions/upload-artifact@v4 with: - name: wheel-${{ matrix.os }}-py${{ matrix.python_version }} + name: wheel-${{ matrix.os }} path: ${{ needs.validate_release.outputs.package_dir }}/dist/*.whl if-no-files-found: error + test_torch_compatibility: + name: Stable ABI / ${{ matrix.os }} / py${{ matrix.compatibility.python_version }} / torch${{ matrix.compatibility.torch_version }} + needs: + - validate_release + - build_cuda_wheels + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04, windows-2022] + compatibility: + - {python_version: "3.10", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.11", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.12", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.13", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.14", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.14", torch_version: "2.11.0", cuda_index: "cu128"} + - {python_version: "3.14", torch_version: "2.12.0", cuda_index: "cu126"} + + steps: + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.compatibility.python_version }} + + - name: Download stable wheel + uses: actions/download-artifact@v4 + with: + name: wheel-${{ matrix.os }} + path: dist + + - name: Install compatibility environment + shell: bash + run: | + python -m pip install --upgrade pip + python -m pip install --no-cache-dir "torch==${{ matrix.compatibility.torch_version }}" \ + --index-url "https://download.pytorch.org/whl/${{ matrix.compatibility.cuda_index }}" + python -m pip install --no-deps dist/*.whl + + - name: Load the same stable binary + shell: bash + env: + EXPECTED_TORCH: ${{ matrix.compatibility.torch_version }} + WITWIN_CORE_MESH_SDF_CUDA_PREBUILT: "1" + WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/must-not-jit-build + run: | + python - <<'PY' + import importlib.util + import os + import site + from pathlib import Path + + import torch + + assert torch.__version__.split("+", 1)[0] == os.environ["EXPECTED_TORCH"] + candidates = [ + Path(root) / "witwin" / "core" / "geometry" / "cuda" / "build.py" + for root in site.getsitepackages() + [site.getusersitepackages()] + ] + build_path = next((path for path in candidates if path.exists()), None) + assert build_path is not None, candidates + spec = importlib.util.spec_from_file_location("witwin_core_mesh_sdf_cuda_build", build_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + native_path = module.prebuilt_extension_path() + assert native_path.exists(), native_path + extension = module.build_extension() + assert Path(extension.__file__).resolve() == native_path.resolve() + assert extension.query_mesh_unsigned_distance is not None + print(f"Loaded {native_path.name} with torch {torch.__version__} on Python {os.sys.version.split()[0]}") + PY + build_sdist: name: Build source distribution needs: validate_release @@ -257,6 +352,7 @@ jobs: if: github.event_name == 'release' needs: - build_cuda_wheels + - test_torch_compatibility - build_sdist runs-on: ubuntu-latest permissions: diff --git a/.gitignore b/.gitignore index 618b9e1..f3a7b5c 100644 --- a/.gitignore +++ b/.gitignore @@ -171,9 +171,6 @@ venv.bak/ # mkdocs documentation /site -.slangtorch_cache/ -.fdtd_runtime_*.slang -.adjoint_kernels_runtime_*.slang *.png *.npy *.npz diff --git a/README.md b/README.md index 3dcf49c..3a98233 100644 --- a/README.md +++ b/README.md @@ -1,20 +1,30 @@ # WiTwin Core -WiTwin Core is the core package of the WiTwin stack. It provides the shared foundations used across WiTwin projects, with a focus on common scene representation, geometry primitives, mesh utilities, and state management that can be reused by different simulation modules. +WiTwin Core is the shared data and geometry package of the WiTwin stack. It defines the geometry and material contracts consumed by solver packages without imposing a common solver-specific scene lifecycle. ## Get Started -Python 3.10+ and an NVIDIA GPU are required for the main WiTwin simulation workflows. +CPython 3.10-3.14, PyTorch 2.10 or newer, and an NVIDIA GPU are supported for the main WiTwin simulation workflows. ```bash pip install witwin ``` +## Prebuilt CUDA Support + +Release wheels are built for Linux x86_64 and Windows x86_64 with CUDA 12.8. Each platform wheel carries one PyTorch Stable ABI native library built against PyTorch 2.10 and reused across supported Python and PyTorch versions; CI currently verifies PyTorch 2.10-2.12. The fat binaries contain native code for compute capabilities 7.0, 7.5, 8.0, 8.6, 8.9, 9.0, 10.0, 10.1, and 12.0, plus compute 12.0 PTX for forward-compatible Blackwell execution. This includes native coverage for RTX 2080-class Turing GPUs and current data-center and RTX/RTX PRO Blackwell families. + +Linux wheels target `manylinux_2_35_x86_64`. The installed NVIDIA driver must support the CUDA 12.x runtime supplied by PyTorch; the CUDA toolkit is only needed for source/JIT builds. + +For full CUDA 12.8 and Blackwell support, use at least driver 570.26 on Linux or 570.65 on Windows. Pre-Blackwell systems can use NVIDIA's CUDA 12.x minor-version compatibility floor (525.60.13 on Linux or 528.33 on Windows), subject to NVIDIA's compatibility-mode feature limits. + ## What It Provides -- Shared scene and geometry primitives in `witwin.core` -- Common state management for simulation-facing scene objects -- Reusable mesh, structure, and material building blocks for WiTwin solvers +- Runtime-checkable `GeometrySpec` and `MaterialSpec` contracts +- Shared analytic geometry primitives and differentiable mesh/SDF utilities +- Reusable `Material` and `Structure` values for solver-owned scenes + +Each solver owns its public `Scene` implementation. Radar-specific SMPL geometry lives in `witwin.radar.geometry`; Maxwell-specific `PolySlab` geometry lives in `witwin.maxwell.geometry`. ## Related Solvers diff --git a/hatch_build.py b/hatch_build.py index b204e2f..cfe0486 100644 --- a/hatch_build.py +++ b/hatch_build.py @@ -1,5 +1,6 @@ from __future__ import annotations +import sysconfig from pathlib import Path from hatchling.builders.hooks.plugin.interface import BuildHookInterface @@ -13,5 +14,6 @@ def initialize(self, version: str, build_data: dict) -> None: prebuilt_dir = Path(self.root) / "witwin" / "core" / "geometry" / "cuda" / "prebuilt" has_prebuilt_extension = any(prebuilt_dir.glob("witwin_core_mesh_sdf_cuda.*")) if has_prebuilt_extension: - build_data["infer_tag"] = True + platform_tag = sysconfig.get_platform().replace("-", "_").replace(".", "_") + build_data["tag"] = f"py3-none-{platform_tag}" build_data["pure_python"] = False diff --git a/pyproject.toml b/pyproject.toml index 186fe70..a6d507f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,27 +4,27 @@ build-backend = "hatchling.build" [project] name = "witwin" -version = "0.0.2" -description = "Witwin - Multi-physics simulation platform" +version = "0.3.0" +description = "Witwin shared geometry, material, and structure contracts" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.10,<3.15" license = "MIT" dependencies = [ "numpy>=1.24", - "torch>=2.0", + "torch>=2.10", ] [project.optional-dependencies] -maxwell = ["witwin-maxwell>=0.0.1"] -radar = ["witwin-radar>=0.0.1"] -channel = ["witwin-channel>=0.0.1"] -genesis = ["witwin-genesis>=0.0.1"] +maxwell = ["witwin-maxwell>=0.2,<0.3"] +radar = ["witwin-radar>=0.3,<0.4"] +channel = ["witwin-channel-native>=0.1,<0.2"] +genesis = ["witwin-genesis>=0.1,<0.2"] all = [ - "witwin-maxwell>=0.0.1", - "witwin-radar>=0.0.1", - "witwin-channel>=0.0.1", - "witwin-genesis>=0.0.1", + "witwin-maxwell>=0.2,<0.3", + "witwin-radar>=0.3,<0.4", + "witwin-channel-native>=0.1,<0.2", + "witwin-genesis>=0.1,<0.2", ] [tool.hatch.build.targets.wheel] @@ -34,8 +34,8 @@ artifacts = [ "witwin/core/geometry/cuda/**/*.cu", "witwin/core/geometry/cuda/**/*.cuh", "witwin/core/geometry/cuda/**/*.h", - "witwin/core/geometry/cuda/prebuilt/*.pyd", - "witwin/core/geometry/cuda/prebuilt/*.so", + "witwin/core/geometry/cuda/prebuilt/**/*.pyd", + "witwin/core/geometry/cuda/prebuilt/**/*.so", ] [tool.hatch.build.hooks.custom] diff --git a/scripts/build_mesh_sdf_cuda_prebuilt.py b/scripts/build_mesh_sdf_cuda_prebuilt.py index 86e3e43..1d2b226 100644 --- a/scripts/build_mesh_sdf_cuda_prebuilt.py +++ b/scripts/build_mesh_sdf_cuda_prebuilt.py @@ -39,24 +39,24 @@ def main() -> None: parser.add_argument("--verbose", action="store_true") args = parser.parse_args() + cuda_build = _load_cuda_build_module() build_dir = Path( os.environ.get( "WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR", - Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel", + Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel" / "stable_abi_v1", ) ) os.environ["WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR"] = str(build_dir) os.environ["WITWIN_CORE_MESH_SDF_CUDA_SKIP_PREBUILT"] = "1" _ensure_current_device_arch() - cuda_build = _load_cuda_build_module() module = cuda_build.build_extension(verbose=args.verbose) module_file = Path(module.__file__).resolve() target_dir = cuda_build.prebuilt_root() target_dir.mkdir(parents=True, exist_ok=True) for suffix in (".pyd", ".so"): - existing = target_dir / f"witwin_core_mesh_sdf_cuda{suffix}" + existing = target_dir / f"{cuda_build.EXTENSION_NAME}{suffix}" if existing.exists(): existing.unlink() target = cuda_build.prebuilt_extension_path() diff --git a/scripts/verify_cuda_binary_arches.py b/scripts/verify_cuda_binary_arches.py new file mode 100644 index 0000000..b7f4bee --- /dev/null +++ b/scripts/verify_cuda_binary_arches.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import argparse +import subprocess +import tempfile +import zipfile +from pathlib import Path + + +EXPECTED_SASS = ("70", "75", "80", "86", "89", "90", "100", "101", "120") +EXPECTED_PTX_TARGET = "sm_120" + + +def _matches(path: Path, stems: tuple[str, ...]) -> bool: + return path.suffix in {".pyd", ".so"} and any(path.name.startswith(stem) for stem in stems) + + +def _collect_binaries(inputs: list[Path], stems: tuple[str, ...], extract_root: Path) -> list[Path]: + binaries: list[Path] = [] + for input_path in inputs: + if input_path.suffix == ".whl": + with zipfile.ZipFile(input_path) as wheel: + for name in wheel.namelist(): + member = Path(name) + if not _matches(member, stems): + continue + wheel.extract(name, extract_root) + binaries.append(extract_root / member) + continue + if input_path.is_dir(): + binaries.extend(path for path in input_path.rglob("*") if _matches(path, stems)) + continue + if _matches(input_path, stems): + binaries.append(input_path) + return sorted(set(path.resolve() for path in binaries)) + + +def _cuobjdump(flag: str, binary: Path) -> str: + result = subprocess.run( + ["cuobjdump", flag, str(binary)], + check=True, + capture_output=True, + text=True, + ) + return f"{result.stdout}\n{result.stderr}" + + +def main() -> None: + parser = argparse.ArgumentParser(description="Verify CUDA SASS and PTX targets in release binaries.") + parser.add_argument("inputs", nargs="+", type=Path) + parser.add_argument("--stem", action="append", required=True) + args = parser.parse_args() + + with tempfile.TemporaryDirectory(prefix="witwin_cuda_arch_verify_") as temp_dir: + binaries = _collect_binaries(args.inputs, tuple(args.stem), Path(temp_dir)) + if not binaries: + raise SystemExit(f"No native binaries matching {args.stem!r} were found.") + + for binary in binaries: + elf_listing = _cuobjdump("--list-elf", binary) + missing_sass = [arch for arch in EXPECTED_SASS if f"sm_{arch}" not in elf_listing] + if missing_sass: + raise SystemExit(f"{binary} is missing SASS targets: {', '.join(missing_sass)}") + + ptx_dump = _cuobjdump("--dump-ptx", binary) + if f".target {EXPECTED_PTX_TARGET}" not in ptx_dump: + raise SystemExit(f"{binary} is missing compute 12.0 PTX.") + + print(f"Verified CUDA architectures in {binary}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_contracts.py b/tests/test_contracts.py new file mode 100644 index 0000000..692f86a --- /dev/null +++ b/tests/test_contracts.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import pytest +import torch + +from witwin.core import ( + FrequencyMaterialSample, + GeometrySpec, + MaterialCapabilities, + MaterialSpec, + StaticMaterialSample, + Structure, +) + + +class DuckGeometry: + kind = "duck" + + def to_mesh(self, segments: int = 16): + del segments + return torch.zeros((3, 3)), torch.tensor([[0, 1, 2]]) + + +class DuckMaterial: + name = "duck" + + def capabilities(self): + return MaterialCapabilities() + + def evaluate_static(self): + return StaticMaterialSample(eps_r=2.0) + + def evaluate_at_frequency(self, frequency: float): + del frequency + return FrequencyMaterialSample(eps_r=2.0) + + +def test_structure_accepts_protocol_conforming_values(): + geometry = DuckGeometry() + material = DuckMaterial() + + assert isinstance(geometry, GeometrySpec) + assert isinstance(material, MaterialSpec) + assert Structure(geometry=geometry, material=material).geometry is geometry + + +@pytest.mark.parametrize( + ("geometry", "material", "message"), + [ + (object(), DuckMaterial(), "GeometrySpec"), + (DuckGeometry(), object(), "MaterialSpec"), + ], +) +def test_structure_rejects_values_outside_shared_contracts(geometry, material, message): + with pytest.raises(TypeError, match=message): + Structure(geometry=geometry, material=material) diff --git a/tests/test_dependency_contract.py b/tests/test_dependency_contract.py new file mode 100644 index 0000000..96c4365 --- /dev/null +++ b/tests/test_dependency_contract.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import tomllib +from pathlib import Path + + +def test_runtime_and_optional_dependencies_do_not_include_slangtorch(): + pyproject = tomllib.loads(Path("pyproject.toml").read_text(encoding="utf-8")) + groups = [pyproject["project"]["dependencies"]] + groups.extend(pyproject["project"].get("optional-dependencies", {}).values()) + assert all( + not dependency.lower().startswith("slangtorch") + for group in groups + for dependency in group + ) diff --git a/tests/test_geometry_sdf.py b/tests/test_geometry_sdf.py index 1cbc0a1..49a255b 100644 --- a/tests/test_geometry_sdf.py +++ b/tests/test_geometry_sdf.py @@ -4,7 +4,6 @@ import torch from witwin.core import Box, Cone, Cylinder, Ellipsoid, HollowBox, Mesh, Prism, Pyramid, Sphere, Torus -from witwin.core.geometry import ComplexPolySlab, PolySlab from witwin.core.geometry.mesh_sdf import ( _mesh_sdf_available, _triangle_mesh_smooth_signed_distance_torch, @@ -152,39 +151,6 @@ def _require_cuda_mesh_sdf(): (0.5, 0.0, 0.0), (0.7, 0.0, 0.0), ), - ( - PolySlab( - vertices=[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.5), (-0.5, 0.5)], - bounds=(-0.5, 0.5), - axis="z", - ), - (0.0, 0.0, 0.0), - (0.5, 0.0, 0.0), - (0.7, 0.0, 0.0), - ), - ( - # Non-convex L-shape: the notch quadrant must be outside. - PolySlab( - vertices=[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.0), (0.0, 0.0), (0.0, 0.5), (-0.5, 0.5)], - bounds=(-0.5, 0.5), - axis="z", - ), - (-0.25, -0.25, 0.0), - (0.5, -0.25, 0.0), - (0.25, 0.25, 0.0), - ), - ( - ComplexPolySlab( - loops=[ - [(-0.5, -0.5), (0.5, -0.5), (0.5, 0.5), (-0.5, 0.5)], - [(-0.2, -0.2), (0.2, -0.2), (0.2, 0.2), (-0.2, 0.2)], - ], - bounds=(-0.5, 0.5), - ), - (0.35, 0.0, 0.0), - (0.5, 0.0, 0.0), - (0.0, 0.0, 0.0), - ), ], ) def test_primitive_signed_distance_has_expected_sign(geometry, inside_point, surface_point, outside_point): @@ -214,16 +180,6 @@ def test_primitive_signed_distance_has_expected_sign(geometry, inside_point, sur (Cone(radius=0.4, height=1.0, axis="z"), (0.0, 0.0, 0.5), (0.2, 0.0, 0.5), (0.3, 0.0, 0.5)), (Pyramid(base_size=1.0, height=1.0, axis="z"), (0.0, 0.0, 0.5), (0.25, 0.0, 0.5), (0.35, 0.0, 0.5)), (Prism(radius=0.4, height=1.0, num_sides=5, axis="z"), (0.0, 0.0, 0.0), (0.4, 0.0, 0.0), (0.6, 0.0, 0.0)), - ( - PolySlab( - vertices=[(-0.4, -0.4), (0.4, -0.4), (0.4, 0.4), (-0.4, 0.4)], - bounds=(-0.5, 0.5), - axis="z", - ), - (0.0, 0.0, 0.0), - (0.4, 0.0, 0.0), - (0.7, 0.0, 0.0), - ), ], ) def test_primitive_occupancy_is_monotone_and_bounded(geometry, inside_point, boundary_point, outside_point): @@ -371,27 +327,6 @@ def test_prism_occupancy_has_gradients_for_position_radius_height_and_rotation() assert abs(grad.item()) > 1.0e-6 -def test_polyslab_occupancy_has_gradients_for_vertices_bounds_and_sidewall_angle(): - xx, yy, zz = _asymmetric_grid() - vertices = torch.nn.Parameter( - torch.tensor([[-0.31, -0.24], [0.29, -0.24], [0.29, 0.26], [-0.31, 0.26]], dtype=torch.float32) - ) - bounds = torch.nn.Parameter(torch.tensor([-0.28, 0.17], dtype=torch.float32)) - sidewall_angle = torch.nn.Parameter(torch.tensor(0.10, dtype=torch.float32)) - - geometry = PolySlab(vertices, bounds, axis="z", sidewall_angle=sidewall_angle) - loss = geometry.to_mask(xx, yy, zz, beta=0.06).sum() - loss.backward() - - for grad in (vertices.grad, bounds.grad): - assert grad is not None - assert torch.all(torch.isfinite(grad)) - assert torch.any(grad.abs() > 1.0e-6) - assert sidewall_angle.grad is not None - assert torch.isfinite(sidewall_angle.grad) - assert abs(sidewall_angle.grad.item()) > 1.0e-6 - - def test_mesh_signed_distance_matches_box_sign_for_closed_cube(): vertices, faces = _cube_mesh(size=1.0) mesh = Mesh(vertices, faces, recenter=False, fill_mode="solid") diff --git a/tests/test_geometry_smoke.py b/tests/test_geometry_smoke.py index d694e8d..63afb37 100644 --- a/tests/test_geometry_smoke.py +++ b/tests/test_geometry_smoke.py @@ -1,10 +1,11 @@ from __future__ import annotations +import math + import pytest import torch -from witwin.core import Box, Cylinder, Mesh, Sphere -from witwin.core.geometry import PolySlab +from witwin.core import Box, Cylinder, Mesh, Sphere, Torus def _grid(): @@ -37,17 +38,6 @@ def _grid(): (0.4, 0.2, -0.1), 64, ), - ( - PolySlab( - vertices=[(-0.3, -0.3), (0.3, -0.3), (0.3, 0.3), (-0.3, 0.3)], - bounds=(-0.35, 0.35), - position=(0.1, -0.2, 0.3), - ), - 16, - (0.1, -0.2, 0.3), - (0.8, -0.2, 0.3), - 12, - ), ], ) def test_geometry_construction_to_mesh_and_to_mask(geometry, segments, inside_point, outside_point, expected_faces): @@ -83,44 +73,24 @@ def test_geometry_construction_to_mesh_and_to_mask(geometry, segments, inside_po assert torch.any(occupancy < 0.5) -def test_polyslab_nonconvex_to_mesh_produces_expected_topology(): - slab = PolySlab( - vertices=[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.0), (0.0, 0.0), (0.0, 0.5), (-0.5, 0.5)], - bounds=(-0.25, 0.25), - sidewall_angle=0.1, - ) - vertices, faces = slab.to_mesh() - - # 6 polygon vertices: 12 side triangles + 2 * (6 - 2) ear-clipped cap triangles. - assert vertices.shape == (12, 3) - assert faces.shape == (20, 3) - assert faces.dtype == torch.int64 - - -def test_polyslab_to_mesh_matches_signed_distance_for_all_axes(): - # Asymmetric polygon so a u/v mirror in the axis remap cannot go unnoticed. - vertices_2d = [(0.1, 0.0), (0.9, 0.0), (0.1, 0.6)] - for axis in ("x", "y", "z"): - slab = PolySlab(vertices=vertices_2d, bounds=(-0.5, 0.5), axis=axis) - mesh_vertices, faces = slab.to_mesh() - - distance = slab.signed_distance(mesh_vertices[:, 0], mesh_vertices[:, 1], mesh_vertices[:, 2]) - assert torch.all(distance.abs() < 1.0e-5), axis - - # Positive divergence-theorem volume implies outward face normals. - tri = mesh_vertices[faces] - volume = torch.sum(torch.sum(tri[:, 0] * torch.linalg.cross(tri[:, 1], tri[:, 2], dim=1), dim=1)) / 6.0 - assert volume.item() == pytest.approx(0.5 * 0.8 * 0.6, rel=1.0e-4), axis - +def test_torus_mesh_matches_analytic_extents_and_volume(): + major_radius = 0.8 + minor_radius = 0.2 + torus = Torus(major_radius=major_radius, minor_radius=minor_radius, axis="z") + vertices, faces = torus.to_mesh(segments=48) -def test_polyslab_to_mesh_accepts_closed_ring_vertices(): - slab = PolySlab( - vertices=[(-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0), (-1.0, -1.0)], - bounds=(-0.5, 0.5), + expected_extents = torch.tensor( + [major_radius + minor_radius, major_radius + minor_radius, minor_radius], + dtype=vertices.dtype, ) - vertices, faces = slab.to_mesh() - assert vertices.shape == (8, 3) - assert faces.shape == (12, 3) + torch.testing.assert_close(vertices.abs().amax(dim=0), expected_extents, rtol=1e-6, atol=1e-6) + + triangles = vertices[faces].to(torch.float64) + signed_volume = torch.sum( + torch.sum(triangles[:, 0] * torch.cross(triangles[:, 1], triangles[:, 2], dim=1), dim=1) + ) / 6.0 + expected_volume = 2.0 * math.pi**2 * major_radius * minor_radius**2 + assert abs(float(signed_volume)) == pytest.approx(expected_volume, rel=1e-2) def test_mesh_roundtrip_preserves_world_vertices_and_faces(): diff --git a/tests/test_scene_base.py b/tests/test_scene_base.py deleted file mode 100644 index 628d496..0000000 --- a/tests/test_scene_base.py +++ /dev/null @@ -1,18 +0,0 @@ -from witwin.core.scene import SceneBase - - -def test_scene_base_uses_add_mutators_without_with_aliases(): - scene = SceneBase(device="cpu") - structure = object() - source = object() - monitor = object() - - returned = scene.add_structure(structure).add_source(source).add_monitor(monitor) - - assert returned is scene - assert scene.structures == [structure] - assert scene.sources == [source] - assert scene.monitors == [monitor] - assert not hasattr(scene, "with_structure") - assert not hasattr(scene, "with_source") - assert not hasattr(scene, "with_monitor") diff --git a/witwin/core/__init__.py b/witwin/core/__init__.py index 0a08391..0b748ed 100644 --- a/witwin/core/__init__.py +++ b/witwin/core/__init__.py @@ -1,9 +1,8 @@ -"""Witwin Core - Shared geometry, materials, and scene utilities.""" +"""Witwin Core - Shared geometry, material, and structure contracts.""" -__version__ = "0.0.1" +__version__ = "0.3.0" from .material import ( - DifferentiableMaterial, FrequencyMaterialSample, Material, MaterialCapabilities, @@ -18,11 +17,11 @@ Ellipsoid, Geometry, GeometryBase, + GeometrySpec, HollowBox, Mesh, Prism, Pyramid, - SMPLBody, Sphere, Torus, ) @@ -33,38 +32,22 @@ quat_to_rotation_matrix, quat_to_rotation_matrix_np, ) -from .scene import SceneBase -from .scene_to_mitsuba import ( - MitsubaRenderable, - MitsubaSceneHandle, - build_mitsuba_scene, - create_mitsuba_mesh, - update_mitsuba_scene_vertices, -) - __all__ = [ "GeometryBase", + "GeometrySpec", "Box", "Sphere", "Cylinder", "Cone", "Ellipsoid", "Pyramid", "Prism", "Torus", "HollowBox", "Material", - "DifferentiableMaterial", "MaterialCapabilities", "MaterialSpec", "Mesh", - "SMPLBody", "FrequencyMaterialSample", "Geometry", - "SceneBase", "StaticMaterialSample", "Structure", - "MitsubaRenderable", - "MitsubaSceneHandle", "quat_from_euler", "quat_identity", "quat_multiply", "quat_to_rotation_matrix", "quat_to_rotation_matrix_np", - "create_mitsuba_mesh", - "build_mitsuba_scene", - "update_mitsuba_scene_vertices", ] diff --git a/witwin/core/geometry/__init__.py b/witwin/core/geometry/__init__.py index 837b5de..e7152f9 100644 --- a/witwin/core/geometry/__init__.py +++ b/witwin/core/geometry/__init__.py @@ -1,14 +1,14 @@ """Shared differentiable geometry package.""" -from .base import GeometryBase +from .base import GeometryBase, GeometrySpec from .mesh import Mesh -from .primitives import Box, ComplexPolySlab, Cone, Cylinder, Ellipsoid, HollowBox, PolySlab, Prism, Pyramid, Sphere, Torus -from .smpl import SMPLBody +from .primitives import Box, Cone, Cylinder, Ellipsoid, HollowBox, Prism, Pyramid, Sphere, Torus -Geometry = Box | Sphere | Cylinder | Cone | Ellipsoid | Pyramid | Prism | PolySlab | ComplexPolySlab | Torus | HollowBox | Mesh | SMPLBody +Geometry = Box | Sphere | Cylinder | Cone | Ellipsoid | Pyramid | Prism | Torus | HollowBox | Mesh __all__ = [ "GeometryBase", + "GeometrySpec", "Box", "Sphere", "Cylinder", @@ -16,11 +16,8 @@ "Ellipsoid", "Pyramid", "Prism", - "PolySlab", - "ComplexPolySlab", "Torus", "HollowBox", "Mesh", - "SMPLBody", "Geometry", ] diff --git a/witwin/core/geometry/base.py b/witwin/core/geometry/base.py index b2cd58a..36a2f75 100644 --- a/witwin/core/geometry/base.py +++ b/witwin/core/geometry/base.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Protocol, runtime_checkable + import numpy as np import torch @@ -136,6 +138,16 @@ def occupancy_from_signed_distance(signed_distance: torch.Tensor, *, xx, yy, zz, return occupancy.clamp(0.0, 1.0) +@runtime_checkable +class GeometrySpec(Protocol): + """Minimal geometry contract consumed by solver-specific scenes.""" + + kind: str + + def to_mesh(self, segments: int = 16) -> tuple: + ... + + class GeometryBase: """Base class for all geometry primitives.""" diff --git a/witwin/core/geometry/cuda/backend.py b/witwin/core/geometry/cuda/backend.py index 63afa2a..c0cff8f 100644 --- a/witwin/core/geometry/cuda/backend.py +++ b/witwin/core/geometry/cuda/backend.py @@ -1,6 +1,6 @@ """Loader for the native CUDA mesh-SDF extension. -Exposes a module object that is call-compatible with the former SlangTorch +Exposes a module object that is call-compatible with the former JIT-kernel module: ``module.(**kwargs).launchRaw(blockSize=..., gridSize=...)``. The kernels compute their own launch grid, so ``launchRaw`` ignores the block / grid hints and simply dispatches the compiled function. diff --git a/witwin/core/geometry/cuda/build.py b/witwin/core/geometry/cuda/build.py index c66e3ca..44aba49 100644 --- a/witwin/core/geometry/cuda/build.py +++ b/witwin/core/geometry/cuda/build.py @@ -8,7 +8,6 @@ from __future__ import annotations -import importlib.util import os import shutil import subprocess @@ -16,6 +15,7 @@ import tempfile from pathlib import Path +import torch from torch.utils.cpp_extension import load @@ -184,19 +184,49 @@ def extension_sources() -> list[Path]: return [root / "extension.cpp", root / "mesh_sdf_kernels.cu"] -def _load_extension_file(module_path: Path): - spec = importlib.util.spec_from_file_location(EXTENSION_NAME, module_path) - if spec is None or spec.loader is None: - raise ImportError(f"Unable to create import spec for {module_path}.") - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module +def _cuda_gencode_flags() -> list[str]: + """Translate the release architecture list directly into nvcc flags.""" + value = os.environ.get("WITWIN_CUDA_GENCODE_ARCHES") + if not value: + return [] + flags: list[str] = [] + for entry in value.split(";"): + entry = entry.strip() + if not entry: + continue + include_ptx = entry.endswith("+PTX") + number = entry.removesuffix("+PTX").replace(".", "") + if not number.isdigit(): + raise ValueError(f"Invalid CUDA architecture {entry!r} in WITWIN_CUDA_GENCODE_ARCHES.") + flags.append(f"-gencode=arch=compute_{number},code=sm_{number}") + if include_ptx: + flags.append(f"-gencode=arch=compute_{number},code=compute_{number}") + return flags + + +class _StableOpsModule: + """Attribute-compatible view of the Stable ABI dispatcher operators.""" + + def __init__(self, library_path: Path) -> None: + self.__file__ = str(library_path) + + def is_available(self) -> bool: + return bool(torch.cuda.is_available()) + + def __getattr__(self, name: str): + return getattr(torch.ops.witwin_core_mesh_sdf_cuda, name) + + +def _load_extension_file(library_path: Path) -> _StableOpsModule: + torch.ops.load_library(str(library_path)) + if not hasattr(torch.ops.witwin_core_mesh_sdf_cuda, "query_mesh_unsigned_distance"): + raise ImportError(f"{library_path} does not register the Stable ABI mesh SDF operators.") + return _StableOpsModule(library_path) def _load_packaged_prebuilt_extension(): - # Wheels ship a compiled extension in ``prebuilt/`` so installs never invoke - # nvcc/MSVC. A build against a mismatched torch ABI raises on import; the - # caller falls back to a JIT rebuild in that case. + # Wheels ship a Stable ABI library in ``prebuilt/`` so installs never invoke + # nvcc/MSVC and the same binary loads across supported Python/Torch versions. module_path = prebuilt_extension_path() if not module_path.exists(): return None @@ -218,19 +248,30 @@ def _jit_build(build_directory: Path, *, verbose: bool): _ensure_windows_build_tools_on_path() _ensure_cuda_home_from_nvcc() build_directory.mkdir(parents=True, exist_ok=True) - return load( + library_path = load( name=EXTENSION_NAME, sources=[str(path) for path in extension_sources()], build_directory=str(build_directory), - extra_cflags=["/O2"] if os.name == "nt" else ["-O3"], - extra_cuda_cflags=["-O3"], + extra_cflags=( + ["/O2", "/DTORCH_TARGET_VERSION=0x020a000000000000"] + if os.name == "nt" + else ["-O3", "-DTORCH_TARGET_VERSION=0x020a000000000000"] + ), + extra_cuda_cflags=[ + "-O3", + "-DTORCH_TARGET_VERSION=0x020a000000000000", + "-DUSE_CUDA", + *_cuda_gencode_flags(), + ], extra_ldflags=_conda_torch_ldflags(), + is_python_module=False, verbose=verbose, ) + return _StableOpsModule(Path(library_path)) def build_extension(*, verbose: bool = False): - default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME + default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME / "stable_abi_v1" build_directory = Path(os.environ.get("WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR", default_build_directory)) if os.environ.get("WITWIN_CORE_MESH_SDF_CUDA_SKIP_PREBUILT") != "1": diff --git a/witwin/core/geometry/cuda/extension.cpp b/witwin/core/geometry/cuda/extension.cpp index e1cad20..802d8f6 100644 --- a/witwin/core/geometry/cuda/extension.cpp +++ b/witwin/core/geometry/cuda/extension.cpp @@ -1,43 +1,29 @@ -#include -#include +#include -bool is_available() { - return torch::cuda::is_available(); -} - -void query_mesh_unsigned_distance_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, at::Tensor closest_triangle_index); -void query_mesh_distance_and_winding_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, - at::Tensor winding_angle, at::Tensor closest_triangle_index); -void query_mesh_unsigned_distance_bvh_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, - at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, - at::Tensor triangle_indices, at::Tensor unsigned_distance, at::Tensor closest_triangle_index); -void query_mesh_parity_sign_bvh_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, - at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, - at::Tensor triangle_indices, double jitter_scale, at::Tensor inside); -void backward_mesh_unsigned_distance_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, - at::Tensor grad_unsigned_distance, at::Tensor grad_triangles, at::Tensor grad_points); -void backward_mesh_distance_and_winding_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, - at::Tensor grad_unsigned_distance, at::Tensor grad_winding_angle, - at::Tensor grad_triangles, at::Tensor grad_points); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("is_available", &is_available, "Return whether CUDA is available to PyTorch."); - m.def("query_mesh_unsigned_distance", &query_mesh_unsigned_distance_cuda, - "Brute-force unsigned point-to-mesh distance and closest triangle index."); - m.def("query_mesh_distance_and_winding", &query_mesh_distance_and_winding_cuda, - "Brute-force unsigned distance plus solid-angle winding sum."); - m.def("query_mesh_unsigned_distance_bvh", &query_mesh_unsigned_distance_bvh_cuda, - "BVH-accelerated unsigned distance and closest triangle index."); - m.def("query_mesh_parity_sign_bvh", &query_mesh_parity_sign_bvh_cuda, - "BVH-accelerated ray-parity inside/outside test."); - m.def("backward_mesh_unsigned_distance", &backward_mesh_unsigned_distance_cuda, - "Backward pass for the unsigned distance query."); - m.def("backward_mesh_distance_and_winding", &backward_mesh_distance_and_winding_cuda, - "Backward pass for the distance-and-winding query."); +STABLE_TORCH_LIBRARY(witwin_core_mesh_sdf_cuda, m) { + m.def( + "query_mesh_unsigned_distance(Tensor triangles, Tensor points, " + "Tensor(a!) unsigned_distance, Tensor(b!) closest_triangle_index) -> ()"); + m.def( + "query_mesh_distance_and_winding(Tensor triangles, Tensor points, " + "Tensor(a!) unsigned_distance, Tensor(b!) winding_angle, " + "Tensor(c!) closest_triangle_index) -> ()"); + m.def( + "query_mesh_unsigned_distance_bvh(Tensor triangles, Tensor points, " + "Tensor node_bbox_min, Tensor node_bbox_max, Tensor node_left, Tensor node_right, " + "Tensor node_start, Tensor node_count, Tensor triangle_indices, " + "Tensor(a!) unsigned_distance, Tensor(b!) closest_triangle_index) -> ()"); + m.def( + "query_mesh_parity_sign_bvh(Tensor triangles, Tensor points, " + "Tensor node_bbox_min, Tensor node_bbox_max, Tensor node_left, Tensor node_right, " + "Tensor node_start, Tensor node_count, Tensor triangle_indices, float jitter_scale, " + "Tensor(a!) inside) -> ()"); + m.def( + "backward_mesh_unsigned_distance(Tensor triangles, Tensor points, " + "Tensor closest_triangle_index, Tensor grad_unsigned_distance, " + "Tensor(a!) grad_triangles, Tensor(b!) grad_points) -> ()"); + m.def( + "backward_mesh_distance_and_winding(Tensor triangles, Tensor points, " + "Tensor closest_triangle_index, Tensor grad_unsigned_distance, Tensor grad_winding_angle, " + "Tensor(a!) grad_triangles, Tensor(b!) grad_points) -> ()"); } diff --git a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu index 12a80df..a558c61 100644 --- a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu +++ b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu @@ -1,17 +1,22 @@ // Native CUDA triangle-mesh distance / winding / parity kernels. // -// Direct port of the former ``mesh_sdf.slang`` module. The forward kernels -// reproduce the Slang closest-point (Ericson) and solid-angle math verbatim; -// the backward kernels replace Slang autodiff with hand-derived analytic +// Direct port of the former mesh-SDF JIT module. The forward kernels +// reproduce its closest-point (Ericson) and solid-angle math verbatim; +// the backward kernels replace generated autodiff with hand-derived analytic // gradients. Every closest-point feature (vertex / edge / face) contributes // gradients to the triangle vertices through the barycentric weights of the // closest point, so grad_p == -(grad_a + grad_b + grad_c) holds by // construction (translation invariance of the distance). -#include -#include -#include -#include +#include +#include +#include +#include +#include +#include +#include + +#include namespace { @@ -57,7 +62,7 @@ __device__ __forceinline__ float3 read_vec3(const float* __restrict__ data, int return make_float3(v[0], v[1], v[2]); } -// Closest point on triangle (a, b, c) to p (Ericson, matching mesh_sdf.slang). +// Closest point on triangle (a, b, c) to p (Ericson, matching the reference implementation). // Returns the squared distance and the barycentric weights (la, lb, lc) of the // closest point so the backward pass can distribute vertex gradients. __device__ __forceinline__ float closest_point_barycentric( @@ -528,16 +533,29 @@ dim3 linear_grid(int elements, int block_size) { return dim3(static_cast((elements + block_size - 1) / block_size), 1, 1); } -void check_f32(const at::Tensor& t, const char* name) { - TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); - TORCH_CHECK(t.scalar_type() == at::kFloat, name, " must be float32"); - TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +void check_f32(const torch::stable::Tensor& tensor, const char* name) { + STD_TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + STD_TORCH_CHECK( + tensor.scalar_type() == torch::headeronly::ScalarType::Float, + name, + " must be float32"); + STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +void check_i32(const torch::stable::Tensor& tensor, const char* name) { + STD_TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + STD_TORCH_CHECK( + tensor.scalar_type() == torch::headeronly::ScalarType::Int, + name, + " must be int32"); + STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); } -void check_i32(const at::Tensor& t, const char* name) { - TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); - TORCH_CHECK(t.scalar_type() == at::kInt, name, " must be int32"); - TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +cudaStream_t current_cuda_stream(const torch::stable::Tensor& tensor) { + void* stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(tensor.get_device_index(), &stream_ptr)); + return static_cast(stream_ptr); } constexpr int kBlock = 256; @@ -545,48 +563,62 @@ constexpr int kBlock = 256; } // namespace void query_mesh_unsigned_distance_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, at::Tensor closest_triangle_index) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + torch::stable::Tensor& unsigned_distance, + torch::stable::Tensor& closest_triangle_index) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_f32(unsigned_distance, "unsigned_distance"); check_i32(closest_triangle_index, "closest_triangle_index"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - query_unsigned_distance_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), + query_unsigned_distance_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), static_cast(triangles.size(0)), point_count, - unsigned_distance.data_ptr(), closest_triangle_index.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + unsigned_distance.mutable_data_ptr(), closest_triangle_index.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } void query_mesh_distance_and_winding_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, - at::Tensor winding_angle, at::Tensor closest_triangle_index) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + torch::stable::Tensor& unsigned_distance, + torch::stable::Tensor& winding_angle, + torch::stable::Tensor& closest_triangle_index) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_f32(unsigned_distance, "unsigned_distance"); check_f32(winding_angle, "winding_angle"); check_i32(closest_triangle_index, "closest_triangle_index"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - query_distance_and_winding_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), + query_distance_and_winding_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), static_cast(triangles.size(0)), point_count, - unsigned_distance.data_ptr(), winding_angle.data_ptr(), - closest_triangle_index.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + unsigned_distance.mutable_data_ptr(), winding_angle.mutable_data_ptr(), + closest_triangle_index.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } void query_mesh_unsigned_distance_bvh_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, - at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, - at::Tensor triangle_indices, at::Tensor unsigned_distance, at::Tensor closest_triangle_index) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& node_bbox_min, + const torch::stable::Tensor& node_bbox_max, + const torch::stable::Tensor& node_left, + const torch::stable::Tensor& node_right, + const torch::stable::Tensor& node_start, + const torch::stable::Tensor& node_count, + const torch::stable::Tensor& triangle_indices, + torch::stable::Tensor& unsigned_distance, + torch::stable::Tensor& closest_triangle_index) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_f32(node_bbox_min, "node_bbox_min"); @@ -598,25 +630,33 @@ void query_mesh_unsigned_distance_bvh_cuda( check_i32(triangle_indices, "triangle_indices"); check_f32(unsigned_distance, "unsigned_distance"); check_i32(closest_triangle_index, "closest_triangle_index"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - query_unsigned_distance_bvh_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), - node_bbox_min.data_ptr(), node_bbox_max.data_ptr(), - node_left.data_ptr(), node_right.data_ptr(), - node_start.data_ptr(), node_count.data_ptr(), - triangle_indices.data_ptr(), point_count, - unsigned_distance.data_ptr(), closest_triangle_index.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + query_unsigned_distance_bvh_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + node_bbox_min.const_data_ptr(), node_bbox_max.const_data_ptr(), + node_left.const_data_ptr(), node_right.const_data_ptr(), + node_start.const_data_ptr(), node_count.const_data_ptr(), + triangle_indices.const_data_ptr(), point_count, + unsigned_distance.mutable_data_ptr(), closest_triangle_index.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } void query_mesh_parity_sign_bvh_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, - at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, - at::Tensor triangle_indices, double jitter_scale, at::Tensor inside) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& node_bbox_min, + const torch::stable::Tensor& node_bbox_max, + const torch::stable::Tensor& node_left, + const torch::stable::Tensor& node_right, + const torch::stable::Tensor& node_start, + const torch::stable::Tensor& node_count, + const torch::stable::Tensor& triangle_indices, + double jitter_scale, + torch::stable::Tensor& inside) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_f32(node_bbox_min, "node_bbox_min"); @@ -627,46 +667,54 @@ void query_mesh_parity_sign_bvh_cuda( check_i32(node_count, "node_count"); check_i32(triangle_indices, "triangle_indices"); check_i32(inside, "inside"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - query_parity_sign_bvh_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), - node_bbox_min.data_ptr(), node_bbox_max.data_ptr(), - node_left.data_ptr(), node_right.data_ptr(), - node_start.data_ptr(), node_count.data_ptr(), - triangle_indices.data_ptr(), static_cast(jitter_scale), point_count, - inside.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + query_parity_sign_bvh_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + node_bbox_min.const_data_ptr(), node_bbox_max.const_data_ptr(), + node_left.const_data_ptr(), node_right.const_data_ptr(), + node_start.const_data_ptr(), node_count.const_data_ptr(), + triangle_indices.const_data_ptr(), static_cast(jitter_scale), point_count, + inside.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } void backward_mesh_unsigned_distance_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, - at::Tensor grad_unsigned_distance, at::Tensor grad_triangles, at::Tensor grad_points) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& closest_triangle_index, + const torch::stable::Tensor& grad_unsigned_distance, + torch::stable::Tensor& grad_triangles, + torch::stable::Tensor& grad_points) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_i32(closest_triangle_index, "closest_triangle_index"); check_f32(grad_unsigned_distance, "grad_unsigned_distance"); check_f32(grad_triangles, "grad_triangles"); check_f32(grad_points, "grad_points"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - backward_unsigned_distance_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), - closest_triangle_index.data_ptr(), grad_unsigned_distance.data_ptr(), - point_count, grad_triangles.data_ptr(), grad_points.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + backward_unsigned_distance_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + closest_triangle_index.const_data_ptr(), grad_unsigned_distance.const_data_ptr(), + point_count, grad_triangles.mutable_data_ptr(), grad_points.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); } void backward_mesh_distance_and_winding_cuda( - at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, - at::Tensor grad_unsigned_distance, at::Tensor grad_winding_angle, - at::Tensor grad_triangles, at::Tensor grad_points) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& closest_triangle_index, + const torch::stable::Tensor& grad_unsigned_distance, + const torch::stable::Tensor& grad_winding_angle, + torch::stable::Tensor& grad_triangles, + torch::stable::Tensor& grad_points) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_i32(closest_triangle_index, "closest_triangle_index"); @@ -674,15 +722,24 @@ void backward_mesh_distance_and_winding_cuda( check_f32(grad_winding_angle, "grad_winding_angle"); check_f32(grad_triangles, "grad_triangles"); check_f32(grad_points, "grad_points"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - backward_distance_and_winding_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), - closest_triangle_index.data_ptr(), grad_unsigned_distance.data_ptr(), - grad_winding_angle.data_ptr(), static_cast(triangles.size(0)), point_count, - grad_triangles.data_ptr(), grad_points.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + backward_distance_and_winding_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + closest_triangle_index.const_data_ptr(), grad_unsigned_distance.const_data_ptr(), + grad_winding_angle.const_data_ptr(), static_cast(triangles.size(0)), point_count, + grad_triangles.mutable_data_ptr(), grad_points.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); +} + +STABLE_TORCH_LIBRARY_IMPL(witwin_core_mesh_sdf_cuda, CUDA, m) { + m.impl("query_mesh_unsigned_distance", TORCH_BOX(&query_mesh_unsigned_distance_cuda)); + m.impl("query_mesh_distance_and_winding", TORCH_BOX(&query_mesh_distance_and_winding_cuda)); + m.impl("query_mesh_unsigned_distance_bvh", TORCH_BOX(&query_mesh_unsigned_distance_bvh_cuda)); + m.impl("query_mesh_parity_sign_bvh", TORCH_BOX(&query_mesh_parity_sign_bvh_cuda)); + m.impl("backward_mesh_unsigned_distance", TORCH_BOX(&backward_mesh_unsigned_distance_cuda)); + m.impl("backward_mesh_distance_and_winding", TORCH_BOX(&backward_mesh_distance_and_winding_cuda)); } diff --git a/witwin/core/geometry/polygon.py b/witwin/core/geometry/polygon.py new file mode 100644 index 0000000..8b56fb8 --- /dev/null +++ b/witwin/core/geometry/polygon.py @@ -0,0 +1,91 @@ +"""Shared differentiable polygon distance helpers.""" + +from __future__ import annotations + +import torch + + +def _cross_2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] + + +def _safe_signed_denom(value: torch.Tensor, eps: float) -> torch.Tensor: + return torch.where(value >= 0.0, value.clamp_min(eps), value.clamp_max(-eps)) + + +def _polygon_edge_distance_sq_2d( + points: torch.Tensor, + polygon: torch.Tensor, + *, + edges: torch.Tensor | None = None, + rel: torch.Tensor | None = None, +) -> torch.Tensor: + """Squared unsigned distance from (P, 2) points to the closed edge loop of an (M, 2) polygon. + + ``edges`` and ``rel`` may be passed in when the caller has already materialized them. + """ + if edges is None: + edges = torch.roll(polygon, shifts=-1, dims=0) - polygon + if rel is None: + rel = points[:, None, :] - polygon[None, :, :] + edge_norm_sq = torch.sum(edges**2, dim=-1).clamp_min(torch.finfo(points.dtype).eps) + t = torch.clamp( + torch.sum(rel * edges[None, :, :], dim=-1) / edge_norm_sq[None, :], 0.0, 1.0 + ) + closest = polygon[None, :, :] + t[..., None] * edges[None, :, :] + return ( + torch.sum((points[:, None, :] - closest) ** 2, dim=-1) + .min(dim=1) + .values.clamp_min(0.0) + ) + + +def convex_polygon_signed_distance_2d( + points: torch.Tensor, polygon: torch.Tensor +) -> torch.Tensor: + next_polygon = torch.roll(polygon, shifts=-1, dims=0) + edges = next_polygon - polygon + rel = points[:, None, :] - polygon[None, :, :] + dist = torch.sqrt( + _polygon_edge_distance_sq_2d(points, polygon, edges=edges, rel=rel) + ) + + polygon_area = 0.5 * torch.sum(_cross_2d(polygon, next_polygon)) + crosses = _cross_2d(edges[None, :, :], rel) + if polygon_area.detach().item() >= 0.0: + inside = torch.all(crosses >= -1.0e-7, dim=1) + else: + inside = torch.all(crosses <= 1.0e-7, dim=1) + return torch.where(inside, -dist, dist) + + +def _polygon_crossing_count_2d( + points: torch.Tensor, polygon: torch.Tensor +) -> torch.Tensor: + """Count +u horizontal-ray crossings of a closed (M, 2) polygon loop for each (P, 2) point.""" + start = polygon + end = torch.roll(polygon, shifts=-1, dims=0) + pu = points[:, 0:1] + pv = points[:, 1:2] + straddles = (start[None, :, 1] <= pv) != (end[None, :, 1] <= pv) + denom = _safe_signed_denom(end[:, 1] - start[:, 1], torch.finfo(points.dtype).eps) + t = (pv - start[None, :, 1]) / denom[None, :] + crossing_u = start[None, :, 0] + t * (end[:, 0] - start[:, 0])[None, :] + return (straddles & (pu < crossing_u)).sum(dim=1) + + +def polygon_loops_signed_distance_2d(points: torch.Tensor, loops) -> torch.Tensor: + """Even-odd signed distance from (P, 2) points to one or more closed 2D polygon loops. + + Negative inside; the interior is defined by the even-odd rule over the combined + crossing count of all loops, so winding direction is irrelevant and additional + loops carve holes. Valid for arbitrary (including non-convex and self-intersecting) + polygons. + """ + distance_sq = torch.stack( + [_polygon_edge_distance_sq_2d(points, loop) for loop in loops], dim=0 + ) + distance = torch.sqrt(distance_sq.min(dim=0).values) + crossings = sum(_polygon_crossing_count_2d(points, loop) for loop in loops) + inside = crossings % 2 == 1 + return torch.where(inside, -distance, distance) diff --git a/witwin/core/geometry/primitives.py b/witwin/core/geometry/primitives.py index 782a30d..fa48b8b 100644 --- a/witwin/core/geometry/primitives.py +++ b/witwin/core/geometry/primitives.py @@ -13,6 +13,7 @@ _constant_tensor, ) from .mesh_sdf import triangle_mesh_unsigned_distance +from .polygon import _safe_signed_denom, convex_polygon_signed_distance_2d def _faces_tensor(data, *, device) -> torch.Tensor: @@ -33,76 +34,6 @@ def _box_signed_distance(dx, dy, dz, half_size: torch.Tensor) -> torch.Tensor: return outside + inside -def _cross_2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - return a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] - - -def _safe_signed_denom(value: torch.Tensor, eps: float) -> torch.Tensor: - return torch.where(value >= 0.0, value.clamp_min(eps), value.clamp_max(-eps)) - - -def _polygon_edge_distance_sq_2d( - points: torch.Tensor, - polygon: torch.Tensor, - *, - edges: torch.Tensor | None = None, - rel: torch.Tensor | None = None, -) -> torch.Tensor: - """Squared unsigned distance from (P, 2) points to the closed edge loop of an (M, 2) polygon. - - ``edges`` and ``rel`` may be passed in when the caller has already materialized them. - """ - if edges is None: - edges = torch.roll(polygon, shifts=-1, dims=0) - polygon - if rel is None: - rel = points[:, None, :] - polygon[None, :, :] - edge_norm_sq = torch.sum(edges ** 2, dim=-1).clamp_min(torch.finfo(points.dtype).eps) - t = torch.clamp(torch.sum(rel * edges[None, :, :], dim=-1) / edge_norm_sq[None, :], 0.0, 1.0) - closest = polygon[None, :, :] + t[..., None] * edges[None, :, :] - return torch.sum((points[:, None, :] - closest) ** 2, dim=-1).min(dim=1).values.clamp_min(0.0) - - -def _convex_polygon_signed_distance_2d(points: torch.Tensor, polygon: torch.Tensor) -> torch.Tensor: - next_polygon = torch.roll(polygon, shifts=-1, dims=0) - edges = next_polygon - polygon - rel = points[:, None, :] - polygon[None, :, :] - dist = torch.sqrt(_polygon_edge_distance_sq_2d(points, polygon, edges=edges, rel=rel)) - - polygon_area = 0.5 * torch.sum(_cross_2d(polygon, next_polygon)) - crosses = _cross_2d(edges[None, :, :], rel) - if polygon_area.detach().item() >= 0.0: - inside = torch.all(crosses >= -1.0e-7, dim=1) - else: - inside = torch.all(crosses <= 1.0e-7, dim=1) - return torch.where(inside, -dist, dist) - - -def _polygon_crossing_count_2d(points: torch.Tensor, polygon: torch.Tensor) -> torch.Tensor: - """Count +u horizontal-ray crossings of a closed (M, 2) polygon loop for each (P, 2) point.""" - start = polygon - end = torch.roll(polygon, shifts=-1, dims=0) - pu = points[:, 0:1] - pv = points[:, 1:2] - straddles = (start[None, :, 1] <= pv) != (end[None, :, 1] <= pv) - denom = _safe_signed_denom(end[:, 1] - start[:, 1], torch.finfo(points.dtype).eps) - t = (pv - start[None, :, 1]) / denom[None, :] - crossing_u = start[None, :, 0] + t * (end[:, 0] - start[:, 0])[None, :] - return (straddles & (pu < crossing_u)).sum(dim=1) - - -def _polygon_loops_signed_distance_2d(points: torch.Tensor, loops) -> torch.Tensor: - """Even-odd signed distance from (P, 2) points to one or more closed 2D polygon loops. - - Negative inside; the interior is defined by the even-odd rule over the combined - crossing count of all loops, so winding direction is irrelevant and additional - loops carve holes. Valid for arbitrary (including non-convex and self-intersecting) - polygons. - """ - distance_sq = torch.stack([_polygon_edge_distance_sq_2d(points, loop) for loop in loops], dim=0) - distance = torch.sqrt(distance_sq.min(dim=0).values) - crossings = sum(_polygon_crossing_count_2d(points, loop) for loop in loops) - inside = crossings % 2 == 1 - return torch.where(inside, -distance, distance) def _canonical_coords(dx, dy, dz, axis: str): @@ -350,7 +281,7 @@ def signed_distance(self, xx, yy, zz): torch.stack([height, radius]), torch.stack([height, -radius]), ]) - return _convex_polygon_signed_distance_2d(profile_points, triangle).reshape(axial.shape) + return convex_polygon_signed_distance_2d(profile_points, triangle).reshape(axial.shape) def to_mesh(self, segments=16): device = self.device @@ -430,7 +361,7 @@ def signed_distance(self, xx, yy, zz): radius = torch.clamp(self.radius.to(device=px.device, dtype=px.dtype), min=torch.finfo(px.dtype).eps) angles = torch.linspace(0, 2 * np.pi, self.num_sides + 1, device=px.device)[:-1] polygon = torch.stack([radius * torch.cos(angles), radius * torch.sin(angles)], dim=1) - polygon_distance = _convex_polygon_signed_distance_2d(torch.stack([px, py], dim=-1).reshape(-1, 2), polygon).reshape(px.shape) + polygon_distance = convex_polygon_signed_distance_2d(torch.stack([px, py], dim=-1).reshape(-1, 2), polygon).reshape(px.shape) axial_distance = torch.abs(axial) - self.height.to(device=axial.device, dtype=axial.dtype) / 2 outside = torch.sqrt(torch.clamp(polygon_distance, min=0.0) ** 2 + torch.clamp(axial_distance, min=0.0) ** 2) inside = torch.clamp(torch.maximum(polygon_distance, axial_distance), max=0.0) @@ -496,7 +427,7 @@ def to_mesh(self, segments=16): for minor in range(minor_segments): phi = 2 * np.pi * minor / minor_segments vertices.append([ - np.cos(phi) * np.cos(theta), + np.cos(phi), np.cos(phi) * np.sin(theta), np.sin(phi), np.cos(theta), @@ -555,230 +486,3 @@ def to_mesh(self, segments=16): torch.cat([outer_vertices, inner_vertices], dim=0), torch.cat([outer_faces, inner_faces + outer_vertices.shape[0]], dim=0), ) - - -def _as_polygon_vertices(value, *, device=None) -> torch.Tensor: - if isinstance(value, torch.Tensor): - vertices = value.to(device=device, dtype=torch.float32) - else: - vertices = torch.tensor([[float(u), float(v)] for u, v in value], dtype=torch.float32, device=device) - if vertices.ndim != 2 or vertices.shape[0] < 3 or vertices.shape[1] != 2: - raise ValueError("vertices must have shape (N, 2) with N >= 3.") - return vertices - - -def _as_axis_bounds(value, *, device=None) -> torch.Tensor: - if isinstance(value, torch.Tensor): - bounds = value.to(device=device, dtype=torch.float32) - else: - bounds = torch.tensor([float(value[0]), float(value[1])], dtype=torch.float32, device=device) - if bounds.shape != (2,): - raise ValueError("bounds must contain (min, max) along the extrusion axis.") - return bounds - - -def _cross_2d_scalar(origin, first, second) -> float: - return float((first[0] - origin[0]) * (second[1] - origin[1]) - (first[1] - origin[1]) * (second[0] - origin[0])) - - -def _point_in_triangle_2d(point, a, b, c, eps=1.0e-12) -> bool: - """Inclusive point-in-triangle test for a CCW triangle (a, b, c).""" - return ( - _cross_2d_scalar(a, b, point) >= -eps - and _cross_2d_scalar(b, c, point) >= -eps - and _cross_2d_scalar(c, a, point) >= -eps - ) - - -def _ear_clip_faces_2d(polygon: np.ndarray) -> list[list[int]]: - """Triangulate a simple CCW (N, 2) polygon into CCW index triangles via ear clipping.""" - remaining = list(range(polygon.shape[0])) - faces: list[list[int]] = [] - while len(remaining) > 3: - clipped = False - for slot in range(len(remaining)): - i0 = remaining[slot - 1] - i1 = remaining[slot] - i2 = remaining[(slot + 1) % len(remaining)] - a, b, c = polygon[i0], polygon[i1], polygon[i2] - if _cross_2d_scalar(a, b, c) <= 1.0e-12: - continue - others = (polygon[j] for j in remaining if j not in (i0, i1, i2)) - if any(_point_in_triangle_2d(p, a, b, c) for p in others): - continue - faces.append([i0, i1, i2]) - remaining.pop(slot) - clipped = True - break - if not clipped: - raise ValueError("Failed to triangulate polygon; vertices must describe a simple polygon.") - faces.append(list(remaining)) - return faces - - -def _offset_polygon_2d(polygon: np.ndarray, offset: float) -> np.ndarray: - """Approximate outward miter offset of a CCW (N, 2) polygon (negative offset shrinks).""" - if offset == 0.0: - return polygon - previous_edges = polygon - np.roll(polygon, 1, axis=0) - next_edges = np.roll(polygon, -1, axis=0) - polygon - - def _outward_normals(edges: np.ndarray) -> np.ndarray: - normals = np.stack([edges[:, 1], -edges[:, 0]], axis=1) - return normals / np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1.0e-12) - - normal_previous = _outward_normals(previous_edges) - normal_next = _outward_normals(next_edges) - denom = np.maximum(1.0 + np.sum(normal_previous * normal_next, axis=1, keepdims=True), 1.0e-3) - return polygon + offset * (normal_previous + normal_next) / denom - - -class PolySlab(GeometryBase): - """Extruded 2D polygon with optional linear sidewall taper (Tidy3D semantics). - - ``vertices`` is the (N, 2) cross-section in the plane perpendicular to ``axis`` - (plane coordinates follow ``_axial_split``: axis z -> (x, y), y -> (x, z), - x -> (y, z)), ``bounds`` is (min, max) along the axis, and a positive - ``sidewall_angle`` (radians) shrinks the cross-section as the axis coordinate - increases. ``reference_plane`` selects the axis position where the drawn - polygon is exact. The interior uses the even-odd rule, so non-convex polygons - are supported and vertex winding is irrelevant. - """ - - kind = "poly_slab" - - def __init__( - self, - vertices, - bounds, - axis="z", - sidewall_angle=0.0, - reference_plane="middle", - position=(0, 0, 0), - rotation=None, - *, - device=None, - ): - super().__init__(position=position, rotation=rotation, device=device) - self.vertices: torch.Tensor = _as_polygon_vertices(vertices, device=device) - self.bounds: torch.Tensor = _as_axis_bounds(bounds, device=device) - if isinstance(axis, int) and not isinstance(axis, bool): - if axis not in (0, 1, 2): - raise ValueError("integer axis must be 0, 1, or 2.") - axis = ("x", "y", "z")[axis] - self.axis: str = self._validate_axis(axis) - self.sidewall_angle: torch.Tensor = _as_scalar(sidewall_angle, device=device) - if reference_plane not in ("bottom", "middle", "top"): - raise ValueError("reference_plane must be 'bottom', 'middle', or 'top'.") - self.reference_plane: str = reference_plane - self.loop_sizes: tuple[int, ...] = (int(self.vertices.shape[0]),) - - def _reference_position(self, bounds: torch.Tensor) -> torch.Tensor: - if self.reference_plane == "bottom": - return bounds[0] - if self.reference_plane == "top": - return bounds[1] - return 0.5 * (bounds[0] + bounds[1]) - - def signed_distance(self, xx, yy, zz): - dx, dy, dz = self._local_coords(xx, yy, zz) - axial, pu, pv = _axial_split(dx, dy, dz, self.axis) - vertices = self.vertices.to(device=pu.device, dtype=pu.dtype) - loops = torch.split(vertices, list(self.loop_sizes), dim=0) - points = torch.stack([pu, pv], dim=-1).reshape(-1, 2) - polygon_distance = _polygon_loops_signed_distance_2d(points, loops).reshape(pu.shape) - bounds = self.bounds.to(device=axial.device, dtype=axial.dtype) - tan_angle = torch.tan(self.sidewall_angle.to(device=axial.device, dtype=axial.dtype)) - # Signed distance minus an offset is exact morphological erosion/dilation, so - # the taper is a z-dependent inward offset of the cross-section. - tapered_distance = polygon_distance + tan_angle * (axial - self._reference_position(bounds)) - axial_distance = torch.abs(axial - 0.5 * (bounds[0] + bounds[1])) - 0.5 * (bounds[1] - bounds[0]) - outside = torch.sqrt(torch.clamp(tapered_distance, min=0.0) ** 2 + torch.clamp(axial_distance, min=0.0) ** 2) - inside = torch.clamp(torch.maximum(tapered_distance, axial_distance), max=0.0) - return outside + inside - - def to_mesh(self, segments=16): - del segments - device = self.device - polygon = self.vertices.detach().cpu().to(torch.float64).numpy() - # Drop consecutive duplicate vertices (including a closed-ring wrap-around - # repeat); coincident points would block every candidate ear below. - keep = np.linalg.norm(polygon - np.roll(polygon, 1, axis=0), axis=1) > 1.0e-12 - polygon = polygon[keep] - if polygon.shape[0] < 3: - raise ValueError("vertices must contain at least 3 distinct points.") - signed_area = 0.5 * float(np.sum(polygon[:, 0] * np.roll(polygon[:, 1], -1) - np.roll(polygon[:, 0], -1) * polygon[:, 1])) - if signed_area < 0.0: - polygon = polygon[::-1].copy() - bounds = self.bounds.detach().cpu().to(torch.float64).numpy() - reference = {"bottom": bounds[0], "top": bounds[1], "middle": 0.5 * (bounds[0] + bounds[1])}[self.reference_plane] - tan_angle = float(np.tan(self.sidewall_angle.detach().cpu().to(torch.float64).item())) - bottom_ring = _offset_polygon_2d(polygon, -tan_angle * (bounds[0] - reference)) - top_ring = _offset_polygon_2d(polygon, -tan_angle * (bounds[1] - reference)) - count = polygon.shape[0] - bottom = np.concatenate([bottom_ring, np.full((count, 1), bounds[0])], axis=1) - top = np.concatenate([top_ring, np.full((count, 1), bounds[1])], axis=1) - vertices = _constant_tensor(np.concatenate([bottom, top], axis=0), device=device) - faces = [] - for index in range(count): - next_index = (index + 1) % count - faces.extend([ - [index, next_index, count + next_index], - [index, count + next_index, count + index], - ]) - for i0, i1, i2 in _ear_clip_faces_2d(polygon): - faces.append([i0, i2, i1]) - faces.append([count + i0, count + i1, count + i2]) - # Map local (u, v, axial) columns to world axes following _axial_split - # (axis z -> (x, y), y -> (x, z), x -> (y, z)); _remap_axis_torch would - # mirror u/v for axis "x" and disagree with signed_distance/to_mask. - vertices = vertices[:, {"z": [0, 1, 2], "y": [0, 2, 1], "x": [2, 0, 1]}[self.axis]] - faces_tensor = _faces_tensor(faces, device=device) - if self.axis == "y": - # Odd column permutation reflects orientation; flip winding to keep - # face normals outward. - faces_tensor = torch.flip(faces_tensor, dims=(1,)) - vertices = self._transform_mesh_verts(vertices) - return vertices, faces_tensor - - -class ComplexPolySlab(PolySlab): - """PolySlab variant for non-simple cross-sections. - - Accepts a list of vertex loops; self-intersecting loops (e.g. bowties) and - multiple loops are allowed. The interior is the even-odd rule over the combined - crossing count of all loops, so loop winding is irrelevant and extra loops carve - holes. The distance magnitude is the minimum distance over all edges of all loops. - """ - - kind = "complex_poly_slab" - - def __init__( - self, - loops, - bounds, - axis="z", - sidewall_angle=0.0, - reference_plane="middle", - position=(0, 0, 0), - rotation=None, - *, - device=None, - ): - loop_tensors = [_as_polygon_vertices(loop, device=device) for loop in loops] - if not loop_tensors: - raise ValueError("loops must contain at least one polygon loop.") - super().__init__( - torch.cat(loop_tensors, dim=0), - bounds, - axis=axis, - sidewall_angle=sidewall_angle, - reference_plane=reference_plane, - position=position, - rotation=rotation, - device=device, - ) - self.loop_sizes = tuple(int(loop.shape[0]) for loop in loop_tensors) - - def to_mesh(self, segments=16): - raise NotImplementedError("ComplexPolySlab does not support mesh export.") diff --git a/witwin/core/geometry/smpl.py b/witwin/core/geometry/smpl.py deleted file mode 100644 index 79d65aa..0000000 --- a/witwin/core/geometry/smpl.py +++ /dev/null @@ -1,229 +0,0 @@ -"""Differentiable SMPL geometry with optional smplpytorch dependency.""" - -from __future__ import annotations - -import pickle -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from .base import GeometryBase - - -class _Arr(np.ndarray): - @property - def r(self): - return np.asarray(self) - - -class _ChRecon: - def __init__(self, *args, **kwargs): - self._data = np.array([]) - for value in args: - if isinstance(value, np.ndarray): - self._data = value - return - if isinstance(value, _ChRecon): - self._data = value._data - return - - def __setstate__(self, state): - if isinstance(state, dict): - for value in state.values(): - if isinstance(value, np.ndarray): - self._data = value - return - if isinstance(value, _ChRecon): - self._data = value._data - return - elif isinstance(state, np.ndarray): - self._data = state - elif isinstance(state, (list, tuple)): - for value in state: - if isinstance(value, np.ndarray): - self._data = value - return - - -class _Unpickler(pickle.Unpickler): - def find_class(self, module, name): - if module.startswith("chumpy"): - return _ChRecon - return super().find_class(module, name) - - -def _ready_arguments_numpy(fname_or_dict): - import cv2 - - if not isinstance(fname_or_dict, dict): - with open(fname_or_dict, "rb") as handle: - data = _Unpickler(handle, encoding="latin1").load() - else: - data = fname_or_dict - - for key, value in list(data.items()): - if isinstance(value, _ChRecon): - data[key] = value._data - - want_shape_model = "shapedirs" in data - num_pose_params = data["kintree_table"].shape[1] * 3 - - if "trans" not in data: - data["trans"] = np.zeros(3) - if "pose" not in data: - data["pose"] = np.zeros(num_pose_params) - if "shapedirs" in data and "betas" not in data: - data["betas"] = np.zeros(data["shapedirs"].shape[-1]) - - if want_shape_model: - data["v_shaped"] = data["shapedirs"].dot(data["betas"]) + data["v_template"] - v_shaped = data["v_shaped"] - joint_regressor = data["J_regressor"] - data["J"] = np.column_stack([ - joint_regressor.dot(v_shaped[:, 0]), - joint_regressor.dot(v_shaped[:, 1]), - joint_regressor.dot(v_shaped[:, 2]), - ]) - pose = data["pose"].ravel()[3:] - rotations = np.concatenate([ - (cv2.Rodrigues(np.array(pp, dtype=np.float64))[0] - np.eye(3)).ravel() - for pp in pose.reshape((-1, 3)) - ]).ravel() - data["v_posed"] = v_shaped + data["posedirs"].dot(rotations) - else: - pose = data["pose"].ravel()[3:] - rotations = np.concatenate([ - (cv2.Rodrigues(np.array(pp, dtype=np.float64))[0] - np.eye(3)).ravel() - for pp in pose.reshape((-1, 3)) - ]).ravel() - data["v_posed"] = data["v_template"] + data["posedirs"].dot(rotations) - - for key, value in data.items(): - if isinstance(value, np.ndarray): - data[key] = value.view(_Arr) - return data - - -def _setup_smpl_compat(): - try: - import chumpy # noqa: F401 - - return - except ImportError: - pass - - try: - import smplpytorch.native.webuser.serialization as serialization - except ImportError: - return - - serialization.ready_arguments = _ready_arguments_numpy - - -try: - _setup_smpl_compat() - from smplpytorch.pytorch.smpl_layer import SMPL_Layer - - _SMPL_AVAILABLE = True -except ImportError: - _SMPL_AVAILABLE = False - - -_SMPL_LAYER_CACHE: dict[tuple[str, str, str], Any] = {} - - -def _resolve_scene_device(device: str | None) -> str: - requested = "cuda" if device is None else device - resolved = torch.device(requested) - if resolved.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError( - "SMPLBody defaults to CUDA, but torch.cuda.is_available() is False. " - "Pass device='cpu' only for scene construction or non-rendering workflows." - ) - return str(resolved) - - -def _default_smpl_model_root() -> str: - return str(Path(__file__).resolve().parents[4] / "radar" / "models" / "smpl_models") - - -def _to_vertex_tensor(value, *, device: str) -> torch.Tensor: - if isinstance(value, torch.Tensor): - return value.to(device=device, dtype=torch.float32) - return torch.as_tensor(value, device=device, dtype=torch.float32) - - -def _get_smpl_layer(*, gender: str, model_root: str, device: str): - if not _SMPL_AVAILABLE: - raise ImportError("smplpytorch is required to instantiate or evaluate SMPLBody.") - key = (str(gender), str(model_root), str(device)) - layer = _SMPL_LAYER_CACHE.get(key) - if layer is None: - layer = SMPL_Layer(center_idx=0, gender=gender, model_root=model_root).to(device) - _SMPL_LAYER_CACHE[key] = layer - return layer - - -class SMPLBody(GeometryBase): - """Differentiable SMPL geometry with position and rotation.""" - - kind = "smpl" - - def __init__( - self, - pose, - shape, - *, - position=(0.0, 0.0, 0.0), - gender: str = "male", - model_root: str | None = None, - rotation=None, - device=None, - ): - super().__init__(position=position, rotation=rotation, device=device) - tensor_device = str(self.position.device) - self.pose = _to_vertex_tensor(pose, device=tensor_device).reshape(-1) - self.shape = _to_vertex_tensor(shape, device=tensor_device).reshape(-1) - self.gender = str(gender) - self.model_root = _default_smpl_model_root() if model_root is None else str(model_root) - - def updated(self, **changes) -> "SMPLBody": - updated = SMPLBody( - pose=changes.pop("pose", self.pose), - shape=changes.pop("shape", self.shape), - position=changes.pop("position", self.position), - gender=changes.pop("gender", self.gender), - model_root=changes.pop("model_root", self.model_root), - rotation=changes.pop("rotation", self.rotation), - device=changes.pop("device", self.position.device), - ) - if changes: - unsupported = ", ".join(sorted(changes)) - raise TypeError(f"Unsupported SMPLBody updates: {unsupported}") - return updated - - def _evaluate(self, *, device: str): - layer = _get_smpl_layer(gender=self.gender, model_root=self.model_root, device=device) - pose_tensor = self.pose.to(device=device, dtype=torch.float32).view(1, -1) - shape_tensor = self.shape.to(device=device, dtype=torch.float32).view(1, -1) - if shape_tensor.requires_grad: - shape_tensor = shape_tensor + 1e-8 - vertices, joints = layer(pose_tensor, th_betas=shape_tensor) - vertices = self._transform_mesh_verts(vertices[0]) - joints = self._transform_mesh_verts(joints[0]) - faces = layer.th_faces.detach().cpu().numpy().astype(np.int32) - return vertices.contiguous(), np.ascontiguousarray(faces), joints.contiguous() - - def to_mesh(self, segments=16, *, device=None): - del segments - resolved_device = _resolve_scene_device(device or self.position.device) - vertices, faces, _ = self._evaluate(device=resolved_device) - face_tensor = torch.as_tensor(faces, device=vertices.device, dtype=torch.int64) - return vertices, face_tensor - - def joints(self, *, device=None) -> torch.Tensor: - resolved_device = _resolve_scene_device(device or self.position.device) - _, _, joints = self._evaluate(device=resolved_device) - return joints diff --git a/witwin/core/material.py b/witwin/core/material.py index 2d7d99f..1324bbf 100644 --- a/witwin/core/material.py +++ b/witwin/core/material.py @@ -6,42 +6,46 @@ import numpy as np -VACUUM_PERMITTIVITY = 8.8541878128e-12 - - -def _coerce_scalar(value: float, *, name: str) -> float: - return float(value) +from .geometry.base import GeometrySpec - -def _coerce_nonnegative(value: float, *, name: str) -> float: - scalar = float(value) - if scalar < 0.0: - raise ValueError(f"{name} must be >= 0.") - return scalar +VACUUM_PERMITTIVITY = 8.8541878128e-12 -def _coerce_differentiable_scalar(value: Any, *, name: str) -> Any: +def _coerce_scalar(value: Any, *, name: str) -> Any: try: return float(value) except (TypeError, ValueError): return value -def _coerce_differentiable_nonnegative(value: Any, *, name: str) -> Any: +def _numeric_scalar(value: Any) -> float | None: + for candidate in ( + lambda item: float(item), + lambda item: float(item.item()), + lambda item: float(item[0]), + ): + try: + return candidate(value) + except (TypeError, ValueError, IndexError, KeyError, AttributeError): + continue + return None + + +def _coerce_nonnegative(value: Any, *, name: str) -> Any: + scalar = _numeric_scalar(value) + if scalar is not None and scalar < 0.0: + raise ValueError(f"{name} must be >= 0.") try: - scalar = float(value) + return float(value) except (TypeError, ValueError): return value - if scalar < 0.0: - raise ValueError(f"{name} must be >= 0.") - return scalar def _is_numeric_close(value: Any, target: float) -> bool: - try: - return bool(np.isclose(float(value), target)) - except (TypeError, ValueError): + scalar = _numeric_scalar(value) + if scalar is None: return False + return bool(np.isclose(scalar, target)) @dataclass(frozen=True) @@ -82,64 +86,6 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: @dataclass(frozen=True, init=False) class Material(MaterialSpec): - eps_r: float - mu_r: float - sigma_e: float - name: str | None - - def __init__( - self, - eps_r: float = 1.0, - mu_r: float = 1.0, - sigma_e: float = 0.0, - name: str | None = None, - ): - object.__setattr__(self, "eps_r", _coerce_scalar(eps_r, name="eps_r")) - object.__setattr__(self, "mu_r", _coerce_scalar(mu_r, name="mu_r")) - object.__setattr__(self, "sigma_e", _coerce_nonnegative(sigma_e, name="sigma_e")) - object.__setattr__(self, "name", None if name is None else str(name)) - - def capabilities(self) -> MaterialCapabilities: - return MaterialCapabilities( - conductive=not np.isclose(self.sigma_e, 0.0), - magnetic=not np.isclose(self.mu_r, 1.0), - anisotropic=False, - dispersive=False, - ) - - def evaluate_static(self) -> StaticMaterialSample: - return StaticMaterialSample( - eps_r=self.eps_r, - mu_r=self.mu_r, - sigma_e=self.sigma_e, - ) - - def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: - resolved_frequency = float(frequency) - if resolved_frequency < 0.0: - raise ValueError("frequency must be >= 0.") - if resolved_frequency == 0.0: - if not np.isclose(self.sigma_e, 0.0): - raise ValueError("Conductive materials require frequency > 0.") - return FrequencyMaterialSample( - eps_r=complex(self.eps_r, 0.0), - mu_r=self.mu_r, - sigma_e=self.sigma_e, - ) - - eps_r = complex( - self.eps_r, - -self.sigma_e / (2.0 * np.pi * resolved_frequency * VACUUM_PERMITTIVITY), - ) - return FrequencyMaterialSample( - eps_r=eps_r, - mu_r=self.mu_r, - sigma_e=self.sigma_e, - ) - - -@dataclass(frozen=True, init=False) -class DifferentiableMaterial(MaterialSpec): eps_r: Any mu_r: Any sigma_e: Any @@ -152,9 +98,9 @@ def __init__( sigma_e: Any = 0.0, name: str | None = None, ): - object.__setattr__(self, "eps_r", _coerce_differentiable_scalar(eps_r, name="eps_r")) - object.__setattr__(self, "mu_r", _coerce_differentiable_scalar(mu_r, name="mu_r")) - object.__setattr__(self, "sigma_e", _coerce_differentiable_nonnegative(sigma_e, name="sigma_e")) + object.__setattr__(self, "eps_r", _coerce_scalar(eps_r, name="eps_r")) + object.__setattr__(self, "mu_r", _coerce_scalar(mu_r, name="mu_r")) + object.__setattr__(self, "sigma_e", _coerce_nonnegative(sigma_e, name="sigma_e")) object.__setattr__(self, "name", None if name is None else str(name)) def capabilities(self) -> MaterialCapabilities: @@ -176,31 +122,31 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: resolved_frequency = float(frequency) if resolved_frequency < 0.0: raise ValueError("frequency must be >= 0.") + sigma_scalar = _numeric_scalar(self.sigma_e) if resolved_frequency == 0.0: - if not _is_numeric_close(self.sigma_e, 0.0): - try: - scalar_sigma = float(self.sigma_e) - except (TypeError, ValueError): - scalar_sigma = None - if scalar_sigma is not None: - raise ValueError("Conductive materials require frequency > 0.") + if sigma_scalar is not None and not np.isclose(sigma_scalar, 0.0): + raise ValueError("Conductive materials require frequency > 0.") + eps_scalar = _numeric_scalar(self.eps_r) + mu_scalar = _numeric_scalar(self.mu_r) return FrequencyMaterialSample( - eps_r=self.eps_r, - mu_r=self.mu_r, - sigma_e=self.sigma_e, + eps_r=self.eps_r if eps_scalar is None else complex(eps_scalar, 0.0), + mu_r=self.mu_r if mu_scalar is None else mu_scalar, + sigma_e=self.sigma_e if sigma_scalar is None else sigma_scalar, ) - try: - eps_r = complex( - float(self.eps_r), - -float(self.sigma_e) / (2.0 * np.pi * resolved_frequency * VACUUM_PERMITTIVITY), - ) - mu_r = float(self.mu_r) - sigma_e = float(self.sigma_e) - except (TypeError, ValueError): + eps_scalar = _numeric_scalar(self.eps_r) + mu_scalar = _numeric_scalar(self.mu_r) + if eps_scalar is None or sigma_scalar is None: eps_r = self.eps_r mu_r = self.mu_r sigma_e = self.sigma_e + else: + eps_r = complex( + eps_scalar, + -sigma_scalar / (2.0 * np.pi * resolved_frequency * VACUUM_PERMITTIVITY), + ) + mu_r = self.mu_r if mu_scalar is None else mu_scalar + sigma_e = self.sigma_e if sigma_scalar is None else sigma_scalar return FrequencyMaterialSample( eps_r=eps_r, mu_r=mu_r, @@ -210,7 +156,7 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: @dataclass(frozen=True, init=False) class Structure: - geometry: Any + geometry: GeometrySpec material: MaterialSpec name: str | None priority: int @@ -220,7 +166,7 @@ class Structure: def __init__( self, - geometry: Any, + geometry: GeometrySpec, material: MaterialSpec, name: str | None = None, priority: int = 0, @@ -228,8 +174,13 @@ def __init__( tags=(), metadata: Mapping[str, Any] | None = None, ): - if material is None: - raise ValueError("Structure requires material=.") + if not isinstance(geometry, GeometrySpec): + raise TypeError("Structure geometry must implement GeometrySpec (kind and to_mesh()).") + if not isinstance(material, MaterialSpec): + raise TypeError( + "Structure material must implement MaterialSpec " + "(name, capabilities(), evaluate_static(), and evaluate_at_frequency())." + ) object.__setattr__(self, "geometry", geometry) object.__setattr__(self, "material", material) object.__setattr__(self, "name", None if name is None else str(name)) diff --git a/witwin/core/scene.py b/witwin/core/scene.py deleted file mode 100644 index 867c5c6..0000000 --- a/witwin/core/scene.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Base Scene class for all solver-specific scenes.""" - - -class SceneBase: - """Base scene for all solver-specific scenes. - - Solver-specific scenes (e.g. maxwell.Scene, radar.Scene) extend this - with domain-specific fields and methods. - """ - - def __init__( - self, - *, - structures=None, - sources=None, - monitors=None, - metadata=None, - device="cuda", - verbose=False, - ): - self.structures = list(structures or []) - self.sources = list(sources or []) - self.monitors = list(monitors or []) - self.metadata = dict(metadata or {}) - self.device = device - self.verbose = verbose - - def add_structure(self, structure): - self.structures.append(structure) - return self - - def add_source(self, source): - self.sources.append(source) - return self - - def add_monitor(self, monitor): - self.monitors.append(monitor) - return self diff --git a/witwin/core/scene_to_mitsuba.py b/witwin/core/scene_to_mitsuba.py deleted file mode 100644 index 90f5b94..0000000 --- a/witwin/core/scene_to_mitsuba.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Shared helpers for lowering mesh-based scenes into Mitsuba.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Mapping - -import numpy as np -import torch - - -def _load_mitsuba(variant: str): - import drjit as dr - import mitsuba as mi - - mi.set_variant(variant) - return mi, dr - - -def _as_vertices_array(vertices) -> np.ndarray: - if isinstance(vertices, torch.Tensor): - vertices = vertices.detach().cpu().numpy() - array = np.asarray(vertices, dtype=np.float32) - if array.ndim != 2 or array.shape[1] != 3: - raise ValueError("Mitsuba vertices must have shape (V, 3).") - return np.ascontiguousarray(array) - - -def _as_faces_array(faces) -> np.ndarray: - if isinstance(faces, torch.Tensor): - faces = faces.detach().cpu().numpy() - array = np.asarray(faces, dtype=np.uint32) - if array.ndim != 2 or array.shape[1] != 3: - raise ValueError("Mitsuba faces must have shape (F, 3).") - return np.ascontiguousarray(array) - - -def _materialize_plugin(mi, plugin): - if isinstance(plugin, Mapping): - return mi.load_dict(dict(plugin)) - return plugin - - -@dataclass(frozen=True) -class MitsubaRenderable: - name: str - vertices: torch.Tensor | np.ndarray - faces: torch.Tensor | np.ndarray - bsdf: dict[str, Any] | None = None - - -@dataclass(frozen=True) -class MitsubaSceneHandle: - scene: Any - params: Any - - -def create_mitsuba_mesh( - renderable: MitsubaRenderable, - *, - default_bsdf: dict[str, Any] | None = None, - variant: str = "cuda_ad_rgb", - has_vertex_texcoords: bool = False, -): - mi, dr = _load_mitsuba(variant) - vertices = _as_vertices_array(renderable.vertices) - faces = _as_faces_array(renderable.faces) - bsdf = default_bsdf or { - "type": "diffuse", - "reflectance": {"type": "rgb", "value": (0.8, 0.8, 0.8)}, - } - - mesh = mi.Mesh( - renderable.name, - vertex_count=vertices.shape[0], - face_count=faces.shape[0], - has_vertex_texcoords=has_vertex_texcoords, - ) - mesh_params = mi.traverse(mesh) - mesh_params["vertex_positions"] = dr.ravel(dr.cuda.ad.TensorXf(vertices)) - mesh_params["faces"] = dr.ravel(dr.cuda.ad.TensorXu(faces)) - mesh_params.update() - mesh.set_bsdf(mi.load_dict(renderable.bsdf or bsdf)) - return mesh - - -def build_mitsuba_scene( - *, - sensor, - renderables: Mapping[str, MitsubaRenderable | Any], - integrator: dict[str, Any] | Any | None = None, - default_bsdf: dict[str, Any] | None = None, - variant: str = "cuda_ad_rgb", -) -> MitsubaSceneHandle: - mi, _ = _load_mitsuba(variant) - scene_dict = { - "type": "scene", - "integrator": _materialize_plugin(mi, integrator or {"type": "direct"}), - "sensor": _materialize_plugin(mi, sensor), - "default_bsdf": default_bsdf - or { - "type": "diffuse", - "reflectance": {"type": "rgb", "value": (0.8, 0.8, 0.8)}, - }, - } - - for name, renderable in renderables.items(): - if not isinstance(renderable, MitsubaRenderable): - renderable = MitsubaRenderable( - name=name, - vertices=renderable.vertices, - faces=renderable.faces, - bsdf=getattr(renderable, "bsdf", None), - ) - scene_dict[name] = create_mitsuba_mesh( - renderable, - default_bsdf=scene_dict["default_bsdf"], - variant=variant, - ) - - scene = mi.load_dict(scene_dict) - return MitsubaSceneHandle(scene=scene, params=mi.traverse(scene)) - - -def update_mitsuba_scene_vertices( - params, - renderables: Mapping[str, MitsubaRenderable | Any], - *, - variant: str = "cuda_ad_rgb", -) -> None: - mi, dr = _load_mitsuba(variant) - for name, renderable in renderables.items(): - key = f"{name}.vertex_positions" - if key not in params: - continue - vertices = renderable.vertices if not isinstance(renderable, MitsubaRenderable) else renderable.vertices - params[key] = dr.ravel(dr.cuda.ad.TensorXf(_as_vertices_array(vertices))) - params.update()