program testplex ! Testing PETSc dmplex from fortran, using a line of blocks in 3-D. ! Create the DMPlex by passing in the cell definitions to DMPlexCreateFromDAG(), ! interpolate the result to form the faces and edges, distribute over the ! processors, and visualize the partitions in VTK. implicit none #include PetscInt, parameter :: num_cells = 20 PetscReal, parameter :: cell_size = 10. PetscErrorCode :: ierr DM :: dm, dist_dm, cell_dm character(5) :: partitioner = "chaco" PetscInt :: overlap = 0 Vec :: partition call PetscInitialize(PETSC_NULL_CHARACTER, ierr); CHKERRQ(ierr) call create_1d_grid(num_cells, cell_size, dm, ierr); CHKERRQ(ierr) call DMPlexDistribute(dm, partitioner, overlap, PETSC_NULL_OBJECT, dist_dm, ierr) CHKERRQ(ierr) if (dist_dm .ne. 0) then call DMDestroy(dm, ierr); CHKERRQ(ierr) dm = dist_dm end if call DMView(dm, PETSC_VIEWER_STDOUT_WORLD, ierr); CHKERRQ(ierr) call create_partition_vector(dm, cell_dm, partition, ierr); CHKERRQ(ierr) ! Visualize partition vector here... call DMDestroy(dm, ierr); CHKERRQ(ierr) call PetscFinalize(ierr); CHKERRQ(ierr) end program testplex !------------------------------------------------------------------------ subroutine create_1d_grid(num_cells, cell_size, dm, ierr) ! Creates grid consisting of a 1D line of cube cells in 3D. implicit none #include PetscInt, intent(in) :: num_cells PetscReal, intent(in):: cell_size DM, intent(out) :: dm PetscErrorCode, intent(out) :: ierr ! Locals: PetscInt, parameter :: dim = 3 PetscMPIInt :: rank PetscInt :: num_vertices_per_face, num_vertices_per_cell PetscInt :: num_vertices, total_num_points PetscInt :: i, i0, i1, l PetscInt :: depth real(8) :: x0, dx PetscInt, allocatable :: num_points(:) PetscInt, allocatable :: cone_size(:) PetscInt, allocatable :: cones(:) PetscInt, allocatable :: cone_orientations(:) PetscScalar, allocatable :: vertex_coords(:) DM :: interp_dm ierr = 0 call MPI_Comm_rank(PETSC_COMM_WORLD, rank, ierr); CHKERRQ(ierr) call DMPlexCreate(PETSC_COMM_WORLD, dm, ierr); CHKERRQ(ierr) call PetscObjectSetName(dm, "testplex", ierr); CHKERRQ(ierr) call DMPlexSetDimension(dm, dim, ierr); CHKERRQ(ierr) if (rank == 0) then num_vertices_per_face = 2 ** (dim-1) num_vertices_per_cell = 2 * num_vertices_per_face num_vertices = num_vertices_per_face * (num_cells + 1) l = num_vertices_per_face * dim total_num_points = num_vertices + num_cells ! # points in DAG allocate(vertex_coords(0:num_vertices*dim-1), cones(0:num_cells*num_vertices_per_cell-1)) ! Set up vertex coordinates dx = real(cell_size) do i = 0, num_cells i0 = i * l i1 = (i+1)*l x0 = real(i * cell_size) vertex_coords(i0:i1) = \ [x0, 0.d0, 0.d0,\ x0, dx, 0.d0,\ x0, 0.d0, dx,\ x0, dx, dx] end do ! Set up cells do i = 0, num_cells-1 i0 = i * num_vertices_per_cell i1 = (i+1) * num_vertices_per_cell cones(i0:i1) = [0, 1, 5, 4, 2, 6, 7, 3] + i * num_vertices_per_face end do cones = cones + num_cells allocate(num_points(0: 1)) num_points = [num_vertices, num_cells] allocate(cone_size(0: total_num_points-1)) allocate(cone_orientations(0:num_cells*num_vertices_per_cell-1)) cone_size = 0 cone_size(0: num_cells-1) = num_vertices_per_cell cone_orientations = 0 depth = 1 call DMPlexCreateFromDAG(dm, depth, num_points, cone_size, cones, & cone_orientations, vertex_coords, ierr); CHKERRQ(ierr) else ! rank > 0: create empty DMPlex num_points = [0, 0, 0]; cone_size = [0] cones = [0]; cone_orientations = [0] vertex_coords = [0] call DMPlexCreateFromDAG(dm, 1, num_points, cone_size, & cones, cone_orientations, vertex_coords, ierr); CHKERRQ(ierr) end if interp_dm = PETSC_NULL_OBJECT call DMPlexInterpolate(dm, interp_dm, ierr); CHKERRQ(ierr) call DMPlexCopyCoordinates(dm, interp_dm, ierr); CHKERRQ(ierr) call DMDestroy(dm, ierr); CHKERRQ(ierr) dm = interp_dm deallocate(num_points, cone_size, cones, cone_orientations, vertex_coords) end subroutine create_1d_grid !------------------------------------------------------------------------ subroutine create_partition_vector(dm, cell_dm, partition, ierr) ! Returns a DM and vector containing partitioning data implicit none #include DM, intent(in) :: dm DM, intent(out) :: cell_dm Vec, intent(out) :: partition PetscErrorCode, intent(out) :: ierr ! Locals: MPI_Comm :: comm PetscMPIInt :: rank PetscSF :: sfPoint PetscSection :: coord_section Vec :: coordinates PetscSection :: cell_section PetscScalar, pointer :: part(:) PetscInt :: cStart, cEnd, c, off call DMGetCoordinateSection(dm, coord_section, ierr); CHKERRQ(ierr) call DMGetCoordinatesLocal(dm, coordinates, ierr); CHKERRQ(ierr) call DMClone(dm, cell_dm, ierr); CHKERRQ(ierr) write (*,*) dm, cell_dm call DMGetPointSF(dm, sfPoint, ierr); CHKERRQ(ierr) call DMSetPointSF(cell_dm, sfPoint, ierr); CHKERRQ(ierr) call DMSetCoordinateSection(cell_dm, coord_section, ierr); CHKERRQ(ierr) call DMSetCoordinatesLocal(cell_dm, coordinates, ierr); CHKERRQ(ierr) call PetscObjectGetComm(dm, comm, ierr); CHKERRQ(ierr) call PetscSectionCreate(comm, cell_section, ierr); CHKERRQ(ierr) call DMPlexGetHeightStratum(cell_dm, 0, cStart, cEnd, ierr); CHKERRQ(ierr) call PetscSectionSetChart(cell_section, cStart, cEnd, ierr); CHKERRQ(ierr) do c = cStart, cEnd-1 call PetscSectionSetDof(cell_section, c, 1, ierr); CHKERRQ(ierr) end do call PetscSectionSetUp(cell_section, ierr); CHKERRQ(ierr) call DMSetDefaultSection(cell_dm, cell_section, ierr); CHKERRQ(ierr) call DMCreateLocalVector(cell_dm, partition, ierr); CHKERRQ(ierr) call PetscObjectSetName(partition, "partition", ierr); CHKERRQ(ierr) call VecGetArrayF90(partition, part, ierr); CHKERRQ(ierr) call MPI_Comm_rank(comm, rank); CHKERRQ(ierr) do c = cStart, cEnd-1 call PetscSectionGetOffset(cell_section, c, off, ierr); CHKERRQ(ierr) part(off) = rank end do call VecRestoreArrayF90(partition, part, ierr); CHKERRQ(ierr) call PetscSectionDestroy(cell_section, ierr); CHKERRQ(ierr) end subroutine create_partition_vector