Add @brm macro

[nsiccha]

Discuss https://github.com/PTWaade/BayesianRegressionModels.jl/blob/ns/macro/scripts/macro.jl, @penelopeysm @PTWaade @seabbs.

[nsiccha]

I'd have thought github would include the code - but apparently not! Here's another attempt:

https://github.com/PTWaade/BayesianRegressionModels.jl/blob/40d3650e26af91fefa3092b0b96b1340cfb193a6/scripts/macro.jl#L5-L22

[PTWaade]

Example inspired by Niko's work and mostly meant to clarify!

@brm begin
    
    #Generating the BMI with a very simple formula
    @linear_combination begin
        BMI = 1    
        @df dBMI
    end
    
    #Checking the BMI against the measured BMI (just a simple likelihood)
    @regression_submodel "BMI_measurement_error" begin
        DistributionLikelihood(Normal, 
                            # Struggling with this. 
                            # It has to know that it should take the output of the operation above as the BMI mean,
                            # and take the BMI_measured from the dBMI dataset
                            dist_args = (μ = BMI, σ = 1), 
                            observations = dBMI.BMI_measured
                        )

        #Something like this could maybe show that it comes from a previous outcome
        @previous_outcome BMI 
    end

    #This linear combination generates the mean of the predicted performance
    @linear_combination begin
        performance_mean = 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter)
        
        #This specifies how Age should be expanded. Can also be a custom type with appropriate dispatch.
        @expansion (Age_first, Age_second) = Age, PolynomialExpansion(order = 2)
        
        @df dpmean
    end

    #This linear combination generates the sd of the predicted performance
    @linear_combination begin
        performance_sd = 1 + Age * BMI + Age:BMI + (1 + Age * BMI | Subject)

        #This specified which operator using for the interaction between Age and BMI. Can be a custom type with appropriate syntax
        @interaction Age_BMI_max (Age,BMI) MaxOperator()

        @df dpsd

    end

    #This evaluates the likelihood of the perf
    @regression_submodel "Performance_likelihood" begin
        DistributionLikelihood(Normal, 
                            dist_args = (μ = performance_mean, σ = performance_sd), 
                            observations = dpmean.Performance
                        )

        @previous_outcome performance_mean
        @previous_outcome performance_sd inv_link = exp
    end

    #Configurations
    @config begin
        
        ## Maybe the configurations for a random effect factor should 
        #Group specification cannot differ between different linear combinations, and is general for a given random effect factor
        @group Subject, by = ClinicalGroup

        #Block specifications cannot differn between linear combinations, and is general for a given random effect factor
        @block Subject Block1 [performance_mean.Treatment, performance_sd.Age_BMI_max]
        @block Subject Block2 [performance_sd.Age, performance_sd.BMI]

    end

    #Settings that are called in every operation (submodel and linear combination) if not overwritten by specific settings in those operations. 
    @defaults begin
        
        #Something with priors
        
    end


    ## ONE THING THAT THIS EXAMPLE DOESN'T HAVE IS A VARIABLE GENERATED BY A SUBMODEL ##

end

[PTWaade]

Primnary conversation currently on Slack; I mostly wanted here to distinguihs the linear combinations from the regression submodels
As you can see, there are still things missing, and how to specify the data is hard
I would imagine that poelpe can still write the simple formulas, and that they just expand into this format, but there might be much better ways to do it!

[nsiccha]

I think these would be the "rules" of the @brm macro:

• Any assignment operations will broadcast the RHS over the rows of the data frame in question and assign to the LHS such that the LHS can (from that point onwards?) be used as any other column of the data frame (e.g. Age_first, Age_second = ploynomial_expand(Age; order=2))
  • That example is actually a bit more complicated than necessary - behind the scenes, what would have to be happening is that for each row/value of Age, ploynomial_expand(Age; order=2) returns a 2-tuple. If we just broadcasted this over the data frame rows, we'd get a vector of 2-tuples, but we actually want a 2-tuple of vectors, to unsplat(?) into Age_first, Age_second
• Any sampling statement (~) will (roughly) do one of the following:
  • If the LHS is in the data frame (or can be computed from columns in it?), it's treated as a "likelihood" statement, meaning
    a) if the topmost call on the RHS does not return a Distribution, it will be wrapped in something like Normal(rhs, something_automatic?) before we proceed to b),
    b) if the topmost call on the RHS does return a Distribution, the joint log likelihood will be computed by summing the rowwise (broadcasted) log likelihoods
  • If the LHS is not in the data frame, it's treated as a "regression" statement, i.e. (from that point onwards?) the LHS can be used as any other data frame column, but the values will depend on model parameters via the functional form implied by the RHS formula. Furthermore,
    a) if the topmost call on the RHS returns a Distribution I don't know yet how to handle this?
    b) if the topmost call on the RHS does not return a distribution, "there's no measurement noise", i.e. the value of the LHS will just be "what comes out of the RHS".

It's currently a bit unclear how we'd reliably know when the topmost call returns a distribution - but maybe that is actually quite easy. Also, at some point we discussed the question of when to parse terms as formulas and when not to. Should e.g. y ~ Normal(x, 1) parse the arguments to Normal as formulas? I think we may need a way to turn formula parsing on/off via some function/macro call.

@seabbs @penelopeysm @PTWaade

[nsiccha]

I think transforming (something like) this

@brm begin 
    Age_first, Age_second = ploynomial_expand(Age; order=2)
    performance_mean ~ 1 + Age_first * Treatment + Age_second + (1 + Treatment | Subject) + (1 + Age_first | Experimenter)
    log(performance_sd) ~ 1 + Age * BMI + max(Age, BMI) + (1 + Age * BMI | Subject)
    Performance ~ Normal(performance_mean, performance_sd)
end

to an expression (somewhat) like this

model(__df__) = begin
    (; Age, Treatment, Subject, Experimenter, BMI) = data(__df__)
    (; performance_mean, performance_sd, Performance) = maybedata(__df__)
    (Age_first, Age_second) = (ensurecols)(ploynomial_expand, Age; order = 2)
    performance_mean = ((maybedists)(; force = isdata(performance_mean)))(1 + Age_first * Treatment + Age_second + ((1 + Treatment) | Subject) + ((1 + Age_first) | Experimenter))
    log(performance_sd) = ((maybedists)(; force = isdata(log(performance_sd))))(1 + Age * BMI + (ensurecols)(max, Age, BMI) + ((1 + Age * BMI) | Subject))
    Performance = ((maybedists)(; force = isdata(Performance)))(Normal, performance_mean, performance_sd)
    BRM(; Age_first, Age_second, Age, performance_mean, Treatment, Subject, Experimenter, performance_sd, BMI, Performance)
end

should allow us to do everything that we need - thoughts, @penelopeysm?

[nsiccha]

log(performance_sd) = ... would obviously not work and would have to be changed to something slightly more clever - but I don't think much more clever.

[PTWaade]

So this is awesome work Niko!
Only osme thoughts on the functionality side per your comments with the rules:

I think y ~ x + z should be shorthand for two lines:
y_mean = x + z (a "linear combination")
y ~ Normal(y_mean, sd) (which I have implemented as happening in a Turing submodel).
in addition, the sd must either be the output of a previous line, be a fixed value, or have a prior.

So with that in mind, the way I implemented it differs a bit from the rules you made (so we will change one of the two!)

I've made it so that every line (I called them "operations") can either be a linear combination, as above, or a Turing submodel. Both linear combinations and Turing submodels can generate outputs (i.e., the lhs). Outputs must be accessible by subsequent linear combinations and Turing submodels. I split that up in two, so that outputs can be used to update the Predictors design matrix (which can be used by subsequent linear combinations) or be stored in an outputs object (which is passed to subsequent Turing submodels).
Note that there are two core uses that the TUring submodels cover: evaluating a likelihood, or generating some new values probabilistically.

To me, splitting things up (I used = and ~ for linear combinations and Turing submodels, but happy to do whatever) simplified things a lot. Then it is a higher-layer functionality that its possible to just write y ~ x + z, which is then split into two parts.

Okay so happy to change everything, but just important to align on this point I think, and to make sure the macro and the implementation are both able to handle the full flexibility!

[PTWaade]

And as I said in Slack, the single distribution likelihood is just one special case of likelihood, but we want to be more general than that, hence the Turing submodel.

Otherwise only exciting. Standing by for being prompted :)

[PTWaade]

Mock model for my specific usecase:

cognitive_parameter_1 = predictor1 + predictor2 + (1|p|gr(subjID, by = clinicalgroup))
cognitive_parameter2 = predictor1 + predictor3 + (1|p|gr(subjID, by = clinicalgroup)) 

This needs to be able to include both random effect groups and correlations across regressions.
Notably, it should not include a likelihood.
Rather, it should be used to generate a Turing model that returns the columns of cognitive_parameter1 and cognitive_parameter2. This will be used as a submodel in a larger Turing model, where the cognitive parameters are used for other things.
An arbitrary number of cognitive parmaters must be allowed, and arbitrary regressions must be allowed.

[PTWaade]

Second mock model for my specific usecase:

physiological_measurement1 ~ cognitive_state1 + other_predictor + (1|subjID)

Here, the cognitive_state1 will be generated by a Turing submodel called before the regression, but should still be used in the regression as a predictor like any other. Must be able to allow as many cognitive_state predictorsd as desired, and also multiple outcomes, and in general any type of regression.

[PTWaade]

Third mock model for my specific usecase (complex, in the future):

Same as the first mock model, but where the clinicalgroup membership is generated in a previous Turing submodel. This is a discrete variable, so a non-gradient sampler must be used for this one (but NUTS or whatever can be used for the regression coefficients)

[PTWaade]

Fourth mock model:

outcome1 ~ monotonic(predictor1) ...

Here predictor 1 is a categorical variable. It would be ncie to use Paul's approach (which we discussed) where there is a single beta giving the full effect size of going from the lowest to the highest level, and then a simplex defining the differences between each step

[PTWaade]

Finally, I would like to combine mock models 1 and 2 with a Turing model of my own in-between (so the outcome of mock model 1 is used in my custom Turing mdoel, which in turn generates predictors for mock model 3).

It is also in general important that the outcomes (I guess all the lhs of the model) are returned by the regression model so they can be accessed when it is used as a submdoel within a larger Turing model.

[PTWaade]

Mock Turing model that the brm model will be used inside (dummy code, not optimised). This covers usecases 1 and 2, so brm could be used twice inside the model.

Note that, importantly, for the first case here I do not want any likelihood inside the brm model at all - I just want to calculate the linear combinations of predictors and coeffiicents, and then handle the likleihood in the larger model.

@model function full_cognitive_model(cognitive_model, behavioural_data, conditioned_regression_model, post_behavioural_model)

   #This should give a set of different cognitive parameters (i.e. columns) for each participant in an experiment (i.e., each row) 
   #If the regression model outputs more columns than needed, there coukld be a tiny wrapper here selecting the appropriate columns
   cognitive_parameters ~ to_submodel(conditioned_regression_model)
   
   #Initialise storage for cognitive state
   cognitive_states = []
   
   #Go through each row (i.e., each subject)
   for (subject_idx, subject_parameters) in enumerate(zip(cognitive_parameters...))
       
       #Cognitive modelling is essentially agent-based modelling
       agent = initialise_agent(cognitive_model, subject_parameters)
       
       #Extract the empirical observations and actions for a participant
       subject_observations, subject_actions = behavioura_data[i]
       
       #For each trial in the experiment
       for (subject_action_t, subject_observation_t)
       
          #Generate a probability disitribution over the action at that timestep
          predicted_action_dist_t = agent.cognitive_model!(subject_observation_t)
          
          #Evaluate the probability of the real action (sample a new one if it is missing)
          subject_action_t ~ predicted_action_dist_t
          
          #Store the action (since some models depend on their previous actions)
          store_action!(agent, subject_action_t)
       
       end
       
       #Store the trajectory of cognitive states (which are updated on each trial)
       push!(cognitive_states, get_cognitive_state_history(agent))
       
   end
   
   #Optionally continue with a different model (such as a regression relating cognitive states to neuroimaging data).
   physiological_data ~ to_submodel(post_behavioural_model(cognitive_states, otherstuff))
   
end
```julia

[PTWaade]

I think we also want to be able to do something like:

a ~ custom_submodel_1()
b ~ a + z
p ~ custom_submodel_2(a, b, z)

i.e., where a custom submodel generates some value (not via a regression) which is used in a different submodel (which is also not a regression). In other words, we might want to allow storing and passing values in another way than through updatecols!, so that we don't waste time updating columns in a design matrix when not necessary,a nd so that we allow for non-tabular values too

[nsiccha]

@PTWaade you can make github add syntax highlighting by having "julia" after the first three backticks:

# Use three ` instead of two
``julia
a=1+1
``

[PTWaade]

@PTWaade you can make github add syntax highlighting by having "julia" after the first three backticks:

# Use three ` instead of two
``julia
a=1+1
``

Ah - thank you :)