Customize Geometry

from witwin.components import MeshComponent, component, float_field, int_field
import numpy as np
 
@component(name="CustomMesh")
class CustomMeshComponent(MeshComponent):
    radius = float_field(1.0, min=0.1, max=10.0)
 
    def _generate_mesh(self):
        """Generate and return (vertices, faces) as numpy arrays."""
        vertices = np.array([...], dtype=np.float32)  # (N, 3)
        faces = np.array([...], dtype=np.uint32)      # (M, 3)
        return vertices, faces

# Each row is [x, y, z]
vertices = np.array([
    [0.0, 0.0, 0.0],    # Vertex 0
    [1.0, 0.0, 0.0],    # Vertex 1
    [0.5, 1.0, 0.0],    # Vertex 2
], dtype=np.float32)

# Each row is [v0, v1, v2] - indices into vertices array
faces = np.array([
    [0, 1, 2],  # Triangle using vertices 0, 1, 2
], dtype=np.uint32)

def _generate_mesh(self):
    # Access component fields via self
    r = float(self.radius)
    n = int(self.segments)
 
    # Generate vertex positions
    vertices = [...]
 
    # Generate face indices
    faces = [...]
 
    return (
        np.array(vertices, dtype=np.float32),
        np.array(faces, dtype=np.uint32)
    )

def on_value_change(self, field_name, old_value, new_value):
    if field_name in ['radius', 'segments']:
        self._mark_dirty()  # Triggers _generate_mesh

from witwin.components import MeshComponent, component, float_field, int_field, bool_field
import numpy as np
 
@component(name="CircleMesh", display_name="Circle Mesh")
class CircleMeshComponent(MeshComponent):
    radius = float_field(1.0, min=0.01, max=100.0, description="Circle radius")
    segments = int_field(32, min=3, max=256, description="Number of segments")
    fill = bool_field(True, description="Fill the circle (disc vs ring)")
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.mesh_type = "Circle"
        self._mark_dirty()
 
    def _generate_mesh(self):
        r = float(self.radius)
        n = int(self.segments)
        angles = np.linspace(0, 2 * np.pi, n, endpoint=False)
 
        if self.fill:
            # Disc: center point + edge points
            vertices = [[0.0, 0.0, 0.0]]  # Center
            for angle in angles:
                x = r * np.cos(angle)
                z = r * np.sin(angle)
                vertices.append([x, 0.0, z])
            vertices = np.array(vertices, dtype=np.float32)
 
            # Fan triangulation from center
            faces = []
            for i in range(n):
                i_curr = i + 1
                i_next = (i + 1) % n + 1
                faces.append([0, i_next, i_curr])
            faces = np.array(faces, dtype=np.uint32)
        else:
            # Ring: just edge points, no faces
            vertices = []
            for angle in angles:
                x = r * np.cos(angle)
                z = r * np.sin(angle)
                vertices.append([x, 0.0, z])
            vertices = np.array(vertices, dtype=np.float32)
            faces = np.zeros((0, 3), dtype=np.uint32)
 
        return vertices, faces
 
    def on_value_change(self, field_name, old_value, new_value):
        if field_name in ['radius', 'segments', 'fill']:
            self._mark_dirty()

@component(name="GridMesh", display_name="Grid Mesh")
class GridMeshComponent(MeshComponent):
    width = float_field(10.0, min=0.1, max=100.0)
    height = float_field(10.0, min=0.1, max=100.0)
    divisions_x = int_field(10, min=1, max=100)
    divisions_y = int_field(10, min=1, max=100)
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.mesh_type = "Grid"
        self._mark_dirty()
 
    def _generate_mesh(self):
        w = float(self.width)
        h = float(self.height)
        dx = int(self.divisions_x)
        dy = int(self.divisions_y)
 
        vertices = []
        faces = []
 
        # Generate vertices
        for j in range(dy + 1):
            for i in range(dx + 1):
                x = (i / dx - 0.5) * w
                z = (j / dy - 0.5) * h
                vertices.append([x, 0.0, z])
 
        # Generate faces (two triangles per cell)
        for j in range(dy):
            for i in range(dx):
                v0 = j * (dx + 1) + i
                v1 = v0 + 1
                v2 = v0 + (dx + 1)
                v3 = v2 + 1
 
                faces.append([v0, v2, v1])
                faces.append([v1, v2, v3])
 
        return (
            np.array(vertices, dtype=np.float32),
            np.array(faces, dtype=np.uint32)
        )
 
    def on_value_change(self, field_name, old_value, new_value):
        if field_name in ['width', 'height', 'divisions_x', 'divisions_y']:
            self._mark_dirty()

@component(name="ParametricSurface", display_name="Parametric Surface")
class ParametricSurfaceComponent(MeshComponent):
    resolution = int_field(32, min=4, max=128)
    amplitude = float_field(1.0, min=0.0, max=5.0)
    frequency = float_field(2.0, min=0.1, max=10.0)
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.mesh_type = "ParametricSurface"
        self._mark_dirty()
 
    def _generate_mesh(self):
        n = int(self.resolution)
        amp = float(self.amplitude)
        freq = float(self.frequency)
 
        vertices = []
        faces = []
 
        # Generate vertices on a parametric surface
        for j in range(n):
            for i in range(n):
                u = i / (n - 1)
                v = j / (n - 1)
 
                x = u * 10 - 5
                z = v * 10 - 5
 
                # Parametric height function
                y = amp * np.sin(freq * x) * np.cos(freq * z)
 
                vertices.append([x, y, z])
 
        # Generate faces
        for j in range(n - 1):
            for i in range(n - 1):
                v0 = j * n + i
                v1 = v0 + 1
                v2 = v0 + n
                v3 = v2 + 1
 
                faces.append([v0, v2, v1])
                faces.append([v1, v2, v3])
 
        return (
            np.array(vertices, dtype=np.float32),
            np.array(faces, dtype=np.uint32)
        )
 
    def on_value_change(self, field_name, old_value, new_value):
        if field_name in ['resolution', 'amplitude', 'frequency']:
            self._mark_dirty()

@component(name="DynamicMesh")
class DynamicMeshComponent(MeshComponent):
 
    @button(display_name="Load From File")
    def load_from_file(self):
        vertices = load_vertices_from_file()
        faces = load_faces_from_file()
        self.set_mesh_data(vertices, faces)
        return "Mesh loaded"

def set_mesh_data(
    self,
    vertices: np.ndarray,  # (N, 3) float32
    faces: np.ndarray,     # (M, 3) uint32
    notify: bool = True    # Sync to frontend
)

class MeshComponent:
    @property
    def vertex_count(self) -> int:
        """Number of vertices."""
        pass
 
    @property
    def face_count(self) -> int:
        """Number of triangular faces."""
        pass
 
    @property
    def is_empty(self) -> bool:
        """Check if mesh has no geometry."""
        pass
 
    @property
    def drjit_vertices(self):
        """Get vertices as drjit Array3f (3, N) for ray tracing."""
        pass
 
    @property
    def drjit_faces(self):
        """Get faces as drjit UInt32 (M*3,) flattened."""
        pass

from witwin.components import MeshComponent, component, float_field, int_field, string_field
from witwin import Console
import numpy as np
 
@component(name="AntennaPattern", display_name="Antenna Pattern Mesh")
class AntennaPatternComponent(MeshComponent):
    """Visualize antenna radiation pattern as a 3D mesh."""
 
    pattern_type = string_field(
        "dipole",
        options=["dipole", "patch", "horn", "isotropic"],
        description="Antenna pattern type"
    )
    scale = float_field(1.0, min=0.1, max=10.0, description="Pattern scale")
    theta_segments = int_field(32, min=8, max=64, description="Elevation segments")
    phi_segments = int_field(64, min=16, max=128, description="Azimuth segments")
 
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.mesh_type = "AntennaPattern"
        self._mark_dirty()
 
    def _generate_mesh(self):
        n_theta = int(self.theta_segments)
        n_phi = int(self.phi_segments)
        scale = float(self.scale)
 
        vertices = []
        faces = []
 
        # Generate vertices based on pattern
        for j in range(n_theta + 1):
            theta = np.pi * j / n_theta  # 0 to pi
 
            for i in range(n_phi):
                phi = 2 * np.pi * i / n_phi  # 0 to 2*pi
 
                # Calculate pattern gain
                gain = self._calculate_pattern(theta, phi)
                r = scale * gain
 
                # Spherical to Cartesian
                x = r * np.sin(theta) * np.cos(phi)
                y = r * np.cos(theta)
                z = r * np.sin(theta) * np.sin(phi)
 
                vertices.append([x, y, z])
 
        # Generate faces
        for j in range(n_theta):
            for i in range(n_phi):
                v0 = j * n_phi + i
                v1 = j * n_phi + (i + 1) % n_phi
                v2 = (j + 1) * n_phi + i
                v3 = (j + 1) * n_phi + (i + 1) % n_phi
 
                faces.append([v0, v2, v1])
                faces.append([v1, v2, v3])
 
        Console.log(f"Generated {self.pattern_type} pattern: {len(vertices)} vertices")
 
        return (
            np.array(vertices, dtype=np.float32),
            np.array(faces, dtype=np.uint32)
        )
 
    def _calculate_pattern(self, theta: float, phi: float) -> float:
        """Calculate antenna gain at given angles."""
        if self.pattern_type == "isotropic":
            return 1.0
 
        elif self.pattern_type == "dipole":
            # Dipole pattern: sin(theta)
            return abs(np.sin(theta)) + 0.01
 
        elif self.pattern_type == "patch":
            # Simple patch antenna approximation
            return max(0.01, np.cos(theta) ** 2)
 
        elif self.pattern_type == "horn":
            # Horn antenna approximation
            return max(0.01, np.cos(theta) ** 4)
 
        return 1.0
 
    def on_value_change(self, field_name, old_value, new_value):
        if field_name in ['pattern_type', 'scale', 'theta_segments', 'phi_segments']:
            self._mark_dirty()