Add homogenization

.github/workflows/test.yml

@@ -13,6 +13,9 @@ jobs:
     env:
       HF_HOME: ${{ github.workspace }}/.cache/huggingface
       HF_HUB_CACHE: ${{ github.workspace }}/.cache/huggingface/hub
+      HF_HUB_DOWNLOAD_TIMEOUT: "60"
+      HF_HUB_ETAG_TIMEOUT: "30"
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
     strategy:
       matrix:
@@ -36,12 +39,9 @@ jobs:
       - name: Cache huggingface
         uses: actions/cache@v4
         with:
-          path: ~/.cache/huggingface
-          key: ${{ runner.os }}-huggingface
-
-      - name: Ensure Hugging Face cache directories exist
-        run: |
-          mkdir -p "$HF_HUB_CACHE"
+          path: ${{ github.workspace }}/.cache/huggingface
+          key: huggingface-${{ hashFiles('pyproject.toml') }}
+          restore-keys: huggingface-
 
       - name: Ensure Hugging Face cache directories exist
         run: |

DeepSDFStruct/SDF.py

@@ -366,6 +366,8 @@ def plot_slice(
         cmap="seismic",
         show_zero_level=True,
         deformation_function=None,
+        xlim=None,
+        ylim=None,
     ):
         """Plot a 2D slice through an SDF as a contour plot.
 
@@ -435,8 +437,15 @@ def plot_slice(
             fig, ax = plt.subplots()
             plt_show = True
         bounds = self._get_domain_bounds()
-        xlim, ylim = project_bounds(origin, normal, bounds=bounds)
-        points = generate_plane_points(origin, normal, res, xlim, ylim)
+
+        if self.geometric_dim == 3:
+            if xlim is None:
+                xlim, ylim = project_bounds(origin, normal, bounds=bounds)
+            points = generate_plane_points(origin, normal, res, xlim, ylim)
+        else:
+            if xlim is None:
+                xlim, ylim = bounds.detach().cpu().numpy().T
+            points = generate_plane_points(origin, normal, res, xlim, ylim)[:, :2]
 
         sdf_device = self.get_device()
         points = torch.from_numpy(points).to(torch.float32).to(sdf_device)
@@ -1636,6 +1645,11 @@ def project_bounds(origin, normal, bounds=None):
 
     if bounds is None:
         bounds = np.array([[0, 0, 0], [1, 1, 1]])
+    elif isinstance(bounds, torch.Tensor):
+        bounds = bounds.detach().cpu().numpy()
+    if bounds.shape[1] == 2:
+        logger.info("cannot project 2D onto 2D, returning input")
+        return bounds[0], bounds[1]
 
     bmin, bmax = bounds
 

DeepSDFStruct/deep_sdf/data.py

@@ -227,9 +227,9 @@ def __init__(
                         neg_tensor[torch.randperm(neg_tensor.shape[0])],
                     ]
                 )
-                if "E" in npz.keys():
+                if "C" in npz.keys():
                     self.loaded_mat_properties.append(
-                        torch.from_numpy(npz["E.npy"]).to(torch.float32)
+                        torch.from_numpy(npz["C.npy"]).to(torch.float32).reshape(-1)
                     )
                 self.filenames.append(filename)
 
@@ -247,12 +247,12 @@ def __getitem__(self, idx):
         if self.load_ram:
             return (
                 unpack_sdf_samples_from_ram(self.loaded_data[idx], self.subsample),
-                mat_prop,
+                mat_prop.unsqueeze(0).expand(self.subsample, -1),
                 idx,
             )
         else:
             return (
                 unpack_sdf_samples(filename, self.geom_dimension, self.subsample),
-                mat_prop,
+                mat_prop.unsqueeze(0).expand(self.subsample, -1),
                 idx,
             )

DeepSDFStruct/deep_sdf/networks/deep_sdf_decoder.py

@@ -29,72 +29,87 @@
 
 
 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):
@@ -130,3 +145,11 @@
                     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/training.py

@@ -485,6 +485,9 @@ def empirical_stat(latent_vecs, indices):
     )
     loss_fun_spec = get_spec_with_default(specs, "LossFunction", "clampedL1")
     loss_fun = get_loss_function(loss_fun_spec)
+    add_homogeniation = get_spec_with_default(specs, "AddHomogenization", False)
+    if add_homogeniation:
+        logger.info("adding homogenization loss")
 
     optimizer_all = torch.optim.Adam(
         [
@@ -561,7 +564,10 @@ def empirical_stat(latent_vecs, indices):
         for sdf_data, properties, indices in sdf_loader:
             # Process the input data
             sdf_data = sdf_data.reshape(-1, geom_dimension + 1).to(device)
-            properties = properties.to(device)
+            properties_expanded = properties.reshape(-1, properties.shape[-1]).to(
+                device
+            )
+
             indices = indices.to(device)
 
             num_sdf_samples = sdf_data.shape[0]
@@ -581,13 +587,16 @@ def empirical_stat(latent_vecs, indices):
             )
 
             sdf_gt = torch.chunk(sdf_gt, batch_split)
+            properties_chunked = torch.chunk(properties_expanded, batch_split)
 
             batch_loss = 0.0
             batch_reg_loss = 0.0
+            batch_hom_loss = 0.0
             optimizer_all.zero_grad()
 
             for i in range(batch_split):
                 batch_lat_vecs = lat_vecs(indices[i])
+                batch_properties = properties_chunked[i]
 
                 input = torch.cat([batch_lat_vecs, xyz[i]], dim=1)
 
@@ -599,11 +608,22 @@ def empirical_stat(latent_vecs, indices):
 
                 chunk_loss = loss_fun(pred_sdf, sdf_gt[i].to(device))
 
-                l2_size_loss = torch.sum(torch.norm(batch_lat_vecs, dim=1))
-                reg_loss = (code_reg_lambda * min(1, epoch / 100) * l2_size_loss) / (
-                    num_sdf_samples / batch_split
-                )
+                if add_homogeniation:
+                    pred_properties = decoder.predict_C_homogenized(batch_lat_vecs)
+
+                    hom_loss = torch.nn.functional.mse_loss(
+                        pred_properties, batch_properties, reduction="mean"
+                    )
+
+                    chunk_loss += hom_loss.to(device)
+                    batch_hom_loss += hom_loss.item()
 
+                # standard l2 latent vector loss
+                reg_loss = (
+                    code_reg_lambda
+                    * min(1, epoch / 100)
+                    * torch.mean(torch.norm(batch_lat_vecs, dim=1))
+                )
                 chunk_loss = chunk_loss + reg_loss.to(device)
                 batch_reg_loss = batch_reg_loss + reg_loss.to(device)
 
@@ -637,9 +657,18 @@ def empirical_stat(latent_vecs, indices):
                 f"Finished {epoch} ({epoch}/{num_epochs}) [{epoch/num_epochs*100:.2f}%] after {total_time}"
             )
         else:
-            pbar.set_postfix(
-                {"Loss": f"{batch_loss:.4f}", "Reg": f"{batch_reg_loss:.4f}"}
-            )
+            if add_homogeniation:
+                pbar.set_postfix(
+                    {
+                        "Loss": f"{batch_loss:.4f}",
+                        "Reg": f"{batch_reg_loss:.4f}",
+                        "Hom": f"{batch_hom_loss:.4f}",
+                    }
+                )
+            else:
+                pbar.set_postfix(
+                    {"Loss": f"{batch_loss:.4f}", "Reg": f"{batch_reg_loss:.4f}"}
+                )
         seconds_elapsed = end - start
         timing_log.append(seconds_elapsed)
 

DeepSDFStruct/sampling.py

@@ -48,6 +48,8 @@
 from functools import partial
 from itertools import starmap
 import multiprocessing
+from dataclasses import dataclass
+
 
 import vtk
 import numpy as np
@@ -114,6 +116,18 @@ class SphereParameters(typing.TypedDict):
     r: float
 
 
+@dataclass
+class GeometryInstance:
+    name: str
+    geom: trimesh.Trimesh
+    C: np.ndarray | None
+
+    def __init__(self, name: str, geom: trimesh.Trimesh, C: np.ndarray | None):
+        self.name = name
+        self.geom = geom
+        self.C = C
+
+
 class SampledSDF:
     """Container for sampled SDF points and their distance values.
 
@@ -229,7 +243,7 @@ def __add__(self, other):
 
 
 def _process_single_geometry_instance(
-    geometry,
+    geom_instance: GeometryInstance,
     file_name: str,
     folder_name: pathlib.Path,
     scale,
@@ -247,16 +261,17 @@ def _process_single_geometry_instance(
         logger.warning(f"File {fname} already exists")
         return
     mesh = None
-    if isinstance(geometry, SDFBase):
-        sdf = geometry
-    elif isinstance(geometry, trimesh.Trimesh):
-        mesh = geometry
+    geom = geom_instance.geom
+    if isinstance(geom, SDFBase):
+        sdf = geom
+    elif isinstance(geom, trimesh.Trimesh):
+        mesh = geom
         sdf = SDFfromMesh(mesh, scale=scale)
         if also_save_mesh:
             mesh.export(fname.with_suffix(".stl"))
     else:
         raise NotImplementedError(
-            f"Geometry must be either trimesh or SDFBase, but not {type(geometry)}."
+            f"Geometry must be either trimesh or SDFBase, but not {type(geom)}."
         )
     sampled_sdf = random_sample_sdf(
         sdf,
@@ -268,16 +283,18 @@ def _process_single_geometry_instance(
         if not isinstance(mesh, trimesh.Trimesh):
             logger.warning(
                 "Add surface samples was specified, but geometry"
-                f"is not given as a mesh but as {type(geometry)}"
+                f"is not given as a mesh but as {type(geom)}"
             )
         else:
             surf_samples = sample_mesh_surface(
                 sdf, mesh, int(n_samples // 2), stds, device="cpu", dtype=torch.float32
             )
             sampled_sdf += surf_samples
     pos, neg = sampled_sdf.split_pos_neg()
-
-    np.savez(fname, neg=neg.stacked, pos=pos.stacked)
+    if geom_instance.C is not None:
+        np.savez(fname, neg=neg.stacked, pos=pos.stacked, C=geom_instance.C)
+    else:
+        np.savez(fname, neg=neg.stacked, pos=pos.stacked)
     if also_save_vtk:
         save_points_to_vtp(fname.with_suffix(".vtp"), neg=neg.stacked, pos=pos.stacked)
 
@@ -294,7 +311,7 @@ def __init__(
         self.outdir = outdir
         self.splitdir = splitdir
         self.dataset_name = dataset_name
-        self.geometries = {}
+        self.geometries: dict[str, dict[str, GeometryInstance]] = {}
         self.stds = stds
         folder_name = pathlib.Path(outdir) / dataset_name
         if os.path.exists(folder_name):
@@ -306,12 +323,22 @@ def __init__(
         else:
             os.makedirs(folder_name)
 
-    def add_class(self, geom_list: list, class_name: str, n_faces=100) -> None:
+    def add_class(
+        self,
+        geom_list: list,
+        class_name: str,
+        n_faces=100,
+        homogenized_c: list[np.ndarray] | None = None,
+    ) -> None:
         """
         Adds a geometry to the sampler object. Tries to transform inputs to
         trimesh data. In case the geometry is a spline object, the n_faces
         parameter determines the accuracy of the extracted mesh
         """
+        if homogenized_c is not None:
+            assert len(geom_list) == len(
+                homogenized_c
+            ), f"Length of geometries ({len(geom_list)}) does not match length of homogenized elasticity tensors ({len(homogenized_c)})"
         instances = {}
         for i, geom in enumerate(geom_list):
             instance_name = f"{class_name}_{i:05}"
@@ -323,7 +350,11 @@ def add_class(self, geom_list: list, class_name: str, n_faces=100) -> None:
             elif isinstance(geom, torchSurfMesh):
                 geom = geom.to_trimesh()
 
-            instances[instance_name] = geom
+            if homogenized_c is not None:
+                C = homogenized_c[i]
+            else:
+                C = None
+            instances[instance_name] = GeometryInstance(instance_name, geom, C)
         self.geometries[class_name] = instances
 
     def process_geometries(