diff --git a/src/HyperGraphNeuralNetworks.jl b/src/HyperGraphNeuralNetworks.jl
index 0f68b6f..309a889 100644
--- a/src/HyperGraphNeuralNetworks.jl
+++ b/src/HyperGraphNeuralNetworks.jl
@@ -35,7 +35,7 @@ export hypergraph_ids, has_multi_hyperedges
 
 include("core/transform.jl")
 
-export add_self_loops, remove_self_loops, remove_multi_hyperedges
+export add_selfloops, remove_selfloops, remove_multihyperedges
 export rewire_hyperedges, to_undirected
 export combine_hypergraphs, get_hypergraph
 export AbstractNegativeSamplingStrategy, UniformSample, SizedSample, MotifSample, CliqueSample, negative_sample_hyperedge
diff --git a/src/core/transform.jl b/src/core/transform.jl
index 9a91413..7c568aa 100644
--- a/src/core/transform.jl
+++ b/src/core/transform.jl
@@ -1,5 +1,5 @@
 """
-    add_self_loops(
+    add_selfloops(
         hg::HGNNHypergraph{T, D};
         add_repeated_hyperedge::Bool = false
     ) where {T<:Real, D<:AbstractDict{Int,T}}
@@ -9,11 +9,11 @@
 
     NOTE: this function will throw an AssertionError if hg.hedata is not empty
 """
-function add_self_loops(hg::HGNNHypergraph{T, D}; add_repeated_hyperedge::Bool = false) where {T<:Real, D<:AbstractDict{Int,T}}
+function add_selfloops(hg::HGNNHypergraph{T, D}; add_repeated_hyperedge::Bool = false) where {T<:Real, D<:AbstractDict{Int,T}}
     
     @assert isempty(hg.hedata)
     
-    vertices = [1:hg.num_vertices]
+    vertices = collect(1:hg.num_vertices)
     v2he = deepcopy(hg.v2he)    
     he2v = deepcopy(hg.he2v)
 
@@ -46,7 +46,7 @@ function add_self_loops(hg::HGNNHypergraph{T, D}; add_repeated_hyperedge::Bool =
 end
 
 """
-    add_self_loops(
+    add_selfloops(
         hg::HGNNDiHypergraph{T, D};
         add_repeated_hyperedge::Bool = false
     ) where {T<:Real, D<:AbstractDict{Int,T}}
@@ -57,15 +57,15 @@ end
 
     NOTE: this function will throw an AssertionError if hg.hedata is not empty
 """
-function add_self_loops(hg::HGNNDiHypergraph{T, D}; add_repeated_hyperedge::Bool = false) where {T<:Real, D<:AbstractDict{Int,T}}
+function add_selfloops(hg::HGNNDiHypergraph{T, D}; add_repeated_hyperedge::Bool = false) where {T<:Real, D<:AbstractDict{Int,T}}
 
     @assert isempty(hg.hedata)
 
-    vertices = [1:hg.num_vertices]
+    vertices = collect(1:hg.num_vertices)
     v2he_tail = deepcopy(hg.hg_tail.v2he)
     v2he_head = deepcopy(hg.hg_head.v2he)
     he2v_tail = deepcopy(hg.hg_tail.he2v)
-    he2v_head = deepcoyp(hg.hg_head.he2v)
+    he2v_head = deepcopy(hg.hg_head.he2v)
 
     all_he_keys = collect(zip(collect.(keys.(he2v_tail)), collect.(keys.(he2v_head))))
     if add_repeated_hyperedge
@@ -88,9 +88,18 @@ function add_self_loops(hg::HGNNDiHypergraph{T, D}; add_repeated_hyperedge::Bool
         end
     end
 
+    # Construct tail and head hypergraphs
+    hg_tail = Hypergraph{T, Nothing, Nothing, D}(length(v2he_tail), length(he2v_tail))
+    hg_tail.v2he .= v2he_tail
+    hg_tail.he2v .= he2v_tail
+
+    hg_head = Hypergraph{T, Nothing, Nothing, D}(length(v2he_head), length(he2v_head))
+    hg_head.v2he .= v2he_head
+    hg_head.he2v .= he2v_head
+
     return HGNNDiHypergraph(
-        Hypergraph(v2he_tail, he2v_tail, Vector{Nothing}(nothing, length(v2he_tail)), Vector{Nothing}(nothing, length(he2v_tail))),
-        Hypergraph(v2he_head, he2v_head, Vector{Nothing}(nothing, length(v2he_head)), Vector{Nothing}(nothing, length(he2v_head))),
+        hg_tail,
+        hg_head,
         hg.num_vertices,
         length(he2v_tail),
         hg.num_hypergraphs,
@@ -102,11 +111,11 @@ function add_self_loops(hg::HGNNDiHypergraph{T, D}; add_repeated_hyperedge::Bool
 end
 
 """
-    remove_self_loops(hg::HGNNHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
+    remove_selfloops(hg::HGNNHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
 
     Remove self-loops (hyperedges containing only one vertex) from an undirected hypergraph `hg`.
 """
-function remove_self_loops(hg::HGNNHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
+function remove_selfloops(hg::HGNNHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
     @assert isempty(hg.hedata)
     
     v2he = deepcopy(hg.v2he)    
@@ -137,18 +146,18 @@ function remove_self_loops(hg::HGNNHypergraph{T, D}) where {T<:Real, D<:Abstract
 end
 
 """
-    remove_self_loops(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
+    remove_selfloops(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
 
     Remove self-loops (hyperedges where the tail and head both contain only a single, shared vertex `v`) from a
     directed hypergraph `hg`.
 """
-function remove_self_loops(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
+function remove_selfloops(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int,T}}
     @assert isempty(hg.hedata)
 
     v2he_tail = deepcopy(hg.hg_tail.v2he)
     v2he_head = deepcopy(hg.hg_head.v2he)
     he2v_tail = deepcopy(hg.hg_tail.he2v)
-    he2v_head = deepcoyp(hg.hg_head.he2v)
+    he2v_head = deepcopy(hg.hg_head.he2v)
 
     to_remove = Int[]
     for (i, (he_tail, he_head)) in enumerate(zip(he2v_tail, he2v_head))
@@ -164,9 +173,18 @@ function remove_self_loops(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:Abstra
     he2v_tail = he2v_tail[Not(to_remove)]
     he2v_head = he2v_head[Not(to_remove)]
 
+    # Construct tail and head hypergraphs
+    hg_tail = Hypergraph{T, Nothing, Nothing, D}(length(v2he_tail), length(he2v_tail))
+    hg_tail.v2he .= v2he_tail
+    hg_tail.he2v .= he2v_tail
+
+    hg_head = Hypergraph{T, Nothing, Nothing, D}(length(v2he_head), length(he2v_head))
+    hg_head.v2he .= v2he_head
+    hg_head.he2v .= he2v_head
+
     return HGNNDiHypergraph(
-        Hypergraph(v2he_tail, he2v_tail, Vector{Nothing}(nothing, length(v2he_tail)), Vector{Nothing}(nothing, length(he2v_tail))),
-        Hypergraph(v2he_head, he2v_head, Vector{Nothing}(nothing, length(v2he_head)), Vector{Nothing}(nothing, length(he2v_head))),
+        hg_tail,
+        hg_head,
         hg.num_vertices,
         length(he2v_tail),
         hg.num_hypergraphs,
@@ -178,11 +196,11 @@ function remove_self_loops(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:Abstra
 end
 
 """
-    remove_multi_hyperedges(hg::HGNNHypergraph)
+    remove_multihyperedges(hg::HGNNHypergraph)
 
     Remove duplicate hyperedges (hyperedges containing identical vertices) from an undirected hypergraph `hg`.
 """
-function remove_multi_hyperedges(hg::HGNNHypergraph)
+function remove_multihyperedges(hg::HGNNHypergraph)
     unique_vs = Set{Set{Int}}()
 
     v2he = deepcopy(hg.v2he)
@@ -219,12 +237,12 @@ function remove_multi_hyperedges(hg::HGNNHypergraph)
 end
 
 """
-    remove_multi_hyperedges(hg::HGNNHypergraph)
+    remove_multihyperedges(hg::HGNNHypergraph)
 
     Remove duplicate hyperedges, i.e., hyperedges ei = (ti, hi), ej = (tj, hj), where t represents a hyperedge tail, h
     represents a hyperedge head, and ti = tj and hi = hj, from a directed hypergraph `hg`.
 """
-function remove_multi_hyperedges(hg::HGNNDiHypergraph)
+function remove_multihyperedges(hg::HGNNDiHypergraph{T, D}) where {T<:Real, D<:AbstractDict{Int, T}}
     unique_vs = Set{Tuple{Set{Int}, Set{Int}}}()
 
     v2he_tail = deepcopy(hg.hg_tail.v2he)
@@ -250,11 +268,20 @@ function remove_multi_hyperedges(hg::HGNNDiHypergraph)
     he2v_tail = hg.hg_tail.he2v[mask_to_keep]
     he2v_head = hg.hg_head.he2v[mask_to_keep]
 
+    # Construct tail and head hypergraphs
+    hg_tail = Hypergraph{T, Nothing, Nothing, D}(length(v2he_tail), length(he2v_tail))
+    hg_tail.v2he .= v2he_tail
+    hg_tail.he2v .= he2v_tail
+
+    hg_head = Hypergraph{T, Nothing, Nothing, D}(length(v2he_head), length(he2v_head))
+    hg_head.v2he .= v2he_head
+    hg_head.he2v .= he2v_head
+
     hedata = getobs(hg.hedata, mask_to_keep)
 
     HGNNDiHypergraph(
-        Hypergraph(v2he_tail, he2v_tail, Vector{Nothing}(nothing, length(v2he_tail)), Vector{Nothing}(nothing, length(he2v_tail))),
-        Hypergraph(v2he_head, he2v_head, Vector{Nothing}(nothing, length(v2he_head)), Vector{Nothing}(nothing, length(he2v_head))),
+        hg_tail,
+        hg_head,
         hg.num_vertices,
         length(he2v_tail),
         hg.num_hypergraphs,
@@ -280,7 +307,7 @@ end
     is in a given hyperedge 
 
 """
-function to_undirected(hg::HGNNDiHypergraph{T,D}) where {T <: Real, D <: AbstractDict{Int, T}}
+function SimpleDirectedHypergraphs.to_undirected(hg::HGNNDiHypergraph{T,D}) where {T <: Real, D <: AbstractDict{Int, T}}
 
     incidence = Matrix{Union{T, Nothing}}(nothing, nhv(hg), nhe(hg))
 
@@ -288,7 +315,7 @@ function to_undirected(hg::HGNNDiHypergraph{T,D}) where {T <: Real, D <: Abstrac
 
     for row in 1:nhv(hg)
         for column in 1:this_nhe
-            tail_val, head_val = h[row, column]
+            tail_val, head_val = hg[row, column]
             if tail_val === nothing && head_val === nothing
                 incidence[row, column] = nothing
             else
@@ -356,11 +383,11 @@ function combine_hypergraphs(
     if isnothing(hg1.hypergraph_ids) && isnothing(hg2.hypergraph_ids)
         hypergraph_ids = ones(num_vertices)
     elseif isnothing(hg1.hypergraph_ids)
-        hypergraph_ids = cat(ones(hg1.num_vertices), (hg2.hypergraph_ids .+ hgid_increment))
+        hypergraph_ids = vcat(ones(hg1.num_vertices), (hg2.hypergraph_ids .+ hgid_increment))
     elseif isnothing(hg2.hypergraph_ids)
-        hypergraph_ids = cat(hg1.hypergraph_ids, (ones(hg2.num_vertices) .+ hgid_increment))
+        hypergraph_ids = vcat(hg1.hypergraph_ids, (ones(hg2.num_vertices) .+ hgid_increment))
     else
-        hypergraph_ids = cat(hg1.hypergraph_ids, (hg2.hypergraph_ids .+ hgid_increment))
+        hypergraph_ids = vcat(hg1.hypergraph_ids, (hg2.hypergraph_ids .+ hgid_increment))
     end
 
     HGNNHypergraph(
@@ -392,54 +419,13 @@ end
 function combine_hypergraphs(
     hgs::AbstractVector{HGNNHypergraph{T,D}}
 ) where {T <: Real, D <: AbstractDict{Int, T}}
-    num_vs = [hg.num_vertices for hg in hgs]
-    num_hes = [hg.num_hyperedges for hg in hgs]
-    num_hgs = [hg.num_hypergraphs for hg in hgs]
-    
-    vsum = cumsum([0; num_vs])[1:(end - 1)]
-    hesum = cumsum([0; num_hes])[1:(end - 1)]
-    hgsum = cumsum([0; num_hgs])[1:(end - 1)]
-
-    v2hes = [hg.v2he for hg in hgs]
-    he2vs = [hg.he2v for hg in hgs]
-
-    v2he = D[]
-    he2v = D[]
-
-    for (i, v) in enumerate(v2hes)
-        new_v = D()
-        for (he, val) in v
-            new_v[he + hesum[i]] = val
-        end
-        push!(v2he, new_v)
-    end
-
-    for (i, he) in enumerate(he2vs)
-        new_he = D()
-        for (v, val) in he
-            new_he[v + vsum[i]] = val
-        end
-        push!(he2v, new_he)
-    end
-
-    function obtain_hg_inds(hg)
-        hg.hypergraph_ids === nothing ? ones(hg.num_vertices) : hg.hypergraph_ids
+    if length(hgs) == 0
+        return nothing
+    elseif length(hgs) == 1
+        return hgs[1]
     end
     
-    hypergraph_id_vecs = obtain_hg_inds.(hgs)
-    hypergraph_ids = cat_features([nhg .+ inc for (nhg, inc) in zip(hgsum, hypergraph_id_vecs)])
-
-    HGNNHypergraph(
-        v2he,
-        he2v,
-        sum(num_vs),
-        sum(num_hes),
-        sum(num_hgs),
-        hypergraph_ids,
-        cat_features([hg.vdata for hg in hgs]),
-        cat_features([hg.hedata for hg in hgs]),
-        cat_features([hg.hgdata for hg in hgs])
-    )
+    combine_hypergraphs(hgs...)
 end
 
 """
@@ -510,16 +496,25 @@ function combine_hypergraphs(
     if isnothing(hg1.hypergraph_ids) && isnothing(hg1.hypergraph_ids)
         hypergraph_ids = ones(num_vertices)
     elseif isnothing(hg1.hypergraph_ids)
-        hypergraph_ids = cat(ones(hg1.num_vertices), (hg2.hypergraph_ids .+ hgid_increment))
+        hypergraph_ids = vcat(ones(hg1.num_vertices), (hg2.hypergraph_ids .+ hgid_increment))
     elseif isnothing(hg2.hypergraph_ids)
-        hypergraph_ids = cat(hg1.hypergraph_ids, (ones(hg2.num_vertices) .+ hgid_increment))
+        hypergraph_ids = vcat(hg1.hypergraph_ids, (ones(hg2.num_vertices) .+ hgid_increment))
     else
-        hypergraph_ids = cat(hg1.hypergraph_ids, (hg2.hypergraph_ids .+ hgid_increment))
+        hypergraph_ids = vcat(hg1.hypergraph_ids, (hg2.hypergraph_ids .+ hgid_increment))
     end
 
+    # Construct tail and head hypergraphs
+    hg_tail = Hypergraph{T, Nothing, Nothing, D}(length(v2he_tail), length(he2v_tail))
+    hg_tail.v2he .= v2he_tail
+    hg_tail.he2v .= he2v_tail
+
+    hg_head = Hypergraph{T, Nothing, Nothing, D}(length(v2he_head), length(he2v_head))
+    hg_head.v2he .= v2he_head
+    hg_head.he2v .= he2v_head
+
     HGNNDiHypergraph(
-        Hypergraph(v2he_tail, he2v_tail, Vector{Nothing}(undef, num_vertices), Vector{Nothing}(undef, num_hyperedges)),
-        Hypergraph(v2he_head, he2v_head, Vector{Nothing}(undef, num_vertices), Vector{Nothing}(undef, num_hyperedges)),
+        hg_tail,
+        hg_head,
         num_vertices,
         num_hyperedges,
         num_hypergraphs,
@@ -546,74 +541,13 @@ end
 function combine_hypergraphs(
     hgs::AbstractVector{HGNNDiHypergraph{T,D}}
 ) where {T <: Real, D <: AbstractDict{Int, T}}
-    num_vs = [hg.num_vertices for hg in hgs]
-    num_hes = [hg.num_hyperedges for hg in hgs]
-    num_hgs = [hg.num_hypergraphs for hg in hgs]
-    
-    vsum = cumsum([0; num_vs])[1:(end - 1)]
-    hesum = cumsum([0; num_hes])[1:(end - 1)]
-    hgsum = cumsum([0; num_hgs])[1:(end - 1)]
-
-    v2he_tails = [hg.hg_tail.v2he for hg in hgs]
-    v2he_heads = [hg.hg_head.v2he for hg in hgs]
-    he2v_tails = [hg.hg_tail.he2v for hg in hgs]
-    he2v_heads = [hg.hg_tail.he2v for hg in hgs]
-
-    v2he_tail = D[]
-    v2he_head = D[]
-    he2v_tail = D[]
-    he2v_head = D[]
-
-    for (i, v) in enumerate(v2he_tails)
-        new_v = D()
-        for (he, val) in v
-            new_v[he + hesum[i]] = val
-        end
-        push!(v2he_tail, new_v)
-    end
-
-    for (i, v) in enumerate(v2he_heads)
-        new_v = D()
-        for (he, val) in v
-            new_v[he + hesum[i]] = val
-        end
-        push!(v2he_head, new_v)
+    if length(hgs) == 0
+        return nothing
+    elseif length(hgs) == 1
+        return hgs[1]
     end
 
-    for (i, he) in enumerate(he2v_tails)
-        new_he = D()
-        for (v, val) in he
-            new_he[v + vsum[i]] = val
-        end
-        push!(he2v_tail, new_he)
-    end
-
-    for (i, he) in enumerate(he2v_heads)
-        new_he = D()
-        for (v, val) in he
-            new_he[v + vsum[i]] = val
-        end
-        push!(he2v_head, new_he)
-    end
-
-    function obtain_hg_inds(hg)
-        hg.hypergraph_ids === nothing ? ones(hg.num_vertices) : hg.hypergraph_ids
-    end
-    
-    hypergraph_id_vecs = obtain_hg_inds.(hgs)
-    hypergraph_ids = cat_features([nhg .+ inc for (nhg, inc) in zip(hgsum, hypergraph_id_vecs)])
-
-    HGNNDiHypergraph(
-        Hypergraph(v2he_tail, he2v_tail, Vector{Nothing}(undef, num_vertices), Vector{Nothing}(undef, num_hyperedges)),
-        Hypergraph(v2he_head, he2v_head, Vector{Nothing}(undef, num_vertices), Vector{Nothing}(undef, num_hyperedges)),
-        sum(num_vs),
-        sum(num_hes),
-        sum(num_hgs),
-        hypergraph_ids,
-        cat_features([hg.vdata for hg in hgs]),
-        cat_features([hg.hedata for hg in hgs]),
-        cat_features([hg.hgdata for hg in hgs])
-    )
+    combine_hypergraphs(hgs...)
 end
 
 function MLUtils.batch(hgs::AbstractVector{HGNNHypergraph{T,D}}) where {T <: Real, D <: AbstractDict{Int, T}}
@@ -633,28 +567,39 @@ end
     hypergraphs. This is used for unbatching (`MLUtils.unbatch`).
 
 """
-get_hypergraph(hg::HGNNHypergraph, i::Int; kws...) = get_hypergraph(hg, [i]; kws...)
+function get_hypergraph(
+    hg::HGNNHypergraph{T,D}, i::Int; kws...
+    ) where {T <: Real, D <: AbstractDict{Int, T}}
+    
+    get_hypergraph(hg, [i]; kws...)
+end
+
+function get_hypergraph(
+    hg::HGNNHypergraph{T,D},
+    i::AbstractVector{Int};
+    map_vertices::Bool = false
+) where {T <: Real, D <: AbstractDict{Int, T}}
 
-function get_hypergraph(hg::HGNNHypergraph, i::AbstractVector{Int}; map_vertices::Bool = false)
     if hg.hypergraph_ids === nothing
         @assert i == [1]
 
         if map_vertices
-            return hg, 1:(hg.num_vertices)
+            return hg, collect(1:(hg.num_vertices))
         else
             return hg
         end
     end
 
     vertex_mask = hg.hypergraph_ids .∈ Ref(i)
-    vertices = (1:(hg.num_vertices))[vertex_mask]
+    vertices = collect(1:(hg.num_vertices))[vertex_mask]
     vertex_map = Dict(v => vnew for (vnew, v) in enumerate(vertices))
 
     hgmap = Dict(i => inew for (inew, i) in enumerate(i))
     hypergraph_ids = [hgmap[i] for i in hg.hypergraph_ids[vertex_mask]]
 
-    he_mask = all.(keys.(hg.he2v) .∈ Ref(vertices))
-    hyperedges = (1:(hg.num_hyperedges))[he_mask]
+    vset = Set(vertices)
+    he_mask = [issubset(Set(keys(he)), vset) for he in hg.he2v]
+    hyperedges = collect(1:(hg.num_hyperedges))[he_mask]
     hyperedge_map = Dict(he => henew for (henew, he) in enumerate(hyperedges))
 
     he2v = hg.he2v[he_mask]
@@ -685,7 +630,7 @@ function get_hypergraph(hg::HGNNHypergraph, i::AbstractVector{Int}; map_vertices
     num_hyperedges = length(hyperedges)
     num_hypergraphs = length(i)
 
-    HGNNHypergraph(
+    hg_new = HGNNHypergraph{T,D}(
         v2he, he2v,
         num_vertices, num_hyperedges, num_hypergraphs,
         hypergraph_ids,
@@ -711,9 +656,20 @@ end
     Extract one or more hypergraphs (based on `hypergraph_ids`) from an `HGNNHypergraph` containing multiple
     hypergraphs. This is used for unbatching (`MLUtils.unbatch`).
 """
-get_hypergraph(hg::HGNNDiHypergraph, i::Int; kws...) = get_hypergraph(hg, [i]; kws...)
+function get_hypergraph(
+    hg::HGNNDiHypergraph{T, D},
+    i::Int;
+    kws...
+) where {T <: Real, D <: AbstractDict{Int, T}}
+    
+    get_hypergraph(hg, [i]; kws...)
+end
 
-function get_hypergraph(hg::HGNNDiHypergraph, i::AbstractVector{Int}; map_vertices::Bool = false)
+function get_hypergraph(
+    hg::HGNNDiHypergraph{T, D},
+    i::AbstractVector{Int};
+    map_vertices::Bool = false
+) where {T <: Real, D <: AbstractDict{Int, T}}
     if hg.hypergraph_ids === nothing
         @assert i == [1]
 
@@ -725,14 +681,18 @@ function get_hypergraph(hg::HGNNDiHypergraph, i::AbstractVector{Int}; map_vertic
     end
 
     vertex_mask = hg.hypergraph_ids .∈ Ref(i)
-    vertices = (1:(hg.num_vertices))[vertex_mask]
+    vertices = collect(1:(hg.num_vertices))[vertex_mask]
     vertex_map = Dict(v => vnew for (vnew, v) in enumerate(vertices))
 
     hgmap = Dict(i => inew for (inew, i) in enumerate(i))
     hypergraph_ids = [hgmap[i] for i in hg.hypergraph_ids[vertex_mask]]
 
-    he_mask = all.(keys.(hg.hg_tail.he2v) .∈ Ref(vertices)) .* all.(keys.(hg.hg_head.he2v) .∈ Ref(vertices))
-    hyperedges = (1:(hg.num_hyperedges))[he_mask]
+    vset = Set(vertices)
+    he_mask = [
+        issubset(Set(keys(he_tail)), vset) && issubset(Set(keys(he_head)), vset)
+        for (he_tail, he_head) in zip(hg.hg_tail.he2v, hg.hg_head.he2v)
+    ]
+    hyperedges = collect(1:(hg.num_hyperedges))[he_mask]
     hyperedge_map = Dict(he => henew for (henew, he) in enumerate(hyperedges))
 
     he2v_tail = hg.hg_tail.he2v[he_mask]
@@ -783,9 +743,18 @@ function get_hypergraph(hg::HGNNDiHypergraph, i::AbstractVector{Int}; map_vertic
     num_hyperedges = length(hyperedges)
     num_hypergraphs = length(i)
 
-    HGNNDiHypergraph(
-        Hypergraph(v2he_tail, he2v_tail, Vector{Nothing}(undef, num_vertices), Vector{Nothing}(undef, num_hyperedges)),
-        Hypergraph(v2he_head, he2v_head, Vector{Nothing}(undef, num_vertices), Vector{Nothing}(undef, num_hyperedges)),
+    # Construct tail and head hypergraphs
+    hg_tail = Hypergraph{T, Nothing, Nothing, D}(length(v2he_tail), length(he2v_tail))
+    hg_tail.v2he .= v2he_tail
+    hg_tail.he2v .= he2v_tail
+
+    hg_head = Hypergraph{T, Nothing, Nothing, D}(length(v2he_head), length(he2v_head))
+    hg_head.v2he .= v2he_head
+    hg_head.he2v .= he2v_head
+
+    hg_new = HGNNDiHypergraph(
+        hg_tail,
+        hg_head,
         num_vertices, num_hyperedges, num_hypergraphs,
         hypergraph_ids,
         vdata, hedata, hgdata
@@ -798,7 +767,7 @@ function get_hypergraph(hg::HGNNDiHypergraph, i::AbstractVector{Int}; map_vertic
     end
 end
 
-function MLUtils.unbatch(hg::HGNNDiHypergraph)
+function MLUtils.unbatch(hg::HGNNDiHypergraph{T, D}) where {T <: Real, D <: AbstractDict{Int, T}}
     return [get_hypergraph(hg, i) for i in 1:(hg.num_hypergraphs)]
 end
 
@@ -1093,9 +1062,15 @@ function motif_negative_sample(
 
     choices = Set{Set{Int}}()
 
-    adjmat = get_twosection_adjacency_mx(hg)
-    g = SimpleGraph(adjmat)
-    edges = [Set([src(e), dst(e)]) for e in edges(g)]
+    # Necessary because of issues that could arise when a weight is 0
+    mx = Matrix(hg)
+    mx[mx .!== nothing] .= 1.0
+    hg_ = Hypergraph{T, D}(mx)
+
+    adj = get_twosection_adjacency_mx(hg_)
+    
+    g = SimpleGraph(adj)
+    edges = [Set([src(e), dst(e)]) for e in Graphs.edges(g)]
 
     for _ in 1:max_trials
         for _ in 1:(n - length(choices))
@@ -1111,7 +1086,8 @@ function motif_negative_sample(
                 e = rand(rng, e_options)
                 union!(he, e)
             end
-            if length(he) >= size
+
+            if length(he) >= size && !((he in he2v) || (he in choices))
                 push!(choices, he)
             end
         end
@@ -1171,6 +1147,8 @@ function motif_negative_sample(
     sources = [src(e) for e in edges(g)]
     destinations = [dst(e) for e in edges(g)]
 
+    he2v = collect(zip(he2v_tail, he2v_head))
+
     choices = Tuple{Set{Int}, Set{Int}}[]
 
     for _ in 1:max_trials
@@ -1181,7 +1159,7 @@ function motif_negative_sample(
             he_tail = Set{Int}(sources[i])
             he_head = Set{Int}(destinations[i])
 
-            while length(he) < total_size
+            while length(he_tail) + length(he_head) < total_size
                 e_options = [
                     i for i in 1:ne(g) if 
                     (sources[i] ∈ he_tail && destinations[i] ∉ he_head) || (destinations[i] ∈ he_head&& sources[i] ∉ he_tail)
@@ -1196,8 +1174,10 @@ function motif_negative_sample(
                 push!(he_head, destinations[e])
             end
 
-            if length(he_tail) + length(he_head) >= total_size
-                push!(choices, (he_tail, he_head))
+            he = (he_tail, he_head)
+            
+            if length(he_tail) + length(he_head) >= total_size && !((he in he2v) || (he in choices))
+                push!(choices, he)
             end
 
         end
diff --git a/test/runtests.jl b/test/runtests.jl
index 189b935..c39ad27 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -4,6 +4,7 @@ using LinearAlgebra
 using Test
 using Graphs
 using GNNGraphs
+import GNNGraphs: getn, getdata, normalize_graphdata, cat_features, shortsummary
 using MLUtils
 using SimpleHypergraphs
 using SimpleDirectedHypergraphs
@@ -660,8 +661,6 @@ end
 end
 
 @testset "HyperGraphNeuralNetworks query" begin
-    # TODO: do this for directed hypergraphs
-
     hgnn = HGNNHypergraph(uh1.v2he, uh1.he2v, 11, 5, 2, uid1, DataStore(), DataStore(), DataStore())
     dhgnn = HGNNDiHypergraph(
         dh1;
@@ -996,6 +995,308 @@ end
 end
 
 @testset "HyperGraphNeuralNetworks transforms" begin
+    # TODO: directed hyperedges
+
+    hgnn = HGNNHypergraph(uh1.v2he, uh1.he2v, 11, 5, 2, uid1, DataStore(), DataStore(), DataStore())
+    
+    dhgnn = HGNNDiHypergraph(
+        dh1;
+        hypergraph_ids = did1,
+        vdata = nothing,
+        hedata = nothing,
+        hgdata = nothing
+    )
+
+    # add_selfloops
+    hgnn0 = add_selfloops(hgnn)
+    @test hgnn0.num_hyperedges == 16
+    for i in 1:hgnn0.num_vertices
+        @test Dict{Int, Float64}(i => 1.0) in hgnn0.he2v
+    end
+
+    hgnn1 = add_hyperedge(hgnn, DataStore(); vertices=Dict{Int, Float64}(1 => 1.0))
+    hgnn2 = add_selfloops(hgnn1)
+    @test hgnn2.num_hyperedges == 16
+    hgnn2 = add_selfloops(hgnn1; add_repeated_hyperedge=true)
+    @test hgnn2.num_hyperedges == 17
+
+    dhgnn0 = add_selfloops(dhgnn)
+    @test dhgnn0.num_hyperedges == 16
+    he_verts = collect(zip(Set.(keys.(dhgnn0.hg_tail.he2v)), Set.(keys.(dhgnn0.hg_head.he2v))))
+    for i in 1:dhgnn0.num_vertices
+        @test (Set(i), Set(i)) in he_verts
+    end
+
+    dhgnn1 = add_hyperedge(
+        dhgnn,
+        DataStore();
+        vertices_tail=Dict{Int, Float64}(1 => 1.0),
+        vertices_head=Dict{Int, Float64}(1 => 1.0)
+    )
+    dhgnn2 = add_selfloops(dhgnn1)
+    @test dhgnn2.num_hyperedges == 16
+    dhgnn2 = add_selfloops(dhgnn1; add_repeated_hyperedge=true)
+    @test dhgnn2.num_hyperedges == 17
+
+    # remove_selfloops
+    hgnn1 = remove_selfloops(hgnn2)
+    @test hgnn1.num_hyperedges == 5
+
+    dhgnn1 = remove_selfloops(dhgnn2)
+    @test dhgnn1.num_hyperedges == 5
+
+    # remove_multihyperedges
+    hgnn2 = add_hyperedge(hgnn, DataStore(); vertices=Dict{Int, Float64}(1 => 1.0, 2 => 2.0, 4 => 4.0))
+    @test remove_multihyperedges(hgnn2).num_hyperedges == 5
+
+    dhgnn2 = add_hyperedge(
+        dhgnn,
+        DataStore();
+        vertices_tail=Dict{Int, Float64}(1 => 1.0, 2 => 2.0),
+        vertices_head=Dict{Int, Float64}(4 => 4.0)
+    )
+    @test remove_multihyperedges(dhgnn2).num_hyperedges == 5
+
+    # to_undirected
+    @test Set.(keys.(to_undirected(dhgnn).he2v)) == Set.(keys.(hgnn.he2v))
+
+    # combine_hypergraphs / MLUtils.batch
+    hg1 = HGNNHypergraph(
+        [
+            1.0     nothing
+            nothing 2.0
+            nothing 3.0
+        ];
+        hypergraph_ids=[1,2,2],
+        vdata = rand(Float64, 5, 3),
+        hedata = rand(Float64, 5, 2),
+        hgdata = rand(Float64, 5, 2)
+    )
+    hg2 = HGNNHypergraph(
+        [
+            1.0     nothing     1.0
+            nothing 2.0         3.0
+        ];
+        hypergraph_ids=[1,1],
+        vdata = rand(Float64, 5, 2),
+        hedata = rand(Float64, 5, 3),
+        hgdata = rand(Float64, 5, 1)
+    )
+
+    hg_comb1 = combine_hypergraphs(hg1, hg2)
+    @test hg_comb1.num_vertices == 5
+    @test hg_comb1.num_hyperedges == 5
+    @test hg_comb1.num_hypergraphs == 3
+    @test hg_comb1.hypergraph_ids == [1, 2, 2, 3, 3]
+    @test hg_comb1.vdata == cat_features(hg1.vdata, hg2.vdata)
+    @test hg_comb1.hedata == cat_features(hg1.hedata, hg2.hedata)
+    @test hg_comb1.hgdata == cat_features(hg1.hgdata, hg2.hgdata)
+
+    hg_comb2 = combine_hypergraphs(hg1, hg1, hg2, hg2)
+    @test hg_comb2.num_vertices == 10
+    @test hg_comb2.num_hyperedges == 10
+    @test hg_comb2.num_hypergraphs == 6
+    @test hg_comb2.hypergraph_ids == [1, 2, 2, 3, 4, 4, 5, 5, 6, 6]
+    @test hg_comb2.vdata == cat_features([hg1.vdata, hg1.vdata, hg2.vdata, hg2.vdata])
+    @test hg_comb2.hedata == cat_features([hg1.hedata, hg1.hedata, hg2.hedata, hg2.hedata])
+    @test hg_comb2.hgdata == cat_features([hg1.hgdata, hg1.hgdata, hg2.hgdata, hg2.hgdata])
+
+    @test isnothing(combine_hypergraphs(HGNNHypergraph{Float64, Dict{Int,Float64}}[]))
+    @test combine_hypergraphs([hg1]) == hg1
+    @test combine_hypergraphs([hg1, hg1, hg2, hg2]) == hg_comb2
+
+    @test combine_hypergraphs([hg1, hg2]) == batch([hg1, hg2])
+
+    dhg1 = HGNNDiHypergraph(
+        [
+            1.0     nothing
+            nothing 2.0
+            nothing 3.0
+        ],
+        [
+            nothing nothing
+            2.0     nothing
+            nothing 6.0
+        ];
+        hypergraph_ids=[1,2,2],
+        vdata = rand(Float64, 5, 3),
+        hedata = rand(Float64, 5, 2),
+        hgdata = rand(Float64, 5, 2)
+    )
+    dhg2 = HGNNDiHypergraph(
+        [
+            1.0     nothing     nothing
+            nothing 2.0         3.0
+        ],
+        [
+            nothing nothing 1.0
+            2.0     4.0     nothing
+        ];
+        hypergraph_ids=[1,1],
+        vdata = rand(Float64, 5, 2),
+        hedata = rand(Float64, 5, 3),
+        hgdata = rand(Float64, 5, 1)
+    )
+
+    dhg_comb1 = combine_hypergraphs(dhg1, dhg2)
+    @test dhg_comb1.num_vertices == 5
+    @test dhg_comb1.num_hyperedges == 5
+    @test dhg_comb1.num_hypergraphs == 3
+    @test dhg_comb1.hypergraph_ids == [1, 2, 2, 3, 3]
+    @test dhg_comb1.vdata == cat_features(dhg1.vdata, dhg2.vdata)
+    @test dhg_comb1.hedata == cat_features(dhg1.hedata, dhg2.hedata)
+    @test dhg_comb1.hgdata == cat_features(dhg1.hgdata, dhg2.hgdata)
+
+    dhg_comb2 = combine_hypergraphs(dhg1, dhg1, dhg2, dhg2)
+    @test dhg_comb2.num_vertices == 10
+    @test dhg_comb2.num_hyperedges == 10
+    @test dhg_comb2.num_hypergraphs == 6
+    @test dhg_comb2.hypergraph_ids == [1, 2, 2, 3, 4, 4, 5, 5, 6, 6]
+    @test dhg_comb2.vdata == cat_features([dhg1.vdata, dhg1.vdata, dhg2.vdata, dhg2.vdata])
+    @test dhg_comb2.hedata == cat_features([dhg1.hedata, dhg1.hedata, dhg2.hedata, dhg2.hedata])
+    @test dhg_comb2.hgdata == cat_features([dhg1.hgdata, dhg1.hgdata, dhg2.hgdata, dhg2.hgdata])
+
+    @test isnothing(combine_hypergraphs(HGNNDiHypergraph{Float64, Dict{Int,Float64}}[]))
+    @test combine_hypergraphs([dhg1]) == dhg1
+    @test combine_hypergraphs([dhg1, dhg1, dhg2, dhg2]) == dhg_comb2
+
+    @test combine_hypergraphs([dhg1, dhg2]) == batch([dhg1, dhg2])
+
+    # get_hypergraph / MLUtils.unbatch
+    hg1_1 = get_hypergraph(hg1, 1)
+    @test hg1_1 == get_hypergraph(hg1, [1])
+    @test get_hypergraph(hg1, [1,2]) == hg1
+    @test hg1_1.num_vertices == 1
+    @test hg1_1.num_hyperedges == 1
+    @test hg1_1.num_hypergraphs == 1
+    @test get_hypergraph(hg1, 2; map_vertices=true)[2] == [2, 3]
+    @test unbatch(hg1) == [get_hypergraph(hg1, 1), get_hypergraph(hg1, 2)]
+
+    dhg1_1 = get_hypergraph(dhg1, 1)
+    @test dhg1_1 == get_hypergraph(dhg1, [1])
+    @test get_hypergraph(dhg1, [1,2]) == dhg1
+    @test dhg1_1.num_vertices == 1
+    @test dhg1_1.num_hyperedges == 0
+    @test dhg1_1.num_hypergraphs == 1
+    @test get_hypergraph(dhg1, 2; map_vertices=true)[2] == [2, 3]
+    @test unbatch(dhg1) == [get_hypergraph(dhg1, 1), get_hypergraph(dhg1, 2)]
+
+    start_he_keys = Set.(keys.(hgnn.he2v))
+
+    # uniform_negative_sample
+    hgnn_u = negative_sample_hyperedge(hgnn, 3, Xoshiro(42), UniformSample(); max_trials=100)
+    @test hgnn_u.num_vertices == 11
+    @test hgnn_u.num_hyperedges == 3
+    # No hyperedge should be in the original hypergraph
+    for he in hgnn_u.he2v
+        @test Set(keys(he)) ∉ start_he_keys
+    end
+    # All hyperedges should be unique
+    @test length(Set(Set.(keys.(hgnn_u.he2v)))) == 3
+
+    start_he_keys = collect(
+        zip(
+            Set.(keys.(dhgnn.hg_tail.he2v)),
+            Set.(keys.(dhgnn.hg_head.he2v))
+        )
+    )
+
+    dhgnn_u = negative_sample_hyperedge(dhgnn, 3, Xoshiro(42), UniformSample(); max_trials=100)
+    @test dhgnn_u.num_vertices == 11
+    @test dhgnn_u.num_hyperedges == 3
+    # No hyperedge should be in the original hypergraph
+    
+    all_he_inds = Set{Tuple{Set{Int}, Set{Int}}}()
+
+    for (he_tail, he_head) in zip(dhgnn_u.hg_tail.he2v, dhgnn_u.hg_head.he2v)
+        he_inds = (Set(keys(he_tail)), Set(keys(he_head)))
+        @test he_inds ∉ start_he_keys
+        push!(all_he_inds, he_inds)
+    end
+    # All hyperedges should be unique
+    @test length(all_he_inds) == 3
+
+
+    # sized_negative_sample
+    hgnn_s = negative_sample_hyperedge(hgnn, 3, Xoshiro(42), SizedSample(); max_trials=100)
+    @test hgnn_s.num_vertices == 11
+    @test hgnn_s.num_hyperedges == 3
+    for he in hgnn_s.he2v
+        @test Set(keys(he)) ∉ start_he_keys
+    end
+    @test length(Set(Set.(keys.(hgnn_s.he2v)))) == 3
+
+    dhgnn_s = negative_sample_hyperedge(dhgnn, 3, Xoshiro(42), SizedSample(); max_trials=100)
+    @test dhgnn_s.num_vertices == 11
+    @test dhgnn_s.num_hyperedges == 3
+    # No hyperedge should be in the original hypergraph
+
+    all_he_inds = Set{Tuple{Set{Int}, Set{Int}}}()
+
+    for (he_tail, he_head) in zip(dhgnn_s.hg_tail.he2v, dhgnn_s.hg_head.he2v)
+        he_inds = (Set(keys(he_tail)), Set(keys(he_head)))
+        @test he_inds ∉ start_he_keys
+        push!(all_he_inds, he_inds)
+    end
+    # All hyperedges should be unique
+    @test length(all_he_inds) == 3
+
+
+    # motif_negative_sample
+    hgnn_m = negative_sample_hyperedge(hgnn, 3, Xoshiro(42), MotifSample(); max_trials=100)
+    @test hgnn_m.num_vertices == 11
+    @test hgnn_m.num_hyperedges == 3
+    for he in hgnn_m.he2v
+        @test Set(keys(he)) ∉ start_he_keys
+    end
+    @test length(Set(Set.(keys.(hgnn_m.he2v)))) == 3
+
+    dhgnn_m = negative_sample_hyperedge(dhgnn, 3, Xoshiro(42), MotifSample(); max_trials=100)
+    @test dhgnn_m.num_vertices == 11
+    @test dhgnn_m.num_hyperedges == 3
+    # No hyperedge should be in the original hypergraph
+
+    all_he_inds = Set{Tuple{Set{Int}, Set{Int}}}()
+
+    for (he_tail, he_head) in zip(dhgnn_m.hg_tail.he2v, dhgnn_m.hg_head.he2v)
+        he_inds = (Set(keys(he_tail)), Set(keys(he_head)))
+        @test he_inds ∉ start_he_keys
+        push!(all_he_inds, he_inds)
+    end
+    # All hyperedges should be unique
+    @test length(all_he_inds) == 3
+
+    # clique_negative_sample
+    hgnn_c = negative_sample_hyperedge(hgnn, 3, Xoshiro(42), CliqueSample(); max_trials=100)
+    @test hgnn_c.num_vertices == 11
+    @test hgnn_c.num_hyperedges == 3
+    for he in hgnn_c.he2v
+        @test Set(keys(he)) ∉ start_he_keys
+    end
+    @test length(Set(Set.(keys.(hgnn_c.he2v)))) == 3
+
+    @test_throws "negative_sample not implemented for strategy of type CliqueSample" negative_sample_hyperedge(
+        dhgnn,
+        1,
+        Xoshiro(42),
+        CliqueSample()
+    )
+
+    # negative_sample_hyperedge
+    struct NewSample <: AbstractNegativeSamplingStrategy end
+    @test_throws "negative_sample not implemented for strategy of type NewSample" negative_sample_hyperedge(
+        hgnn,
+        1,
+        Xoshiro(42),
+        NewSample()
+    )
+
+    @test_throws "negative_sample not implemented for strategy of type NewSample" negative_sample_hyperedge(
+        dhgnn,
+        1,
+        Xoshiro(42),
+        NewSample()
+    )
 
 end