Feature Request: `flatten_only`

[jonniedie]

One of the biggest problems I have when handling parameters (whether I'm using ComponentArrays, Functors, ParameterHandling, etc.) is that I often only want to optimize (or do some other analysis) over a certain subset of the parameters. And, critically, that subset is going to be constantly changing. It would be nice to have a way to tag parameters for certain operations.

One way to accomplish this with ParameterHandling would be a flatten_only function that digs through the structure and only pulls out elements wrapped in a specified type. The idea is that a user would create a wrapper type and method for value that unwraps the object. I'm not 100% sure this is the best way to handle it (maybe it would be better for ParameterHandling to have a Tagged type that allows users to tag with symbols?), but here is an idea of what it would look like to use:

struct Tunable{T}
    val::T
end

ParameterHandling.value(x::Tunable) = x.val

julia> raw_params = (
       a = 1,
       b = 2.5,
       c = Tunable(3.1),
       d = (
           a = Tunable(4),
           b = 5.1,
       ),
   );

julia> tunable_params, unflatten = flatten_only(Tunable, raw_params);

julia> tunable_params
2-element Vector{Float64}:
 3.1
 4.0

[willtebbutt]

We have what is (I believe) the exact opposite of what you're after! ParameterHandling.fixed allows you to specify something as not being tunable. We really need better docs...

Is it going to be much more convenient for you to specify what's tunable than what's not tunable?

[jonniedie]

Oh cool! That's good to know.

For me in controls engineering work, it tends to be more convenient to specify what's tunable rather than what isn't because simulations will have 1000s of internal parameters but I only need to optimize over a handful of gains or filter coefficients at a time. And sometimes I might want to run a sensitivity analysis over a subset of the parameters that aren't tunables, so it would be nice to pick and choose what I'd like to pull out.

[willtebbutt]

Ah interesting. Do you have a real-world example with that kind of number of parameters? In intruiged to know what your parameter container looks like.

[jonniedie]

I don't have anything on me that I can share, but in general it's going to be a bunch of nested NamedTuples or similar. For our flight simulations, we model vehicles down pretty fine detail, so each vehicle is going to be made up of a pretty complex structure of nested subsystems. The parameter structure basically reflects the nesting of the vehicle subsystems. For example, you might have something like vehicle.first_stage.engine[1].tvc_actuator[1].friction_coefficient and you'd like to run a sensitivity analysis with that and some other select parameters to see what is contributing to some certain flight characteristic. Generally, 100s or 1000s of parameters is pretty small for high-fidelity engineering simulations. More often it's in the 10,000s. Some people I've talked to in the automotive industry have 100,000s of parameters in their models. Most of the analyses or optimizations they're running only need to deal with a small handful of those at a time, though.

[jonniedie]

The more I think about it, tagging with a symbol seems like a better idea than tagging with a wrapper type because sometimes you might want certain values to participate in more than one operation.

[willtebbutt]

The more I think about it, tagging with a symbol seems like a better idea than tagging with a wrapper type because sometimes you might want certain values to participate in more than one operation.

Could you give an example? I'm not quite sure what you mean.

[jonniedie]

Yeah, the more I'm thinking about it, the more I'm not sure it's the right call either. The idea would be that you can have different variables tagged for different reasons. So sometimes you might want to flatten a certain group of parameters to, say, optimize over them. Another time you might want to flatten a different group (that might possibly share some elements with the first group) for a different task.

julia> raw_params = (
       a = 1,
       b = Tagged{(:group1,)}(2.5),
       c = Tagged{(:group1, :group2)}(3.1),
       d = (
           a = Tagged{(:group2,)}(4),
           b = Tagged{(:group2,)}(5.1),
       ),
    );

julia> params, unflatten = flatten_only(:group1, raw_params);

julia> params
2-element Vector{Float64}:
 2.5
 3.1

julia> unflatten(params)
(a = 1, b = 2.5, c = 3.1, d = (a = 4, b = 5.1))

[jonniedie]

Here Tagged would be a type that lives in ParameterHandling.jl and would look something like:

struct Tagged{Syms, Eltype}
    val::Eltype
    Tagged{Syms}(val::Eltype) where {Syms, Eltype} = new{Syms, Eltype}(val)
end
Tagged(val) = Tagged{()}(val)

[willtebbutt]

Hmm this is really interesting.

One way I could imagine implementing your flatten_only as the composition of a function which strips out the tags, and calls fixed on everything else, and flatten.

i.e. something like

remove_tags(x::Tagged, group) = is_in_group(x, :group) ? x.val : fixed(x.val)
remove_tags(x::Real, group) = fixed(x)
remove_tags(x::NamedTuple, group) = map(val -> remove_tags(val, group), x)

etc. Then flatten(remove_tags(raw_params)) would give you what you need I believe. Do you agree?

[jonniedie]

Yeah, that would definitely work.

[willtebbutt]

Excellent. If you've got the time to make a PR, I would be happy to review. I might get around to implementing this at some point, but probably not in the near term.