Skip to content

Conversation

JakobAsslaender
Copy link
Member

@nHackel: I drafted some code to close #107. I hacked the tests to print outputs and the calculated metrics look right.

However, I did not check if I was able to remove all allocations. And the code certainly did not become more readable, as I am certainly missusing several temp variables while they are not in use.

Let me know what you think! I am sure this can be improved upon! Or maybe we can add the readable code in comments?

@nHackel
Copy link
Member

nHackel commented Sep 28, 2025

Hmm good question about the readability of the code. We could introduce additional temporary vectors, but those would be "unecessary" and cost more memory.

I think we should have the "clean" calculations as comments above with an explanation of how/why we are splitting things and when the temp variables are used as partial storage and when their intended value. We could potentially rename them to be more general as well but that might add more mental load in the places where they are properly used.

A quick profile showed me that a lot of allocations are also happening in:

  AHA = solver.AHA
  for i  eachindex(solver.reg)
    mul!(state.β, adjoint(solver.regTrafo[i]), state.z[i],  state.ρ[i], 1)
    mul!(state.β, adjoint(solver.regTrafo[i]), state.u[i], -state.ρ[i], 1)
    AHA += state.ρ[i] * adjoint(solver.regTrafo[i]) * solver.regTrafo[i] # here 
  end

In particular, the composite operator formed by the multiplications and addition always allocates temporary result vectors.

Unless I am missing something elsewhere in the code, the only thing changing for AHA is state.ρ[i]. Which means we could use the normal operator from LinearOperatorCollection and precompute AHA by first computing and storing references to the individual solver.AHA_components = normalOperator(regTrafo[i], state.p[i]) and then performing the sum and storing/using that in the iteration.

That way we could do:

  AHA = solver.AHA
  for i  eachindex(solver.reg)
    mul!(state.β, adjoint(solver.regTrafo[i]), state.z[i],  state.ρ[i], 1)
    mul!(state.β, adjoint(solver.regTrafo[i]), state.u[i], -state.ρ[i], 1)
    solver.AHA_components[i].weights = state.p[i]
end

At the moment the normal operator does not accept numbers as weights, but we can add that by changing the ::Nothing here to a ::Union{Nothing, Number}

@JakobAsslaender
Copy link
Member Author

Thanks for checking all this, @nHackel! I followed your suggestion and added the old code back as comments with some explanatory words. I think renaming the variables would hurt the readability of the actual algorithm, and I think it is good enough to spell out what we are doing.

I like your idea about the composite operator. It seems you have a good grasp of what to do, and I do not. Do you mind adding these changes to the PR?

@nHackel
Copy link
Member

nHackel commented Sep 29, 2025

@JakobAsslaender yeah I'll try adding the changes during this or next week and I'll ping you again once things look okay on basic data and then you can try with proper data?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Allocations in ADMM
2 participants