DeepSDFStruct 1.7.2__py3-none-any.whl

tests/test_structural_optimization.py

@@ -0,0 +1,178 @@
+from DeepSDFStruct.pretrained_models import get_model, PretrainedModels
+from DeepSDFStruct.SDF import SDFfromDeepSDF, CappedBorderSDF
+from DeepSDFStruct.lattice_structure import LatticeSDFStruct
+from DeepSDFStruct.torch_spline import TorchSpline
+from DeepSDFStruct.mesh import create_3D_mesh, torchVolumeMesh
+from DeepSDFStruct.parametrization import SplineParametrization
+from DeepSDFStruct.optimization import MMA, tet_signed_vol, get_mesh_from_torchfem
+from DeepSDFStruct.utils import configure_logging
+import DeepSDFStruct
+import torchfem.materials
+import torchfem.solid
+import splinepy
+import torch
+import logging
+import numpy as np
+
+logger = logging.getLogger(DeepSDFStruct.__name__)
+configure_logging()
+
+
+def test_structural_optimization(num_iter=1):
+    # torch.set_default_device("cuda")
+
+    logger.info("Loading Model")
+    # Load a pretrained DeepSDF model
+    model = get_model(PretrainedModels.AnalyticRoundCross)
+    sdf = SDFfromDeepSDF(model)
+
+    cap_border_dict = {
+        "x0": {"cap": 1, "measure": 0.05},
+        "z1": {"cap": 1, "measure": 0.1},
+    }
+
+    param_spline_sp = splinepy.BSpline(
+        [1, 0, 1],
+        [[0, 0, 0.25, 0.5, 0.75, 1, 1], [0, 1], [0, 0, 0.5, 1, 1]],
+        [[0.3] * model._trained_latent_vectors[0].shape[0]] * 3 * 5,
+    )
+    logger.info("using spline with knot vectors:")
+    logger.info(param_spline_sp.knot_vectors)
+    logger.info(f"and control points: {param_spline_sp.control_points}")
+    tiling = [2, 1, 1]
+
+    param_spline = SplineParametrization(param_spline_sp, device=model.device)
+    # Define a spline-based deformation field
+    deformation_spline = TorchSpline(
+        splinepy.helpme.create.box(2, 1, 1).bspline, device=model.device
+    )
+
+    # Create the lattice structure with deformation and microtile
+    lattice_struct_uncapped = LatticeSDFStruct(
+        tiling=tiling, microtile=sdf, parametrization=param_spline
+    )
+    lattice_struct = CappedBorderSDF(
+        CappedBorderSDF(lattice_struct_uncapped, cap_border_dict)
+    )
+
+    lr = 1e-2
+    param = next(lattice_struct_uncapped.parametrization.parameters())
+    # optimizer = torch.optim.Adam([param], lr=lr)
+    init_vol = None
+    init_compl = None
+    target_vol = 0.5
+
+    bounds = np.zeros(param.shape) + np.array([0.15, 0.75])
+    optimizer = MMA(param, bounds)
+
+    for e in range(num_iter):
+        # log(
+        #     f"Starting iteration with parameters: {lattice_struct.parametrization.parameters.T}"
+        # )
+        # torch.set_default_device("cuda")
+        torch.set_default_dtype(torch.float32)
+        mesh, _ = create_3D_mesh(
+            lattice_struct,
+            10,
+            mesh_type="volume",
+            differentiate=False,
+            device=model.device,
+            deformation_function=deformation_spline,
+        )
+        surf_mesh, _ = create_3D_mesh(
+            lattice_struct,
+            10,
+            mesh_type="surface",
+            differentiate=False,
+            device=model.device,
+            deformation_function=deformation_spline,
+        )
+
+        surf_trimesh = surf_mesh.to_trimesh()
+        assert surf_trimesh.is_watertight
+        if isinstance(mesh, torchVolumeMesh):
+            tets = mesh.volumes
+            verts = mesh.vertices
+        else:
+            raise RuntimeError("Resulting mesh should be volume mesh.")
+
+        # Ensure consistent positive orientation for each tetrahedron.
+        tets_oriented = tets.clone()
+        vols = tet_signed_vol(verts, tets_oriented)
+        neg_mask = vols < 0
+        if neg_mask.any():
+            tets_oriented_neg = tets_oriented[neg_mask]
+            tets_oriented_neg = tets_oriented_neg[:, [0, 2, 1, 3]]
+            tets_oriented[neg_mask] = tets_oriented_neg
+
+        vols = tet_signed_vol(verts, tets_oriented)
+        if init_vol is None:
+            init_vol = vols.sum().item()
+            logger.info(f"Initial volume: {init_vol} on {len(vols)} elements.")
+        eps = 1e-12
+        # Use float64 for the mask to match torchfem's internal precision;
+        # tets with tiny float32-positive volumes can become negative in float64.
+        vols_f64 = tet_signed_vol(verts.to(torch.float64), tets_oriented)
+        mask = vols_f64 > eps
+        vol = vols[mask].sum()
+
+        # keep only the good tets
+        tets_clean = tets_oriented[mask]
+
+        # check how many were removed
+        removed = (~mask).sum()
+        # print(f"Removed {removed} negative volume tets")
+
+        # Create model
+        # torch.set_default_device("cpu")
+        torch.set_default_dtype(torch.float64)
+        verts_tf = verts.to("cpu").to(torch.float64)
+        tets_tf = tets_clean.to("cpu")
+
+        # Material
+        material = torchfem.materials.IsotropicElasticity3D(E=1000.0, nu=0.3)
+
+        cantilever = torchfem.solid.Solid(verts_tf, tets_tf, material)
+
+        # Constrained displacement at left end [Node_IDs, DOFs]
+        left_const_mask = cantilever.nodes[:, 0] < 1e-5
+        cantilever.constraints[left_const_mask, :] = True
+
+        # Load at tip [Node_ID, DOF]
+        top_mask = cantilever.nodes[:, 2] > (1 - 1e-1)
+        top_mask_no_side = top_mask & (~left_const_mask)
+        num_nodes = top_mask_no_side.sum().item()
+        cantilever.forces[top_mask_no_side, 2] = -100.0 / num_nodes
+
+        # log("Starting Simulation")
+        u, f, _, _, _ = cantilever.solve(
+            rtol=1e-2, atol=1e-2, device="cpu", method="spsolve"
+        )
+
+        # Compute sensitivity of compliance w.r.t. element thicknesses
+        compliance = torch.inner(f.ravel(), u.ravel())
+
+        F = compliance
+        G = vol - target_vol
+        dF = torch.autograd.grad(F, param, retain_graph=True)[0]
+        dG = torch.autograd.grad(G, param, retain_graph=True)[0]
+        optimizer.step(F, dF, G, dG)
+
+        # Reset default dtype to float32 for other tests
+        torch.set_default_dtype(torch.float32)
+
+    # torch.autograd.grad(compliance, cantilever.thickness)[0]
+    mesh = get_mesh_from_torchfem(cantilever)
+    mesh.point_data["u"] = u.detach().cpu().numpy()
+    out_file_name = "sim_out.vtk"
+    logger.info(f"Writing Output to {out_file_name}")
+    mesh.save(out_file_name)
+    assert abs(dF.sum().item()) > 1e-8, "Derivative of objective is zero"
+    assert abs(dG.sum().item()) > 1e-8, "Derivative of constraint is zero"
+
+
+if __name__ == "__main__":
+    import warnings
+
+    warnings.filterwarnings("error")
+    test_structural_optimization()

tests/test_torch_spline.py

@@ -0,0 +1,216 @@
+import torch
+import splinepy
+from DeepSDFStruct.torch_spline import (
+    # TorchSpline,
+    generate_bbox_spline,
+    torch_spline_1D,
+    torch_spline_3D,
+)
+from DeepSDFStruct.torch_spline import TorchSpline
+import numpy as np
+import pytest
+
+
+@pytest.fixture
+def np_rng():
+    return np.random.default_rng(0)
+
+
+@pytest.fixture
+def bspline():
+    control_points = np.random.rand(64, 3)  # 4x4x4 Bezier -> 64 control points
+
+    knot_vecs = [
+        [0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0],
+    ]
+
+    return splinepy.BSpline(
+        degrees=[2, 2, 2], knot_vectors=knot_vecs, control_points=control_points
+    )
+
+
+def test_torchspline_evaluation(np_rng, bspline, device="cpu"):
+    """Check that TorchSpline evaluation matches splinepy evaluation"""
+
+    spline_module = TorchSpline(bspline, device=device)
+
+    queries = torch.tensor(np_rng.random((5, 3)), dtype=torch.float32, device=device)
+
+    # Torch output
+    output_torch = spline_module(queries).detach().cpu().numpy()
+
+    # Numpy output
+    output_np = bspline.evaluate(queries.detach().cpu().numpy())
+
+    assert np.allclose(output_torch, output_np, atol=1e-6)
+
+
+def test_torchspline_derivative(np_rng, bspline, device="cpu", eps=1e-3):
+    """Check that TorchSpline forward derivatives w.r.t input are correct via finite differences"""
+
+    spline_module = TorchSpline(bspline, device=device)
+
+    queries = torch.tensor(
+        np_rng.random((5, 3)), dtype=torch.float32, device=device, requires_grad=True
+    )
+
+    # Torch output
+    output = spline_module(queries)
+
+    # Finite differences approximation
+    fd_jacobian = np.zeros((queries.shape[0], output.shape[1], queries.shape[1]))
+    for i in range(queries.shape[0]):
+        for j in range(queries.shape[1]):
+            q_plus = queries.clone()
+            q_plus[i, j] += eps
+            q_minus = queries.clone()
+            q_minus[i, j] -= eps
+            f_plus = spline_module(q_plus).detach().cpu().numpy()[i]
+            f_minus = spline_module(q_minus).detach().cpu().numpy()[i]
+            fd_jacobian[i, :, j] = (f_plus - f_minus) / (2 * eps)
+
+    # Torch autograd jacobian
+    jacobian_torch = []
+    for i in range(output.shape[1]):
+        grad = (
+            torch.autograd.grad(
+                output[:, i].sum(), queries, create_graph=False, retain_graph=True
+            )[0]
+            .detach()
+            .cpu()
+            .numpy()
+        )
+        jacobian_torch.append(grad)
+    jacobian_torch = np.stack(jacobian_torch, axis=1)
+
+    fd_flat = fd_jacobian.flatten()
+    jac_flat = jacobian_torch.flatten()
+
+    np.testing.assert_allclose(
+        fd_jacobian,
+        jacobian_torch,
+        rtol=1e-5,
+        atol=1e-4,
+        err_msg="Finite difference Jacobian does not match autograd Jacobian.\n"
+        f"First 5 entries of FD Jacobian:\n{fd_flat[:5]}\n"
+        f"First 5 entries of Autograd Jacobian:\n{jac_flat[:5]}",
+    )
+
+
+def test_generate_bbox_spline():
+    # Define bounding box
+    bbox = np.array([[-1.0, -2.0, -3.0], [4.0, 5.0, 6.0]])
+
+    spline = generate_bbox_spline(bbox)
+
+    # Check that min/max corners match bbox
+    cp = spline.control_points
+    mins, maxs = cp.min(axis=0), cp.max(axis=0)
+    np.testing.assert_allclose(mins, bbox[0])
+    np.testing.assert_allclose(maxs, bbox[1])
+
+    params = np.array(np.meshgrid([0, 1], [0, 1], [0, 1])).T.reshape(-1, 3)
+
+    evals = spline.evaluate(params)
+
+    expected = np.array(
+        np.meshgrid(
+            [bbox[0, 0], bbox[1, 0]], [bbox[0, 1], bbox[1, 1]], [bbox[0, 2], bbox[1, 2]]
+        )
+    ).T.reshape(-1, 3)
+
+    # Order might differ, so compare sets
+    assert set(map(tuple, np.round(evals, 8))) == set(map(tuple, np.round(expected, 8)))
+
+
+def test_custom_torchspline_3D(bspline):
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    n_queries = 1000
+
+    queries = torch.rand((n_queries, 3), dtype=torch.float32, device=device)
+
+    splinepy_eval = bspline.evaluate(queries.detach().cpu().numpy())
+
+    splinepy_eval_as_tensor = torch.tensor(
+        splinepy_eval, dtype=torch.float32, device=device
+    )
+
+    output_ip = torch_spline_3D(
+        torch.tensor(bspline.knot_vectors, device=device, dtype=torch.float32),
+        control_points=torch.tensor(
+            bspline.control_points, device=device, dtype=torch.float32
+        ),
+        degrees=bspline.degrees,
+        queries=queries,
+    )
+
+    diff = torch.linalg.norm(splinepy_eval_as_tensor - output_ip, axis=1)
+
+    print("Max L2 difference:", diff.max())
+    torch.testing.assert_close(
+        diff.max(),
+        torch.tensor(0, device=diff.device, dtype=diff.dtype),
+        atol=1e-5,
+        rtol=1e-5,
+    )
+
+
+def test_custom_torchspline_1D():
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    n_queries = 1000
+
+    queries = torch.rand((n_queries, 1), dtype=torch.float32, device=device)
+    control_points = np.random.rand(4, 1)  # 4x4x4 Bezier -> 64 control points
+
+    knot_vecs = [[0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0]]
+    bspline = splinepy.BSpline(
+        degrees=[2], knot_vectors=knot_vecs, control_points=control_points
+    )
+    splinepy_eval = bspline.evaluate(queries.detach().cpu().numpy())
+
+    splinepy_eval_as_tensor = torch.tensor(
+        splinepy_eval, dtype=torch.float32, device=device
+    )
+
+    output_ip = torch_spline_1D(
+        torch.tensor(knot_vecs, device=device, dtype=torch.float32),
+        control_points=torch.tensor(control_points, device=device, dtype=torch.float32),
+        degrees=[2],
+        queries=queries,
+    )
+
+    diff = torch.linalg.norm(splinepy_eval_as_tensor - output_ip, axis=1)
+
+    print("Max L2 difference:", diff.max())
+    torch.testing.assert_close(
+        diff.max(),
+        torch.tensor(0, device=diff.device, dtype=diff.dtype),
+        atol=1e-5,
+        rtol=1e-5,
+    )
+
+
+if __name__ == "__main__":
+    control_points = np.random.rand(64, 3)  # 4x4x4 Bezier -> 64 control points
+
+    knot_vecs = [
+        [0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0],
+        [0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0],
+    ]
+
+    bspline_test = splinepy.BSpline(
+        degrees=[2, 2, 2], knot_vectors=knot_vecs, control_points=control_points
+    )
+    test_custom_torchspline_3D(bspline_test)
+    test_custom_torchspline_1D()
+    test_generate_bbox_spline()
+    np_rng = np.random.default_rng(0)
+    test_torchspline_evaluation(np_rng, bspline_test)
+    test_torchspline_derivative(np_rng, bspline_test)

tests/test_train_model.py

@@ -0,0 +1,110 @@
+from DeepSDFStruct.deep_sdf.training import (
+    train_deep_sdf,
+    reconstruct_meshs_from_latent,
+    create_interpolated_meshes_from_latent,
+)
+from DeepSDFStruct.pretrained_models import get_model
+from huggingface_hub import snapshot_download
+from huggingface_hub.utils import HfHubHTTPError
+import pytest
+import torch
+import time
+
+REVISION = "dbe58ebaa00057d5f15096c2b253c7efa91e19d3"
+
+
+def snapshot_download_with_retry(*args, max_retries=3, **kwargs):
+    """Download with retry on 429 rate limit errors."""
+    for attempt in range(max_retries):
+        try:
+            return snapshot_download(*args, **kwargs)
+        except HfHubHTTPError as e:
+            if e.response and e.response.status_code == 429:
+                if attempt < max_retries - 1:
+                    wait_time = 60 * (attempt + 1)
+                    print(
+                        f"Rate limited (429). Waiting {wait_time}s before retry {attempt + 2}/{max_retries}"
+                    )
+                    time.sleep(wait_time)
+                else:
+                    raise
+            raise
+
+
+@pytest.fixture(scope="module", autouse=True)
+def set_float32_dtype():
+    torch.set_default_dtype(torch.float32)
+    torch.set_default_device("cpu")
+    yield
+
+
+@pytest.fixture(scope="module")
+def data_dir():
+    return snapshot_download_with_retry(
+        "mkofler/lattice_structure_unit_cells",
+        repo_type="dataset",
+        revision=REVISION,
+        ignore_patterns=["*.stl", "**/*.stl"],
+    )
+
+
+def test_train_homogenization_model(data_dir):
+    exp_dir = "DeepSDFStruct/trained_models/test_experiment_homogenization"
+
+    device = "cpu"
+    train_deep_sdf(exp_dir, data_dir, device=device)
+
+
+def test_train_hierarchical_model(data_dir):
+    exp_dir = "DeepSDFStruct/trained_models/test_experiment_hierarchical"
+
+    device = "cpu"
+    train_deep_sdf(exp_dir, data_dir, device=device)
+
+
+def test_train_model(data_dir):
+    exp_dir = "DeepSDFStruct/trained_models/test_experiment"
+
+    device = "cpu"
+    train_deep_sdf(exp_dir, data_dir, device=device)
+
+
+def test_continue_from(data_dir):
+    exp_dir = "DeepSDFStruct/trained_models/test_experiment"
+
+    device = "cpu"
+    train_deep_sdf(exp_dir, data_dir, device=device, continue_from="1")
+
+
+def test_latent_recon():
+    exp_dir = "DeepSDFStruct/trained_models/analytic_round_cross"
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = "cpu"
+    reconstruct_meshs_from_latent(exp_dir, filetype="obj", device=device)
+    create_interpolated_meshes_from_latent(exp_dir, [1, 2, 3], 4, device=device)
+
+
+def test_cpp_file_export():
+    exp_dir = "DeepSDFStruct/trained_models/analytic_round_cross"
+    model = get_model(exp_dir)
+    model.export_libtorch_executable("tests/tmp_outputs/test_cpp_model.pt")
+
+
+if __name__ == "__main__":
+    import warnings
+
+    warnings.filterwarnings("error")
+    torch.set_default_dtype(torch.float32)
+    torch.set_default_device("cpu")
+    data_dir = snapshot_download_with_retry(
+        "mkofler/lattice_structure_unit_cells",
+        repo_type="dataset",
+        revision=REVISION,
+        ignore_patterns=["*.stl", "**/*.stl"],
+    )
+    test_train_homogenization_model(data_dir)
+    test_train_hierarchical_model(data_dir)
+    test_train_model(data_dir)
+    test_continue_from(data_dir)
+    test_latent_recon()
+    test_cpp_file_export()

tests/test_training_data_ids.py

@@ -0,0 +1,41 @@
+import json
+import os
+
+from DeepSDFStruct.deep_sdf.training import save_latent_code_data_map
+import DeepSDFStruct.deep_sdf.workspace as ws
+
+
+def test_save_latent_code_data_map(tmp_path):
+    experiment_directory = str(tmp_path / "experiment")
+    data_source = str(tmp_path / "data_source")
+    npz_filenames = [
+        "dataset_a/class_1/instance_1.npz",
+        "dataset_b/class_2/instance_2.npz",
+    ]
+
+    save_latent_code_data_map(experiment_directory, data_source, npz_filenames)
+
+    mapping_filename = ws.get_latent_code_data_map_filename(experiment_directory)
+    assert os.path.isfile(mapping_filename)
+
+    with open(mapping_filename, "r", encoding="utf-8") as f:
+        latent_code_data_map = json.load(f)
+
+    assert latent_code_data_map["data_source"] == data_source
+    assert latent_code_data_map["sdf_samples_subdir"] == ws.sdf_samples_subdir
+    assert latent_code_data_map["latent_codes"] == [
+        {
+            "latent_index": 0,
+            "relative_npz_filename": npz_filenames[0],
+            "npz_filename": os.path.join(
+                data_source, ws.sdf_samples_subdir, npz_filenames[0]
+            ),
+        },
+        {
+            "latent_index": 1,
+            "relative_npz_filename": npz_filenames[1],
+            "npz_filename": os.path.join(
+                data_source, ws.sdf_samples_subdir, npz_filenames[1]
+            ),
+        },
+    ]