add onehot_batch

src/bit_str.jl

@@ -1,4 +1,4 @@
-export BitStr, @bit_str, bcat, bit_literal, bit, to_location
+export BitStr, @bit_str, bcat, bit_literal, bit, to_location, onehot, onehot_batch
 
 """
     BitStr{T}
@@ -288,6 +288,9 @@ Returns an onehot vector of type `Vector{T}`, where the `bit_str`-th element is
 onehot(::Type{T}, n::BitStr) where T = onehot(T, length(n), n.val)
 onehot(n::BitStr) = onehot(Float64, n)
 
+onehot_batch(::Type{T}, n::BitStr, nbatch::Int) where T = onehot_batch(T, length(n), n.val, nbatch)
+onehot_batch(n::BitStr, nbatch::Int) = onehot_batch(Float64, n, nbatch)
+
 # conversions
 for IntType in [:Int8, :Int16, :Int32, :Int64, :Int128, :BigInt]
     @eval Base.convert(::Type{$IntType}, x::BitStr) = $IntType(x.val)

src/utils.jl

@@ -1,4 +1,4 @@
-export bsizeof, bdistance, bit_length, onehot, log2i, hypercubic, log2dim1, indices_with
+export bsizeof, bdistance, bit_length, onehot, onehot_batch, log2i, hypercubic, log2dim1, indices_with
 # NOTE: all binary specified operations begin with b
 
 """
@@ -29,6 +29,20 @@ end
 
 onehot(nbits::Int, x::Integer) = onehot(Float64, nbits, x)
 
+"""
+    onehot_batch([T=Float64], nbits, x::Integer, nbatch::Int)
+
+Create a batch of onehot vector in type `Matrix{T}`, where the last dimension is
+the batch dimension, and the index `x + 1` is `1`.
+"""
+function onehot_batch(::Type{T}, nbits::Int, x::Integer, nbatch::Int) where T
+    v = zeros(T, 1 << nbits, nbatch)
+    v[x + 1, :] .= 1
+    return v
+end
+
+onehot_batch(nbits::Int, x::Integer, nbatch::Int) = onehot_batch(Float64, nbits, x, nbatch)
+
 """
     unsafe_sub(a::UnitRange, b::NTuple{N}) -> NTuple{N}
 

test/utils.jl

@@ -3,6 +3,7 @@ using Test, BitBasis
 @test bsizeof(ind) == sizeof(Int) * 8
 @test onehot(ComplexF64, 2, 2) == [0, 0, 1, 0]
 @test bdistance(1,7) == 2
+@test onehot_batch(ComplexF64, bit"01", 2) == transpose(ComplexF64[0 1 0 0;0 1 0 0])
 @test log2dim1(rand(4, 4)) == log2i(4)
 
 @testset "log2i" begin