Add components method to Abelian group (#1632)

* Add another getindex method To be able to call `x[begin:end]`. Co-authored-by: Tommy Hofmann <[email protected]>
thofma · Oct 18, 2024 · 542c23e · 542c23e
1 parent 6b7d0e0
commit 542c23e
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/docs/src/manual/abelian/elements.md b/docs/src/manual/abelian/elements.md
@@ -25,6 +25,7 @@ parent(x::FinGenAbGroupElem)
 ## Access
 
 ```@docs
+getindex(x::FinGenAbGroupElem, v::AbstractVector{Int})
 getindex(x::FinGenAbGroupElem, i::Int)
 ```
 

diff --git a/src/GrpAb/Elem.jl b/src/GrpAb/Elem.jl
@@ -146,6 +146,26 @@ function getindex(x::FinGenAbGroupElem, i::Int)
   return x.coeff[1, i]
 end
 
+@doc raw"""
+    getindex(x::FinGenAbGroupElem, v::AbstractVector{Int}) -> Vector{ZZRingElem}
+
+Returns the $i$-th components of the element $x$ where $i \in v$.
+
+!!! note
+    This function is inefficient since the elements are internally stored using ZZMatrix but this function outputs a vector.
+"""
+function getindex(x::FinGenAbGroupElem, v::AbstractVector{Int})
+  return [x.coeff[1, i] for i in v]
+end
+
+function Base.firstindex(x::FinGenAbGroupElem)
+  return Int(1)
+end
+
+function Base.lastindex(x::FinGenAbGroupElem)
+  return ngens(parent(x))
+end
+
 ################################################################################
 #
 #  Comparison

diff --git a/test/GrpAb/Elem.jl b/test/GrpAb/Elem.jl
@@ -6,12 +6,14 @@
     a = @inferred FinGenAbGroupElem(G, N)
     @test parent(a) == G
     @test a.coeff == N
+    @test a[begin:end] == [0, 0, 0]
 
     G = @inferred abelian_group([3, 0])
     N = FlintZZ[1 1]
     a = @inferred FinGenAbGroupElem(G, N)
     @test @inferred parent(a) == G
     @test a.coeff == N
+    @test a[begin:end] == [1, 1]
 
     N = matrix(FlintZZ, 1, 2, [ 1, 1 ])
     a = @inferred G(N)
@@ -21,6 +23,7 @@
     a = @inferred G(N)
     @test @inferred parent(a) == G
     @test a.coeff == transpose(N)
+    @test a[begin:end] == [1, 1]
   end
 
   @testset "Generators" begin