corrected-e-number-clean.py

import numpy as np
from mpi4py import MPI
from pathlib import Path
import dolfinx
import dolfinx.fem.petsc
import dolfinx.graph
from dolfinx import mesh, fem, default_real_type
from dolfinx.fem.petsc import LinearProblem
from dolfinx.io import XDMFFile
from dolfinx.geometry import (bb_tree, compute_collisions_points, compute_colliding_cells)
from dolfinx.mesh import (CellType, GhostMode, compute_midpoints, create_box,
                          create_cell_partitioner, create_mesh, refine)
import ufl
from ufl import TrialFunction, TestFunction, dx, grad, inner, dot
from basix.ufl import element
from functools import partial
from petsc4py import PETSc
from slepc4py import SLEPc

def create_periodic_mesh(mesh, indicator, indicator_dual, mapping_function, mapping_function_inv):

    geometry = mesh.geometry._cpp_object
    topology = mesh.topology

    mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim-1)
    mesh.topology.create_connectivity(mesh.topology.dim-1, mesh.topology.dim)

    # Map left side to right side
    left_vertices = dolfinx.mesh.locate_entities_boundary(mesh, 0, indicator)
    left_midpoints = dolfinx.mesh.compute_midpoints(mesh, 0, left_vertices)

    right_midpoints = mapping_function(left_midpoints.T).T
    right_vertices = dolfinx.mesh.locate_entities_boundary(mesh, 0, indicator_dual)

    # Find closest vertex on right side
    bb_tree = dolfinx.geometry.bb_tree(mesh,0, right_vertices)
    mid_tree = dolfinx.geometry.create_midpoint_tree(mesh, 0, right_vertices)
    closest_vertex = dolfinx.geometry.compute_closest_entity(bb_tree, mid_tree, mesh, right_midpoints)

    # Keep only left side vertices
    num_vertices_local = mesh.topology.index_map(0).size_local + mesh.topology.index_map(0).num_ghosts
    keep_vertices = np.ones(num_vertices_local, dtype=np.bool_)
    keep_vertices[right_vertices] = False
    # Create submap
    new_vertices = np.flatnonzero(keep_vertices)
    new_vertex_map, sub_to_parent = dolfinx.cpp.common.create_sub_index_map(mesh.topology.index_map(0), new_vertices, allow_owner_change=True)

    # Invert map
    num_vertices_local = mesh.topology.index_map(0).size_local + mesh.topology.index_map(0).num_ghosts
    parent_to_sub = np.full(num_vertices_local, -1, dtype=np.int32)
    parent_to_sub[sub_to_parent] = np.arange(sub_to_parent.size, dtype=np.int32)

    # Create map from right to left vertices for replacement
    replace_map = np.arange(num_vertices_local, dtype=np.int32)
    replace_map[closest_vertex] = mapping_function_inv(left_midpoints.T, left_vertices)

    # First map vertices from right to left vertices
    c_to_v = mesh.topology.connectivity(mesh.topology.dim, 0)
    new_c = replace_map[c_to_v.array]
    # Then map vertices to new indices
    new_c = parent_to_sub[new_c]
    new_o = c_to_v.offsets.copy()
    new_c_to_v = dolfinx.graph.adjacencylist(new_c, new_o)
    new_v_to_v = dolfinx.graph.adjacencylist(np.arange(len(sub_to_parent), dtype=np.int32))

    topology = dolfinx.cpp.mesh.Topology(MPI.COMM_WORLD, mesh.topology.cell_type)
    topology.set_index_map(0, new_vertex_map)
    topology.set_index_map(mesh.topology.dim, mesh.topology.index_map(mesh.topology.dim))
    topology.set_connectivity(new_v_to_v, 0,0)
    topology.set_connectivity(new_c_to_v, mesh.topology.dim, 0)
    c_el = dolfinx.fem.coordinate_element(mesh._ufl_domain.ufl_coordinate_element().basix_element)
    geometry = dolfinx.mesh.create_geometry(mesh.geometry.index_map(), mesh.geometry.dofmap, c_el._cpp_object, mesh.geometry.x[:, :mesh.geometry.dim].copy(),  mesh.geometry.input_global_indices)
    if mesh.geometry.x.dtype == np.float64:
        cpp_mesh = dolfinx.cpp.mesh.Mesh_float64(mesh.comm, topology, geometry._cpp_object)
    elif mesh.geometry.x.dtype == np.float32:
        cpp_mesh = dolfinx.cpp.mesh.Mesh_float32(mesh.comm, topology, geometry._cpp_object)
    else:
        raise RuntimeError(f"Unsupported dtype for mesh {mesh.geometry.x.dtype}")

    new_mesh = dolfinx.mesh.Mesh(cpp_mesh, domain = ufl.Mesh(mesh._ufl_domain.ufl_coordinate_element()))

    return new_mesh
    

def indicator(x):
    tol = 1e-10

    is_corner = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol)

    is_z_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_y_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)
    is_x_axis = np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[0, :], 0.0, atol=tol) & ~np.isclose(x[0, :], Lx, atol=tol)

    is_yz_face = np.isclose(x[0, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[1, :], Ly, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_zx_face = np.isclose(x[1, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_xy_face = np.isclose(x[2, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)

    return is_corner | is_x_axis | is_y_axis | is_z_axis | is_xy_face | is_yz_face | is_zx_face
    
def indicator_dual(x):
    tol = 1e-10

    is_corner = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol)

    is_z_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_y_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)
    is_x_axis = np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[0, :], 0.0, atol=tol) & ~np.isclose(x[0, :], Lx, atol=tol)

    is_yz_face = np.isclose(x[0, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[1, :], Ly, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_zx_face = np.isclose(x[1, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_xy_face = np.isclose(x[2, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)

    x_condition = np.isclose(x[0, :], 0.0, atol=tol) | np.isclose(x[0, :], Lx, atol=tol)
    
    y_condition = np.isclose(x[1, :], 0.0, atol=tol) | np.isclose(x[1, :], Ly, atol=tol)
    
    z_condition = np.isclose(x[2, :], 0.0, atol=tol) | np.isclose(x[2, :], Lz, atol=tol)

    combined_condition = x_condition | y_condition | z_condition

    return ~(is_corner | is_x_axis | is_y_axis | is_z_axis | is_xy_face | is_yz_face | is_zx_face) & combined_condition
    
    
def mapping_function(x):

    tol = 1e-10

    is_corner = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol)

    is_z_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_y_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)
    is_x_axis = np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[0, :], 0.0, atol=tol) & ~np.isclose(x[0, :], Lx, atol=tol)

    is_yz_face = np.isclose(x[0, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[1, :], Ly, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_zx_face = np.isclose(x[1, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_xy_face = np.isclose(x[2, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)
    
    ncol = x.shape[1] + 6 + (nx + ny + nz - 3) * 2
    new_array = np.zeros((3, ncol), dtype=x.dtype)
    
    array1 = x[:, is_yz_face] 
    array1[0, :] += Lx
    array2 = x[:, is_zx_face] 
    array2[1, :] += Ly
    array3 = x[:, is_xy_face] 
    array3[2, :] += Lz
    
    
    array4 = x[:, is_x_axis] 
    array4[1, :] += Ly
    array5 = x[:, is_x_axis] 
    array5[2, :] += Lz
    array6 = x[:, is_x_axis] 
    array6[2, :] += Lz
    array6[1, :] += Ly
    
    array7 = x[:, is_y_axis] 
    array7[0, :] += Lx
    array8 = x[:, is_y_axis] 
    array8[2, :] += Lz
    array9 = x[:, is_y_axis] 
    array9[0, :] += Lx
    array9[2, :] += Lz
    
    array10 = x[:, is_z_axis] 
    array10[1, :] += Ly
    array11 = x[:, is_z_axis] 
    array11[0, :] += Lx
    array12 = x[:, is_z_axis] 
    array12[0, :] += Lx
    array12[1, :] += Ly    
    
    # Conditionally define array13 based on is_corner
    if np.any(is_corner):
        array13 = np.array([
            [Lx, 0.0, 0.0, Lx, 0.0, Lx, Lx],  # x-coordinates
            [0.0, Ly, 0.0, Ly, Ly, 0.0, Ly],  # y-coordinates
            [0.0, 0.0, Lz, 0.0, Lz, Lz, Lz]   # z-coordinates
        ])
    else:
        array13 = np.zeros((3, 0), dtype=x.dtype)    
    
    new_array = np.hstack((array1, array2, array3, array4, array5, array6, array7, array8, array9, array10, array11, array12, array13))
    
    return new_array
    
def mapping_function_inv(x, left_vertices):

    tol = 1e-10

    is_corner = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol)

    is_z_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_y_axis = np.isclose(x[0, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[1, :], 0.0, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)
    is_x_axis = np.isclose(x[1, :], 0.0, atol=tol) & np.isclose(x[2, :], 0.0, atol=tol) & ~np.isclose(x[0, :], 0.0, atol=tol) & ~np.isclose(x[0, :], Lx, atol=tol)

    is_yz_face = np.isclose(x[0, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[1, :], Ly, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_zx_face = np.isclose(x[1, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[2, :], Lz, atol=tol)
    is_xy_face = np.isclose(x[2, :], 0, atol=tol) & ~is_corner & ~is_x_axis & ~is_y_axis & ~is_z_axis & ~np.isclose(x[0, :], Lx, atol=tol) & ~np.isclose(x[1, :], Ly, atol=tol)
    
    array1 = left_vertices[is_yz_face] 
    array2 = left_vertices[is_zx_face] 
    array3 = left_vertices[is_xy_face] 
    
    
    array4 = left_vertices[is_x_axis] 
    array5 = left_vertices[is_x_axis] 
    array6 = left_vertices[is_x_axis] 
    
    array7 = left_vertices[is_y_axis] 
    array8 = left_vertices[is_y_axis] 
    array9 = left_vertices[is_y_axis] 
    
    array10 = left_vertices[is_z_axis] 
    array11 = left_vertices[is_z_axis] 
    array12 = left_vertices[is_z_axis] 
    
    corner_column = left_vertices[is_corner]
    array13 = np.repeat(corner_column, 7)
    
    replace_map = np.concatenate((array1, array2, array3, array4, array5, array6, array7, array8, array9, array10, array11, array12, array13))
    
    return replace_map
    
def read_cube(filename):
    with open(filename, "r") as f:
        lines = f.readlines()

    # Parse origin (line 3)
    origin = np.array(list(map(float, lines[2].split()[1:4])))

    # Determine the number of atoms (line 7)
    num_atoms = int(lines[2].split()[0])

    # Parse grid dimensions (lines 3-5)
    nx = int(lines[3].split()[0])
    ny = int(lines[4].split()[0])
    nz = int(lines[5].split()[0])

    # Parse grid spacing (lines 3-5)
    dx = float(lines[3].split()[1])
    dy = float(lines[4].split()[2])
    dz = float(lines[5].split()[3])
    grid_spacing = np.array([dx, dy, dz])

    atom_data = []
    start_line = 6  # Line index starts at 0, so line 7 is index 6
    for i in range(start_line, start_line + num_atoms):
        values = lines[i].split()
        atom_data.append((int(values[0]), *map(float, values[1:5])))

    data_start_line = 6 + num_atoms + 1
    # Parse data
    data = np.array([float(value) for line in lines[data_start_line-1:] for value in line.split()])
    data = data.reshape((nx, ny, nz), order='C')   # should be F?????

    return origin, grid_spacing, data, atom_data, num_atoms
    
origin, grid_spacing, rho_elec_np, atom_data, num_atoms = read_cube("rho_elec.cube")
nx, ny, nz = rho_elec_np.shape
#_, _, vhartree_np, _, _ = read_cube("vhartree.cube")
_, _, vxc_np, _, _ = read_cube("vxc.cube")
#vxc_np = -(3 / np.pi) ** (1 / 3)*np.cbrt(rho_elec_np)
_, _, total_density_np, _, _ = read_cube("total_density.cube")

Lx, Ly, Lz = grid_spacing * np.array([nx, ny, nz])
domain = mesh.create_box(
    MPI.COMM_WORLD,
    [[0.0, 0.0, 0.0], [Lx, Ly, Lz]],
    [nx, ny, nz],
    mesh.CellType.hexahedron,
)

mpi_comm = MPI.COMM_WORLD
filename = "tmp.xdmf"
with XDMFFile(mpi_comm, filename, "w") as file:
    file.write_mesh(domain)

# Read all geometry data on all processes
with XDMFFile(MPI.COMM_SELF, filename, "r") as file:
    x_global = file.read_geometry_data()

# Read topology data
with XDMFFile(MPI.COMM_WORLD, filename, "r") as file:
    cell_shape, cell_degree = file.read_cell_type()
    x = file.read_geometry_data()
    topo = file.read_topology_data()
    
num_local_coor = x.shape[0]
all_sizes = mpi_comm.allgather(num_local_coor)
all_sizes.insert(0, 0)
all_ranges = np.cumsum(all_sizes)

# Testing the premise: coordinates are read contiguously in chunks
rank = mpi_comm.rank
assert np.all(x_global[all_ranges[rank]:all_ranges[rank + 1]] == x)

domain1 = ufl.Mesh(element("Lagrange", cell_shape.name, cell_degree, shape=(3, )))

# Partition mesh in layers, capture geometrical data and topological
# data from outer scope
def partitioner(*args):
    midpoints = np.mean(x_global[topo], axis=1)
    ###################################################################################################
    # Initialize the target array dest to store the target rank for each cell
    dest = np.zeros(len(midpoints), dtype=np.int32)

    # Identify indices that satisfy the conditions
    max_x, min_x = float(Lx) - 2.0 * float(Lx) / float(nx), 2.0 * float(Lx) / float(nx)
    max_y, min_y = float(Ly) - 2.0 * float(Ly) / float(ny), 2.0 * float(Ly) / float(ny)
    max_z, min_z = float(Lz) - 2.0 * float(Lz) / float(nz), 2.0 * float(Lz) / float(nz)

    # Assign cells that satisfy the boundary condition to rank=0
    boundary_indices = np.where(
        ((midpoints[:, 0] - max_x) > 0.0) | ((midpoints[:, 0] - min_x) < 0.0) |
        ((midpoints[:, 1] - max_y) > 0.0) | ((midpoints[:, 1] - min_y) < 0.0) |
        ((midpoints[:, 2] - max_z) > 0.0) | ((midpoints[:, 2] - min_z) < 0.0)
    )[0]
    dest[boundary_indices] = 0

    # Find the remaining cells to distribute
    if mpi_comm.size > 1:
        to_distribute = np.setdiff1d(np.arange(len(midpoints)), boundary_indices)

        # Distribute the remaining cells evenly across other ranks
        remaining_ranks = np.arange(1, mpi_comm.size)
        num_remaining_points = len(to_distribute)
        num_ranks = len(remaining_ranks)

        # Create an array of the same length as to_distribute and replace the beginning with repeated_ranks
        repeated_ranks = np.repeat(remaining_ranks, num_remaining_points // num_ranks)
        final_dest = np.zeros(len(to_distribute), dtype=np.int32)  # Initialize with 0
        final_dest[:len(repeated_ranks)] = repeated_ranks  # Replace the beginning values
        dest[to_distribute] = final_dest
    #####################################################################################################
    return dolfinx.cpp.graph.AdjacencyList_int32(dest)

new_domain = create_mesh(mpi_comm, topo, x, domain1, partitioner)

tdim = new_domain.topology.dim
assert domain.topology.index_map(tdim).size_global == new_domain.topology.index_map(tdim).size_global
num_cells = new_domain.topology.index_map(tdim).size_local
cell_midpoints = compute_midpoints(new_domain, tdim, np.arange(num_cells))

#new_mesh = create_mesh(mpi_comm, topo, x, domain1, partitioner)
new_mesh = create_periodic_mesh(new_domain, indicator, indicator_dual, mapping_function, mapping_function_inv)

V = dolfinx.fem.functionspace(new_mesh, ("CG", 1))
#print("weiye niubi!!!", V.dofmap.index_map.size_global)
rho_elec = fem.Function(V)
vhartree = fem.Function(V)
vxc = fem.Function(V)
total_density = fem.Function(V)

rho_elec_np = MPI.COMM_WORLD.bcast(rho_elec_np if MPI.COMM_WORLD.rank == 0 else None, root=0)
#vhartree_np = MPI.COMM_WORLD.bcast(vhartree_np if MPI.COMM_WORLD.rank == 0 else None, root=0)
vxc_np = MPI.COMM_WORLD.bcast(vxc_np if MPI.COMM_WORLD.rank == 0 else None, root=0)
total_density_np = MPI.COMM_WORLD.bcast(total_density_np if MPI.COMM_WORLD.rank == 0 else None, root=0)

def interpolate_to_mesh(data, grid_spacing, nx, ny, nz):
    def interpolator(x):
        return data[
            (np.round(x[0] / grid_spacing[0]).astype(int)) % nx,
            (np.round(x[1] / grid_spacing[1]).astype(int)) % ny,
            (np.round(x[2] / grid_spacing[2]).astype(int)) % nz
        ]
    return interpolator

#def interpolate_to_mesh(data, grid_spacing, nx, ny, nz):
#    def interpolator(x):
#        return data[
#            np.clip(np.round(x[0] / grid_spacing[0]).astype(int), 0, nx-1),
#            np.clip(np.round(x[1] / grid_spacing[1]).astype(int), 0, ny-1),
#            np.clip(np.round(x[2] / grid_spacing[2]).astype(int), 0, nz-1),
#        ]
#    return interpolator

rho_elec.interpolate(interpolate_to_mesh(rho_elec_np, grid_spacing, nx, ny, nz))
#vhartree.interpolate(interpolate_to_mesh(vhartree_np, grid_spacing, nx, ny, nz))
vxc.interpolate(interpolate_to_mesh(vxc_np, grid_spacing, nx, ny, nz))
total_density.interpolate(interpolate_to_mesh(total_density_np, grid_spacing, nx, ny, nz))

##################### correct electron number and core charge number ######################
total_density.x.array[:] = total_density.x.array[:] - rho_elec.x.array[:]
# 计算周期性胞元体积 measure
dxx = ufl.dx(domain=new_mesh)
vol = fem.assemble_scalar(fem.form(1.0*dxx))

# 净电荷
Q_local = fem.assemble_scalar(fem.form(rho_elec*dxx))
Q = MPI.COMM_WORLD.allreduce(Q_local, op=MPI.SUM)
print("old total electron number", Q)
rho_elec.x.array[:] *= 10.0/Q
Q_local = fem.assemble_scalar(fem.form(rho_elec*dxx))
Q = MPI.COMM_WORLD.allreduce(Q_local, op=MPI.SUM)
print("new total electron number", Q)

# 净电荷
Q_local = fem.assemble_scalar(fem.form(total_density*dxx))
Q = MPI.COMM_WORLD.allreduce(Q_local, op=MPI.SUM)
print("old total core charge number", Q)
total_density.x.array[:] *= -10.0/Q
Q_local = fem.assemble_scalar(fem.form(total_density*dxx))
Q = MPI.COMM_WORLD.allreduce(Q_local, op=MPI.SUM)
print("new total core charge number", Q)
#
#core_density = fem.Function(V)
#core_density.x.array[:] = total_density.x.array[:]
#
total_density.x.array[:] = total_density.x.array[:] + rho_elec.x.array[:]
##################### correct electron number and core charge number ######################

##################### background counter charge #######################
# 计算周期性胞元体积 measure
dxx = ufl.dx(domain=new_mesh)
vol = fem.assemble_scalar(fem.form(1.0*dxx))

# 净电荷
Q_local = fem.assemble_scalar(fem.form(total_density*dxx))
Q = MPI.COMM_WORLD.allreduce(Q_local, op=MPI.SUM)
print("netcharge:",Q)
# 若想令 net charge = 0，则对 rho 做平移
average = Q/vol   # 这是 rho 在网格上的平均值
total_density.x.array[:] -= average  # 每个自由度都减去 average

# 净电荷
Q_local = fem.assemble_scalar(fem.form(total_density*dxx))
Q = MPI.COMM_WORLD.allreduce(Q_local, op=MPI.SUM)
print("netcharge_new:",Q)
##################### background counter charge #######################

kai = TrialFunction(V)
lai = TestFunction(V)
a = (
      grad(kai)[0] * grad(lai)[0]
    + grad(kai)[1] * grad(lai)[1]
    + grad(kai)[2] * grad(lai)[2]
) * dx
L = 4.0 * np.pi * total_density * lai * dx

petsc_options = {
    "ksp_type": "minres", #"preonly", #"minres",
    "ksp_rtol": 1e-8,
    "ksp_atol": 1e-10,
    "ksp_max_it": 100000,
    "ksp_monitor": None,
    "ksp_converged_reason": None,
    "pc_type": "jacobi",
#    "pc_factor_mat_solver_type": "mumps",
#    "mat_scale": "symmetric"
#    "mat_symmetric": True
}

#petsc_options = {
#    "ksp_type": "preonly", #"minres",
#    "ksp_rtol": 1e-6,
#    "ksp_atol": 1e-10,
#    "ksp_max_it": 100000,
#    "ksp_monitor": None,
#    "ksp_converged_reason": None,
#    "pc_type": "jacobi",
##    "pc_factor_mat_solver_type": "mumps",
##    "mat_scale": "symmetric"
##    "mat_symmetric": True
#}

problem = dolfinx.fem.petsc.LinearProblem(a, L, petsc_options=petsc_options)
null_space = PETSc.NullSpace().create(constant=True)
problem.A.setNullSpace(null_space)
vhartree = problem.solve()

########################## 最外层迭代，为了优化vxc ########################
for i in range(1):

########################## 构建KS方程并且解方程 ###########################
    psi = ufl.TrialFunction(V)
    phi = ufl.TestFunction(V)
    
    # Kohn-Sham Hamiltonian form
    H_form = (0.5*ufl.inner(ufl.grad(psi), ufl.grad(phi)) + (vhartree + vxc)*psi*phi)*dx
    # Mass form
    M_form = psi*phi*dx
    
    # Assemble
    A = fem.petsc.assemble_matrix(fem.form(H_form))
    A.assemble()
    B = fem.petsc.assemble_matrix(fem.form(M_form))
    B.assemble()
   
    num_states = 5
 
    # Create SLEPc eigenvalue solver
    solver = SLEPc.EPS().create()
    solver.setOperators(A, B)  # <-- 重点：传递(A, B)
    solver.setProblemType(SLEPc.EPS.ProblemType.GHEP)  # 广义Hermitian
    solver.setWhichEigenpairs(SLEPc.EPS.Which.SMALLEST_REAL)
    solver.setDimensions(num_states, PETSc.DECIDE)
    
    # 设置求解器和预条件
    solver.setTolerances(tol=5e-7, max_it=5000)
    solver.setConvergenceTest(SLEPc.EPS.Conv.ABS)
    solver.setType(SLEPc.EPS.Type.LOBPCG) #KRYLOVSCHUR)
    pc = solver.getST().getKSP().getPC()
    pc.setType("jacobi") #hypre")
    pc.setFactorSolverType("mumps") #superlu_dist")
    # 其它 PETSc/SLEPc 选项
    opts = PETSc.Options()
    opts.setValue("eps_monitor_conv", None)
    opts.setValue("eps_converged_reason", None)
    opts.setValue("mat_mumps_icntl_14", 200)
    solver.setFromOptions()
    
    # Solve
    solver.solve()
#    psi = TrialFunction(V)
#    phi = TestFunction(V)
#
#    hbar2_over_2m = 0.5  #0.5  #0.5
#    H = (
#        hbar2_over_2m * inner(grad(psi), grad(phi)) * dx
#        +  (vhartree + vxc) * psi * phi * dx
#    )
#    
#    A = fem.petsc.assemble_matrix(fem.form(H))
#    A.assemble()
#    
#    num_states = 5
#
#    solver = SLEPc.EPS().create()
#    solver.setOperators(A)
#    solver.setProblemType(SLEPc.EPS.ProblemType.HEP)  
#    solver.setWhichEigenpairs(SLEPc.EPS.Which.SMALLEST_REAL)  
#    solver.setDimensions(num_states, PETSc.DECIDE)
#    solver.setTolerances(tol=1e-6, max_it=5000)
#    solver.setConvergenceTest(SLEPc.EPS.Conv.ABS)  # 绝对误差准则
##    solver.setType(SLEPc.EPS.Type.GD)  # 幂法
#    solver.setType(SLEPc.EPS.Type.KRYLOVSCHUR)  # 幂法
#    pc = solver.getST().getKSP().getPC()
#
#    pc.setType("lu")
#    pc.setFactorSolverType("mumps")
#    
#    # Set additional options
#    opts = PETSc.Options()
#    opts.setValue("eps_monitor_conv", None)
#    opts.setValue("eps_monitor_all", None)
#    opts.setValue("eps_converged_reason", None)
#    opts.setValue("mat_mumps_icntl_14", 200)  # Increase workspace for MUMPS
#    
#    # Apply options
#    solver.setFromOptions()
#
#    solver.solve()
    
    psi_list = []
    epsilon_list = []
    
    nconv = solver.getConverged()
    print(f"Number of converged eigenvalues: {nconv}")
    
    for i in range(nconv):
        vr, _ = A.createVecs()
        solver.getEigenpair(i, vr)
        psi_i = fem.Function(V)
        tmp = fem.Function(V) 
        psi_i.x.array[:len(vr.array_r)] = vr.array_r
        dolfinx.la.Vector.scatter_forward(psi_i.x)
        tmp.x.array[:] = psi_i.x.array[:]**2
        local_norm = fem.assemble_scalar(fem.form(tmp * dx))
        global_norm = np.sqrt(MPI.COMM_WORLD.allreduce(local_norm, op=MPI.SUM))
        print("global_norm",i,global_norm)
        psi_i.x.array[:] /= global_norm
        tmp.x.array[:] = psi_i.x.array[:]**2
        local_norm = fem.assemble_scalar(fem.form(tmp * dx))
        global_norm = np.sqrt(MPI.COMM_WORLD.allreduce(local_norm, op=MPI.SUM))
        print("global_norm after",i,global_norm)
        psi_list.append(psi_i)
        epsilon_list.append(solver.getEigenvalue(i).real) 

########################## 计算电子密度 ###########################
    rho = fem.Function(V)
    rho.x.array[:] = 0.0  # 将电子密度初始化为零

    # 遍历每个波函数 psi_i，累加 |psi_i|^2
    for psi in psi_list:
       rho.x.array[:] += 2.0*(psi.x.array[:]**2)

    rho_expression = rho
    local_norm = fem.assemble_scalar(fem.form(rho_expression * dx))
    global_norm = MPI.COMM_WORLD.allreduce(local_norm, op=MPI.SUM)
    print("global_norm rho",global_norm)
    
########################## 计算一系列gi ###########################
    g_list = []
    for i, psi_i in enumerate(psi_list):
        g_i = fem.Function(V)
        g_i.x.array[:] = 8.0 * (rho_elec.x.array[:] - rho.x.array[:]) * psi_i.x.array[:]
        g_list.append(g_i)
   
    solver.destroy()
    del solver
         
########################## 解后两个方程获得pi函数 ###########################
    p_list = []  
    mu_list = []
    
    p_var = TrialFunction(V)  
    q_var = TestFunction(V)
   
    hbar2_over_2m = 0.5
 
    for i, (psi_i, e_i, g_i) in enumerate(zip(psi_list, epsilon_list, g_list)):
        a_form = (
            hbar2_over_2m * inner(grad(p_var), grad(q_var)) * dx
            + (vhartree + vxc) * p_var * q_var * dx
            - e_i * p_var * q_var * dx
        )
        A = fem.petsc.assemble_matrix(fem.form(a_form))
        A.assemble()
    
        # 修正方程右端项 b
        L_form = g_i * q_var * dx
        b = fem.petsc.assemble_vector(fem.form(L_form))
        b.assemble()
     
        # 正交约束部分矩阵 C
        C_form = 2.0 * psi_i * q_var * dx
        C = fem.petsc.assemble_vector(fem.form(C_form))
        C.assemble()
    
        # 当前进程 MPI 信息
        rank = MPI.COMM_WORLD.rank
        size = MPI.COMM_WORLD.size    
    
        # 扩展矩阵和右端项
        n = A.size[0]
        last_row = n  # “扩展”矩阵多出来的行(列)索引就是 n
    
        # 4) 创建扩展矩阵 M (形状: (n+1) x (n+1))
        M = PETSc.Mat().create()
        M.setSizes([n + 1, n + 1])#, [n + 1, n + 1])  
        M.setType("aij")
        M.setUp()
    
        row_start, row_end = A.getOwnershipRange()
        
        # 填充 A 的主块
        for row in range(row_start, row_end):
            col_indices, values = A.getRow(row)
            M.setValues(row, col_indices, values, addv=False)
    
        global_C = np.zeros(C.getSize(), dtype=float)  # 在 root 上准备全局数组
        global_C[row_start:row_end] = C.getArray(readonly=True)
        global_C = MPI.COMM_WORLD.allreduce(global_C, op=MPI.SUM)
    
        row_start, row_end = M.getOwnershipRange()
    
        # 填充正交性条件
        if row_start <= last_row < row_end:
            # 最后一行
            for col in range(n):
#                print("col", col)
                M.setValue(last_row, col, global_C[col], addv=False)
            # 最右下角设置为 0
            M.setValue(last_row, last_row, 0.0, addv=False)
            print("row_start, row_end, last_row", row_start, row_end, last_row)
       
        row_start, row_end = A.getOwnershipRange()
       
        # 填充最后一列 (C^T)
        for row_loc in range(row_start, row_end):
            M.setValue(row_loc, last_row, global_C[row_loc], addv=False)

        # 装配矩阵
        M.assemble()
   
        # 创建扩展右端项 b_ext
        b_ext = PETSc.Vec().create()
        b_ext.setSizes(n + 1)
        b_ext.setUp()
       
        # 5.1) 复制 b 的本地部分到 b_ext
        b_local_start, b_local_end = b.getOwnershipRange()
        for row in range(b_local_start, b_local_end):
            val = b.getValue(row)
            b_ext.setValue(row, val, addv=False)

        # 5.2) 给最后一个分量(对应拉格朗日乘子行)设为0
        if row_start <= last_row < row_end:
            b_ext.setValue(last_row, 0.0, addv=False)
    
        b_ext.assemble()
   
        # 6) 创建并初始化扩展解向量 p_ext
        p_ext = PETSc.Vec().create()
        p_ext.setSizes(n + 1)
        p_ext.setUp()
        p_ext.assemble()
    
        # 求解扩展系统
        solver = PETSc.KSP().create(MPI.COMM_WORLD)
        solver.setOperators(M)
########################################
        solver.setType("preonly")
        solver.getPC().setType("lu")

#        solver.setType("cg") #preonly")
#        solver.getPC().setType("HYPRE_BoomerAMG") #lu")
########################################
        solver.solve(b_ext, p_ext)
    
        converged_reason = solver.getConvergedReason()
        
        if converged_reason > 0:
            print("Solver converged successfully.")
        elif converged_reason == 0:
            print("Solver is still iterating (did not converge or diverge).")
        else:
            print(f"Solver diverged with reason code {converged_reason}.")
    
        # 8) 提取修正波函数 p_i (取前 n 个分量)
        p_i = fem.Function(V)
        row_start, row_end = M.getOwnershipRange()
    
        global_p_ext = np.zeros(p_ext.getSize(), dtype=float)  # 在 root 上准备全局数组
        global_p_ext[row_start:row_end] = p_ext.getArray(readonly=True)
        global_p_ext = MPI.COMM_WORLD.allreduce(global_p_ext, op=MPI.SUM)

        local_indices = np.arange(V.dofmap.index_map.size_local, dtype=np.int32)
        local_to_global_dofs = V.dofmap.index_map.local_to_global(local_indices)   
        p_i.x.array[:len(local_to_global_dofs)] = global_p_ext[local_to_global_dofs]
        dolfinx.la.Vector.scatter_forward(p_i.x)

        local_norm = fem.assemble_scalar(fem.form(psi_list[0] * psi_i * dx))
        global_norm = MPI.COMM_WORLD.allreduce(local_norm, op=MPI.SUM)
        print("integral of psi_i * pi * dx", global_norm)

        p_list.append(p_i)
    
        # 9) 提取拉格朗日乘子 mu_i (最后一个分量, 全局 index = n)
        #    只有拥有 last_row 的进程能直接读 p_ext[last_row]
        mu_i = 0.0
        if row_start <= last_row < row_end:
            mu_i = p_ext.getValue(last_row)
    
        # 如果需要让所有进程都拿到 mu_i，可做一次MPI广播:
        mu_i = MPI.COMM_WORLD.bcast(mu_i, root=(rank if row_start <= last_row < row_end else 0))
   
        print("mu_i,",mu_i)
 
        mu_list.append(mu_i)
       
        solver.destroy()
        del solver
        A.destroy()
        M.destroy()
 
    grad_vxc = fem.Function(V)
    grad_vxc.x.array[:] = 0.0
    
    for psi, p in zip(psi_list, p_list):
        psi_values = psi.x.array
        p_values = p.x.array

        grad_vxc.x.array[:] += psi_values * p_values  # 按元素相乘并累加

    learning_rate = 0.3

    vxc.x.array[:] -= learning_rate * grad_vxc.x.array[:]
    
    ##################### Vxc 计算完成 #######################
    print("Updated vxc computed.")
    
vertices = new_mesh.geometry.x

x_min, x_max = vertices[:, 0].min(), vertices[:, 0].max()
y_min, y_max = vertices[:, 1].min(), vertices[:, 1].max()
z_min, z_max = vertices[:, 2].min(), vertices[:, 2].max()

Lx, Ly, Lz = x_max - x_min, y_max - y_min, z_max - z_min
nx_new = len(np.unique(vertices[:, 0]))
ny_new = len(np.unique(vertices[:, 1]))
nz_new = len(np.unique(vertices[:, 2]))

grid_spacing_new = [Lx / nx_new, Ly / ny_new, Lz / nz_new]

vertices_coords = new_mesh.geometry.x[:new_mesh.geometry.index_map().size_local]

tree = bb_tree(new_mesh, new_mesh.topology.dim)
cell_candidates = compute_collisions_points(tree, vertices_coords)
colliding_cells = compute_colliding_cells(new_mesh, cell_candidates, vertices_coords)

cells = []
for i in range(len(vertices_coords)):
    candidates = colliding_cells.links(i)
    if len(candidates) == 0:
        raise ValueError(f"Point {vertices_coords[i]} is outside the mesh.")
    cells.append(candidates[0])
cells = np.array(cells, dtype=np.int32)

for idx, psi in enumerate(psi_list):
     
    values_1d=psi.eval(vertices_coords, cells)
#    values_1d=psi_list[4].eval(vertices_coords, cells)
    #values_1d=grad_vxc.eval(vertices_coords, cells)
    #values_1d=vhartree.eval(vertices_coords, cells)
    #values_1d=rho_elec.eval(vertices_coords, cells)
    #values_1d=rho.eval(vertices_coords, cells)
    #values_1d=vxc.eval(vertices_coords, cells)
    #values_1d=p_list[1].eval(vertices_coords, cells)
    #values_1d=g_list[0].eval(vertices_coords, cells)
    
    print("epsilon_list:", epsilon_list)
    
    global_coords = None
    global_values_1d = None
    
    local_coord_size = vertices_coords.shape[0]*vertices_coords.shape[1]
    local_value_size = values_1d.size
    print(f"lallaal {new_mesh.geometry.index_map().size_local}")
    recv_counts_coords = MPI.COMM_WORLD.gather(local_coord_size, root=0)
    recv_counts_values = MPI.COMM_WORLD.gather(local_value_size, root=0)
    
    if rank == 0:
        global_size = new_mesh.geometry.index_map().size_global
        global_coords = np.zeros((global_size, vertices_coords.shape[1]), dtype=vertices_coords.dtype)
        global_values_1d = np.zeros(global_size, dtype=values_1d.dtype)
    
        displs_coords = np.cumsum([0] + recv_counts_coords[:-1])
        displs_values = np.cumsum([0] + recv_counts_values[:-1])
    else:
        displs_coords = None
        displs_values = None
    
    MPI.COMM_WORLD.Gatherv(
        sendbuf=vertices_coords,
        recvbuf=(global_coords, recv_counts_coords, displs_coords, MPI.DOUBLE),
        root=0
    )
    
    MPI.COMM_WORLD.Gatherv(
        sendbuf=values_1d,
        recvbuf=(global_values_1d, recv_counts_values, displs_values, MPI.DOUBLE),
        root=0
    )
    values_3d = None
    if rank == 0:
       sorted_indices = np.lexsort((global_coords[:, 2], global_coords[:, 1], global_coords[:, 0]))
       sorted_values = global_values_1d[sorted_indices]
    
       values_3d = sorted_values.reshape((nx+1, ny+1, nz+1), order="C")
    
    def write_cube(values_3d, origin, spacing, atom_data, num_atoms, output_filename, comm=MPI.COMM_WORLD):
        rank = comm.Get_rank()
    
        if rank == 0:
            nx, ny, nz = values_3d.shape
            nx -= 1
            ny -= 1
            nz -= 1
    
            with open(output_filename, "w") as f:
                header = [
                    "CUBE FILE\n",
                    "OUTER LOOP: X, MIDDLE LOOP: Y, INNER LOOP: Z\n",
                    f"{num_atoms:5d} {origin[0]:12.6f} {origin[1]:12.6f} {origin[2]:12.6f}\n",
                    f"{nx:5d} {spacing[0]:12.6f} {0.0:12.6f} {0.0:12.6f}\n",
                    f"{ny:5d} {0.0:12.6f} {spacing[1]:12.6f} {0.0:12.6f}\n",
                    f"{nz:5d} {0.0:12.6f} {0.0:12.6f} {spacing[2]:12.6f}\n"
                ]
                f.writelines(header)
    
                for atom in atom_data:
                    line = f"{atom[0]:d} {atom[1]:.6f} {atom[2]:.6f} {atom[3]:.6f} {atom[4]:.6f}\n"
                    f.write(line)
    
                for i in range(nx):
                    for j in range(ny):
                        row_values = values_3d[i, j, :nz]
                        row_str = " ".join(f"{v:13.5e}" for v in row_values)
                        f.write(row_str + "\n")
    
            print(f".cube file '{output_filename}' written successfully!")
    
        comm.Barrier()
#        exit()
    
    if __name__ == "__main__":
        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()
    
        values_3d = comm.bcast(values_3d, root=0)
        output_filename = f"grad_vxc_{idx}.cube"
        write_cube(values_3d, origin, grid_spacing, atom_data, num_atoms, output_filename, comm)