DeepSDFStruct 1.7.2__py3-none-any.whl

DeepSDFStruct/optimization.py

@@ -0,0 +1,357 @@
+"""
+Structural Optimization Utilities
+==================================
+
+This module provides tools for gradient-based optimization of SDF-based geometries,
+with a focus on structural design problems. It integrates with TorchFEM for
+finite element analysis and provides optimization algorithms suitable for
+constrained design problems.
+
+Key Features
+------------
+
+MMA Optimizer
+    Implementation of the Method of Moving Asymptotes (MMA), a gradient-based
+    algorithm well-suited for structural optimization with nonlinear constraints.
+    MMA is particularly effective for:
+    - Topology optimization
+    - Shape optimization with constraints
+    - Problems with expensive objective evaluations
+    - Highly nonlinear design spaces
+
+Finite Element Integration
+    - Conversion between TorchFEM and PyVista mesh formats
+    - Support for tetrahedral and hexahedral elements
+    - Linear and quadratic element types
+    - Integration with gradient computation
+
+Mesh Quality Utilities
+    - Signed volume computation for tetrahedra
+    - Mesh quality metrics
+    - Degeneracy detection
+
+The module is designed to work seamlessly with differentiable SDF representations,
+enabling gradient-based optimization of complex 3D structures.
+"""
+
+import torchfem.materials
+import torchfem.solid
+from torchfem.elements import Hexa1, Hexa2, Tetra1, Tetra2
+import torch
+import numpy as np
+from mmapy import mmasub
+import pyvista
+import logging
+import DeepSDFStruct
+
+logger = logging.getLogger(DeepSDFStruct.__name__)
+
+
+def get_mesh_from_torchfem(Solid: torchfem.Solid) -> pyvista.UnstructuredGrid:
+    """Convert a TorchFEM Solid mesh to PyVista UnstructuredGrid.
+
+    This function enables visualization and export of TorchFEM finite element
+    meshes using PyVista. It supports both tetrahedral and hexahedral elements
+    with linear and quadratic shape functions.
+
+    Parameters
+    ----------
+    Solid : torchfem.Solid
+        TorchFEM solid mesh object containing nodes, elements, and element type.
+
+    Returns
+    -------
+    pyvista.UnstructuredGrid
+        PyVista mesh representation suitable for visualization and I/O.
+
+    Raises
+    ------
+    NotImplementedError
+        If input is not a torchfem.Solid object.
+
+    Notes
+    -----
+    Supported element types:
+    - Tetra1: 4-node linear tetrahedron
+    - Tetra2: 10-node quadratic tetrahedron
+    - Hexa1: 8-node linear hexahedron
+    - Hexa2: 20-node quadratic hexahedron
+
+    Examples
+    --------
+    >>> from DeepSDFStruct.optimization import get_mesh_from_torchfem
+    >>> import torchfem
+    >>>
+    >>> # Assume we have a TorchFEM solid mesh
+    >>> # solid = torchfem.Solid(...)
+    >>>
+    >>> # Convert to PyVista for visualization
+    >>> pv_mesh = get_mesh_from_torchfem(solid)
+    >>> pv_mesh.plot()
+    """
+    if not isinstance(Solid, torchfem.Solid):
+        raise NotImplementedError("Currently only solid mesh is supported.")
+    # VTK cell types
+    etype = Solid.etype
+
+    if etype is Tetra1 or isinstance(etype, Tetra1):
+        cell_types = [pyvista.CellType.TETRA] * Solid.n_elem
+    elif etype is Tetra2 or isinstance(etype, Tetra2):
+        cell_types = [pyvista.CellType.QUADRATIC_TETRA] * Solid.n_elem
+    elif etype is Hexa1 or isinstance(etype, Hexa1):
+        cell_types = [pyvista.CellType.HEXAHEDRON] * Solid.n_elem
+    elif etype is Hexa2 or isinstance(etype, Hexa2):
+        cell_types = [pyvista.CellType.QUADRATIC_HEXAHEDRON] * Solid.n_elem
+    else:
+        raise TypeError(f"Unsupported element type: {etype} ({type(etype)})")
+
+    # VTK element list
+    el = len(Solid.elements[0]) * torch.ones(Solid.n_elem, dtype=Solid.elements.dtype)
+    elements = torch.cat([el[:, None], Solid.elements], dim=1).view(-1).tolist()
+
+    # Deformed node positions
+    pos = Solid.nodes
+
+    # Create unstructured mesh
+    mesh = pyvista.UnstructuredGrid(elements, cell_types, pos.tolist())
+    return mesh
+
+
+def tet_signed_vol(vertices, tets):
+    """Compute signed volumes of tetrahedral elements.
+
+    Calculates the signed volume of each tetrahedron, which is positive for
+    correctly oriented elements and negative for inverted elements. This is
+    useful for detecting mesh degeneracies and enforcing mesh quality constraints.
+
+    Parameters
+    ----------
+    vertices : torch.Tensor
+        Vertex coordinates of shape (N, 3).
+    tets : torch.Tensor
+        Tetrahedral connectivity of shape (M, 4), where each row contains
+        vertex indices [v0, v1, v2, v3].
+
+    Returns
+    -------
+    torch.Tensor
+        Signed volumes of shape (M,), one per tetrahedron. Positive volumes
+        indicate correctly oriented elements.
+
+    Notes
+    -----
+    The signed volume is computed as:
+        V = (1/6) * ((v1-v0) × (v2-v0)) · (v3-v0)
+
+    Examples
+    --------
+    >>> import torch
+    >>> from DeepSDFStruct.optimization import tet_signed_vol
+    >>>
+    >>> # Define a simple tetrahedron
+    >>> vertices = torch.tensor([
+    ...     [0.0, 0.0, 0.0],
+    ...     [1.0, 0.0, 0.0],
+    ...     [0.0, 1.0, 0.0],
+    ...     [0.0, 0.0, 1.0]
+    ... ])
+    >>> tets = torch.tensor([[0, 1, 2, 3]])
+    >>> volumes = tet_signed_vol(vertices, tets)
+    >>> print(f"Volume: {volumes[0]:.3f}")  # Should be 1/6 ≈ 0.167
+    """
+    v0 = vertices[tets[:, 0]]
+    v1 = vertices[tets[:, 1]]
+    v2 = vertices[tets[:, 2]]
+    v3 = vertices[tets[:, 3]]
+    vols = torch.einsum("ij,ij->i", torch.cross(v1 - v0, v2 - v0, dim=1), v3 - v0) / 6.0
+    return vols
+
+
+class MMA:
+    """Method of Moving Asymptotes (MMA) optimizer for constrained problems.
+
+    MMA is a gradient-based optimization algorithm designed for nonlinear
+    constrained problems. It constructs convex subproblems using moving
+    asymptotes and is particularly effective for structural optimization.
+
+    The optimizer handles a single objective function and a single constraint,
+    with box bounds on design variables. It automatically normalizes the
+    objective by its initial value for better numerical behavior.
+
+    Parameters
+    ----------
+    parameters : torch.Tensor
+        Initial design variables (will be optimized in-place).
+    bounds : array-like of shape (n, 2)
+        Box constraints [[lower_1, upper_1], ..., [lower_n, upper_n]]
+        for each design variable.
+    max_step : float, default 0.1
+        Maximum allowed change in design variables per iteration,
+        as a fraction of the bound range.
+
+    Attributes
+    ----------
+    parameters : torch.Tensor
+        Current design variables (updated in-place each iteration).
+    loop : int
+        Current iteration number.
+    x : ndarray
+        Current design variables in numpy format.
+    xold1, xold2 : ndarray
+        Design variables from previous two iterations (for MMA history).
+
+    Methods
+    -------
+    step(F, dF, G, dG)
+        Perform one MMA optimization step given objective, constraint,
+        and their gradients.
+
+    Notes
+    -----
+    MMA was developed by Krister Svanberg and is widely used in topology
+    optimization. It is particularly effective for problems where:
+    - The objective and constraints are expensive to evaluate
+    - Gradients are available (via automatic differentiation)
+    - The design space is high-dimensional
+    - Strong nonlinearity is present
+
+    The implementation uses the mmapy package for the core MMA algorithm.
+
+    Examples
+    --------
+    >>> import torch
+    >>> from DeepSDFStruct.optimization import MMA
+    >>>
+    >>> # Define design variables
+    >>> params = torch.ones(10, requires_grad=True)
+    >>> bounds = [[0.0, 2.0]] * 10
+    >>>
+    >>> # Create optimizer
+    >>> optimizer = MMA(params, bounds, max_step=0.1)
+    >>>
+    >>> # Optimization loop
+    >>> for i in range(100):
+    ...     # Compute objective and constraint
+    ...     objective = (params ** 2).sum()
+    ...     constraint = params.sum() - 5.0
+    ...
+    ...     # Compute gradients
+    ...     dF = torch.autograd.grad(objective, params, create_graph=True)[0]
+    ...     dG = torch.autograd.grad(constraint, params, create_graph=True)[0]
+    ...
+    ...     # MMA step
+    ...     optimizer.step(objective, dF, constraint, dG)
+    ...
+    ...     if optimizer.ch < 1e-3:
+    ...         break
+
+    References
+    ----------
+    .. [1] Svanberg, K. (1987). "The method of moving asymptotes—a new method
+           for structural optimization." International Journal for Numerical
+           Methods in Engineering, 24(2), 359-373.
+    .. [2] mmapy: Python implementation of MMA
+           https://github.com/arjendeetman/mmapy
+    """
+
+    def __init__(self, parameters, bounds, max_step=0.1):
+        self.max_step = max_step
+        self.bounds = np.array(bounds)
+        self.parameters = parameters
+        self.m = 1
+        self.n = len(parameters)
+        self.x = parameters.detach().cpu().numpy()
+        self.xold1 = parameters.detach().cpu().numpy()
+        self.xold2 = parameters.detach().cpu().numpy()
+        self.low = []
+        self.upp = []
+        self.a0_MMA = 1
+        self.a_MMA = np.zeros((self.m, 1))
+        self.c_MMA = 10000 * np.ones((self.m, 1))
+        self.d_MMA = np.zeros((self.m, 1))
+
+        self.loop = 0
+        self.ch = 1.0
+        self.F0 = None
+
+    def step(self, F, dF, G, dG):
+        """Perform one MMA optimization step.
+
+        Updates design variables by solving a convex subproblem constructed
+        from the objective, constraint, and their gradients.
+
+        Parameters
+        ----------
+        F : torch.Tensor or float
+            Objective function value at current design.
+        dF : torch.Tensor
+            Gradient of objective w.r.t. design variables, shape (n,).
+        G : torch.Tensor or float
+            Constraint function value at current design (≤ 0 is feasible).
+        dG : torch.Tensor
+            Gradient of constraint w.r.t. design variables, shape (n,).
+
+        Notes
+        -----
+        The method automatically:
+        - Normalizes the objective by its initial value
+        - Enforces move limits based on max_step
+        - Updates MMA history (xold1, xold2)
+        - Computes and logs convergence metric (ch)
+        - Updates self.parameters in-place
+
+        The convergence metric ch is the relative change in design variables.
+        """
+        orig_shape = dF.shape
+        F_np = F.detach().cpu().numpy().reshape(-1, 1)
+        dFdx_np = dF.detach().cpu().numpy().reshape(-1, 1)
+        G_np = G.detach().cpu().numpy().reshape(-1, 1)
+        dGdx_np = dG.detach().cpu().numpy().reshape(1, -1)
+        if self.loop == 0:
+            self.F0 = F_np
+        F_np = F_np / self.F0
+        dFdx_np = dFdx_np / self.F0
+
+        xmin = np.maximum(self.x - self.max_step, self.bounds[:, 0:1])
+        xmax = np.minimum(self.x + self.max_step, self.bounds[:, 1:2])
+        move = 0.1
+        self.loop = self.loop + 1
+        xmma, ymma, zmma, lam, xsi, eta, muMMA, zet, s, low, upp = mmasub(
+            self.m,
+            self.n,
+            self.loop,
+            self.x,
+            xmin,
+            xmax,
+            self.xold1,
+            self.xold2,
+            F_np,
+            dFdx_np,
+            G_np,
+            dGdx_np,
+            self.low,
+            self.upp,
+            self.a0_MMA,
+            self.a_MMA,
+            self.c_MMA,
+            self.d_MMA,
+        )
+
+        self.xold2 = self.xold1.copy()
+        self.xold1 = self.x.copy()
+        self.x = xmma
+        self.upp = upp
+        self.low = low
+
+        ch = np.abs(np.mean(self.x.T - self.xold1.T) / np.mean(self.x.T))
+        with torch.no_grad():
+            self.parameters.copy_(
+                torch.tensor(
+                    xmma, dtype=self.parameters.dtype, device=self.parameters.device
+                )
+            )
+        logger.info(
+            "It.: {0:4} | J.: {1:1.3e} | Constr.:  {2:1.3e} | ch.: {3:1.3e}".format(
+                self.loop, F_np[0][0], G_np[0][0], ch
+            )
+        )

DeepSDFStruct/parametrization.py

@@ -0,0 +1,215 @@
+"""
+Parametrization Functions for Spatially-Varying Properties
+===========================================================
+
+This module provides parametrization functions that map from spatial coordinates
+to parameter values, enabling SDFs with spatially-varying properties. This is
+particularly useful for lattice structures where properties like strut thickness
+or material density vary across the structure.
+
+Classes
+-------
+Constant
+    Spatially-constant parameter that is independent of position.
+    Useful as a baseline or for uniform structures.
+
+SplineParametrization
+    B-spline-based parametrization providing smooth spatial variation.
+    Parameters vary smoothly according to a spline function.
+
+Both classes inherit from torch.nn.Module, making their parameters
+automatically discoverable for gradient-based optimization.
+
+Examples
+--------
+Create a lattice with varying thickness::
+
+    from DeepSDFStruct.parametrization import SplineParametrization
+    from DeepSDFStruct.torch_spline import TorchSpline
+    import splinepy
+
+    # Create a B-spline for thickness variation
+    spline = splinepy.BSpline(
+        degrees=[2, 2, 2],
+        control_points=...,
+        knot_vectors=...
+    )
+
+    # Wrap in parametrization module
+    param_func = SplineParametrization(spline)
+
+    # Use in lattice structure
+    lattice = LatticeSDFStruct(
+        tiling=[3, 3, 3],
+        microtile=unit_cell,
+        parametrization=param_func
+    )
+"""
+
+import torch
+import torch.nn as nn
+import splinepy as sp
+from DeepSDFStruct.torch_spline import TorchSpline
+
+
+class Constant(nn.Module):
+    """Spatially-constant parameter value.
+
+    This parametrization returns the same parameter value(s) for all query
+    points, independent of position. Useful for uniform structures or as
+    a simple baseline.
+
+    The parameter is stored as a torch.nn.Parameter, making it automatically
+    discoverable for optimization.
+
+    Parameters
+    ----------
+    value : float, list, or torch.Tensor
+        The constant parameter value(s). Can be scalar or vector.
+    device : str or torch.device, optional
+        Device for computation ('cpu' or 'cuda').
+    dtype : torch.dtype, optional
+        Data type for the parameter.
+
+    Attributes
+    ----------
+    param : torch.nn.Parameter
+        The learnable constant parameter.
+
+    Examples
+    --------
+    >>> from DeepSDFStruct.parametrization import Constant
+    >>> import torch
+    >>>
+    >>> # Scalar constant
+    >>> const = Constant(0.5)
+    >>> queries = torch.rand(100, 3)
+    >>> params = const(queries)
+    >>> print(params.shape)  # (100, 1)
+    >>> print(params[0])  # tensor([0.5])
+    >>>
+    >>> # Vector constant
+    >>> const_vec = Constant([0.5, 1.0, 1.5])
+    >>> params = const_vec(queries)
+    >>> print(params.shape)  # (100, 3)
+    """
+
+    def __init__(self, value, device=None, dtype=None):
+        super().__init__()
+        if not isinstance(value, torch.Tensor):
+            value = torch.tensor(value, device=device, dtype=dtype)
+        self.param = nn.Parameter(value)
+
+    def forward(self, queries: torch.Tensor) -> torch.Tensor:
+        """Evaluate constant parameter at query points.
+
+        Parameters
+        ----------
+        queries : torch.Tensor
+            Query point coordinates of shape (N, d). The spatial
+            coordinates are ignored.
+
+        Returns
+        -------
+        torch.Tensor
+            Parameter values of shape (N, m) where m is the
+            parameter dimension. All rows are identical.
+        """
+        N = queries.shape[0]
+        return self.param.expand(N, -1)
+
+    def set_param(self, new_value: torch.Tensor):
+        """Update the constant parameter value.
+
+        Parameters
+        ----------
+        new_value : torch.Tensor
+            New parameter value(s).
+        """
+        with torch.no_grad():
+            self.param.copy_(new_value.to(self.param.device))
+
+
+class SplineParametrization(nn.Module):
+    """B-spline-based spatially-varying parametrization.
+
+    Uses a B-spline, Bezier, or NURBS function to provide smoothly-varying
+    parameters across space. The spline control points are learnable parameters
+    that can be optimized.
+
+    This is useful for creating gradual transitions in material properties,
+    thickness variations, or other spatially-dependent design variables.
+
+    Parameters
+    ----------
+    spline : splinepy.BSpline, splinepy.Bezier, or splinepy.NURBS
+        The spline function defining parameter variation. The spline's
+        output dimension determines the parameter dimension.
+    device : str or torch.device, optional
+        Device for computation ('cpu' or 'cuda').
+
+    Attributes
+    ----------
+    torch_spline : TorchSpline
+        PyTorch-compatible wrapper for the spline.
+
+    Examples
+    --------
+    >>> from DeepSDFStruct.parametrization import SplineParametrization
+    >>> import splinepy
+    >>> import torch
+    >>>
+    >>> # Create a 3D B-spline for thickness variation
+    >>> spline = splinepy.BSpline(
+    ...     degrees=[2, 2, 2],
+    ...     control_points=torch.rand(27, 1),  # 3x3x3 grid, 1D output
+    ...     knot_vectors=[
+    ...         [0, 0, 0, 1, 1, 1],
+    ...         [0, 0, 0, 1, 1, 1],
+    ...         [0, 0, 0, 1, 1, 1]
+    ...     ]
+    ... )
+    >>>
+    >>> # Create parametrization
+    >>> param = SplineParametrization(spline)
+    >>>
+    >>> # Evaluate at query points
+    >>> queries = torch.rand(100, 3)
+    >>> thickness = param(queries)
+    >>> print(thickness.shape)  # (100, 1)
+    """
+
+    def __init__(self, spline: sp.BSpline | sp.Bezier | sp.NURBS, device=None):
+        super().__init__()
+        self.torch_spline = TorchSpline(spline, device=device)
+
+    def forward(self, queries: torch.Tensor) -> torch.Tensor:
+        """Evaluate spline parametrization at query points.
+
+        Parameters
+        ----------
+        queries : torch.Tensor
+            Query point coordinates of shape (N, d) where d matches
+            the spline's parametric dimension.
+
+        Returns
+        -------
+        torch.Tensor
+            Parameter values of shape (N, m) where m is the spline's
+            output dimension (number of parameter values per point).
+        """
+        return self.torch_spline(queries)
+
+    def set_param(self, new_value: torch.Tensor):
+        """Update the spline control points.
+
+        Parameters
+        ----------
+        new_value : torch.Tensor
+            New control point values. Shape must match the spline's
+            control point array.
+        """
+        with torch.no_grad():
+            self.torch_spline.control_points.copy_(
+                new_value.to(self.torch_spline.control_points.device)
+            )

DeepSDFStruct/pretrained_models.py

@@ -0,0 +1,128 @@
+"""
+Pretrained DeepSDF Models
+========================
+
+This module provides access to pretrained DeepSDF neural network models for
+common microstructure geometries. These models can be used directly for
+geometry generation or as starting points for transfer learning.
+
+Available Models
+----------------
+
+ChiAndCross
+    Chi-shaped and cross-shaped lattice structures, commonly used in
+    mechanical metamaterials.
+
+AnalyticRoundCross
+    Round cross-section variations with analytical parameterization,
+    useful for smooth stress distribution.
+
+RoundCross
+    Standard round cross structures with various connectivity patterns.
+
+Primitives
+    Basic 3D geometric primitives (spheres, cylinders, cubes, etc.)
+    for building more complex structures.
+
+Primitives2D
+    2D geometric primitives for planar structures and cross-sections.
+
+Functions
+---------
+
+get_model(model, checkpoint='latest', device=None)
+    Load a pretrained model by name or enum value. Returns a DeepSDFModel
+    ready for inference.
+
+list_available_models()
+    Get a list of all available pretrained models.
+
+The pretrained models are stored within the package and loaded on demand,
+enabling quick prototyping and exploration without requiring training.
+
+Examples
+--------
+Load and use a pretrained model::
+
+    from DeepSDFStruct.pretrained_models import get_model, PretrainedModels
+    import torch
+
+    # Load a model
+    model = get_model(PretrainedModels.RoundCross)
+
+    # Generate geometry with latent code
+    latent_code = torch.zeros(256)  # Use learned latent vector
+    points = torch.rand(1000, 3)
+    sdf_values = model(points, latent_code)
+
+List available models::
+
+    from DeepSDFStruct.pretrained_models import list_available_models
+
+    models = list_available_models()
+    for model in models:
+        print(f"Available model: {model.value}")
+"""
+
+from enum import Enum
+import importlib.resources
+from DeepSDFStruct.deep_sdf.workspace import load_trained_model, load_latent_vectors
+from DeepSDFStruct.deep_sdf.models import DeepSDFModel
+import torch
+
+
+class PretrainedModels(Enum):
+    ChiAndCross = "chi_and_cross"
+    AnalyticRoundCross = "analytic_round_cross"
+    RoundCross = "round_cross"
+    Primitives = "primitives"
+    Primitives2D = "primitives_2d"
+
+
+main_dir = importlib.resources.files("DeepSDFStruct")
+with importlib.resources.as_file(main_dir) as path:
+    PRETRAINED_MODELS_DIR = path / "trained_models"
+
+# Maps enum entries to file paths
+_MODEL_REGISTRY = {
+    PretrainedModels.ChiAndCross: PRETRAINED_MODELS_DIR / "chi_and_cross",
+    PretrainedModels.AnalyticRoundCross: PRETRAINED_MODELS_DIR / "analytic_round_cross",
+    PretrainedModels.RoundCross: PRETRAINED_MODELS_DIR / "round_cross",
+    PretrainedModels.Primitives: PRETRAINED_MODELS_DIR / "primitives",
+    PretrainedModels.Primitives2D: PRETRAINED_MODELS_DIR / "primitives_2d",
+}
+
+
+def get_model(
+    model: str | PretrainedModels, checkpoint: str = "latest", device=None
+) -> DeepSDFModel:
+    """
+    Load a pretrained model by name or enum.
+
+    Args:
+        model (str | PretrainedModels): model identifier
+        checkpoint (str): checkpoint file name (default: 'latest')
+
+    Returns:
+        Trained PyTorch model
+    """
+    if isinstance(model, str):
+        path = model
+    else:
+        model_enum = model
+        path = _MODEL_REGISTRY.get(model_enum)
+
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    if not path:
+        raise ValueError(f"Model path {path} not found.")
+    decoder = load_trained_model(path, checkpoint, device=device)
+    latent_vectors = load_latent_vectors(path, checkpoint, device=device)
+    decoder.eval()
+    deep_sdf_model = DeepSDFModel(decoder, latent_vectors, device=device)
+    return deep_sdf_model
+
+
+def list_available_models():
+    return list(PretrainedModels)