gri-plot 0.2.0.post1__py3-none-any.whl

gri_plot/shapes/ellipsoid.py

@@ -0,0 +1,227 @@
+"""Ellipsoid shape for 3D visualization."""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from gri_plot.shapes.meshgen import ellipsoid_mesh
+from gri_plot.surfaces import ImplicitShape
+from gri_plot.surfaces.mesh import vertices_to_mesh3d
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import plotly.graph_objects as go
+    from numpy.typing import NDArray
+
+
+class Ellipsoid(ImplicitShape):
+    """A 3D ellipsoid defined by center and covariance matrix.
+
+    The ellipsoid represents a 1-sigma contour of a 3D Gaussian distribution,
+    or can be specified by semi-axes and rotation.
+
+    The residual function returns Mahalanobis distance squared - 1, so:
+    - 0 on the 1-sigma surface
+    - Positive outside
+    - Negative inside
+
+    Attributes:
+        center_xyz: Center of the ellipsoid in XYZ coordinates.
+        covariance: Covariance matrix (3x3).
+        semi_axes: Semi-axis lengths (a, b, c).
+        rotation: Rotation matrix (3x3).
+        label: Optional label for legends.
+        is_volume: Whether this shape represents a volume or surface.
+    """
+
+    def __init__(  # noqa: PLR0913 - supports covariance or semi-axes + rotation
+        self,
+        center_xyz: NDArray[np.floating],
+        covariance: NDArray[np.floating] | None = None,
+        semi_axes: NDArray[np.floating] | None = None,
+        rotation: NDArray[np.floating] | None = None,
+        label: str | None = None,
+        *,
+        as_volume: bool = True,
+    ) -> None:
+        """Initialize the ellipsoid.
+
+        The ellipsoid can be defined either by a covariance matrix or by
+        semi-axes and rotation. If covariance is provided, semi_axes and
+        rotation are derived from its eigendecomposition.
+
+        Args:
+            center_xyz: Center of the ellipsoid, shape (3,).
+            covariance: Covariance matrix, shape (3, 3). Mutually exclusive
+                with semi_axes/rotation.
+            semi_axes: Semi-axis lengths (a, b, c) in descending order.
+                Required if covariance is None.
+            rotation: Rotation matrix, shape (3, 3). Columns are the principal
+                axes. Defaults to identity if not provided.
+            label: Optional label for legends.
+            as_volume: Keyword-only. If True, interior is solution region.
+                If False, boundary is solution locus. Default True.
+
+        Raises:
+            ValueError: If neither covariance nor semi_axes is provided, or
+                if both are provided.
+        """
+        self._center = np.asarray(center_xyz, dtype=np.float64)
+        self._label = label
+        self._is_volume = as_volume
+
+        if covariance is not None and semi_axes is not None:
+            raise ValueError("Provide either covariance or semi_axes, not both")
+
+        if covariance is not None:
+            self._covariance = np.asarray(covariance, dtype=np.float64)
+            # Eigendecomposition
+            eigenvalues, eigenvectors = np.linalg.eigh(self._covariance)
+            # Sort in descending order
+            idx = np.argsort(eigenvalues)[::-1]
+            eigenvalues = eigenvalues[idx]
+            eigenvectors = eigenvectors[:, idx]
+
+            self._semi_axes = np.sqrt(np.abs(eigenvalues))
+            self._rotation = eigenvectors
+            # Compute inverse covariance for residual
+            self._cov_inv = np.linalg.inv(self._covariance)
+
+        elif semi_axes is not None:
+            self._semi_axes = np.asarray(semi_axes, dtype=np.float64)
+            if rotation is None:
+                self._rotation = np.eye(3)
+            else:
+                self._rotation = np.asarray(rotation, dtype=np.float64)
+
+            # Reconstruct covariance: C = R @ diag(s^2) @ R.T
+            s_squared = np.diag(self._semi_axes**2)
+            self._covariance = self._rotation @ s_squared @ self._rotation.T
+            self._cov_inv = np.linalg.inv(self._covariance)
+
+        else:
+            raise ValueError("Either covariance or semi_axes must be provided")
+
+    @property
+    def center_xyz(self) -> NDArray[np.floating]:
+        """Get the center in XYZ coordinates."""
+        return self._center
+
+    @property
+    def covariance(self) -> NDArray[np.floating]:
+        """Get the covariance matrix."""
+        return self._covariance
+
+    @property
+    def semi_axes(self) -> NDArray[np.floating]:
+        """Get the semi-axis lengths."""
+        return self._semi_axes
+
+    @property
+    def rotation(self) -> NDArray[np.floating]:
+        """Get the rotation matrix."""
+        return self._rotation
+
+    @property
+    def label(self) -> str | None:
+        """Get the label."""
+        return self._label
+
+    @property
+    def is_volume(self) -> bool:
+        """Whether this shape represents a volume (True) or surface (False)."""
+        return self._is_volume
+
+    def residual_fn(self, xyz: NDArray[np.floating]) -> NDArray[np.floating]:
+        """Compute the residual function (Mahalanobis distance squared - 1).
+
+        Args:
+            xyz: Points to evaluate, shape (..., 3).
+
+        Returns:
+            Residual values, shape (...). Zero on surface, positive outside,
+            negative inside.
+        """
+        xyz = np.asarray(xyz)
+        diff = xyz - self._center
+
+        # Mahalanobis distance squared: d^2 = (x-mu)^T @ C^-1 @ (x-mu)
+        # Using einsum for batch computation
+        temp = np.einsum("...i,ij->...j", diff, self._cov_inv)
+        mahal_sq = np.einsum("...i,...i->...", diff, temp)
+
+        return mahal_sq - 1.0
+
+    def to_mesh(
+        self,
+        resolution: int | None = None,
+    ) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+        """Generate mesh vertices and faces.
+
+        Args:
+            resolution: Number of samples per dimension. If None, uses 50.
+
+        Returns:
+            Tuple of (vertices, faces).
+        """
+        if resolution is None:
+            resolution = 50
+        return ellipsoid_mesh(
+            self._center,
+            semi_axes=self._semi_axes,
+            rotation=self._rotation,
+            resolution=resolution,
+        )
+
+    def get_bounds_xyz(
+        self,
+    ) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
+        """Get axis-aligned bounding box in XYZ coordinates.
+
+        The bounding box accounts for the ellipsoid's rotation.
+
+        Returns:
+            Tuple of (min_corner, max_corner).
+        """
+        # The extent in each axis direction
+        # For an ellipsoid with semi-axes s and rotation R, the extent along
+        # axis e_i is sqrt(sum_j (R[i,j] * s[j])^2)
+        extents = np.zeros(3)
+        for i in range(3):
+            extents[i] = np.sqrt(np.sum((self._rotation[i, :] * self._semi_axes) ** 2))
+
+        return self._center - extents, self._center + extents
+
+    def to_trace(
+        self,
+        resolution: int | None = None,
+        intensity_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]]
+        | None = None,
+        **kwargs,
+    ) -> go.Mesh3d:
+        """Generate a Plotly Mesh3d trace.
+
+        Args:
+            resolution: Mesh resolution.
+            intensity_fn: Optional intensity function for coloring.
+            **kwargs: Additional arguments for vertices_to_mesh3d.
+
+        Returns:
+            Plotly Mesh3d trace.
+        """
+        vertices, faces = self.to_mesh(resolution)
+        return vertices_to_mesh3d(
+            vertices,
+            faces,
+            intensity_fn=intensity_fn,
+            name=self._label,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return f"Ellipsoid(center={self._center}, semi_axes={self._semi_axes})"
+
+
+__all__ = ["Ellipsoid"]

gri_plot/shapes/meshgen.py

@@ -0,0 +1,245 @@
+"""Mesh generation for geometric shapes.
+
+This module provides functions to generate vertices and faces for common
+3D shapes. The output can be passed to vertices_to_mesh3d() for rendering.
+"""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from gri_plot.surfaces.mesh import grid_to_mesh
+
+# Threshold for choosing perpendicular vector when building orthonormal basis.
+# When |axis[2]| < this value, cross with Z-axis; otherwise cross with X-axis.
+AXIS_ALIGNMENT_THRESHOLD = 0.9
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+
+def sphere_mesh(
+    center: NDArray[np.floating],
+    radius: float,
+    resolution: int = 50,
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Generate vertices and faces for a sphere.
+
+    Args:
+        center: Center of the sphere, shape (3,).
+        radius: Radius of the sphere.
+        resolution: Number of samples per dimension.
+
+    Returns:
+        Tuple of (vertices, faces).
+    """
+    center = np.asarray(center)
+
+    u = np.linspace(0, 2 * np.pi, resolution)
+    v = np.linspace(0, np.pi, resolution)
+    u_grid, v_grid = np.meshgrid(u, v)
+
+    x = center[0] + radius * np.cos(u_grid) * np.sin(v_grid)
+    y = center[1] + radius * np.sin(u_grid) * np.sin(v_grid)
+    z = center[2] + radius * np.cos(v_grid)
+
+    return grid_to_mesh(x, y, z)
+
+
+def ellipsoid_mesh(
+    center: NDArray[np.floating],
+    covariance: NDArray[np.floating] | None = None,
+    semi_axes: NDArray[np.floating] | None = None,
+    rotation: NDArray[np.floating] | None = None,
+    resolution: int = 50,
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Generate vertices and faces for an ellipsoid.
+
+    The ellipsoid can be specified either by a covariance matrix or by
+    semi-axes and rotation matrix.
+
+    Args:
+        center: Center of the ellipsoid, shape (3,).
+        covariance: Covariance matrix, shape (3, 3). If provided, semi_axes
+            and rotation are computed from it.
+        semi_axes: Semi-axis lengths (a, b, c). Required if covariance is None.
+        rotation: Rotation matrix, shape (3, 3). Defaults to identity.
+        resolution: Number of samples per dimension.
+
+    Returns:
+        Tuple of (vertices, faces).
+
+    Raises:
+        ValueError: If neither covariance nor semi_axes is provided.
+    """
+    center = np.asarray(center)
+
+    if covariance is not None:
+        covariance = np.asarray(covariance)
+        eigenvalues, eigenvectors = np.linalg.eigh(covariance)
+        # Sort in descending order
+        idx = np.argsort(eigenvalues)[::-1]
+        eigenvalues = eigenvalues[idx]
+        eigenvectors = eigenvectors[:, idx]
+
+        semi_axes = np.sqrt(np.abs(eigenvalues))
+        rotation = eigenvectors
+
+    elif semi_axes is not None:
+        semi_axes = np.asarray(semi_axes)
+        rotation = np.eye(3) if rotation is None else np.asarray(rotation)
+    else:
+        raise ValueError("Either covariance or semi_axes must be provided")
+
+    # Generate unit sphere
+    u = np.linspace(0, 2 * np.pi, resolution)
+    v = np.linspace(0, np.pi, resolution)
+    u_grid, v_grid = np.meshgrid(u, v)
+
+    x_sphere = np.cos(u_grid) * np.sin(v_grid)
+    y_sphere = np.sin(u_grid) * np.sin(v_grid)
+    z_sphere = np.cos(v_grid)
+
+    # Scale by semi-axes
+    x_ellipsoid = semi_axes[0] * x_sphere
+    y_ellipsoid = semi_axes[1] * y_sphere
+    z_ellipsoid = semi_axes[2] * z_sphere
+
+    # Stack and rotate
+    points = np.stack([x_ellipsoid, y_ellipsoid, z_ellipsoid], axis=-1)
+    rotated = np.einsum("ij,...j->...i", rotation, points)
+
+    # Translate to center
+    x = rotated[..., 0] + center[0]
+    y = rotated[..., 1] + center[1]
+    z = rotated[..., 2] + center[2]
+
+    return grid_to_mesh(x, y, z)
+
+
+def cone_mesh(
+    apex: NDArray[np.floating],
+    axis: NDArray[np.floating],
+    half_angle: float,
+    height: float,
+    resolution: int = 50,
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Generate vertices and faces for a cone.
+
+    The cone extends from the apex along the axis direction.
+
+    Args:
+        apex: Apex (tip) of the cone, shape (3,).
+        axis: Axis direction (normalized or not), shape (3,).
+        half_angle: Half-angle of the cone in radians.
+        height: Height (length along axis) of the cone.
+        resolution: Number of samples per dimension.
+
+    Returns:
+        Tuple of (vertices, faces).
+    """
+    apex = np.asarray(apex)
+    axis = np.asarray(axis)
+    axis = axis / np.linalg.norm(axis)
+
+    # Create orthonormal basis
+    if abs(axis[2]) < AXIS_ALIGNMENT_THRESHOLD:
+        perp1 = np.cross(axis, np.array([0.0, 0.0, 1.0]))
+    else:
+        perp1 = np.cross(axis, np.array([1.0, 0.0, 0.0]))
+    perp1 = perp1 / np.linalg.norm(perp1)
+    perp2 = np.cross(axis, perp1)
+
+    # Parametric surface: h from 0 to height, theta from 0 to 2*pi
+    h = np.linspace(0, height, resolution)
+    theta = np.linspace(0, 2 * np.pi, resolution)
+    h_grid, theta_grid = np.meshgrid(h, theta)
+
+    # Radius at height h
+    r = h_grid * np.tan(half_angle)
+
+    # Position along cone surface
+    x = (
+        apex[0]
+        + h_grid * axis[0]
+        + r * (np.cos(theta_grid) * perp1[0] + np.sin(theta_grid) * perp2[0])
+    )
+    y = (
+        apex[1]
+        + h_grid * axis[1]
+        + r * (np.cos(theta_grid) * perp1[1] + np.sin(theta_grid) * perp2[1])
+    )
+    z = (
+        apex[2]
+        + h_grid * axis[2]
+        + r * (np.cos(theta_grid) * perp1[2] + np.sin(theta_grid) * perp2[2])
+    )
+
+    return grid_to_mesh(x, y, z)
+
+
+def cylinder_mesh(
+    center: NDArray[np.floating],
+    axis: NDArray[np.floating],
+    radius: float,
+    height: float,
+    resolution: int = 50,
+) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+    """Generate vertices and faces for a cylinder.
+
+    The cylinder is centered at center and extends height/2 in each
+    direction along the axis.
+
+    Args:
+        center: Center of the cylinder, shape (3,).
+        axis: Axis direction (normalized or not), shape (3,).
+        radius: Radius of the cylinder.
+        height: Total height of the cylinder.
+        resolution: Number of samples per dimension.
+
+    Returns:
+        Tuple of (vertices, faces).
+    """
+    center = np.asarray(center)
+    axis = np.asarray(axis)
+    axis = axis / np.linalg.norm(axis)
+
+    # Create orthonormal basis
+    if abs(axis[2]) < AXIS_ALIGNMENT_THRESHOLD:
+        perp1 = np.cross(axis, np.array([0.0, 0.0, 1.0]))
+    else:
+        perp1 = np.cross(axis, np.array([1.0, 0.0, 0.0]))
+    perp1 = perp1 / np.linalg.norm(perp1)
+    perp2 = np.cross(axis, perp1)
+
+    # Parametric surface: h from -height/2 to height/2, theta from 0 to 2*pi
+    h = np.linspace(-height / 2, height / 2, resolution)
+    theta = np.linspace(0, 2 * np.pi, resolution)
+    h_grid, theta_grid = np.meshgrid(h, theta)
+
+    # Position along cylinder surface
+    x = (
+        center[0]
+        + h_grid * axis[0]
+        + radius * (np.cos(theta_grid) * perp1[0] + np.sin(theta_grid) * perp2[0])
+    )
+    y = (
+        center[1]
+        + h_grid * axis[1]
+        + radius * (np.cos(theta_grid) * perp1[1] + np.sin(theta_grid) * perp2[1])
+    )
+    z = (
+        center[2]
+        + h_grid * axis[2]
+        + radius * (np.cos(theta_grid) * perp1[2] + np.sin(theta_grid) * perp2[2])
+    )
+
+    return grid_to_mesh(x, y, z)
+
+
+__all__ = [
+    "cone_mesh",
+    "cylinder_mesh",
+    "ellipsoid_mesh",
+    "sphere_mesh",
+]

gri_plot/shapes/sphere.py

@@ -0,0 +1,152 @@
+"""Sphere shape for 3D visualization."""
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from gri_plot.shapes.meshgen import sphere_mesh
+from gri_plot.surfaces import ImplicitShape
+from gri_plot.surfaces.mesh import vertices_to_mesh3d
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import plotly.graph_objects as go
+    from numpy.typing import NDArray
+
+
+class Sphere(ImplicitShape):
+    """A 3D sphere defined by center and radius.
+
+    Attributes:
+        center_xyz: Center of the sphere in XYZ coordinates.
+        radius: Radius of the sphere.
+        label: Optional label for legends.
+        is_volume: Whether this shape represents a volume or surface.
+    """
+
+    def __init__(
+        self,
+        center_xyz: NDArray[np.floating],
+        radius: float,
+        label: str | None = None,
+        *,
+        as_volume: bool = True,
+    ) -> None:
+        """Initialize the sphere.
+
+        Args:
+            center_xyz: Center of the sphere, shape (3,).
+            radius: Radius of the sphere (positive).
+            label: Optional label for legends.
+            as_volume: Keyword-only. If True, interior is solution region.
+                If False, boundary is solution locus. Default True.
+
+        Raises:
+            ValueError: If radius is not positive.
+        """
+        self._center = np.asarray(center_xyz, dtype=np.float64)
+        if radius <= 0:
+            raise ValueError("Radius must be positive")
+        self._radius = float(radius)
+        self._label = label
+        self._is_volume = as_volume
+
+    @property
+    def center_xyz(self) -> NDArray[np.floating]:
+        """Get the center in XYZ coordinates."""
+        return self._center
+
+    @property
+    def radius(self) -> float:
+        """Get the radius."""
+        return self._radius
+
+    @property
+    def label(self) -> str | None:
+        """Get the label."""
+        return self._label
+
+    @property
+    def is_volume(self) -> bool:
+        """Whether this shape represents a volume (True) or surface (False)."""
+        return self._is_volume
+
+    def residual_fn(self, xyz: NDArray[np.floating]) -> NDArray[np.floating]:
+        """Compute the residual function for the sphere.
+
+        The residual is distance_from_center - radius, so:
+        - 0 on the surface
+        - Positive outside
+        - Negative inside
+
+        Args:
+            xyz: Points to evaluate, shape (..., 3).
+
+        Returns:
+            Residual values, shape (...).
+        """
+        xyz = np.asarray(xyz)
+        diff = xyz - self._center
+        distance = np.linalg.norm(diff, axis=-1)
+        return distance - self._radius
+
+    def to_mesh(
+        self,
+        resolution: int | None = None,
+    ) -> tuple[NDArray[np.floating], NDArray[np.integer]]:
+        """Generate mesh vertices and faces.
+
+        Args:
+            resolution: Number of samples per dimension. If None, uses 50.
+
+        Returns:
+            Tuple of (vertices, faces).
+        """
+        if resolution is None:
+            resolution = 50
+        return sphere_mesh(self._center, self._radius, resolution)
+
+    def get_bounds_xyz(
+        self,
+    ) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
+        """Get axis-aligned bounding box in XYZ coordinates.
+
+        Returns:
+            Tuple of (min_corner, max_corner).
+        """
+        offset = np.array([self._radius, self._radius, self._radius])
+        return self._center - offset, self._center + offset
+
+    def to_trace(
+        self,
+        resolution: int | None = None,
+        intensity_fn: Callable[[NDArray[np.floating]], NDArray[np.floating]]
+        | None = None,
+        **kwargs,
+    ) -> go.Mesh3d:
+        """Generate a Plotly Mesh3d trace.
+
+        Args:
+            resolution: Mesh resolution. If None, uses default (50).
+            intensity_fn: Optional intensity function for coloring.
+            **kwargs: Additional arguments for vertices_to_mesh3d.
+
+        Returns:
+            Plotly Mesh3d trace.
+        """
+        vertices, faces = self.to_mesh(resolution)
+        return vertices_to_mesh3d(
+            vertices,
+            faces,
+            intensity_fn=intensity_fn,
+            name=self._label,
+            **kwargs,
+        )
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return f"Sphere(center={self._center}, radius={self._radius})"
+
+
+__all__ = ["Sphere"]