Add axes support for SnakePattern (#55)

* Add `axes` for `SnakePattern`

* Bump version

* More fixes

* SnakePattern tests

Author: lkdvos

Project.toml

@@ -1,7 +1,7 @@
 name = "MPSKitModels"
 uuid = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
 authors = ["Maarten Van Damme", "Lukas Devos", "Gertian Roose", "Klaas Gunst"]
-version = "0.4.2"
+version = "0.4.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/lattices/snakepattern.jl

@@ -10,12 +10,21 @@ end
 
 SnakePattern(lattice) = SnakePattern(lattice, identity)
 
+Base.axes(lattice::SnakePattern) = axes(lattice.lattice)
+Base.isfinite(::Type{SnakePattern{N,G}}) where {N,G} = isfinite(G)
+
 function linearize_index(snake::SnakePattern, i...)
     return snake.pattern(linearize_index(snake.lattice, i...))
 end
 
-vertices(snake::SnakePattern) = vertices(snake.lattice)
-nearest_neighbours(snake::SnakePattern) = nearest_neighbours(snake.lattice)
+function vertices(snake::SnakePattern)
+    return map(x -> LatticePoint(x.coordinates, snake), vertices(snake.lattice))
+end
+function nearest_neighbours(snake::SnakePattern)
+    return map(nearest_neighbours(snake.lattice)) do (x, y)
+        return LatticePoint(x.coordinates, snake), LatticePoint(y.coordinates, snake)
+    end
+end
 bipartition(snake::SnakePattern) = bipartition(snake.lattice)
 
 """

test/lattices.jl

@@ -201,3 +201,43 @@ end
         @test allunique(NNN)
     end
 end
+
+@testset "SnakePattern" begin
+    base_lattice = FiniteCylinder(4, 16)
+    base_V = vertices(base_lattice)
+    base_NN = nearest_neighbours(base_lattice)
+
+    lattice = SnakePattern(base_lattice, identity)
+
+    V = vertices(lattice)
+    @test all(zip(V, base_V)) do (x, y)
+        return linearize_index(x) == linearize_index(y)
+    end
+
+    @test length(base_lattice) == length(lattice)
+    @test lattice[1, 1] == first(V)
+
+    NN = nearest_neighbours(lattice)
+    @test all(zip(NN, base_NN)) do (x, y)
+        return linearize_index.(Tuple(x)) ==
+               linearize_index.(Tuple(y))
+    end
+
+    pattern(i) = length(lattice) - i + 1
+    lattice = SnakePattern(base_lattice, pattern)
+
+    V = vertices(lattice)
+    @test all(zip(V, base_V)) do (x, y)
+        return linearize_index(x) == pattern(linearize_index(y))
+    end
+
+    @test length(base_lattice) == length(lattice)
+    @test lattice[1, 1] == first(V)
+
+    base_NN = nearest_neighbours(base_lattice)
+    NN = nearest_neighbours(lattice)
+    @test all(zip(NN, base_NN)) do (x, y)
+        return linearize_index.(Tuple(x)) ==
+               pattern.(linearize_index.(Tuple(y)))
+    end
+end