Creating ics on grid for higher dimensions

src/mapping/attractor_mapping.jl

@@ -11,7 +11,8 @@ export AttractorMapper,
     automatic_Δt_basins,
     extract_attractors,
     subdivision_based_grid,
-    SubdivisionBasedGrid
+    SubdivisionBasedGrid,
+    generate_ics_on_grid
 
 #########################################################################################
 # AttractorMapper structure definition
@@ -166,14 +167,41 @@ corresponding to the state space partitioning indicated by `grid`.
 """
 function basins_of_attraction(mapper::AttractorMapper, grid::Tuple; kwargs...)
     basins = zeros(Int32, map(length, grid))
-    I = CartesianIndices(basins)
-    A = StateSpaceSet([generate_ic_on_grid(grid, i) for i in vec(I)])
+    A = generate_ics_on_grid(grid, basins)
     fs, labels = basins_fractions(mapper, A; kwargs...)
     attractors = extract_attractors(mapper)
     vec(basins) .= vec(labels)
     return basins, attractors
 end
 
+"""
+Generates initial conditions on `grid`, used typically in `basins_of_attraction`. A common
+use case is studying basins in small-dimensional systems. In this case,
+`generate_ic_on_grid`, using CartesianIndices on all dimensions, works well. A particular
+case occurs when studying basins in high-dimensional systems, in which `grid` is
+high-dimensional but only has varying values in two dimensions. The first function will
+not work, so a special function `_generate_ics_on_grid_vary_few_dims` is used to avoid the
+problem. This separation is controlled by:
+
+* `min_num_fixed_dims`: minimum number of fixed dimensions above which it becomes worth it
+  to use the special function to generate ics.
+"""
+function generate_ics_on_grid(grid, basins; min_num_fixed_dims=4)
+    grid_lengths = length.(grid)
+    number_fixed_dims = length(findall(len->len<=1, grid_lengths))
+    if number_fixed_dims >= min_num_fixed_dims
+        return _generate_ics_on_grid_vary_few_dims(grid)
+    else
+        return _generate_ics_on_grid(grid, basins)
+    end
+end
+
+function _generate_ics_on_grid(grid, basins)
+    I = CartesianIndices(basins)
+    A = StateSpaceSet([generate_ic_on_grid(grid, i) for i in vec(I)])
+    return A
+end
+
 # Type-stable generation of an initial condition given a grid array index
 @generated function generate_ic_on_grid(grid::NTuple{B, T}, ind) where {B, T}
     gens = [:(grid[$k][ind[$k]]) for k=1:B]
@@ -183,6 +211,30 @@ end
     end
 end
 
+"""
+    Generates ics on grid for a reduced grid, which contains only dimensions that are varying.
+    Then construct the full-dimension vector from that. 
+"""
+function _generate_ics_on_grid_vary_few_dims(grid::NTuple{B, T}) where {B, T} #only worth it if only a few varyign dims
+    grid_lengths = length.(grid)
+    idxs_varying_dims = findall(len->len>1, grid_lengths)
+    reduced_grid = grid[idxs_varying_dims]
+    I_reduced = CartesianIndices(length.(reduced_grid))
+    A_reduced = [generate_ic_on_grid(reduced_grid, i) for i in vec(I_reduced)]
+
+    ics_fixed = [grid_dim[1] for grid_dim in grid]
+    A = _expand_A(A_reduced, ics_fixed, idxs_varying_dims)
+    return Dataset(A)
+end
+
+function _expand_A(vec_reduced, ic_fixed, idxs_varying_dims) where {A}
+    vec = [deepcopy(ic_fixed) for _ in vec_reduced]
+    for i in eachindex(vec)
+        vec[i][idxs_varying_dims] .= vec_reduced[i]
+    end
+    return vec
+end
+
 #########################################################################################
 # Includes
 #########################################################################################

test/mapping/ics_on_grid.jl

@@ -0,0 +1,46 @@
+using Attractors, Test
+
+@testset "Generating ics" begin
+    function _construct_grid_varying_2_dims(idx_dim_1, idx_dim_2, ics_dim_1, ics_dim_2, ics_base_dims)
+        _grid = [ic:ic for ic in ics_base_dims]
+        _grid[idx_dim_1] = ics_dim_1
+        _grid[idx_dim_2] = ics_dim_2
+        return Tuple(_grid)
+    end
+
+    function construct_grid_varying_2_dims(size_grid, num_ics_1, num_ics_2, idx_dim_1, idx_dim_2) #silly function just to avoid repeating code
+        ics_dim_1 = range(1, 10; length=num_ics_1)
+        ics_dim_2 = range(1, 10; length=num_ics_2)
+        ics_base = Float64.(collect(1:size_grid))
+        grid = _construct_grid_varying_2_dims(idx_dim_1, idx_dim_2, ics_dim_1, ics_dim_2, ics_base)
+        return grid, ics_dim_1, ics_dim_2
+    end
+
+   @testset "Varying in 10 dimensions" begin
+        xg = range(1, 10; length=3)
+        grid = ntuple(x->xg, 10)
+        A = @time Attractors._generate_ics_on_grid_vary_few_dims(grid) 
+        basins = zeros(Int32, map(length, grid))
+        A_previous = @time Attractors._generate_ics_on_grid(grid, basins)
+        @test A == A_previous
+    end
+
+    @testset "Varying in 2 dimensions - small" begin
+        grid, _, _ = construct_grid_varying_2_dims(10, 3, 4, 1, 2)
+        A = @time Attractors._generate_ics_on_grid_vary_few_dims(grid)
+        basins = zeros(Int32, map(length, grid))
+        A_previous = @time Attractors._generate_ics_on_grid(grid, basins) 
+    end
+
+    @testset "Varying in 2 dimensions - big" begin
+        num_ics_1 = 100; num_ics_2 = 10; idx_dim_1 = 1; idx_dim_2 = 2
+        grid, ics_dim_1, ics_dim_2 = construct_grid_varying_2_dims(100, num_ics_1, num_ics_2, idx_dim_1, idx_dim_2)
+        A = @time Attractors._generate_ics_on_grid_vary_few_dims(grid)
+        @test length(A) == num_ics_1 * num_ics_2
+        @test all(ics_dim_1 .== [A[idx][idx_dim_1] for idx in 1:num_ics_1])
+        @test all(ics_dim_2[1] .== [A[idx][idx_dim_2] for idx in 1:num_ics_1])
+        @test all(ics_dim_2[2] .== [A[idx][idx_dim_2] for idx in (num_ics_1+1):2*num_ics_1])
+    end
+
+end
+