Skip to content

Commit

Permalink
refactor get_influx for performance (#1753)
Browse files Browse the repository at this point in the history
This is a follow-up of #1606.

Hopefully fixes the runtime of the HWS model.
  • Loading branch information
SouthEndMusic authored Aug 27, 2024
1 parent d0494c1 commit d8fe0a1
Showing 1 changed file with 15 additions and 6 deletions.
21 changes: 15 additions & 6 deletions core/src/util.jl
Original file line number Diff line number Diff line change
Expand Up @@ -540,12 +540,21 @@ function get_influx(basin::Basin, node_id::NodeID)::Float64
end

function get_influx(basin::Basin, basin_idx::Int; prev::Bool = false)::Float64
(; vertical_flux, vertical_flux_prev, vertical_flux_from_input) = basin
vertical_flux = wrap_forcing(vertical_flux[parent(vertical_flux_from_input)])
flux_vector = prev ? vertical_flux_prev : vertical_flux
(; precipitation, evaporation, drainage, infiltration) = flux_vector
return precipitation[basin_idx] - evaporation[basin_idx] + drainage[basin_idx] -
infiltration[basin_idx]
(; node_id, vertical_flux, vertical_flux_prev, vertical_flux_from_input) = basin
influx = if prev
vertical_flux_prev.precipitation[basin_idx] -
vertical_flux_prev.evaporation[basin_idx] +
vertical_flux_prev.drainage[basin_idx] -
vertical_flux_prev.infiltration[basin_idx]
else
n = length(node_id)
vertical_flux = vertical_flux[parent(vertical_flux_from_input)]
vertical_flux[basin_idx] - # precipitation
vertical_flux[n + basin_idx] + # evaporation
vertical_flux[2n + basin_idx] - # drainage
vertical_flux[3n + basin_idx] # infiltration
end
return influx
end

inflow_edge(graph, node_id)::EdgeMetadata = graph[inflow_id(graph, node_id), node_id]
Expand Down

0 comments on commit d8fe0a1

Please sign in to comment.