Automatic differentiation with subset of parameters in ModelingToolkit

Thanks Chris, that’s helpful!

For others who are looking for a potential solution; here’s what I ended up doing, inspired by Chris’s comment.

After marking parameters as ‘tunable’ you can extract their indices easily:

function split_parameters(params)
    indexmap = [p => i for (p,i) in zip(params, 1:length(params))]
    subset_indexmap = eltype(indexmap)[]
    for (p, i) in indexmap
        istunable(p) && push!(subset_indexmap, p => i)
    end
    return indexmap, subset_indexmap
end

Then one can rebuild the p vector as follows, replacing the relevant subset parameter values:

p_idxs_map, p_subset_idx_map = split_params(parameters(sys))
p_subset_indices = last.(p_subset_idx_map)
p = prob.p
p_subset = ...
idx_to_subset_idx(i) = findfirst(isequal(i), p_subset_indices)
new_p = [i in p_subset_indices ? p_subset[idx_to_subset_idx(i)] : p[i] for i in 1:length(p)]

which should (hopefully) work fine with Zygote also.

1 Like