vtk output for solutions about pymfem HOT 14 OPEN

import mfem.ser as mfem
import gmsh
import sys
import os
import io


create_gmsh = True
create_mfem = True
gmsh_mesh = "cube.msh"
lc = 0.1
if create_gmsh:
	gmsh.initialize()
	gmsh.option.setNumber("Mesh.MshFileVersion", 2.2)	
	gmsh.model.add("gmsh to mfem VTK")
	iterX = [0,1,1,0]
	iterY = [0,0,1,1]
	cA = [1, 2, 3, 4, 1]
	pID = 1 ; eID = 1 ; cID = 1; sID = 1; 
	a6 = range(6)
	b_names = {1: "bottom", 2: "top", 3: "left", 4: "front", 5: "right", 6: "back"}
	for zi in range(2):
		for k in range(4):
			print(" ADD POINT ",iterX[k], iterY[k], zi, lc, pID)
			gmsh.model.geo.addPoint(iterX[k],iterY[k],zi, lc, pID)
			pID +=1
	gmsh.model.geo.synchronize()
	for zi in range(2):
		offL = zi * 4
		for k in range(len(cA)-1):
			print(" ADD LINE ", cA[k] + offL, cA[k+1]+ offL, eID)
			gmsh.model.geo.addLine(cA[k] + offL, cA[k+1]+ offL, eID)
			eID += 1
	for k in range(len(cA)-1):
		print(" ADD LINE ", cA[k], cA[k] + 4, eID)
		gmsh.model.geo.addLine(cA[k], cA[k] + 4, eID)
		eID += 1

	## THE CURVES
	print(" CURVE LOOP 1 ",[1,2,3,4],  cID)  # bottom
	gmsh.model.geo.addCurveLoop([1,2,3,4],  cID)  # bottom
	print(" CURVE LOOP 2", [5,6,7,8],  cID+1)  # top
	gmsh.model.geo.addCurveLoop([5,6,7,8],  cID+1)  # top
	print(" CURVE LOOP 3", [9,-8,-12,4],  cID+2)  # left
	gmsh.model.geo.addCurveLoop([9,-8,-12,4],  cID+2)  # left
	print(" CURVE LOOP 4", [1,10,-5, -9], cID+3)  # front
	gmsh.model.geo.addCurveLoop([1,10,-5, -9], cID+3)  # front
	print(" CURVE LOOP 5", [2,11,-6, -10], cID+4)  # right
	gmsh.model.geo.addCurveLoop([2,11,-6, -10], cID+4)  # right
	print(" CURVE LOOP 6", [-3, 11,7,-12], cID+5)  # back
	gmsh.model.geo.addCurveLoop([-3, 11,7,-12], cID+5)  # back
	cID += 6

	for j in range(6):
		print(" ADD SURFACE ", [sID + j], sID + j)
		sj = gmsh.model.geo.addPlaneSurface([sID + j], sID + j)

	print(" SURFACE LOOP ", [a + sID for a in a6])
	loop = gmsh.model.geo.addSurfaceLoop([a + sID for a in a6])
	vol = gmsh.model.geo.addVolume([loop], tag=1)
	gmsh.model.geo.synchronize()

	for j in range(6):
		print(" Adding boundary ",  [sID + j], "box_"+b_names[j+1] + "_" + str(j+1))
		gmsh.model.addPhysicalGroup(2, [sID + j], name=b_names[j+1])
	sID += 6
	gmsh.model.addPhysicalGroup(3, [vol], name="volume1", tag = vol)
	# We can then generate a 2D mesh...
	gmsh.model.mesh.generate(3)
	gmsh.model.geo.synchronize()
	gmsh.write(gmsh_mesh)
	gmsh.write(gmsh_mesh.replace(".msh", ".vtk"))
	gmsh.finalize()

if create_mfem:
	
	mesh = mfem.Mesh(gmsh_mesh)
	device = mfem.Device("cpu")
	device.Print()
	dim = mesh.Dimension()
	order = 1
	fec = mfem.H1_FECollection(order, dim)
	fespace = mfem.FiniteElementSpace(mesh, fec, dim)

	# 6. Determine the list of true (i.e. conforming) essential boundary dofs.
	#    In this example, the boundary conditions are defined by marking only
	#    boundary attribute 1 from the mesh as essential and converting it to a
	#    list of true dofs.
	ess_tdof_list = mfem.intArray()
	ess_bdr = mfem.intArray([0] * (mesh.bdr_attributes.Max()))

	# impose Dirichlet BCs in the bottom box to fix objects in space
	ess_bdr[0] = 1

	fespace.GetEssentialTrueDofs(ess_bdr, ess_tdof_list)

	# 7. Set up the linear form b(.) which corresponds to the right-hand side of
	#    the FEM linear system. In this case, b_i equals the boundary integral
	#    of f*phi_i where f represents a "pull down" force on the Neumann part
	#    of the boundary and phi_i are the basis functions in the finite element
	#    fespace. The force is defined by the VectorArrayCoefficient object f,
	#    which is a vector of Coefficient objects. The fact that f is non-zero
	#    on boundary attribute 2 is indicated by the use of piece-wise constants
	#    coefficient for its last component.
	f = mfem.VectorArrayCoefficient(dim)
	for i in range(dim - 1):
		f.Set(i, mfem.ConstantCoefficient(0.0))

	# this is a boundary vector. we are applying a vertical load (?)
	pull_force = mfem.Vector([0] * mesh.bdr_attributes.Max())
	pull_force[1] = -2e5

	f.Set(dim - 1, mfem.PWConstCoefficient(pull_force))
	b = mfem.LinearForm(fespace)
	b.AddBoundaryIntegrator(mfem.VectorBoundaryLFIntegrator(f))
	b.Assemble()
	# 8. Define the solution vector x as a finite element grid function
	#    corresponding to fespace. Initialize x with initial guess of zero,
	#    which satisfies the boundary conditions.
	x = mfem.GridFunction(fespace)
	x.Assign(0.0)

	# 9. Set up the material properties: top box will be rigid, bottom will be very elastic

	lamb = mfem.Vector(mesh.attributes.Max())
	mu = mfem.Vector(mesh.attributes.Max())
	# CREATE AN ELASTIC MATERIAL
	E_0 = 5e5
	nu_0 = 0.4
	lamb[0] = E_0 * nu_0 / ((1 + nu_0) * (1 - 2 * nu_0))
	mu[0] = E_0 / (2 * (1 + nu_0))
	lambda_func = mfem.PWConstCoefficient(lamb)
	mu_func = mfem.PWConstCoefficient(mu)
	a = mfem.BilinearForm(fespace)
	a.AddDomainIntegrator(mfem.ElasticityIntegrator(lambda_func, mu_func))

	# 10. Assemble the bilinear form and the corresponding linear system,
	#     applying any necessary transformations such as: eliminating boundary
	#     conditions, applying conforming constraints for non-conforming AMR,
	#     static condensation, etc.
	print('matrix...')
	a.Assemble()

	A = mfem.OperatorPtr()
	B = mfem.Vector()
	X = mfem.Vector()
	a.FormLinearSystem(ess_tdof_list, x, b, A, X, B)
	print('...done')  # Here, original version calls heigth, which is not defined in the header...!?

	# 10. Solve
	AA = mfem.OperatorHandle2SparseMatrix(A)
	M = mfem.GSSmoother(AA)
	mfem.PCG(AA, M, B, X, 1, 5000, 1e-8, 1e-8)

	# 11. Recover the solution as a finite element grid function.
	a.RecoverFEMSolution(X, b, x)

	# 13. For non-NURBS meshes, make the mesh curved based on the finite element
	#     space. This means that we define the mesh elements through a fespace
	#     based transformation of the reference element. This allows us to save
	#     the displaced mesh as a curved mesh when using high-order finite
	#     element displacement field. We assume that the initial mesh (read from
	#     the file) is not higher order curved mesh compared to the chosen FE
	#     space.
	if not mesh.NURBSext:
		mesh.SetNodalFESpace(fespace)

	# 14. Save the displaced mesh and the inverted solution (which gives the
	#     backward displacements to the original grid). This output can be
	#     viewed later using GLVis: "glvis -m mesh.mesh -g data.gf".

	nodes = mesh.GetNodes()
	nodes += x
	x.Neg()

	output = io.StringIO()
	output.precision = 8
	ref = 1
	mesh.PrintVTK(output)

	x.SaveVTK(output, 'U', ref)
	fid = open("converted_mesh.vtk", "w")
	fid.write(output.getvalue())
	fid.close()

def grid_function_save_vtk(gf: mfem.GridFunction, os: io, field_name: str):

	fes = gf.FESpace()
	mesh = fes.GetMesh()
	vec_dim = gf.VectorDim()
	assert gf.Size()/vec_dim == mesh.GetNV()
	os.write("POINT_DATA "+str(fes.GetNDofs())+"\n")
	if vec_dim == 1:
		print('Scalar field', field_name)
		os.write("SCALARS "+field_name+ " double 1\n LOOKUP_TABLE default\n")
		for i in range(fes.GetNDofs()):
			os.write(gf[i]+"\n")
	elif (vec_dim == 2 or vec_dim == 3) and mesh.SpaceDimension() > 1:
		print('Vector field', field_name)
		os.write("VECTORS "+field_name+" double\n")
		vdofs = mfem.intArray(vec_dim)
		for i in range(fes.GetNDofs()):
			vdofs.SetSize(1)
			vdofs[0] = i
			fes.DofsToVDofs(vdofs)
			os.write(str(gf[vdofs[0]]) +' '+str(gf[vdofs[1]]) +' ')
			if vec_dim == 2:
				os.write("0.0\n")
			else:
				os.write(str(gf[vdofs[2]])+"\n")
	else:
		pass

#       // other data: save the components as separate scalars
#       for (int vd = 0; vd < vec_dim; vd++)
#       {
#          out << "SCALARS " << field_name << vd << " double 1\n"
#              << "LOOKUP_TABLE default\n";
#          Array<int> vdofs(vec_dim);
#          for (int i=0; i < fes->GetNDofs(); ++i)
#          {
#             vdofs.SetSize(1);
#             vdofs[0] = i;
#             fes->DofsToVDofs(vdofs);
#             out << gf[vdofs[vd]] << '\n';
#          }
#       }
#    }

# }

	output = io.StringIO()
	output.precision = 8
	ref = 1
	mesh.PrintVTK(output)
	grid_function_save_vtk(x, output, 'U')
	fid = open(os.path.join("vtkdata.vtk"), "w")
	fid.write(output.getvalue())