New non-invasive interface for specifying parameters on arbitrary types

[bgroenks96]

@rafaqz Since we're talking about making breaking changes to the package anyway, I would like to submit another slightly more fundamental change (or possibly just extension) to the ModelParameters interface.

As you might recall, I recently incorporated ModelParameters into SpeedyWeather. However, one longstanding drawback of the Param interface is the need to intrusively modify existing struct types to assign separate type parameters for each field that may optionally be assigned a Param. This has the unfortunate effect of greatly increasing type complexity.

We wanted to avoid this in SpeedyWeather but still specify Params for all relevant fields on model component types. To achieve this, I developed an extension of ModelParameters that defines a parameters(obj; kwargs...) method which recursively constructs a nested named tuple of parameters defined on each type, e.g:

parameters(obj::MyType) = (
    fieldA = Param(obj.fieldA, desc="Field A"),
    fieldB = Param(obj.fieldB, desc="Field B"),
    ...
)

This solution has the benefit of being both flexible and non-invasive, since it does not require the Params to be set directly as field values but rather constructed only on demand in a separate type. However, it had the disadvantage of requiring the separate implementation of the parameters method for each type.

To ameliorate this, I developed a macro @parameterized which looks something like this (taking an example here from SpeedyWeather):

@parameterized @kwdef mutable struct Earth{NF<:AbstractFloat} <: AbstractPlanet

    "angular frequency of Earth's rotation [rad/s]"
    @param rotation::NF = DEFAULT_ROTATION
    
    "gravitational acceleration [m/s^2]"
    @param gravity::NF = DEFAULT_GRAVITY (bounds=Nonnegative,)
    
    "switch on/off daily cycle"
    daily_cycle::Bool = true
    
    "Seconds in a daily rotation"
    length_of_day::Second = Hour(24)

    "switch on/off seasonal cycle"
    seasonal_cycle::Bool = true

    "Seconds in an orbit around the sun"
    length_of_year::Second = Day(365.25)
    
    "time of spring equinox (year irrelevant)"
    equinox::DateTime = DateTime(2000, 3, 20) 

    "angle [˚] rotation axis tilt wrt to orbit"
    @param axial_tilt::NF = 23.4 (bounds=-90..90,)

    "Total solar irradiance at the distance of 1 AU [W/m²]"
    @param solar_constant::NF = 1365 (bounds=Nonnegative,)
end

All fields marked with @param (which is only a syntax marker not a real macro) will have parameters defined in an automatically generated dispatch of parameters. The named tuple following the field still permits the inclusion of any arbitrary number of additional properties for each parameter (e.g. bounds, units, etc.).

As an added bonus, a description field is automatically populated with the docstring for the struct field, thereby reducing duplication.

The parameters can then be handled completely separately from the original type itself, which is actually kind of nice. The separate parameter structure can be modified either via a Model type, with a ComponentArray or Setfield or however the user wants to handle it. The updated parameters are then applied via a reconstruct method (confusingly not the same as Flatten.reconstruct since its a nested data structure) which recursively updates the values using ConstructionBase.setproperties. So far this is type stable with the help of a generated function dispatch.

I actually quite like this system, and I think it has some major advantages over the current ModelParameters system. It also could easily accommodate #66 since we could add @const and/or @var tags to distinguish between Var and Const.

If you also like it, then I think there are three options here:

1. We migrate the code from SpeedyWeather into ModelParameters and have the @parameterized/parameters interface as a complimentary alternative to the usual direct Param field assignment and params approach. The downside to this is that we would have two different interfaces, which is maybe confusing, and the overlap between method names (e.g. params and parameters) which perform different functions is also not great.
2. We simply adopt this approach as the new primary interface for ModelParameters, possibly maintaining some backwards compatibility via a deprecated params method or something like that.
3. We keep ModelParameters as it is and I can instead put this code into another package that explicitly extends ModelParameters.

Looking forward to your feedback!

[rafaqz]

Yes I was thinking the same reading over your PRs, I could use something like this too in some cases.

But how are you actually defining the models and parameter objects, and using them? Some complete example code would help me understand it a bit more.

Also, one of the reasons I use objects rather than all NamedTuples in models is for constructors that validate parameter values at model definition rather than during model run. Doesn't this approach lose that?

[bgroenks96]

But how are you actually defining the models and parameter objects, and using them?

Do you want just the struct definitions? Or also the generated parameters methods?

Also, one of the reasons I use objects rather than all NamedTuples in models is for constructors that validate parameter values at model definition rather than during model run. Doesn't this approach lose that?

Yes and no. You're still free to define constructors on the struct types themselves which do this validation. Since it uses ConstructionBase to reconstruct the model types, it would also be possible to modify the constructor interface via constructorof. However, for @kwdef types, I currently auto-generate a dispatch for setproperties! that just uses the default keyword constructor. This is mostly because sometimes we have default values that are defined based on the preceding fields.

The trickier part would be if you wanted to validate parameter values already in the separate NamedTuple-like structure. I don't currently have a way to do this. I am of course open to suggestions!

Also, parameters actually returns an AbstractModel type SpeedyParams (obviously the name could be generalized) which wraps the NamedTuple and implements the Model interface, which hopefully we can try again to expand and provide actual table-like behavior. In principle, we could also support some kind of pre-validation using this type? But I'll have to think about exactly how to do that.

[bgroenks96]

I guess actually it would be pretty trivial to implement basic validation like checking bounds in the parameter container.

[rafaqz]

I just wanted to see a miny version of a whole workflow to understand what it's for.

I have a similar situation from other peoples code but I was intending to change it, like I'm still not sure that it's ever good design to have the model and params separate. But I'm open to it.

[bgroenks96]

Hey @rafaqz sorry for forgetting about this.

Here is an excerpt from the Speedy development docs. Does this give you a sufficient idea of the workflow?

Parameter handling

SpeedyWeather also provides automated parameter handling for all models and subcomponents via an extension of ModelParameters.jl. Parameters can be automatically collected via the parameters method:

spectral_grid = SpectralGrid(trunc=23, nlayers=1) 
model = Barotropic(; spectral_grid)
params = parameters(model)

# output (truncated)
Parameters:
┌───────┬──────────────────────────┬──────────┬────────────┬──────────────────────────────┬─────────────────────────────────────┬──────────────────────────────────────────────────────────────────────────────────────────┐       
│   idx │                fieldname │      val │  component │               componentttype │                              bounds │                                                                                     desc │       
│ Int64 │                   Symbol │  Float32 │     Symbol │                     DataType │ IntervalSets.TypedEndpointsInterval │                                                                                   String │       
├───────┼──────────────────────────┼──────────┼────────────┼──────────────────────────────┼─────────────────────────────────────┼──────────────────────────────────────────────────────────────────────────────────────────┤       
│     1 │                 rotation │  7.29e-5 │     planet │               Earth{Float32} │       -Inf .. Inf (open) (RealLine) │                                            angular frequency of Earth's rotation [rad/s] │       
│     2 │                  gravity │     9.81 │     planet │               Earth{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                                       gravitational acceleration [m/s^2] │       
│     3 │               axial_tilt │     23.4 │     planet │               Earth{Float32} │                           -90 .. 90 │                                                angle [˚] rotation axis tilt wrt to orbit │       
│     4 │           solar_constant │   1365.0 │     planet │               Earth{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                    Total solar irradiance at the distance of 1 AU [W/m²] │       
│     5 │         mol_mass_dry_air │  28.9649 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                                                            molar mass of dry air [g/mol] │       
│     6 │          mol_mass_vapour │  18.0153 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                                                       molar mass of water vapour [g/mol] │       
│     7 │            heat_capacity │   1004.0 │ atmosphere │     EarthAtmosphere{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                           specific heat at constant pressure cₚ [J/K/kg] │       
│     8 │                 R_vapour │  461.524 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                                          specific gas constant for water vapour [J/kg/K] │       
│     9 │                mol_ratio │  0.62197 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                       Ratio of gas constants: dry air / water vapour, often called ε [1] │       
│    10 │              μ_virt_temp │ 0.607794 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │ Virtual temperature Tᵥ calculation, Tᵥ = T(1 + μ*q), humidity q, absolute tempereature T │       
│    11 │                        κ │ 0.285911 │ atmosphere │     EarthAtmosphere{Float32} │       -Inf .. Inf (open) (RealLine) │                                         = R_dry/cₚ, gas const for air over heat capacity │       
│    12 │            water_density │   1000.0 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                                                                    water density [kg/m³] │       
│    13 │ latent_heat_condensation │  2.501e6 │ atmosphere │     EarthAtmosphere{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                                       latent heat of condensation [J/kg] │       
│    14 │  latent_heat_sublimation │  2.801e6 │ atmosphere │     EarthAtmosphere{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                                        latent heat of sublimation [J/kg] │       
│    15 │                 pres_ref │ 100000.0 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                                                          surface reference pressure [Pa] │       
│    16 │                 temp_ref │    288.0 │ atmosphere │     EarthAtmosphere{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                                        surface reference temperature [K] │       
│    17 │         moist_lapse_rate │    0.005 │ atmosphere │     EarthAtmosphere{Float32} │       -Inf .. Inf (open) (RealLine) │                                   reference moist-adiabatic temperature lapse rate [K/m] │       
│    18 │           dry_lapse_rate │   0.0098 │ atmosphere │     EarthAtmosphere{Float32} │       -Inf .. Inf (open) (RealLine) │                                     reference dry-adiabatic temperature lapse rate [K/m] │       
│    19 │          layer_thickness │   8500.0 │ atmosphere │     EarthAtmosphere{Float32} │        0.0 .. Inf (open) (HalfLine) │                                          layer thickness for the shallow water model [m] │       
│    20 │                 strength │  3.0e-12 │    forcing │      KolmogorovFlow{Float32} │       -Inf .. Inf (open) (RealLine) │                                                      [OPTION] Strength of forcing [1/s²] │       
│    21 │               wavenumber │      8.0 │    forcing │      KolmogorovFlow{Float32} │        0.0 .. Inf (open) (HalfLine) │                    [OPTION] Wavenumber of forcing in meridional direction (pole to pole) │       
│    22 │                        c │   1.0e-7 │       drag │ LinearVorticityDrag{Float32} │ 0.0 .. Inf (closed-open) (HalfLine) │                                                          [OPTION] drag coefficient [1/s] │       
└───────┴──────────────────────────┴──────────┴────────────┴──────────────────────────────┴─────────────────────────────────────┴──────────────────────────────────────────────────────────────────────────────────────────┘

The returned SpeedyParams object implements the Model interface from ModelParmaeters.jl which allows you to interact with the parameter metadata in tablar form. For example, we could extract the values of the parameters with params[:,:val] or the bounds with params[:,:bounds]. Subsets of parameters can also be extracted by indexing params with one or more String variable names (or prefxes), e.g:

param_subset = params[["planet.gravity", "atmosphere.heat_capacity"]]

# output (truncated)

Parameters:
┌───────┬───────────────┬─────────┬────────────┬──────────────────────────┬───────────────────────────────────────┬────────────────────────────────────────────────┐
│   idx │     fieldname │     val │  component │           componentttype │                                bounds │                                           desc │
│ Int64 │        Symbol │ Float32 │     Symbol │                 DataType │ DomainSets.HalfLine{Float64, :closed} │                                         String │
├───────┼───────────────┼─────────┼────────────┼──────────────────────────┼───────────────────────────────────────┼────────────────────────────────────────────────┤
│     1 │       gravity │    9.81 │     planet │           Earth{Float32} │   0.0 .. Inf (closed-open) (HalfLine) │             gravitational acceleration [m/s^2] │
│     2 │ heat_capacity │  1004.0 │ atmosphere │ EarthAtmosphere{Float32} │   0.0 .. Inf (closed-open) (HalfLine) │ specific heat at constant pressure cₚ [J/K/kg] │
└───────┴───────────────┴─────────┴────────────┴──────────────────────────┴───────────────────────────────────────┴────────────────────────────────────────────────┘

Vectorizing parameters

Many sensitivity analysis, optimization, or uncertainty quantification algorithms require the parameters to be supplied as one or more vectors of values. SpeedyParams provides a dispatch for Base.vec that flattens the model parameters into a ComponentVector:

param_vec = vec(params)

# output

ComponentVector{Float32}(planet = (rotation = 7.29f-5, gravity = 9.81f0, axial_tilt = 23.4f0, solar_constant = 1365.0f0), atmosphere = (mol_mass_dry_air = 28.9649f0, mol_mass_vapour = 18.0153f0, heat_capacity = 1004.0f0, R_vapour = 461.52438f0, mol_ratio = 0.62197006f0, μ_virt_temp = 0.60779446f0, κ = 0.2859107f0, water_density = 1000.0f0, latent_heat_condensation = 2.501f6, latent_heat_sublimation = 2.801f6, pres_ref = 100000.0f0, temp_ref = 288.0f0, moist_lapse_rate = 0.005f0, dry_lapse_rate = 0.0098f0, layer_thickness = 8500.0f0), forcing = (strength = 3.0f-12, wavenumber = 8.0f0), drag = (c = 1.0f-7))

ComponentVectors behave like normal Arrays but additionally allow you to access the components following the original nested structure in the model, e.g. param_vec.planet.solar_constant will extract the solar constant parameter from the Earth component.

We can use the resulting parameter vector to calculate sensitivities over a single time step:

initialize!(simulation)
run!(simulation, period=Day(10))
(; Δt, Δt_sec) = simulation.model.time_stepping
ps = parameters(model)
pvec = vec(ps)
dp = zero(pvec)
dprogn = one(progn) # shadow for the prognostic variabels
ddiagn = make_zero(diagn) # shadow for the diagnostic variables

function timestep_with_new_params!(progn, diagn, dt, model, p)
    new_model = SpeedyWeather.reconstruct(model, p)
    SpeedyWeather.timestep!(progn, diagn, dt, new_model)
    return nothing
end

autodiff(Reverse, timestep_with_new_params!, Const, Duplicated(progn, dprogn), Duplicated(diagn, ddiagn), Const(dt), Duplicated(model, make_zero(model)))

[bgroenks96]

We would like to move this parameter handling code out of Speedy soon, so if you could let me know whether or not you want it to live in ModelParameters that would be great. Otherwise, I would create a separate package for it.

[rafaqz]

I love this syntax, have been meaning to add it too:

... params[:,:val] or the bounds with params[:,:bounds]

This is a bit clunkier but still cool:

param_subset = params[["planet.gravity", "atmosphere.heat_capacity"]]

If there was more nesting would it be params["outer.planet.gravity"] ?

The component vectors compatability is also good to have, but I think it should use ComponentVector(model) rather than vec as its easier to read and understand in my opinion. ComponentVector has some performance and type stability issues in my experience (That make me hesitant to use it for vec).

Also, ComponentVector(obj) is clearer about what you are getting than vec, and can go in an extension so its not a hard dependency.

All the above is good stuff that we want here.

Now the @param macros are also fine and may help some people. Jonestly I cant see many character savings, and macros add magic so you really want to justify using them - this package was partly a response to the magic overkill I wrote into FieldMetadata.jl and this is a pretty strong return to that syntax. But I'm fine to have them.

What I'm worried about is added complexity from duplicating Flatten.jl internals here without significant gains. Flatten has been very little work for many years now, and this sounds like some work to understand, significant moving parts that can break, and a requirement to use macros. (Keeping things simple is really high priority to me now as I've got a lot more maintenance than I can practically do)

To understand your reasoning - I think I need to understand this statement:

However, one longstanding drawback of the Param interface is the need to intrusively modify existing struct types to assign separate type parameters for each field that may optionally be assigned a Param. This has the unfortunate effect of greatly increasing type complexity

You mean because Param types can vary? Could you just standardise Param types like Model does in the constructor when you define it with the @parametrizee macro? Then there can be one T. That should be relatively trivial, and we don't have to rewrite internals. Of course, you will have to swap <:AbstractFloat to <:AbstractNumber too.

Is that a possible compromise?

(otherwise maybe you are talking about a new package, its really totally different? Like it doens't even use Param? It might even be the perfect sweet spot between this and FieldMetdata.jl, I dont know. But its certainly giving "lisp curse"...)

[bgroenks96]

If there was more nesting would it be params["outer.planet.gravity"] ?

Yes.

The component vectors compatability is also good to have, but I think it should use ComponentVector(model) rather than vec as its easier to read and understand in my opinion. ComponentVector has some performance and type stability issues in my experience (That make me hesitant to use it for vec).

Fair enough.

Now the @param macros are also fine and may help some people. Jonestly I cant see many character savings, and macros add magic so you really want to justify using them - this package was partly a response to the magic overkill I wrote into FieldMetadata.jl and this is a pretty strong return to that syntax

So firstly, let me clarify that @param is technically not a macro, but rather a syntax marker for the @parameterized macro. It just looks like a macro so as not to be to be too weird...

Secondly, it's important to note that the purpose of @param is not to save characters. Rather, it's to allow the user to specify which fields should be treated as parameters without changing the actual type/value of the parameter in the struct. Without @param, this would need to be done manually in a dispatch of parameters for the struct type. The @parameterized macro automates adding that method dispatch and also allows us to use the docstring directly as metadata, which I think is really nice.

I do agree though that macro magic should be avoided as much as possible, I just didn't really see an alternative here. I also agree that this is similar to FieldMetadata, though I think it's a bit better :)

What I'm worried about is added complexity from duplicating Flatten.jl internals here without significant gains.

I don't understand? What duplication of Flatten.jl internals are there here? You mean SpeedyWeather.reconstruct?

Flatten has been very little work for many years now, and this sounds like some work to understand, significant moving parts that can break, and a requirement to use macros.

It's actually not too bad, it uses all of the same basic functionality as ModelParameters. The @parameterized macro is the worst part.

You mean because Param types can vary? Could you just standardise Param types like Model does in the constructor when you define it with the @parametrizee macro? Then there can be one T. That should be relatively trivial, and we don't have to rewrite internals. Of course, you will have to swap <:AbstractFloat to <:AbstractNumber too.

Homogenizing the Param fields partially solves this problem, but not fully, since even with homogenized Param fields they can still take on heterogeneous types (e.g. if you use IntervalSets to specify bounds or distribution types for a prior, etc).

However, I think the bigger disadvantage to the ModelParameters approach is that it results in an unintuitive constructor interface. Let me try to explain.

As a user, I would usually like to be able to do something like Earth(rotation=1.0) when creating the type, or @set earth.rotation = 1.0 if modifying the (immutable) object. Param breaks this because now this field is no longer a Param and will not show up as a parameter anymore when collecting with params. It also means that all metadata associated with that parameter is lost (for that instance), and of course you can't expect the user to know all of that information by heart unless they are already very experienced with the model.

In this way, ModelParameters kind of forces you into always modifying parameters ex post facto via Model or update. That's fine for many use cases, but it is a trade-off, I think. This is why in my practical applications of ModelParameters I always ended up inverting the workflow and manually marking fields as Param on construction that I wanted to optimize/vary. This is also OK for some use cases, but I think we would ideally also like to be able to clearly communicate to the user an overview of all model parameters and their associated metadata.

So it's really a problem of top-down vs. bottom-up parameter specification. I think @parameterized basically solves this problem. You get a non-invasive interface that lets you interact with the original struct normally while also having all of the parameter information available via parameters.

I hope that makes it more clear!

[rafaqz]

You can still homogenise mixed Interval or distribution types, they will just have to be abstract and boxed somewhere. Like the type parameter of Param can be a Union.

As a user, I would usually like to be able to do something like Earth(rotation=1.0) when creating the type, or @set earth.rotation = 1.0 if modifying the (immutable) object. Param breaks this because now this field is no longer a Param and will not show up as a parameter anymore when collecting with params

Yeah, seems that can be annoying in some contexts, and this is really the key part of the difference, and where what you're doing is more similar to FieldMetadata.jl. But from my perspective with ModelParameters I made this specific change intentionally! Because making some field not be a param, or even just changing the bounds or priors after it was defined was pretty annoying with FieldMetadata.jl. It was possible but kind of ugly and hard to understand for newbies, and involved redefining existing methods (behind a magic macro of course).

• I work on a lot of models where what is a free parameter or constant is not fixed by the designer, but remains somewhat flexible for the user to define. We need to be able to easily stop the optimiser looking at a bunch of parameters in objects that we don't have access to the code or macro for.
• Another case is where some Params are thrown loosely in a NamedTuple that doesn't even have a constructor to put your macros on.
• Part of the original design was that you can put Param objects in Tuples, NamedTupes, and existing structs that were not intended to hold them when they were designed. Like you can put RealParam in a distribution.

[bgroenks96]

You can still homogenise mixed Interval or distribution types, they will just have to be abstract and boxed somewhere. Like the type parameter of Param can be a Union.

Yeah but then you will lose type stability; this would almost certainly break an optimization routine using Enzyme.

Because making some field not be a param, or even just changing the bounds or priors after it was defined was pretty annoying with FieldMetadata.jl.

I think in this system it's not really that hard? You either need to redefine or add an additional dispatch for the parameters method, or just take a subset of the comprehensive parameter set. I have come to really prefer the latter approach.

I work on a lot of models where what is a free parameter or constant is not fixed by the designer, but remains somewhat flexible for the user to define.

So do I!

We need to be able to easily stop the optimiser looking at a bunch of parameters in objects that we don't have access to the code or macro for.

That's where I think I would prefer the subset approach over modifying field types.

Another case is where some Params are thrown loosely in a NamedTuple that doesn't even have a constructor to put your macros on.

I guess this can be solved either by defining a wrapper type for the named tuple or by choosing to implement the parameters dispatch directly instead of using the macro.

Regardless, I actually don't think it would be all that crazy to support both workflows. There are ways in which they can complement each other; e.g. we could have a withparams method that reconstructs a model with the parameters returned by the parameters method, and the parameters method could have a default dispatch that returns the nested parameter structure of a model type with Param fields.

We just need a coherent naming scheme for the methods. params and parameters are too close.

[bgroenks96]

One idea would be parameters -> parametersof and the current params -> flatparams or flatten_params.

[rafaqz]

You either need to redefine or add an additional dispatch for the parameters method

This is FieldMetadata.jl... to me its a pretty hefty requirement for a new user.

or just take a subset of the comprehensive parameter set

Can you explain this option? where would it happen?