Review Treatment of TimestepArrays

src/core/time_arrays.jl

@@ -452,9 +452,22 @@ end
 #
 # TimestepArray methods
 #
+
+# _dotview_helper TODOs - we should add options for if the arg is a TimestepValue
+# or an array of TimestepIndexes or TimestepValues
+function _dotview_helper(arg)
+	if arg isa AbstractTimestep
+		return arg.t
+	elseif arg isa TimestepIndex
+		return arg.index
+	else
+		return arg
+	end
+end
+
 function Base.dotview(v::Mimi.TimestepArray, args...)
 	# convert any timesteps to their underlying index
-	args = map(arg -> (arg isa AbstractTimestep ? arg.t : arg), args)
+	args = map(_dotview_helper, args)
 	Base.dotview(v.data, args...)
 end
 

test/test_timesteparrays.jl

@@ -663,11 +663,6 @@ set_param!(m, :first, :par1, 1:length(years))
 
 @test_throws MissingException run(m)
 
-# Check broadcast assignment to underlying array
-x = Mimi.TimestepVector{Mimi.FixedTimestep{2005,10}, Float64}(zeros(10))
-x[:] .= 10
-@test all(x.data .== 10)
-
 #------------------------------------------------------------------------------
 # 7. Test TimestepArrays with time not as the first dimension
 #------------------------------------------------------------------------------
@@ -718,4 +713,33 @@ run(m)
 @test all(!ismissing, m[:gdp, :pop])
 @test all(!ismissing, m[:gdp, :mat2])
 
+#------------------------------------------------------------------------------
+# 8. Check broadcast assignment to underlying array
+#------------------------------------------------------------------------------
+
+x_arr = zeros(10)
+y_arr = collect(reshape(zeros(8), 4, 2))
+
+x = Mimi.TimestepVector{Mimi.FixedTimestep{2005,10}, Float64}(x_arr)
+y = TimestepMatrix{FixedTimestep{2000, 1}, Float64, 1}(y_arr)
+
+# colon and ints
+x[:] .= 10
+y[:] .= 10
+@test all(y.data .== 10)
+@test all(x.data .== 10)
+
+y[:,1] .= 20
+@test all(y.data[:,1] .== 20)
+
+reset_time_val(x, x_arr)
+reset_time_val(y, y_arr)
+
+# TimestepIndex
+y[TimestepIndex(2),:] .= 10
+@test all(y.data[2,:] .== 10)
+
+reset_time_val(x, x_arr)
+reset_time_val(y, y_arr)
+
 end #module