@@ -75,8 +75,16 @@ function Base.length(Ψ::GradVector)
7575end
7676
7777
78- function Base. size (G:: GradgenOperator )
79- return Base. size (G. G)
78+ function Base. size (O:: GradgenOperator{num_controls,GT,CGT} ) where {num_controls,GT,CGT}
79+ return (num_controls + 1 ) .* size (O. G)
80+ end
81+
82+
83+ function Base. size (
84+ O:: GradgenOperator{num_controls,GT,CGT} ,
85+ dim:: Integer
86+ ) where {num_controls,GT,CGT}
87+ return (num_controls + 1 ) * size (O. G, dim)
8088end
8189
8290
@@ -88,6 +96,9 @@ function Base.similar(G::GradgenOperator{num_controls,GT,CGT}) where {num_contro
8896 return GradgenOperator {num_controls,GT,CGT} (similar (G. G), similar (G. control_deriv_ops))
8997end
9098
99+ function Base. eltype (O:: GradgenOperator{num_controls,GT,CGT} ) where {num_controls,GT,CGT}
100+ return promote_type (eltype (GT), eltype (CGT))
101+ end
91102
92103function Base. copyto! (dest:: GradgenOperator , src:: GradgenOperator )
93104 copyto! (dest. G, src. G)
@@ -161,23 +172,23 @@ function *(
161172end
162173
163174
164- @inline function convert_gradgen_to_dense (G)
175+ @inline function convert_gradgen_to_dense (G:: GradGenerator )
165176 N = size (G. G)[1 ]
166177 L = length (G. control_derivs)
167178 G_full = zeros (eltype (G. G), N * (L + 1 ), N * (L + 1 ))
168179 convert_gradgen_to_dense! (G_full, G)
169180end
170181
171182
172- @inline function convert_gradgen_to_dense! (G_full, G)
183+ @inline function convert_gradgen_to_dense! (G_full, G:: GradGenerator )
173184 N = size (G. G)[1 ]
174185 L = length (G. control_derivs)
175- @inbounds for i = 1 : L+ 1
176- G_full[(i- 1 )* N+ 1 : i* N, (i- 1 )* N+ 1 : i* N] .= G. G
186+ @inbounds for i = 1 : ( L+ 1 )
187+ G_full[(( i- 1 )* N+ 1 ) : ( i* N) , (( i- 1 )* N+ 1 ) : ( i* N) ] .= G. G
177188 end
178189 # Set the control-derivatives in the last (block-)column
179190 @inbounds for i = 1 : L
180- G_full[(i- 1 )* N+ 1 : i* N, L* N+ 1 : ( L+ 1 )* N] .= G. control_derivs[i]
191+ G_full[(( i- 1 )* N+ 1 ) : ( i* N), ( L* N+ 1 ) : (( L+ 1 )* N) ] .= G. control_derivs[i]
181192 end
182193 return G_full
183194end
195206 N = length (Ψ. state)
196207 L = length (Ψ. grad_states)
197208 @inbounds for i = 1 : L
198- Ψ_full[(i- 1 )* N+ 1 : i* N] .= Ψ. grad_states[i]
209+ Ψ_full[(( i- 1 )* N+ 1 ) : ( i* N) ] .= Ψ. grad_states[i]
199210 end
200- @inbounds Ψ_full[L* N+ 1 : ( L+ 1 )* N] .= Ψ. state
211+ @inbounds Ψ_full[( L* N+ 1 ) : (( L+ 1 )* N) ] .= Ψ. state
201212 return Ψ_full
202213end
203214
206217 N = length (Ψ. state)
207218 L = length (Ψ. grad_states)
208219 @inbounds for i = 1 : L
209- Ψ. grad_states[i] .= Ψ_full[(i- 1 )* N+ 1 : i* N]
220+ Ψ. grad_states[i] .= Ψ_full[(( i- 1 )* N+ 1 ) : ( i* N) ]
210221 end
211- @inbounds Ψ. state .= Ψ_full[L* N+ 1 : ( L+ 1 )* N]
222+ @inbounds Ψ. state .= Ψ_full[( L* N+ 1 ) : (( L+ 1 )* N) ]
212223 return Ψ
213224end
214225
@@ -225,8 +236,8 @@ function Base.convert(
225236 L = num_controls
226237 N = length (vec) ÷ (L + 1 ) # dimension of state
227238 @assert length (vec) == (L + 1 ) * N
228- grad_states = [convert (T, vec[(i- 1 )* N+ 1 : i* N]) for i = 1 : L]
229- state = convert (T, vec[L* N+ 1 : ( L+ 1 )* N])
239+ grad_states = [convert (T, vec[(( i- 1 )* N+ 1 ) : ( i* N) ]) for i = 1 : L]
240+ state = convert (T, vec[( L* N+ 1 ) : (( L+ 1 )* N) ])
230241 return GradVector {num_controls,T} (state, grad_states)
231242end
232243
@@ -236,8 +247,7 @@ function Base.Array{T}(G::GradgenOperator) where {T}
236247 𝟘 = zeros (T, N, M)
237248 μ = G. control_deriv_ops
238249 block_rows = [
239- hcat ([𝟘 for j = 1 : i- 1 ]. .. , Array {T} (G. G), [𝟘 for j = i+ 1 : L]. .. , Array {T} (μ[i]))
240- for i = 1 : L
250+ hcat ([𝟘 for j = 1 : (i- 1 )]. .. , Array {T} (G. G), [𝟘 for j = (i+ 1 ): L]. .. , Array {T} (μ[i])) for i = 1 : L
241251 ]
242252 last_block_row = hcat ([𝟘 for j = 1 : L]. .. , Array {T} (G. G))
243253 return Base. Array {T} (vcat (block_rows... , last_block_row))
0 commit comments