Avoid Array allocations in concentration code

Deltares · Dec 20, 2024 · ae2589c · ae2589c
1 parent 3cd6fcf
commit ae2589c
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/core/src/callback.jl b/core/src/callback.jl
@@ -190,15 +190,15 @@ function update_concentrations!(u, t, integrator)::Nothing
     # of the basins after processing inflows only
     cumulative_in .= 0.0
 
-    mass .+= concentration[1, :, :] .* vertical_flux.drainage * dt
+    @views mass .+= concentration[1, :, :] .* vertical_flux.drainage * dt
     basin.concentration_data.cumulative_in .= vertical_flux.drainage * dt
 
     # Precipitation depends on fixed area
     for node_id in basin.node_id
         fixed_area = basin_areas(basin, node_id.idx)[end]
         added_precipitation = fixed_area * vertical_flux.precipitation[node_id.idx] * dt
-
-        mass[node_id.idx, :] .+= concentration[2, node_id.idx, :] .* added_precipitation
+        @views mass[node_id.idx, :] .+=
+            concentration[2, node_id.idx, :] .* added_precipitation
         cumulative_in[node_id.idx] += added_precipitation
     end
 
@@ -212,7 +212,8 @@ function update_concentrations!(u, t, integrator)::Nothing
         if active
             outflow_id = edge[1].edge[2]
             volume = integral(flow_rate, tprev, t)
-            mass[outflow_id.idx, :] .+= flow_boundary.concentration[id.idx, :] .* volume
+            @views mass[outflow_id.idx, :] .+=
+                flow_boundary.concentration[id.idx, :] .* volume
             cumulative_in[outflow_id.idx] += volume
         end
     end

diff --git a/core/src/concentration.jl b/core/src/concentration.jl
@@ -6,8 +6,8 @@ function mass_updates_user_demand!(integrator::DEIntegrator)::Nothing
     (; basin, user_demand) = integrator.p
     (; concentration_state, mass) = basin.concentration_data
 
-    for (inflow_edge, outflow_edge) in
-        zip(user_demand.inflow_edge, user_demand.outflow_edge)
+    @views for (inflow_edge, outflow_edge) in
+               zip(user_demand.inflow_edge, user_demand.outflow_edge)
         from_node = inflow_edge.edge[1]
         to_node = outflow_edge.edge[2]
         userdemand_idx = outflow_edge.edge[1].idx
@@ -41,7 +41,7 @@ function mass_inflows_basin!(integrator::DEIntegrator)::Nothing
     for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge)
         from_node = inflow_edge.edge[1]
         to_node = outflow_edge.edge[2]
-        if from_node.type == NodeType.Basin
+        @views if from_node.type == NodeType.Basin
             flow = flow_update_on_edge(integrator, inflow_edge.edge)
             if flow < 0
                 cumulative_in[from_node.idx] -= flow
@@ -67,7 +67,7 @@ function mass_inflows_basin!(integrator::DEIntegrator)::Nothing
             flow = flow_update_on_edge(integrator, outflow_edge.edge)
             if flow > 0
                 cumulative_in[to_node.idx] += flow
-                if from_node.type == NodeType.Basin
+                @views if from_node.type == NodeType.Basin
                     mass[to_node.idx, :] .+= concentration_state[from_node.idx, :] .* flow
                 elseif from_node.type == NodeType.LevelBoundary
                     mass[to_node.idx, :] .+=
@@ -95,7 +95,7 @@ function mass_outflows_basin!(integrator::DEIntegrator)::Nothing
     (; state_inflow_edge, state_outflow_edge, basin) = integrator.p
     (; mass, concentration_state) = basin.concentration_data
 
-    for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge)
+    @views for (inflow_edge, outflow_edge) in zip(state_inflow_edge, state_outflow_edge)
         from_node = inflow_edge.edge[1]
         to_node = outflow_edge.edge[2]
         if from_node.type == NodeType.Basin