improve performance of IndexAtom a little

src/VectorAffineFunctionAsMatrix.jl

@@ -44,6 +44,26 @@ function to_vaf(vaf_as_matrix::VectorAffineFunctionAsMatrix{T}) where {T}
     return MOI.VectorAffineFunction{T}(vats, vaf_as_matrix.aff.vector)
 end
 
+function to_saf(tape::SparseTape{T}, output_index) where {T}
+    A = tape.operation.matrix
+    m, n = size(A)
+    sats = MOI.ScalarAffineTerm{T}[]
+    rows = SparseArrays.rowvals(A)
+    vals = SparseArrays.nonzeros(A)
+    for j in 1:n # for each column
+        for i in SparseArrays.nzrange(A, j)
+            row = rows[i]
+            row == output_index || continue
+            val = vals[i]
+            push!(sats, MOI.ScalarAffineTerm{T}(val, tape.variables[j]))
+        end
+    end
+    return MOI.ScalarAffineFunction{T}(
+        sats,
+        tape.operation.vector[output_index],
+    )
+end
+
 # method for adding constraints and coverting to standard VAFs as needed
 function MOI_add_constraint(model, f, set)
     return MOI.add_constraint(model, f, set)

src/atoms/IndexAtom.jl

@@ -42,33 +42,45 @@ function evaluate(x::IndexAtom)
     return output(result)
 end
 
-function new_conic_form!(context::Context{T}, x::IndexAtom) where {T}
-    obj = conic_form!(context, only(AbstractTrees.children(x)))
-    m = length(x)
-    n = length(x.children[1])
+function _index(tape::SparseTape{T}, keep_rows::Vector{Int}) where {T}
+    @assert issorted(keep_rows)
+    A = tape.operation.matrix
+    indexed = A[keep_rows, :]
+    af = SparseAffineOperation{T}(indexed, tape.operation.vector[keep_rows])
+    return SparseTape{T}(af, tape.variables)
+end
+_index(tape::Vector, keep_rows::Vector{Int}) = tape[keep_rows]
+
+function _index_real!(
+    obj_size::Tuple,
+    obj_tape::Union{SparseTape{T},SPARSE_VECTOR{T}},
+    x::IndexAtom,
+) where {T}
     if x.inds === nothing
-        sz = length(x.cols) * length(x.rows)
-        J = Vector{Int}(undef, sz)
-        k = 1
-        num_rows = x.children[1].size[1]
-        for c in x.cols
-            for r in x.rows
-                J[k] = num_rows * (convert(Int, c) - 1) + convert(Int, r)
-                k += 1
-            end
-        end
-        index_matrix = create_sparse(T, collect(1:sz), J, one(T), m, n)
+        linear_indices =
+            LinearIndices(CartesianIndices(obj_size))[x.rows, x.cols]
     else
-        index_matrix = create_sparse(
-            T,
-            collect(1:length(x.inds)),
-            collect(x.inds),
-            one(T),
-            m,
-            n,
-        )
+        linear_indices = collect(x.inds)
+    end
+    return _index(obj_tape, vec(linear_indices))
+end
+
+function new_conic_form!(context::Context{T}, x::IndexAtom) where {T}
+    input = x.children[1]
+    if !iscomplex(x) # real case
+        input_tape = conic_form!(context, input)
+        return _index_real!(size(input), input_tape, x)
+    else # complex case
+        input_tape = conic_form!(context, input)
+        re = _index_real!(size(input), real(input_tape), x)
+        im = _index_real!(size(input), imag(input_tape), x)
+        if re isa SPARSE_VECTOR
+            @assert im isa SPARSE_VECTOR
+            return ComplexStructOfVec(re, im)
+        else
+            return ComplexTape(re, im)
+        end
     end
-    return operate(add_operation, T, sign(x), index_matrix, obj)
 end
 
 function Base.getindex(