diff --git a/src/Graphs/partitionedgraphs/abstractpartitionedge.jl b/src/Graphs/partitionedgraphs/abstractpartitionedge.jl
index e907238..3470ed3 100644
--- a/src/Graphs/partitionedgraphs/abstractpartitionedge.jl
+++ b/src/Graphs/partitionedgraphs/abstractpartitionedge.jl
@@ -3,5 +3,6 @@ abstract type AbstractPartitionEdge{V} <: AbstractNamedEdge{V} end
 parent(pe::AbstractPartitionEdge) = not_implemented()
 src(pe::AbstractPartitionEdge) = not_implemented()
 dst(pe::AbstractPartitionEdge) = not_implemented()
+reverse(pe::AbstractPartitionEdge) = not_implemented()
 
 #Don't have the vertices wrapped. But wrap them with source and edge.
diff --git a/src/Graphs/partitionedgraphs/abstractpartitionedgraph.jl b/src/Graphs/partitionedgraphs/abstractpartitionedgraph.jl
index 5ef0692..a536306 100644
--- a/src/Graphs/partitionedgraphs/abstractpartitionedgraph.jl
+++ b/src/Graphs/partitionedgraphs/abstractpartitionedgraph.jl
@@ -3,18 +3,23 @@ abstract type AbstractPartitionedGraph{V,PV} <: AbstractNamedGraph{V} end
 #Needed for interface
 partitioned_graph(pg::AbstractPartitionedGraph) = not_implemented()
 unpartitioned_graph(pg::AbstractPartitionedGraph) = not_implemented()
-which_partition(pg::AbstractPartitionedGraph, vertex) = not_implemented()
-partitioned_vertices(pg::AbstractPartitionedGraph) = not_implemented()
+partitionvertex(pg::AbstractPartitionedGraph, vertex) = not_implemented()
+partitionvertices(pg::AbstractPartitionedGraph, verts::Vector) = not_implemented()
+partitionvertices(pg::AbstractPartitionedGraph) = not_implemented()
 copy(pg::AbstractPartitionedGraph) = not_implemented()
 delete_from_vertex_map!(pg::AbstractPartitionedGraph, vertex) = not_implemented()
 insert_to_vertex_map!(pg::AbstractPartitionedGraph, vertex) = not_implemented()
-partition_edge(pg::AbstractPartitionedGraph, edge) = not_implemented()
-function edges(pg::AbstractPartitionedGraph, partition_edge::AbstractPartitionEdge)
+partitionedge(pg::AbstractPartitionedGraph, edge) = not_implemented()
+function partitionedges(pg::AbstractPartitionedGraph, edges::Vector{<:AbstractEdge})
+  return not_implemented()
+end
+partitionedges(pg::AbstractPartitionedGraph) = not_implemented()
+function edges(pg::AbstractPartitionedGraph, partitionedge::AbstractPartitionEdge)
   return not_implemented()
 end
 vertices(pg::AbstractPartitionedGraph, pv::AbstractPartitionVertex) = not_implemented()
 function vertices(
-  pg::AbstractPartitionedGraph, partition_verts::Vector{V}
+  pg::AbstractPartitionedGraph, partitionverts::Vector{V}
 ) where {V<:AbstractPartitionVertex}
   return not_implemented()
 end
@@ -31,22 +36,22 @@ end
 edgetype(pg::AbstractPartitionedGraph) = edgetype(unpartitioned_graph(pg))
 parent_graph_type(pg::AbstractPartitionedGraph) = parent_graph_type(unpartitioned_graph(pg))
 nv(pg::AbstractPartitionedGraph, pv::AbstractPartitionVertex) = length(vertices(pg, pv))
-function has_vertex(pg::AbstractPartitionedGraph, partition_vertex::AbstractPartitionVertex)
-  return has_vertex(partitioned_graph(pg), parent(partition_vertex))
+function has_vertex(pg::AbstractPartitionedGraph, partitionvertex::AbstractPartitionVertex)
+  return has_vertex(partitioned_graph(pg), parent(partitionvertex))
 end
 
-function has_edge(pg::AbstractPartitionedGraph, edge::AbstractPartitionEdge)
-  return has_edge(partitioned_graph(pg), parent(partition_edge))
+function has_edge(pg::AbstractPartitionedGraph, partitionedge::AbstractPartitionEdge)
+  return has_edge(partitioned_graph(pg), parent(partitionedge))
 end
 
 function is_boundary_edge(pg::AbstractPartitionedGraph, edge::AbstractEdge)
-  p_edge = partition_edge(pg, edge)
+  p_edge = partitionedge(pg, edge)
   return src(p_edge) == dst(p_edge)
 end
 
 function add_edge!(pg::AbstractPartitionedGraph, edge::AbstractEdge)
   add_edge!(unpartitioned_graph(pg), edge)
-  pg_edge = parent(partition_edge(pg, edge))
+  pg_edge = parent(partitionedge(pg, edge))
   if src(pg_edge) != dst(pg_edge)
     add_edge!(partitioned_graph(pg), pg_edge)
   end
@@ -55,7 +60,7 @@ function add_edge!(pg::AbstractPartitionedGraph, edge::AbstractEdge)
 end
 
 function rem_edge!(pg::AbstractPartitionedGraph, edge::AbstractEdge)
-  pg_edge = partition_edge(pg, edge)
+  pg_edge = partitionedge(pg, edge)
   if has_edge(partitioned_graph(pg), pg_edge)
     g_edges = edges(pg, pg_edge)
     if length(g_edges) == 1
@@ -65,50 +70,50 @@ function rem_edge!(pg::AbstractPartitionedGraph, edge::AbstractEdge)
   return rem_edge!(unpartitioned_graph(pg), edge)
 end
 
-function rem_edge!(pg::AbstractPartitionedGraph, partition_edge::AbstractPartitionEdge)
-  return rem_edges!(pg, edges(pg, parent(partition_edge)))
+function rem_edge!(pg::AbstractPartitionedGraph, partitionedge::AbstractPartitionEdge)
+  return rem_edges!(pg, edges(pg, parent(partitionedge)))
 end
 
-function rem_edge(pg::AbstractPartitionedGraph, partition_edge::AbstractPartitionEdge)
+function rem_edge(pg::AbstractPartitionedGraph, partitionedge::AbstractPartitionEdge)
   pg_new = copy(pg)
-  rem_edge!(pg_new, partition_edge)
+  rem_edge!(pg_new, partitionedge)
   return pg_new
 end
 
 function rem_edges!(
-  pg::AbstractPartitionedGraph, partition_edges::Vector{<:AbstractPartitionEdge}
+  pg::AbstractPartitionedGraph, partitionedges::Vector{<:AbstractPartitionEdge}
 )
-  for pe in partition_edges
+  for pe in partitionedges
     rem_edge!(pg, pe)
   end
   return pg
 end
 
 function rem_edges(
-  pg::AbstractPartitionedGraph, partition_edges::Vector{<:AbstractPartitionEdge}
+  pg::AbstractPartitionedGraph, partitionedges::Vector{<:AbstractPartitionEdge}
 )
   pg_new = copy(pg)
-  rem_edges!(pg_new, partition_edges)
+  rem_edges!(pg_new, partitionedges)
   return pg_new
 end
 
 #Vertex addition and removal. I think it's important not to allow addition of a vertex without specification of PV
 function add_vertex!(
-  pg::AbstractPartitionedGraph, vertex, partition_vertex::AbstractPartitionVertex
+  pg::AbstractPartitionedGraph, vertex, partitionvertex::AbstractPartitionVertex
 )
   add_vertex!(unpartitioned_graph(pg), vertex)
-  add_vertex!(partitioned_graph(pg), parent(partition_vertex))
-  insert_to_vertex_map!(pg, vertex, partition_vertex)
+  add_vertex!(partitioned_graph(pg), parent(partitionvertex))
+  insert_to_vertex_map!(pg, vertex, partitionvertex)
   return pg
 end
 
 function add_vertices!(
   pg::AbstractPartitionedGraph,
   vertices::Vector,
-  partition_vertices::Vector{<:AbstractPartitionVertex},
+  partitionvertices::Vector{<:AbstractPartitionVertex},
 )
-  @assert length(vertices) == length(partition_vertices)
-  for (v, pv) in zip(vertices, partition_vertices)
+  @assert length(vertices) == length(partitionvertices)
+  for (v, pv) in zip(vertices, partitionvertices)
     add_vertex!(pg, v, pv)
   end
 
@@ -116,14 +121,14 @@ function add_vertices!(
 end
 
 function add_vertices!(
-  pg::AbstractPartitionedGraph, vertices::Vector, partition_vertex::AbstractPartitionVertex
+  pg::AbstractPartitionedGraph, vertices::Vector, partitionvertex::AbstractPartitionVertex
 )
-  add_vertices!(pg, vertices, fill(partition_vertex, length(vertices)))
+  add_vertices!(pg, vertices, fill(partitionvertex, length(vertices)))
   return pg
 end
 
 function rem_vertex!(pg::AbstractPartitionedGraph, vertex)
-  pv = which_partition(pg, vertex)
+  pv = partitionvertex(pg, vertex)
   delete_from_vertex_map!(pg, pv, vertex)
   rem_vertex!(unpartitioned_graph(pg), vertex)
   if !haskey(partitioned_vertices(pg), parent(pv))
@@ -132,16 +137,14 @@ function rem_vertex!(pg::AbstractPartitionedGraph, vertex)
   return pg
 end
 
-function rem_vertex!(
-  pg::AbstractPartitionedGraph, partition_vertex::AbstractPartitionVertex
-)
-  rem_vertices!(pg, vertices(pg, partition_vertex))
+function rem_vertex!(pg::AbstractPartitionedGraph, partitionvertex::AbstractPartitionVertex)
+  rem_vertices!(pg, vertices(pg, partitionvertex))
   return pg
 end
 
-function rem_vertex(pg::AbstractPartitionedGraph, partition_vertex::AbstractPartitionVertex)
+function rem_vertex(pg::AbstractPartitionedGraph, partitionvertex::AbstractPartitionVertex)
   pg_new = copy(pg)
-  rem_vertex!(pg_new, partition_vertex)
+  rem_vertex!(pg_new, partitionvertex)
   return pg_new
 end
 
@@ -155,19 +158,19 @@ function (pg1::AbstractPartitionedGraph == pg2::AbstractPartitionedGraph)
     return false
   end
   for v in vertices(pg1)
-    if which_partition(pg1, v) != which_partition(pg2, v)
+    if partitionvertex(pg1, v) != partitionvertex(pg2, v)
       return false
     end
   end
   return true
 end
 
-function subgraph(pg::AbstractPartitionedGraph, partition_vertex::AbstractPartitionVertex)
-  return first(induced_subgraph(unpartitioned_graph(pg), vertices(pg, [partition_vertex])))
+function subgraph(pg::AbstractPartitionedGraph, partitionvertex::AbstractPartitionVertex)
+  return first(induced_subgraph(unpartitioned_graph(pg), vertices(pg, [partitionvertex])))
 end
 
 function induced_subgraph(
-  pg::AbstractPartitionedGraph, partition_vertex::AbstractPartitionVertex
+  pg::AbstractPartitionedGraph, partitionvertex::AbstractPartitionVertex
 )
-  return subgraph(pg, partition_vertex), nothing
+  return subgraph(pg, partitionvertex), nothing
 end
diff --git a/src/Graphs/partitionedgraphs/partitionedge.jl b/src/Graphs/partitionedgraphs/partitionedge.jl
index 4607b9d..b344e0b 100644
--- a/src/Graphs/partitionedgraphs/partitionedge.jl
+++ b/src/Graphs/partitionedgraphs/partitionedge.jl
@@ -1,4 +1,4 @@
-struct PartitionEdge{V,E<:AbstractEdge{V}} <: AbstractPartitionEdge{V}
+struct PartitionEdge{V,E<:AbstractEdge{<:V}} <: AbstractPartitionEdge{V}
   edge::E
 end
 
@@ -6,3 +6,5 @@ parent(pe::PartitionEdge) = getfield(pe, :edge)
 src(pe::PartitionEdge) = PartitionVertex(src(parent(pe)))
 dst(pe::PartitionEdge) = PartitionVertex(dst(parent(pe)))
 PartitionEdge(p::Pair) = PartitionEdge(NamedEdge(first(p) => last(p)))
+PartitionEdge(vsrc, vdst) = PartitionEdge(vsrc => vdst)
+reverse(pe::PartitionEdge) = PartitionEdge(reverse(parent(pe)))
diff --git a/src/Graphs/partitionedgraphs/partitionedgraph.jl b/src/Graphs/partitionedgraphs/partitionedgraph.jl
index 33c71a2..27ad0fe 100644
--- a/src/Graphs/partitionedgraphs/partitionedgraph.jl
+++ b/src/Graphs/partitionedgraphs/partitionedgraph.jl
@@ -33,8 +33,8 @@ function PartitionedGraph(partitioned_vertices)
 end
 
 function PartitionedGraph(g::AbstractGraph; kwargs...)
-  partitioned_vertices = partition_vertices(g; kwargs...)
-  return PartitionedGraph(g, partitioned_vertices)
+  partitioned_verts = partitioned_vertices(g; kwargs...)
+  return PartitionedGraph(g, partitioned_verts)
 end
 
 #Needed for interface
@@ -43,25 +43,42 @@ unpartitioned_graph(pg::PartitionedGraph) = getfield(pg, :graph)
 partitioned_vertices(pg::PartitionedGraph) = getfield(pg, :partitioned_vertices)
 which_partition(pg::PartitionedGraph) = getfield(pg, :which_partition)
 parent_graph_type(PG::Type{<:PartitionedGraph}) = fieldtype(PG, :graph)
-function vertices(pg::PartitionedGraph, partition_vert::PartitionVertex)
-  return partitioned_vertices(pg)[parent(partition_vert)]
+function vertices(pg::PartitionedGraph, partitionvert::PartitionVertex)
+  return partitioned_vertices(pg)[parent(partitionvert)]
 end
-function vertices(pg::PartitionedGraph, partition_verts::Vector{<:PartitionVertex})
-  return unique(reduce(vcat, [vertices(pg, pv) for pv in partition_verts]))
+function vertices(pg::PartitionedGraph, partitionverts::Vector{<:PartitionVertex})
+  return unique(reduce(vcat, [vertices(pg, pv) for pv in partitionverts]))
 end
-function which_partition(pg::PartitionedGraph, vertex)
+function partitionvertex(pg::PartitionedGraph, vertex)
   return PartitionVertex(which_partition(pg)[vertex])
 end
 
-function partition_edge(pg::PartitionedGraph, edge::AbstractEdge)
+function partitionvertices(pg::PartitionedGraph, verts::Vector)
+  return unique(partitionvertex(pg, v) for v in verts)
+end
+
+function partitionvertices(pg::PartitionedGraph)
+  return PartitionVertex.(vertices(partitioned_graph(pg)))
+end
+
+function partitionedge(pg::PartitionedGraph, edge::AbstractEdge)
   return PartitionEdge(
-    parent(which_partition(pg, src(edge))) => parent(which_partition(pg, dst(edge)))
+    parent(partitionvertex(pg, src(edge))) => parent(partitionvertex(pg, dst(edge)))
   )
 end
 
-function partition_edges(pg::PartitionedGraph, partition_edge::PartitionEdge)
-  psrc_vs = vertices(pg, PartitionVertex(src(partition_edge)))
-  pdst_vs = vertices(pg, PartitionVertex(dst(partition_edge)))
+#Lets filter out any self-edges from this. Although this makes it a bit consistent with partitionedge
+function partitionedges(pg::PartitionedGraph, edges::Vector{<:AbstractEdge})
+  return filter(e -> src(e) != dst(e), unique([partitionedge(pg, e) for e in edges]))
+end
+
+function partitionedges(pg::PartitionedGraph)
+  return PartitionEdge.(edges(partitioned_graph(pg)))
+end
+
+function edges(pg::PartitionedGraph, partitionedge::PartitionEdge)
+  psrc_vs = vertices(pg, PartitionVertex(src(partitionedge)))
+  pdst_vs = vertices(pg, PartitionVertex(dst(partitionedge)))
   psrc_subgraph = subgraph(unpartitioned_graph(pg), psrc_vs)
   pdst_subgraph = subgraph(pg, pdst_vs)
   full_subgraph = subgraph(pg, vcat(psrc_vs, pdst_vs))
@@ -69,23 +86,27 @@ function partition_edges(pg::PartitionedGraph, partition_edge::PartitionEdge)
   return setdiff(edges(full_subgraph), vcat(edges(psrc_subgraph), edges(pdst_subgraph)))
 end
 
+function edges(pg::PartitionedGraph, partitionedges::Vector{<:PartitionEdge})
+  return unique(reduce(vcat, [edges(pg, pe) for pe in partitionedges]))
+end
+
 function copy(pg::PartitionedGraph)
   return PartitionedGraph(
     copy(unpartitioned_graph(pg)),
     copy(partitioned_graph(pg)),
     copy_keys_values(partitioned_vertices(pg)),
-    copy_keys_values(which_partition(pg)),
+    copy_keys_values(partitionvertex(pg)),
   )
 end
 
 function insert_to_vertex_map!(
-  pg::PartitionedGraph, vertex, partition_vertex::PartitionVertex
+  pg::PartitionedGraph, vertex, partitionvertex::PartitionVertex
 )
-  pv = parent(partition_vertex)
+  pv = parent(partitionvertex)
   if pv ∉ keys(partitioned_vertices(pg))
     insert!(partitioned_vertices(pg), pv, [vertex])
   else
-    partitioned_vertices(pg)[pv] = unique(vcat(vertices(pg, partition_vertex), [vertex]))
+    partitioned_vertices(pg)[pv] = unique(vcat(vertices(pg, partitionvertex), [vertex]))
   end
 
   insert!(which_partition(pg), vertex, pv)
@@ -93,7 +114,7 @@ function insert_to_vertex_map!(
 end
 
 function delete_from_vertex_map!(pg::PartitionedGraph, vertex)
-  pv = which_partition(pg, vertex)
+  pv = partitionvertex(pg, vertex)
   return delete_from_vertex_map!(pg, pv, vertex)
 end
 
@@ -127,7 +148,7 @@ function induced_subgraph(pg::PartitionedGraph, vertices::Vector)
 end
 
 function induced_subgraph(
-  pg::PartitionedGraph, partition_verts::Vector{V}
+  pg::PartitionedGraph, partitionverts::Vector{V}
 ) where {V<:PartitionVertex}
-  return induced_subgraph(pg, vertices(pg, partition_verts))
+  return induced_subgraph(pg, vertices(pg, partitionverts))
 end
diff --git a/src/Graphs/partitionedgraphs/partitioning.jl b/src/Graphs/partitionedgraphs/partitioning.jl
index 05551b5..3074fc4 100644
--- a/src/Graphs/partitionedgraphs/partitioning.jl
+++ b/src/Graphs/partitionedgraphs/partitioning.jl
@@ -71,7 +71,7 @@ function _npartitions(
   return error("Must specify either `npartitions` or `nvertices_per_partition`")
 end
 
-function partition_vertices(
+function partitioned_vertices(
   g::Graph;
   npartitions=nothing,
   nvertices_per_partition=nothing,
@@ -84,15 +84,15 @@ function partition_vertices(
     return group(v -> 1, collect(vertices(g)))
   end
 
-  return partition_vertices(
+  return partitioned_vertices(
     Backend(backend), g, _npartitions(g, npartitions, nvertices_per_partition); kwargs...
   )
 end
 
-function partition_vertices(
+function partitioned_vertices(
   g::AbstractNamedGraph; npartitions=nothing, nvertices_per_partition=nothing, kwargs...
 )
-  vertex_partitions = partition_vertices(
+  vertex_partitions = partitioned_vertices(
     parent_graph(g); npartitions, nvertices_per_partition, kwargs...
   )
   #[inv(vertex_to_parent_vertex(g))[v] for v in partitions]
diff --git a/src/NamedGraphs.jl b/src/NamedGraphs.jl
index e3f1e0a..62a337b 100644
--- a/src/NamedGraphs.jl
+++ b/src/NamedGraphs.jl
@@ -181,6 +181,11 @@ export NamedGraph,
   neighbors,
   nv,
   partitioned_graph,
+  partitionedge,
+  partitionedges,
+  partitionvertex,
+  partitionvertices,
+  partitioned_vertices,
   path_digraph,
   path_graph,
   periphery,
diff --git a/src/requires/kahypar.jl b/src/requires/kahypar.jl
index eecb0ef..d7709d3 100644
--- a/src/requires/kahypar.jl
+++ b/src/requires/kahypar.jl
@@ -4,14 +4,14 @@ set_partitioning_backend!(Backend"KaHyPar"())
 KaHyPar.HyperGraph(g::SimpleGraph) = incidence_matrix(g)
 
 """
-    partition_vertices(::Backend"KaHyPar", g::Graph, npartiations::Integer; objective="edge_cut", alg="kway", kwargs...)
+partitioned_vertices(::Backend"KaHyPar", g::Graph, npartiations::Integer; objective="edge_cut", alg="kway", kwargs...)
 
 - default_configuration => "cut_kKaHyPar_sea20.ini"
 - :edge_cut => "cut_kKaHyPar_sea20.ini"
 - :connectivity => "km1_kKaHyPar_sea20.ini"
 - imbalance::Number=0.03
 """
-function partition_vertices(
+function partitioned_vertices(
   ::Backend"KaHyPar",
   g::SimpleGraph,
   npartitions::Integer;
@@ -28,6 +28,6 @@ function partition_vertices(
       kahypar_configurations[(objective=objective, alg=alg)],
     )
   end
-  partitions = @suppress KaHyPar.partition(g, npartitions; configuration, kwargs...)
-  return groupfind(partitions .+ 1)
+  partitioned_verts = @suppress KaHyPar.partition(g, npartitions; configuration, kwargs...)
+  return groupfind(partitioned_verts .+ 1)
 end
diff --git a/src/requires/metis.jl b/src/requires/metis.jl
index d88562d..2c8a319 100644
--- a/src/requires/metis.jl
+++ b/src/requires/metis.jl
@@ -1,33 +1,17 @@
 set_partitioning_backend!(Backend"Metis"())
 
 """
-    partition_vertices(::Backend"Metis", g::AbstractGraph, npartitions::Integer; alg="recursive")
+    partitioned_vertices(::Backend"Metis", g::AbstractGraph, npartitions::Integer; alg="recursive")
 
 Partition the graph `G` in `n` parts.
 The partition algorithm is defined by the `alg` keyword:
  - :KWAY: multilevel k-way partitioning
  - :RECURSIVE: multilevel recursive bisection
 """
-function partition_vertices(
+function partitioned_vertices(
   ::Backend"Metis", g::SimpleGraph, npartitions::Integer; alg="recursive", kwargs...
 )
   metis_alg = metis_algs[alg]
-  partitions = Metis.partition(g, npartitions; alg=metis_alg, kwargs...)
-  return groupfind(Int.(partitions))
+  partitioned_verts = Metis.partition(g, npartitions; alg=metis_alg, kwargs...)
+  return groupfind(Int.(partitioned_verts))
 end
-
-## #=
-##     Metis.partition(G, n; alg = :KWAY)
-## 
-## Partition the graph `G` in `n` parts.
-## The partition algorithm is defined by the `alg` keyword:
-##  - :KWAY: multilevel k-way partitioning
-##  - :RECURSIVE: multilevel recursive bisection
-## =#
-## function partition(g::Metis.Graph, npartitions::Integer)
-##   return Metis.partition(g, npartitions; alg=:KWAY)
-## end
-## 
-## function partition(g::Graph, npartitions::Integer)
-##   return partition(Metis.graph(adjacency_matrix(g)), npartitions)
-## end
diff --git a/test/test_partitionedgraph.jl b/test/test_partitionedgraph.jl
index 35faf72..25bd881 100644
--- a/test/test_partitionedgraph.jl
+++ b/test/test_partitionedgraph.jl
@@ -17,22 +17,32 @@ using Graphs
   nx, ny = 10, 10
   g = named_grid((nx, ny))
 
+  #Partition it column-wise (into a 1D chain)
   partitions = [[(i, j) for j in 1:ny] for i in 1:nx]
   pg = PartitionedGraph(g, partitions)
   @test vertextype(partitioned_graph(pg)) == Int64
   @test vertextype(unpartitioned_graph(pg)) == vertextype(g)
+  @test isa(partitionvertices(pg), Vector{PartitionVertex{Int64}})
+  @test isa(partitionedges(pg), Vector{PartitionEdge{Int64,NamedEdge{Int64}}})
   @test is_tree(partitioned_graph(pg))
   @test nv(pg) == nx * ny
   @test nv(partitioned_graph(pg)) == nx
 
+  #Same partitioning but with a dictionary constructor
   partition_dict = Dictionary([first(partition) for partition in partitions], partitions)
   pg = PartitionedGraph(g, partition_dict)
   @test vertextype(partitioned_graph(pg)) == vertextype(g)
   @test vertextype(unpartitioned_graph(pg)) == vertextype(g)
+  @test isa(partitionvertices(pg), Vector{PartitionVertex{Tuple{Int64,Int64}}})
+  @test isa(
+    partitionedges(pg),
+    Vector{PartitionEdge{Tuple{Int64,Int64},NamedEdge{Tuple{Int64,Int64}}}},
+  )
   @test is_tree(partitioned_graph(pg))
   @test nv(pg) == nx * ny
   @test nv(partitioned_graph(pg)) == nx
 
+  #Partition the whole thing into just 1 vertex
   pg = PartitionedGraph([i for i in 1:nx])
   @test unpartitioned_graph(pg) == partitioned_graph(pg)
   @test nv(pg) == nx
@@ -41,6 +51,25 @@ using Graphs
   @test ne(partitioned_graph(pg)) == 0
 end
 
+@testset "Test Partitioned Graph Partition Edge and Vertex Finding" begin
+  nx, ny, nz = 4, 4, 4
+  g = named_grid((nx, ny, nz))
+
+  #Partition it column-wise (into a square grid)
+  partitions = [[(i, j, k) for k in 1:nz] for i in 1:nx for j in 1:ny]
+  pg = PartitionedGraph(g, partitions)
+  @test Set(partitionvertices(pg)) == Set(partitionvertices(pg, vertices(g)))
+  @test Set(partitionedges(pg)) == Set(partitionedges(pg, edges(g)))
+  @test partitionvertex(pg, (1, 1, 1)) == partitionvertex(pg, (1, 1, nz))
+  @test partitionvertex(pg, (2, 1, 1)) != partitionvertex(pg, (1, 1, nz))
+
+  @test partitionedge(pg, NamedEdge((1, 1, 1) => (2, 1, 1))) ==
+    partitionedge(pg, NamedEdge((1, 1, 2) => (2, 1, 2)))
+  inter_column_edges = NamedEdge.([(1, 1, i) => (2, 1, i) for i in 1:nz])
+  @test length(partitionedges(pg, inter_column_edges)) == 1
+  @test length(partitionvertices(pg, [(1, 2, i) for i in 1:nz])) == 1
+end
+
 @testset "Test Partitioned Graph Vertex/Edge Addition and Removal" begin
   nx, ny = 10, 10
   g = named_grid((nx, ny))