DeepSDFStruct 1.7.2__py3-none-any.whl

DeepSDFStruct/deep_sdf/networks/__init__.py

@@ -0,0 +1,8 @@
+"""
+DeepSDF Neural Network Architectures
+====================================
+
+This module provides neural network decoder architectures for DeepSDF models.
+Each architecture implements a different approach to decoding latent codes
+into SDF values.
+"""

DeepSDFStruct/deep_sdf/networks/analytic_round_cross.py

@@ -0,0 +1,105 @@
+"""
+Analytic Round Cross Decoder Network
+====================================
+
+This module implements a specialized decoder architecture for round cross
+lattice structures with analytical parametrization. The network incorporates
+geometric priors specific to cross-shaped lattice unit cells.
+
+The architecture is similar to the standard DeepSDF decoder but may include
+specialized layers or constraints to better represent the round cross geometry
+family, which is common in mechanical metamaterials.
+"""
+
+#!/usr/bin/env python3
+# Copyright 2004-present Facebook. All Rights Reserved.
+
+import torch.nn as nn
+import torch
+
+
+class RoundCrossDecoder(nn.Module):
+    def __init__(
+        self,
+        latent_size,
+        dims,
+        geom_dimension,
+        dropout=None,
+        dropout_prob=0.0,
+        norm_layers=(),
+        latent_in=(),
+        weight_norm=False,
+        xyz_in_all=None,
+        use_tanh=False,
+        latent_dropout=False,
+    ):
+        super(RoundCrossDecoder, self).__init__()
+
+        def make_sequence():
+            return []
+
+        dims = [latent_size + geom_dimension] + dims + [1]
+
+        self.num_layers = len(dims)
+        self.geom_dimension = geom_dimension
+        self.norm_layers = norm_layers
+        self.latent_in = latent_in
+        self.latent_dropout = latent_dropout
+        if self.latent_dropout:
+            self.lat_dp = nn.Dropout(0.2)
+
+        self.xyz_in_all = xyz_in_all
+        self.weight_norm = weight_norm
+
+        for layer in range(0, self.num_layers - 1):
+            if layer + 1 in latent_in:
+                out_dim = dims[layer + 1] - dims[0]
+            else:
+                out_dim = dims[layer + 1]
+                if self.xyz_in_all and layer != self.num_layers - 2:
+                    out_dim -= geom_dimension
+
+            if weight_norm and layer in self.norm_layers:
+                setattr(
+                    self,
+                    "lin" + str(layer),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Linear(dims[layer], out_dim)
+                    ),
+                )
+            else:
+                setattr(self, "lin" + str(layer), nn.Linear(dims[layer], out_dim))
+
+            if (
+                (not weight_norm)
+                and self.norm_layers is not None
+                and layer in self.norm_layers
+            ):
+                setattr(self, "bn" + str(layer), nn.LayerNorm(out_dim))
+
+        self.use_tanh = use_tanh
+        if use_tanh:
+            self.tanh = nn.Tanh()
+        self.relu = nn.ReLU()
+
+        self.dropout_prob = dropout_prob
+        self.dropout = dropout
+        self.th = nn.Tanh()
+
+    # input: N x (L+3), or N x (L+geom_dimension)
+    def forward(self, input):
+        xyz = input[:, -self.geom_dimension :]
+        r = input[:, 0]
+        output = torch.linalg.norm(xyz, dim=1, ord=torch.inf)
+
+        # add x cylinder
+        cylinder = torch.sqrt(xyz[:, 1] ** 2 + xyz[:, 2] ** 2) - r
+        output = torch.minimum(output, cylinder)
+        # add y cylinder
+        cylinder = torch.sqrt(xyz[:, 0] ** 2 + xyz[:, 2] ** 2) - r
+        output = torch.minimum(output, cylinder)
+        # add z cylinder
+        cylinder = torch.sqrt(xyz[:, 0] ** 2 + xyz[:, 1] ** 2) - r
+        output = torch.minimum(output, cylinder)
+
+        return output.reshape(-1, 1)

DeepSDFStruct/deep_sdf/networks/deep_sdf_decoder.py

@@ -0,0 +1,155 @@
+"""
+Standard DeepSDF Decoder Network
+================================
+
+This module implements the standard DeepSDF decoder architecture from
+Park et al. (2019). The network takes a latent code and 3D coordinates
+as input and outputs a signed distance value.
+
+Architecture
+------------
+The decoder is a multi-layer perceptron (MLP) with:
+- Concatenation of latent code and spatial coordinates at input
+- Optional skip connections for latent code injection
+- Optional coordinate injection at multiple layers
+- Batch normalization or weight normalization
+- ReLU activations (or tanh for final layer)
+- Dropout for regularization
+
+The network learns to map from a latent space to implicit SDF representations,
+enabling compact encoding of complex 3D geometries.
+"""
+
+#!/usr/bin/env python3
+# Copyright 2004-present Facebook. All Rights Reserved.
+
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+
+class DeepSDFDecoder(nn.Module):
+    def __init__(
+        self,
+        latent_size,
+        dims,
+        geom_dimension,
+        dropout=None,
+        dropout_prob=0.0,
+        norm_layers=(),
+        latent_in=(),
+        weight_norm=False,
+        xyz_in_all=None,
+        use_tanh=False,
+        latent_dropout=False,
+        homogen_predictor_dims=None,
+    ):
+        super(DeepSDFDecoder, self).__init__()
+
+        def make_sequence():
+            return []
+
+        dims = [latent_size + geom_dimension] + dims + [1]
+
+        self.num_layers = len(dims)
+        self.geom_dimension = geom_dimension
+        self.norm_layers = norm_layers
+        self.latent_in = latent_in
+        self.latent_dropout = latent_dropout
+        if self.latent_dropout:
+            self.lat_dp = nn.Dropout(0.2)
+
+        self.xyz_in_all = xyz_in_all
+        self.weight_norm = weight_norm
+
+        for layer in range(0, self.num_layers - 1):
+            if layer + 1 in latent_in:
+                out_dim = dims[layer + 1] - dims[0]
+            else:
+                out_dim = dims[layer + 1]
+                if self.xyz_in_all and layer != self.num_layers - 2:
+                    out_dim -= geom_dimension
+
+            if weight_norm and layer in self.norm_layers:
+                setattr(
+                    self,
+                    "lin" + str(layer),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Linear(dims[layer], out_dim)
+                    ),
+                )
+            else:
+                setattr(self, "lin" + str(layer), nn.Linear(dims[layer], out_dim))
+
+            if (
+                (not weight_norm)
+                and self.norm_layers is not None
+                and layer in self.norm_layers
+            ):
+                setattr(self, "bn" + str(layer), nn.LayerNorm(out_dim))
+
+        self.use_tanh = use_tanh
+        if use_tanh:
+            self.tanh = nn.Tanh()
+        self.relu = nn.ReLU()
+
+        self.dropout_prob = dropout_prob
+        self.dropout = dropout
+        self.th = nn.Tanh()
+        if homogen_predictor_dims is not None:
+            hp_dims = [latent_size] + homogen_predictor_dims
+            homogen_layers = []
+
+            for i in range(len(hp_dims) - 1):
+                homogen_layers.append(nn.Linear(hp_dims[i], hp_dims[i + 1]))
+
+                # No activation after final layer
+                if i < len(hp_dims) - 2:
+                    homogen_layers.append(nn.ReLU())
+
+            self.homogen_network = nn.Sequential(*homogen_layers)
+        else:
+            self.homogen_network = None
+
+    # input: N x (L+3), or N x (L+geom_dimension)
+    def forward(self, input):
+        xyz = input[:, -self.geom_dimension :]
+
+        if input.shape[1] > self.geom_dimension and self.latent_dropout:
+            latent_vecs = input[:, : -self.geom_dimension]
+            latent_vecs = F.dropout(latent_vecs, p=0.2, training=self.training)
+            x = torch.cat([latent_vecs, xyz], 1)
+        else:
+            x = input
+
+        for layer in range(0, self.num_layers - 1):
+            lin = getattr(self, "lin" + str(layer))
+            if layer in self.latent_in:
+                x = torch.cat([x, input], 1)
+            elif layer != 0 and self.xyz_in_all:
+                x = torch.cat([x, xyz], 1)
+            x = lin(x)
+            # last layer Tanh
+            if layer == self.num_layers - 2 and self.use_tanh:
+                x = self.tanh(x)
+            if layer < self.num_layers - 2:
+                if (
+                    self.norm_layers is not None
+                    and layer in self.norm_layers
+                    and not self.weight_norm
+                ):
+                    bn = getattr(self, "bn" + str(layer))
+                    x = bn(x)
+                x = self.relu(x)
+                if self.dropout is not None and layer in self.dropout:
+                    x = F.dropout(x, p=self.dropout_prob, training=self.training)
+
+        return x
+
+    def predict_C_homogenized(self, latent_vec):
+        if self.homogen_network is not None:
+            return self.homogen_network(latent_vec)
+        else:
+            raise RuntimeError(
+                "Cannot predict homogenized elasticity tensor, because homogenization network is not defined"
+            )

DeepSDFStruct/deep_sdf/networks/hierarchical_deep_sdf_decoder.py

@@ -0,0 +1,139 @@
+"""
+Hierarchical DeepSDF Decoder Network
+====================================
+
+This module implements a hierarchical decoder architecture that injects
+latent codes at multiple layers rather than only at the input. This enables
+multi-scale geometric control and better representation of complex features.
+
+Architecture
+------------
+Unlike the standard DeepSDF decoder, this architecture:
+- Accepts multiple latent vectors of potentially different sizes
+- Injects each latent at a specified layer depth
+- Enables hierarchical control from coarse to fine details
+- Supports various activation functions (ReLU, GELU, SiLU, etc.)
+
+The hierarchical structure is particularly useful for:
+- Multi-resolution geometric features
+- Compositional shape generation
+- Transfer learning across scales
+- Fine-tuning specific detail levels
+
+This architecture is inspired by progressive training and multi-scale
+representations in generative models.
+"""
+
+#!/usr/bin/env python3
+# Copyright 2004-present Facebook. All Rights Reserved.
+
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from DeepSDFStruct.deep_sdf.nn_utils import activations
+
+
+class HierachicalDeepSDFDecoder(nn.Module):
+    def __init__(
+        self,
+        latent_size,
+        dims,
+        geom_dimension,
+        dropout=None,
+        dropout_prob=0.0,
+        latent_in=[],
+        weight_norm=False,
+        norm_layers=[],
+        xyz_in_all=None,
+        activation_fun="relu",
+        latent_dropout=False,
+    ):
+        super(HierachicalDeepSDFDecoder, self).__init__()
+
+        self.layer_dims = [geom_dimension] + dims + [1]
+        assert len(latent_size) == len(
+            latent_in
+        ), "latent shape must match latent in shape"
+        # add latent dimensions to inputs
+        self.latent_dims = {}
+        for lat_size, index in zip(latent_size, latent_in):
+            self.layer_dims[index] = self.layer_dims[index] + lat_size
+            self.latent_dims[index] = lat_size
+
+        self.num_layers = len(dims)
+        self.geom_dimension = geom_dimension
+        self.latent_size = latent_size
+        self.latent_in = latent_in
+        self.latent_dropout = latent_dropout
+        if self.latent_dropout:
+            self.lat_dp = nn.Dropout(0.2)
+
+        self.xyz_in_all = xyz_in_all
+        self.weight_norm = weight_norm
+        self.norm_layers = norm_layers
+
+        for layer in range(0, len(self.layer_dims) - 1):
+            if layer + 1 in latent_in:
+                out_dim = self.layer_dims[layer + 1] - self.latent_dims[layer + 1]
+            else:
+                out_dim = self.layer_dims[layer + 1]
+
+            if weight_norm and layer in self.norm_layers:
+                setattr(
+                    self,
+                    "lin" + str(layer),
+                    nn.utils.parametrizations.weight_norm(
+                        nn.Linear(self.layer_dims[layer], out_dim)
+                    ),
+                )
+            else:
+                setattr(
+                    self, "lin" + str(layer), nn.Linear(self.layer_dims[layer], out_dim)
+                )
+
+        if activation_fun not in activations:
+            raise ValueError(
+                f"Unsupported activation function '{activation_fun}'. "
+                f"Available options are: {list(activations.keys())}"
+            )
+
+        self.activation = activations[activation_fun]
+
+        self.dropout_prob = dropout_prob
+        self.dropout = dropout
+        self.th = nn.Tanh()
+
+    # input: N x (L+3), or N x (L+geom_dimension)
+    def forward(self, input):
+        xyz = input[:, -self.geom_dimension :]
+        latent_vecs = input[:, : -self.geom_dimension]
+        latent_chunks = torch.split(latent_vecs, self.latent_size, dim=1)
+
+        if self.latent_dropout:
+            latent_vecs = F.dropout(latent_vecs, p=0.2, training=self.training)
+        x = xyz
+
+        for layer in range(0, len(self.layer_dims) - 1):
+            lin = getattr(self, "lin" + str(layer))
+            if layer in self.latent_in:
+                latent_idx = self.latent_in.index(layer)
+                latent_chunk = latent_chunks[latent_idx]
+                x = torch.cat([x, latent_chunk], 1)
+                # print(f"Latent chunk in layer {layer}: {latent_chunk}")
+            x = lin(x)
+
+            if layer < self.num_layers - 2:
+                if (
+                    self.norm_layers is not None
+                    and layer in self.norm_layers
+                    and not self.weight_norm
+                ):
+                    bn = getattr(self, "bn" + str(layer))
+                    x = bn(x)
+                if self.dropout is not None and layer in self.dropout:
+                    x = F.dropout(x, p=self.dropout_prob, training=self.training)
+
+            if layer <= (self.num_layers - 2):
+                x = self.activation(x)
+
+        return x

DeepSDFStruct/deep_sdf/networks/hierarchical_positional_sdf_decoder.py

@@ -0,0 +1,190 @@
+import torch
+import torch.nn as nn
+import logging
+import torch.nn.functional as F
+from DeepSDFStruct.deep_sdf.nn_utils import activations
+
+logger = logging.getLogger("DeepSDFStruct")
+
+
+class HierachicalPositionalDeepSDFDecoder(nn.Module):
+    def __init__(
+        self,
+        latent_size,
+        dims,
+        geom_dimension,
+        dropout=None,
+        dropout_prob=0.0,
+        latent_in=[],
+        weight_norm=False,
+        norm_layers=[],
+        xyz_in_all=None,
+        positional_encoding=False,
+        pe_levels=6,
+        split_latent=False,
+        activation_fun="relu",
+        latent_dropout=False,
+        use_tanh=None,
+    ):
+        super(HierachicalPositionalDeepSDFDecoder, self).__init__()
+        if use_tanh is not None:
+            logger.warning("Use tanh is deprecated, use the activation_fun argument")
+
+            # -----------------------------------------------------
+        # Positional Encoding
+        # -----------------------------------------------------
+        self.use_pe = positional_encoding
+        self.pe_levels = pe_levels
+
+        if positional_encoding:
+            # PE increases dimensionality from D -> D * (2*L + 1)
+            self.xyz_dim = geom_dimension * (2 * pe_levels + 1)
+        else:
+            self.xyz_dim = geom_dimension
+
+        self.geom_dimension = geom_dimension
+
+        # -----------------------------------------------------
+        # Base MLP structure
+        # layer_dims[i] = input dimension of layer i
+        # -----------------------------------------------------
+        self.layer_dims = [
+            self.xyz_dim
+        ] + dims  # final layer out_dim is handled separately
+        self.num_layers = len(dims)
+
+        # -----------------------------------------------------
+        # Latent injection logic
+        # -----------------------------------------------------
+        self.split_latent = split_latent
+        self.latent_in = latent_in
+        self.latent_size = latent_size
+        self.latent_dropout = latent_dropout
+
+        if self.split_latent:
+            assert (
+                latent_size % len(latent_in) == 0
+            ), "latent_size must be divisible by number of latent injection layers"
+            chunk_size = latent_size // len(latent_in)
+            latent_chunk_sizes = [chunk_size] * len(latent_in)
+        else:
+            latent_chunk_sizes = [latent_size] * len(latent_in)
+
+        self.latent_dims = {}  # map layer → latent_chunk_dim
+        for size, layer_idx in zip(latent_chunk_sizes, latent_in):
+            self.layer_dims[layer_idx] += size
+            self.latent_dims[layer_idx] = size
+
+        # -----------------------------------------------------
+        # XYZ skip connection logic (xyz_in_all)
+        # xyz is added BEFORE the linear layer, after activation,
+        # but only for layers > 0
+        # -----------------------------------------------------
+        self.xyz_in_all = xyz_in_all
+        if xyz_in_all:
+            for layer in range(1, self.num_layers):
+                self.layer_dims[layer] += self.xyz_dim
+
+        # -----------------------------------------------------
+        # Linear layer construction
+        # -----------------------------------------------------
+        self.weight_norm = weight_norm
+        self.norm_layers = norm_layers
+
+        for layer in range(self.num_layers):
+            in_dim = self.layer_dims[layer]
+            out_dim = dims[layer]
+
+            linear = nn.Linear(in_dim, out_dim)
+            if weight_norm and layer in norm_layers:
+                linear = nn.utils.parametrizations.weight_norm(linear)
+
+            setattr(self, f"lin{layer}", linear)
+
+        # Output layer (final layer: size dims[-1] → 1)
+        self.output_layer = nn.Linear(dims[-1], 1)
+
+        # -----------------------------------------------------
+        # Normalization + Dropout
+        # -----------------------------------------------------
+        self.dropout_prob = dropout_prob
+        self.dropout = dropout
+        if dropout is not None:
+            self.dropout_layers = set(dropout)
+        else:
+            self.dropout_layers = set()
+
+        # BatchNorm for chosen layers
+        for layer in norm_layers:
+            setattr(self, f"bn{layer}", nn.BatchNorm1d(dims[layer]))
+
+        # Latent dropout
+        if latent_dropout:
+            self.lat_dp = nn.Dropout(0.2)
+
+        # -----------------------------------------------------
+        # Activation
+        # -----------------------------------------------------
+        if activation_fun not in activations:
+            raise ValueError(
+                f"Unsupported activation function {activation_fun}. "
+                f"Choose from: {list(activations.keys())}"
+            )
+        self.activation = activations[activation_fun]
+
+    # ----------- Positional Encoding -----------
+    def positional_encoding(self, xyz):
+        """
+        Applies sinusoidal positional encoding to each of the xyz dims.
+        Output shape: [B, 3*(2*L+1)]
+        """
+        enc = [xyz]
+        for i in range(self.pe_levels):
+            freq = 2**i
+            enc.append(torch.sin(freq * xyz))
+            enc.append(torch.cos(freq * xyz))
+        return torch.cat(enc, dim=-1)
+
+    # input: N x (L+3), or N x (L+geom_dimension)
+    def forward(self, input: torch.Tensor):
+        xyz = input[:, -self.geom_dimension :]
+        latent_vecs = input[:, : -self.geom_dimension]
+        if self.latent_dropout:
+            latent_vecs = F.dropout(latent_vecs, p=0.2, training=self.training)
+        if self.split_latent:
+            latent_chunks = torch.split(
+                latent_vecs, int(self.latent_size / len(self.latent_in)), dim=1
+            )
+        else:
+            latent_chunks = [latent_vecs for _ in self.latent_in]
+
+        if self.use_pe:
+            xyz = self.positional_encoding(xyz)
+        x = xyz
+
+        for layer in range(0, len(self.layer_dims) - 1):
+            lin = getattr(self, "lin" + str(layer))
+            if layer in self.latent_in:
+                latent_idx = self.latent_in.index(layer)
+                latent_chunk = latent_chunks[latent_idx]
+                x = torch.cat([x, latent_chunk], 1)
+
+            if layer != 0 and self.xyz_in_all:
+                x = torch.cat([x, xyz])
+            x = lin(x)
+
+            if layer < self.num_layers - 2:
+                if (
+                    self.norm_layers is not None
+                    and layer in self.norm_layers
+                    and not self.weight_norm
+                ):
+                    bn = getattr(self, "bn" + str(layer))
+                    x = bn(x)
+                if self.dropout is not None and layer in self.dropout:
+                    x = F.dropout(x, p=self.dropout_prob, training=self.training)
+
+            if layer <= (self.num_layers - 2):
+                x = self.activation(x)
+
+        return self.output_layer(x)

DeepSDFStruct/deep_sdf/networks/resnet_positional_sdf_decoder.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import logging
+
+logger = logging.getLogger("DeepSDFStruct")
+
+
+class ResNetPositionalDeepSDFDecoder(nn.Module):
+    def __init__(
+        self,
+        latent_size,
+        dims,
+        dropout=None,
+        dropout_prob=None,
+        norm_layers=None,
+        latent_in=None,
+        xyz_in_all=False,
+        latent_dropout=False,
+        weight_norm=True,
+        use_tanh=None,
+        geom_dimension=3,
+        activation_fun="tanh",
+        positional_encoding=False,
+        pe_levels=6,  # number of frequency bands if PE enabled
+        resnet=False,
+        resnet_every=2,  # insert residual connection every N layers
+    ):
+        super().__init__()
+        if use_tanh is not None:
+            logger.warning("Use tanh is deprecated, use the activation_fun argument")
+        self.geom_dimension = geom_dimension
+        self.use_pe = positional_encoding
+        self.pe_levels = pe_levels
+        self.use_resnet = resnet
+        self.resnet_every = resnet_every
+
+        # Optional positional encoding increases xyz dimension
+        input_geom_dim = geom_dimension
+        if positional_encoding:
+            input_geom_dim = geom_dimension * (2 * pe_levels + 1)
+
+        # MLP layer dims
+        dims = [latent_size + input_geom_dim] + dims + [1]
+        self.num_layers = len(dims) - 1
+
+        if activation_fun == "tanh":
+            self.act = nn.Tanh()
+        elif activation_fun == "relu":
+            self.act = nn.ReLU()
+        elif activation_fun == "leaky_relu":
+            self.act = nn.LeakyReLU(0.2)
+        elif activation_fun == "silu":
+            self.act = nn.SiLU()
+        elif activation_fun == "sigmoid":
+            self.act = nn.Sigmoid()
+        else:
+            raise ValueError(f"Unknown activation function: {activation_fun}")
+
+        # Construct layers
+        self.layers = nn.ModuleList()
+        for i in range(self.num_layers):
+            self.layers.append(nn.Linear(dims[i], dims[i + 1]))
+
+    # ----------- Positional Encoding -----------
+    def positional_encoding(self, xyz):
+        """
+        Applies sinusoidal positional encoding to each of the xyz dims.
+        Output shape: [B, 3*(2*L+1)]
+        """
+        enc = [xyz]
+        for i in range(self.pe_levels):
+            freq = 2**i
+            enc.append(torch.sin(freq * xyz))
+            enc.append(torch.cos(freq * xyz))
+        return torch.cat(enc, dim=-1)
+
+    # ---------------- Forward ------------------
+    def forward(self, input):
+        """
+        input: [B, latent + xyz]
+        """
+        xyz = input[:, -self.geom_dimension :]
+        latent = input[:, : -self.geom_dimension]
+
+        if self.use_pe:
+            xyz = self.positional_encoding(xyz)
+
+        x = torch.cat([latent, xyz], dim=-1)
+
+        residual = None
+        for i, layer in enumerate(self.layers):
+            x_in = x
+            x = layer(x)
+            if i < self.num_layers - 1:
+                x = self.act(x)
+
+            if self.use_resnet and i % self.resnet_every == 0:
+
+                if residual is None:
+                    residual = x_in
+                else:
+
+                    if i < self.num_layers - 1:
+                        if residual.shape == x.shape:
+                            x = x + residual
+                    residual = x_in
+
+        return x