Load @assigned_heading in addition to loading @heading

app/controllers/budgets/investments_controller.rb

@@ -89,7 +89,10 @@ module Budgets
       end
 
       def load_heading
-        @heading = @budget.headings.find(params[:heading_id]) if params[:heading_id].present?
+        if params[:heading_id].present?
+          @heading = @budget.headings.find(params[:heading_id])
+          @assigned_heading = @ballot.try(:heading_for_group, @heading.try(:group))
+        end
       end
 
   end