more util functions

polyfem · Huangzizhou · Jun 30, 2024 · Jun 30, 2024 · Jun 30, 2024 · Jul 1, 2024
commit 670378e3cfc0eb7fa8aea1cecdaa9a386c88c3f9
diff --git a/src/mesh/mesh.cpp b/src/mesh/mesh.cpp
@@ -12,6 +12,8 @@ void define_mesh(py::module_ &m)
 {
   py::class_<Mesh>(m, "Mesh", "Mesh")
 
+      .def("dimension", &Mesh::dimension, "Dimension of the mesh (2 or 3)")
+
       .def("n_elements", &Mesh::n_elements, "Number of elements")
 
       .def("n_boundary_elements", &Mesh::n_boundary_elements,
@@ -92,6 +94,18 @@ void define_mesh(py::module_ &m)
           },
           "Get all vertex positions")
 
+      .def(
+          "set_vertices",
+          [](Mesh &mesh, const Eigen::VectorXi &ids,
+             const Eigen::MatrixXd &points) {
+            for (int i = 0; i < ids.size(); i++)
+            {
+              assert(ids(i) < mesh.n_vertices());
+              mesh.set_point(ids(i), points.row(i));
+            }
+          },
+          "Set a subset of vertex positions given a list of vertex indices")
+
       .def(
           "set_vertices",
           [](Mesh &mesh, const Eigen::MatrixXd &points) {
@@ -103,12 +117,13 @@ void define_mesh(py::module_ &m)
       .def(
           "elements",
           [](const Mesh &mesh) {
-            Eigen::MatrixXi elements(mesh.n_elements(), mesh.n_cell_vertices(0));
+            Eigen::MatrixXi elements(mesh.n_elements(),
+                                     mesh.n_cell_vertices(0));
             for (int e = 0; e < mesh.n_elements(); e++)
             {
-               assert(mesh.n_cell_vertices(e) == elements.cols());
-               for (int i = 0; i < elements.cols(); i++)
-                    elements(e, i) = mesh.element_vertex(e, i);
+              assert(mesh.n_cell_vertices(e) == elements.cols());
+              for (int i = 0; i < elements.cols(); i++)
+                elements(e, i) = mesh.element_vertex(e, i);
             }
             return elements;
           },

diff --git a/src/state/state.cpp b/src/state/state.cpp
@@ -156,6 +156,9 @@ void define_solver(py::module_ &m)
            "load PDE and problem parameters from the settings", py::arg("json"),
            py::arg("strict_validation") = false)
 
+      .def("set_max_threads", &State::set_max_threads, 
+      "set maximum number of threads", py::arg("nthreads"))
+
       .def("ndof", &State::ndof, "Dimension of the solution")
 
       .def("n_bases", [](const State &s) { return s.n_bases; }, "Number of basis")
@@ -398,6 +401,28 @@ void define_solver(py::module_ &m)
           "set_cache_level",
           [](State &s, solver::CacheLevel level) {
             s.optimization_enabled = level;
+            if (level == solver::CacheLevel::Derivatives)
+            {
+              if (s.is_contact_enabled())
+              {
+                if (!s.args["contact"]["use_convergent_formulation"])
+                {
+                  s.args["contact"]["use_convergent_formulation"] = true;
+                  logger().info("Use convergent formulation for differentiable contact...");
+                }
+                if (s.args["/solver/contact/barrier_stiffness"_json_pointer].is_string())
+                {
+                  logger().error("Only constant barrier stiffness is supported in differentiable contact!");
+                }
+              }
+
+              if (s.args.contains("boundary_conditions") && s.args["boundary_conditions"].contains("rhs"))
+              {
+                json rhs = s.args["boundary_conditions"]["rhs"];
+                if ((rhs.is_array() && rhs.size() > 0 && rhs[0].is_string()) || rhs.is_string())
+                  logger().error("Only constant rhs over space is supported in differentiable code!");
+              }
+            }
           },
           "Set solution caching level", py::arg("cache_level"))