Add a link arg for broadcast_dayofweek

CDCgov · Jun 5, 2024 · 24c2707 · 24c2707
1 parent f269dd1
commit 24c2707
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/EpiAware/src/EpiLatentModels/broadcast/helpers.jl b/EpiAware/src/EpiLatentModels/broadcast/helpers.jl
@@ -3,12 +3,15 @@ Constructs a `BroadcastLatentModel` appropriate for modelling the day of the wee
 
 # Arguments
 - `model::AbstractTuringLatentModel`: The latent model to be repeated.
+- `link::Function`: The link function to transform the latent model _before_ broadcasting
+to periodic weekly. Default is `x -> 7 * softmax(x)` which implements constraint
+of the sum week effects to be 7.
 
 # Returns
 - `BroadcastLatentModel`: The broadcast latent model.
 "
-function broadcast_dayofweek(model::AbstractTuringLatentModel)
-    transformed_model = TransformLatentModel(model, x -> 7 * softmax(x))
+function broadcast_dayofweek(model::AbstractTuringLatentModel; link = x -> 7 * softmax(x))
+    transformed_model = TransformLatentModel(model, link)
     return BroadcastLatentModel(transformed_model, 7, RepeatEach())
 end