diff --git a/src/TO.jl b/src/TO.jl
index e0eaa05e..5266048b 100644
--- a/src/TO.jl
+++ b/src/TO.jl
@@ -145,13 +145,10 @@ function adaptiveTOCollocationStiffnessMatrix(ctx::GridContext{2}, flow_maps, ti
     else
         N = length(times)
     end
-    num_real_points = ctx.n
-    if on_torus || on_cylinder
-        num_real_points = ctx.n - length(bdata.periodic_dofs_from)
-    end
-    flow_map_images = zeros(Vec{2}, (N,num_real_points))
+
+    flow_map_images = zeros(Vec{2}, (N,ctx.n))
     if times === nothing
-        for i in 1:num_real_points
+        for i in 1:ctx.n
             if flow_map_mode == 0
                 flow_map_images[1,i] = Vec{2}(flow_maps(ctx.grid.nodes[i].x))
             else
@@ -159,7 +156,7 @@ function adaptiveTOCollocationStiffnessMatrix(ctx::GridContext{2}, flow_maps, ti
             end
         end
     else
-        for i in 1:num_real_points
+        for i in 1:ctx.n
             if flow_map_mode == 0
                 flow_map_images[:,i] = Vec{2}.(flow_maps(ctx.grid.nodes[i].x, times))
             else
diff --git a/test/test_fem.jl b/test/test_fem.jl
index 5811202b..337e466c 100644
--- a/test/test_fem.jl
+++ b/test/test_fem.jl
@@ -155,4 +155,15 @@ end
 
     λ, = get_smallest_eigenpairs(D, M, 3)
     @test all(<(sqrt(eps())), λ)
+
+    LL, UR = (0., 0.), (1., 1.)
+    gs = 10
+    ctx, _ = regularTriangularGrid((gs, gs), LL, UR);
+    predicate = (p1, p2) -> peuclidean(p1, p2, [1.0,Inf]) < 2e-10
+    bdata = BoundaryData(ctx, predicate, [])
+
+    @test !isnothing(
+        adaptiveTOCollocationStiffnessMatrix(
+            ctx, identity; on_cylinder=true, bdata)
+       )
 end