test contact

test_branch/test_contact.py

@@ -3,9 +3,23 @@
 from dolfinx.io import XDMFFile, gmshio
 import gmsh
 import numpy as np
-from dolfinx import mesh, fem
+from dolfinx import fem
 from mpi4py import MPI
 import dolfinx.io
+import basix
+from irrevolutions.solvers import SNESSolver
+from dolfinx.fem import (
+    Function,
+    locate_dofs_topological,
+    dirichletbc,
+    assemble_scalar,
+    form,
+)
+from irrevolutions.models import default_model_parameters
+import os
+import yaml
+import ufl
+from petsc4py import PETSc
 
 
 def create_arc_ring_mesh(inner_radius, outer_radius, angle=180, lc=0.1):
@@ -85,12 +99,140 @@ def create_arc_ring_mesh(inner_radius, outer_radius, angle=180, lc=0.1):
         "R_inner": 0.3,  # Inner hole radius (set to 0.0 for no hole)
         "lc": 0.05,  # Mesh element size
         "a": 0.1,  # Half-width of the refined region (-a < x < a)
+        "geometric_dimension": 2,
     }
-    # create_disk_with_hole(comm, geom_params)
+    model, _ = create_arc_ring_mesh(0.5, 1.0, angle=180, lc=0.05)
+    mesh, cell_tags, facet_tags = gmshio.model_to_mesh(
+        model, comm, 0, gdim=geom_params["geometric_dimension"]
+    )
     # gmsh_model, tdim = create_disk_with_hole(comm, geom_params)
     # mesh, cell_tags, facet_tags = gmshio.model_to_mesh(gmsh_model, comm, 0, gdim=2)
-    # __import__("pdb").set_trace()
+    dx = ufl.Measure("dx", domain=mesh)
+    with XDMFFile(MPI.COMM_WORLD, f"output/contact.xdmf", "w") as xdmf:
+        xdmf.write_mesh(mesh)
 
-    model, _ = create_arc_ring_mesh(0.5, 1.0, angle=180, lc=0.05)
-    mesh, cell_tags, facet_tags = gmshio.model_to_mesh(model, comm, 0, gdim=2)
-    __import__("pdb").set_trace()
+    element_u = basix.ufl.element(
+        "Lagrange",
+        mesh.basix_cell(),
+        degree=1,
+        shape=((geom_params["geometric_dimension"]),),
+    )
+
+    V_u = dolfinx.fem.functionspace(mesh, element_u)
+
+    u = dolfinx.fem.Function(V_u, name="Displacement")
+    u_lb = dolfinx.fem.Function(V_u, name="Displacement_lb")
+    u_ub = dolfinx.fem.Function(V_u, name="Displacement_ub")
+    u_lb.interpolate(
+        # lambda x: np.full((x.shape[1], tdim), -np.inf)
+        lambda x: np.stack([np.full_like(x[0], -np.inf), np.full_like(x[1], -x[1])])
+    )
+    g_t = 1.1
+
+    u_ub.interpolate(
+        lambda x: np.stack([np.full_like(x[0], np.inf), np.full_like(x[1], g_t - x[1])])
+    )
+    body_f = dolfinx.fem.Constant(mesh, np.array([0.0, 1.0], dtype=PETSc.ScalarType))
+
+    def bottom_boundary(x):
+        return np.isclose(x[1], 0, atol=1e-2)
+
+    bottom_dofs = locate_dofs_topological(
+        V_u,
+        mesh.topology.dim - 1,
+        dolfinx.mesh.locate_entities_boundary(mesh, 1, bottom_boundary),
+    )
+
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    yaml_file = os.path.join(dir_path, "default_parameters.yml")
+    with open(yaml_file, "r") as f:
+        parameters = yaml.load(f, Loader=yaml.FullLoader)
+    # solver_parameters = parameters["solver"]
+    solver_parameters = {
+        "contact": {
+            "type": "SNES",
+            "prefix": "contact",
+            "snes": {
+                # "snes_type": "vinewtonrsls",
+                "snes_type": "vinewtonssls",
+                "snes_linesearch_type": "basic",
+                "ksp_type": "preonly",
+                "pc_type": "lu",
+                "pc_factor_mat_solver_type": "mumps",
+                "snes_atol": 1.0e-8,
+                "snes_rtol": 1.0e-8,
+                "snes_max_it": 50,
+                "snes_vi_monitor": True,
+                "snes_monitor": "",
+            },
+        }
+    }
+    E, nu = parameters["model"]["E"], parameters["model"]["nu"]
+
+    def epsilon(v):
+        return ufl.sym(ufl.grad(v))
+
+    def sigma(v):
+        return (
+            E
+            / (1 + nu)
+            * (epsilon(v) + nu / (1 - nu) * ufl.tr(epsilon(v)) * ufl.Identity(2))
+        )
+
+    energy = 0.5 * ufl.inner(sigma(u), epsilon(u)) * dx - ufl.dot(body_f, u) * dx
+    # F = ufl.derivative(energy, u)
+    bottom_disp = dolfinx.fem.Constant(
+        mesh, np.array([0.0, 0.0], dtype=PETSc.ScalarType)
+    )
+    bcs_u = [
+        # dirichletbc(top_disp, top_dofs, V_u),
+        dirichletbc(bottom_disp, bottom_dofs, V_u),
+    ]
+    bcs = {"bcs_u": bcs_u}
+    energy_u = ufl.derivative(energy, u, ufl.TestFunction(V_u))
+    contact = SNESSolver(
+        energy_u,
+        u,
+        bcs.get("bcs_u"),
+        petsc_options=solver_parameters.get("contact").get("snes"),
+        bounds=(u_lb, u_ub),
+        prefix=solver_parameters.get("contact").get("prefix"),
+    )
+    loads = np.linspace(-0.1, 1, 30)
+
+    for it, t in enumerate(loads):
+        # u_ub.interpolate(
+        #     lambda x: np.stack(
+        #         [np.full_like(x[0], np.inf), np.full_like(x[1], gap - t)]
+        #     )
+        # )
+        body_f.value = np.array([0.0, -t], dtype=PETSc.ScalarType)
+        contact.solve()
+        # inactive_set = contact.solver.getVIInactiveSet()
+        # inactive_set.setType("general")
+        # __import__("pdb").set_trace()
+        # inactive_set.setOptionsPrefix("inactive_set_")
+        # inactive_set.setFromOptions()
+        # inactive_set.getComm()
+        # if inactive_set.getSize() > 0:
+        #     inactive_set.view()
+        # else:
+        #     print("No inactive constraints detected!")
+
+        # print(f"Number of inactive DOFs: {inactive_set.size}")
+        # inactive_set.setComm(PETSc.COMM_WORLD)
+        # inactive_set.view()
+
+        rnorm = contact.solver.getFunctionNorm()
+        print(f"Final residual norm: {rnorm}")
+        diff = np.linalg.norm(u.x.array - u_ub.x.array)
+        print(f"Max violation of contact constraint: {diff}")
+        print("Lower bound values:", contact.lb.x.array)
+        print("Upper bound values:", contact.ub.x.array)
+        print("SNES Solver Communicator:", contact.solver.getComm())
+        print("MPI Rank:", MPI.COMM_WORLD.rank)
+        with XDMFFile(MPI.COMM_WORLD, f"output/contact.xdmf", "a") as xdmf:
+            xdmf.write_function(u, t)
+            # xdmf.write_function(u_lb, t)
+            # xdmf.write_function(u_ub, t)
+        print(f"Load step {it} done.")