Merge pull request #780 from mimiframework/broadcast

Broadcast

Author: lrennels

docs/src/tutorials/tutorial_3.md

@@ -54,13 +54,10 @@ You have now successfully downloaded MimiDICE2010 to your local machine.
 
 The next step is to run DICE using the provided API for the package:
 
-```jldoctest tutorial2; output = false, filter = r".*"s
+```julia
 using MimiDICE2010
 m = MimiDICE2010.get_model()
 run(m)
-
-# output
-
 ```
 
 These steps should be relatively consistent across models, where a repository for `ModelX` should contain a primary file `ModelX.jl` which exports, at minimum, a function named something like `get_model` or `construct_model` which returns a version of the model, and can allow for model customization within the call.
@@ -77,35 +74,26 @@ Thus there are no required arguments, although the user can input `params`, a di
 
 In the case that you wish to alter an exogenous parameter, you may use the [`update_param!`](@ref) function.  Per usual, you will start by importing the Mimi package to your space with 
 
-```jldoctest tutorial2; output = false
+```julia
 using Mimi
-
-# output
-
 ```
 
 In DICE the parameter `fco22x` is the forcings of equilibrium CO2 doubling in watts per square meter, and exists in the components `climatedynamics` and `radiativeforcing`.  We can change this value from its default value of `3.200` to `3.000` in both components, using the following code:
 
-```jldoctest tutorial2; output = false, filter = r".*"s 
+```julia
 update_param!(m, :fco22x, 3.000)
 run(m)
-
-# output
-
 ```
 
 A more complex example may be a situation where you want to update several parameters, including some with a `:time` dimension, in conjunction with altering the time index of the model itself. DICE uses a default time horizon of 2005 to 2595 with 10 year increment timesteps.  If you wish to change this, say, to 2000 to 2500 by 10 year increment timesteps and use parameters that match this time, you could use the following code:
 
 First you upate the `time` dimension of the model as follows:
 
-```jldoctest tutorial2; output = false, filter = r".*"s
+```julia
 const ts = 10
 const years = collect(2000:ts:2500)
 nyears = length(years)
 set_dimension!(m, :time, years)
-
-# output
-
 ```
 
 Now that you have changed the time dimension, you have a mismatch between the time labels attached to your parameters and the time labels used by the model.  Thus, **you must update at least all parameters with a `:time`** dimension and use the explicit `update_timesteps=true` flag to get the time labels on the parameters to match those in the model. This is required even in cases where you do not want to change the parameter values themselves (see the forum question [here](https://forum.mimiframework.org/t/update-time-index/134/5)) for an in-depth explanation of this case. You may of course also update parameters without a `:time` dimension as desired. 

src/mcs/montecarlo.jl

@@ -326,8 +326,11 @@ function _perturb_param!(param::ScalarModelParameter{T}, md::ModelDef, trans::Tr
     end
 end
 
+# rvalue is an Array so we expect the dims to match and don't need to worry about
+# broadcasting
 function _perturb_param!(param::ArrayModelParameter{T}, md::ModelDef, 
-                         trans::TransformSpec, rvalue::Union{Number, Array{<: Number, N}}) where {T, N}
+    trans::TransformSpec, rvalue::Array{<: Number, N}) where {T, N}
+    
     op = trans.op
     pvalue = value(param)
     indices = _param_indices(param, md, trans)
@@ -340,6 +343,49 @@ function _perturb_param!(param::ArrayModelParameter{T}, md::ModelDef,
 
     else
         pvalue[indices...] += rvalue
+
+    end
+end
+
+# rvalue is a Number so we might need to deal with broadcasting
+function _perturb_param!(param::ArrayModelParameter{T}, md::ModelDef, 
+                         trans::TransformSpec, rvalue::Number) where {T, N}
+    op = trans.op
+    pvalue = value(param)
+    indices = _param_indices(param, md, trans)
+
+    if op == :(=)
+        
+        # first we check for a time index
+        ti = get_time_index_position(param)
+
+        # If there is no time index we have all methods needed to broadcast normally
+        if isnothing(ti)
+            broadcast_flag = sum(map(x -> length(x) > 1, indices)) > 0
+            broadcast_flag ? pvalue[indices...] .= rvalue : pvalue[indices...] = rvalue
+        
+        else
+            indices1, ts, indices2 = split_indices(indices, ti)
+            non_ts_indices = [indices1..., indices2...]
+            broadcast_flag = isempty(non_ts_indices) ? false : sum(map(x -> length(x) > 1, non_ts_indices)) > 0
+            
+            # Loop over the Array of TimestepIndex 
+            if isa(ts, Array) 
+                for el in ts
+                    broadcast_flag ? pvalue[indices1..., el, indices2...] .= rvalue : pvalue[indices1..., el, indices2...] = rvalue
+                end
+
+            # The time is just a single TimestepIndex and we can proceed with broadcast 
+            else     
+                broadcast_flag ? pvalue[indices...] .= rvalue : pvalue[indices...] = rvalue
+            end
+        end
+
+    elseif op == :(*=)
+        pvalue[indices...] *= rvalue
+
+    else 
+        pvalue[indices...] += rvalue
     end
 end
 

test/mcs/test_defmcs.jl

@@ -269,3 +269,20 @@ trial2 = copy(si2.sim_def.rvdict[:name1].dist.values)
 
 @test length(trial1) == length(trial2)
 @test trial1 != trial2
+
+# test broadcasting examples
+sd3 = @defsim begin
+
+    # 1 dimension
+    depk[:] = Uniform(0.1, 0.2)
+    k0[(Region2, Region3)] = Uniform(20, 30)
+    
+    # 2 dimensions
+    tfp[:, Region1] = Uniform(0.75, 1.25)
+    sigma[2020:5:2050, (Region2, Region3)] = Uniform(0.8, 1.2)
+    s[2020, Region1] = Uniform(0.2, 0.3)
+     
+end
+
+N = 5
+si3 = run(sd3, m, N)