Merge pull request #2 from arvigj/main

python binding fixes

Author: Huangzizhou

src/differentiable/adjoint.cpp

@@ -30,9 +30,11 @@ void define_adjoint(py::module_ &m)
       [](State &state) {
         Eigen::VectorXd term;
         if (state.problem->is_time_dependent())
-          AdjointTools::dJ_material_transient_adjoint_term(state, state.get_adjoint_mat(1), state.get_adjoint_mat(0), term);
+          AdjointTools::dJ_material_transient_adjoint_term(
+              state, state.get_adjoint_mat(1), state.get_adjoint_mat(0), term);
         else
-          AdjointTools::dJ_material_static_adjoint_term(state, state.diff_cached.u(0), state.get_adjoint_mat(0), term);
+          AdjointTools::dJ_material_static_adjoint_term(
+              state, state.diff_cached.u(0), state.get_adjoint_mat(0), term);
 
         return utils::unflatten(term, state.bases.size());
       },
@@ -43,7 +45,8 @@ void define_adjoint(py::module_ &m)
       [](State &state) {
         Eigen::VectorXd term;
         if (state.problem->is_time_dependent())
-          AdjointTools::dJ_friction_transient_adjoint_term(state, state.get_adjoint_mat(1), state.get_adjoint_mat(0), term);
+          AdjointTools::dJ_friction_transient_adjoint_term(
+              state, state.get_adjoint_mat(1), state.get_adjoint_mat(0), term);
         else
           log_and_throw_adjoint_error(
               "Friction coefficient derivative is only supported for transient problems!");
@@ -88,7 +91,7 @@ void define_adjoint(py::module_ &m)
               vec += term.segment(state.ndof() + g.index * dim, dim);
             }
         }
-        
+
         return map;
       },
       py::arg("solver"));
@@ -129,8 +132,24 @@ void define_adjoint(py::module_ &m)
               vec += term.segment(g.index * dim, dim);
             }
         }
-        
+
         return map;
       },
       py::arg("solver"));
+
+  m.def(
+      "dirichlet_derivative",
+      [](State &state) {
+        const int dim = state.mesh->dimension();
+
+        Eigen::VectorXd term;
+        if (state.problem->is_time_dependent())
+          log_and_throw_adjoint_error(
+              "Dirichlet derivative is only supported for static problems!");
+
+        AdjointTools::dJ_dirichlet_static_adjoint_term(
+            state, state.get_adjoint_mat(0), term);
+        return utils::unflatten(term, state.mesh->dimension());
+      },
+      py::arg("solver"));
 }

src/differentiable/objective.cpp

@@ -15,26 +15,38 @@ using namespace polyfem::solver;
 
 void define_objective(py::module_ &m)
 {
-    py::class_<AdjointForm, std::shared_ptr<AdjointForm>>(m, "Objective")
-        .def("name", &AdjointForm::name)
+  py::class_<AdjointForm, std::shared_ptr<AdjointForm>>(m, "Objective")
+      .def("name", &AdjointForm::name)
 
-        .def("value", &AdjointForm::value, py::arg("x"))
+      .def("value", &AdjointForm::value, py::arg("x"))
 
-        .def("solution_changed", &AdjointForm::solution_changed, py::arg("x"))
+      .def("solution_changed", &AdjointForm::solution_changed, py::arg("x"))
 
-        .def("derivative", [](AdjointForm &obj, State &solver, const Eigen::VectorXd &x, const std::string &wrt) -> Eigen::VectorXd {
+      .def("is_step_collision_free", &AdjointForm::is_step_collision_free,
+           py::arg("x0"), py::arg("x1"))
+
+      .def("max_step_size", &AdjointForm::max_step_size, py::arg("x0"),
+           py::arg("x1"))
+
+      .def(
+          "derivative",
+          [](AdjointForm &obj, State &solver, const Eigen::VectorXd &x,
+             const std::string &wrt) -> Eigen::VectorXd {
             if (wrt == "solution")
-                return obj.compute_adjoint_rhs(x, solver);
+              return obj.compute_adjoint_rhs(x, solver);
             else if (wrt == obj.get_variable_to_simulations()[0]->name())
             {
-                Eigen::VectorXd grad;
-                obj.compute_partial_gradient(x, grad);
-                return grad;
+              Eigen::VectorXd grad;
+              obj.compute_partial_gradient(x, grad);
+              return grad;
             }
             else
-                throw std::runtime_error("Input type does not match objective derivative type!");
-        }, py::arg("solver"), py::arg("x"), py::arg("wrt"));
+              throw std::runtime_error(
+                  "Input type does not match objective derivative type!");
+          },
+          py::arg("solver"), py::arg("x"), py::arg("wrt"));
 
-    m.def("create_objective", &AdjointOptUtils::create_simple_form,
-        py::arg("obj_type"), py::arg("param_type"), py::arg("solver"), py::arg("parameters"));
+  m.def("create_objective", &AdjointOptUtils::create_simple_form,
+        py::arg("obj_type"), py::arg("param_type"), py::arg("solver"),
+        py::arg("parameters"));
 }

src/differentiable/utils.cpp

@@ -1,4 +1,5 @@
 #include <polyfem/mesh/SlimSmooth.hpp>
+#include <polyfem/solver/AdjointTools.hpp>
 #include <polyfem/utils/MatrixUtils.hpp>
 #include <polyfem/State.hpp>
 #include "binding.hpp"
@@ -8,18 +9,27 @@
 namespace py = pybind11;
 using namespace polyfem;
 using namespace polyfem::mesh;
+using namespace polyfem::solver;
 
 void define_opt_utils(py::module_ &m)
 {
   m.def(
-      "apply_slim",
-      [](const Eigen::MatrixXd &V, const Eigen::MatrixXi &F,
-         const Eigen::MatrixXd &Vnew) {
-        Eigen::MatrixXd Vsmooth;
-        bool succeed = apply_slim(V, F, Vnew, Vsmooth, 1000);
-        if (!succeed)
-          throw std::runtime_error("SLIM failed to converge!");
-        return Vsmooth;
-      },
-      py::arg("Vold"), py::arg("faces"), py::arg("Vnew"));
+       "apply_slim",
+       [](const Eigen::MatrixXd &V, const Eigen::MatrixXi &F,
+          const Eigen::MatrixXd &Vnew) {
+         Eigen::MatrixXd Vsmooth;
+         bool succeed = apply_slim(V, F, Vnew, Vsmooth, 1000);
+         if (!succeed)
+           throw std::runtime_error("SLIM failed to converge!");
+         return Vsmooth;
+       },
+       py::arg("Vold"), py::arg("faces"), py::arg("Vnew"))
+
+      .def("map_primitive_to_node_order",
+           &AdjointTools::map_primitive_to_node_order, py::arg("state"),
+           py::arg("primitives"))
+
+      .def("map_node_to_primitive_order",
+           &AdjointTools::map_node_to_primitive_order, py::arg("state"),
+           py::arg("nodes"));
 }

src/mesh/mesh.cpp

@@ -3,6 +3,7 @@
 #include <polyfem/mesh/mesh2D/CMesh2D.hpp>
 #include "binding.hpp"
 #include <pybind11/eigen.h>
+#include <pybind11/stl.h>
 
 namespace py = pybind11;
 using namespace polyfem;
@@ -50,31 +51,37 @@ void define_mesh(py::module_ &m)
       .def("get_boundary_id", &Mesh::get_boundary_id,
            "Get boundary ID of one boundary primitive", py::arg("primitive"))
 
+      .def("get_body_ids", &Mesh::get_body_ids, "Get body IDs")
+
       // .def(
       //     "set_boundary_side_set_from_bary",
       //     [](Mesh &mesh,
-      //        const std::function<int(const RowVectorNd &)> &boundary_marker) {
+      //        const std::function<int(const RowVectorNd &)> &boundary_marker)
+      //        {
       //       mesh.compute_boundary_ids(boundary_marker);
       //     },
-      //     "Sets the side set for the boundary conditions, the functions takes the barycenter of the boundary (edge or face)",
+      //     "Sets the side set for the boundary conditions, the functions takes
+      //     the barycenter of the boundary (edge or face)",
       //     py::arg("boundary_marker"))
       // .def(
       //     "set_boundary_side_set_from_bary_and_boundary",
       //     [](Mesh &mesh, const std::function<int(const RowVectorNd &, bool)>
       //                        &boundary_marker) {
       //       mesh.compute_boundary_ids(boundary_marker);
       //     },
-      //     "Sets the side set for the boundary conditions, the functions takes the barycenter of the boundary (edge or face) and a flag that says if the element is boundary",
-      //     py::arg("boundary_marker"))
+      //     "Sets the side set for the boundary conditions, the functions takes
+      //     the barycenter of the boundary (edge or face) and a flag that says
+      //     if the element is boundary", py::arg("boundary_marker"))
       // .def(
       //     "set_boundary_side_set_from_v_ids",
       //     [](Mesh &mesh,
       //        const std::function<int(const std::vector<int> &, bool)>
       //            &boundary_marker) {
       //       mesh.compute_boundary_ids(boundary_marker);
       //     },
-      //     "Sets the side set for the boundary conditions, the functions takes the sorted list of vertex id and a flag that says if the element is boundary",
-      //     py::arg("boundary_marker"))
+      //     "Sets the side set for the boundary conditions, the functions takes
+      //     the sorted list of vertex id and a flag that says if the element is
+      //     boundary", py::arg("boundary_marker"))
 
       .def("set_body_ids", &Mesh::set_body_ids, "Set body IDs with an array",
            py::arg("ids"))