Don't call factorize_linear_constraint every time in _update_equilibrium of ProximalProjection?

[dpanici]

DESC/desc/optimize/_constraint_wrappers.py

Line 734 in 555a959

def _update_equilibrium(self, x, store=False):

• profile an optimization and see if this is actually a bottleneck

[dpanici]

This is not a true profile, but I threw "verbose=3" in the solve_options for an L=M=14 N=12 free bdry solve I am doing on della (with 8 threads requested so it should be multi-threaded), and the linear constraint projection of each equilibrum subproblem takes like 12.6 seconds (and seems like each time it takes about that long). Each subproblem solution time itself is on the same order of magnitude (about 12-14 seconds)

So I think this could be a significant timesaver, if we can just pass in the A matrix of the factorization to the _update_equilibrium steps and have it internally, if the A was passed in, just do the particular solution and not the factorization (will have to account for the different b but this should be doable)

[YigitElma]

Can you share your script so I can run it on nsys?

[f0uriest]

I think we should add an update_b method to LinearConstraintProjection, similar to how we have an update_target method for most of the linear objectives.

Then, in proximal, we would call eq.solve with the already wrapped and updated linear constraint projection rather than passing in the constraints each time

[YigitElma]

I will send some profiles but it is hard to see them probably. So, maybe I can show you on Zoom if you want.
Each script is here and I ran them on Della A100 with 32 core cpu.
proximal-lsq-exact-script.txt
lsq-auglag-script.txt

Optimization with lsq-auglag

It builds LinearConstraintProjection.build twice at the beginning, and there is no more factorize_linear_constraints operation. Taking jacobian and evaluating the function are very small portions of each iteration. Yes, there are multiple function evaluations inside "tr_subproblem" but it is on the order of couple hundred microseconds.

solve continuation_automatic

This part is even more complicated to see.

Basically, it goes like perturb, solve, perturb .... Every perturb calls factorize linear constraints twice (at the beginning and at the end). solve calls it once. If we zoom to one of them, it looks like this,

proximal-lsq-exact optimization

This is even more complicated to see, but here it is,

Here I zoom the second iteration,

Zooming even further,

I guess the the first to factorize_linear_constraints call is from perturb, and the other one that is under LinearConstraintProjection is for eq.solve.

[YigitElma]

Btw, I have noticed that perturb still uses SVD and tr_step_exact_svd. Was that deliberate?

The second factorization in perturbation can be made faster since Z is the same but I don't know if it is easy or worth working on. D and b should change, so the new method might end up looking like the same function but with a faster particular solution finder without Z finder. I think #1374 should be enough for this.

I am not too familiar with that portion of the code, so sorry if there is any mistake.

[f0uriest]

In perturb and in ProximalProjection we need to call factorize_linear_constraints multiple times because the vector b is changing, so we need to update the particular solution. If we stored Ainv (maybe as part of the project/recover operators) we could get around that.

[YigitElma]

Something like,

(
    xp,
    A,
    b,
    Z,
    D,
    unfixed_idx,
    project,
    recover,
    Ainv,
    A_full_nondegenerate,
    degenerate_idx, # maybe we need those for b_new
) = factorize_linear_constraints(
    objective,
    constraint,
)
b_new = -constraint.compute_scaled_error(0)
b_new = np.delete(b_new, degenerate_idx)
xp_new = jnp.zeros_like(xp)
fixed_idx = np.setdiff1d(np.arange(xp.size), unfixed_idx)
xp_new[fixed_idx] = b_new[fixed_idx]
xp_new[unfixed_idx] = Ainv @ (b_new-A_full_nondegenerate[:,fixed_idx] @ xp_new[fixed_idx])
recover = lambda x: xp_new + Z @ x

[YigitElma]

Maybe in the build method of ProximalProjection, we can create and build a LinearConstraintProjection and use it for the rest. We also need to add a update_constraint_target method to LinearConstraintProjection such that ForceBalance objective, self-consistency constraints and null-space etc stays the same, but the target for the FixedParameter constraints and project/recover methods of the LinearConstraintProjection get updated.

Then pass this LinearConstraintProjection to perturb and solve in update_equilibrium method (we need some type checking in those though), and only update the targets along the way.

[dpanici]

Pass in empty tuple for constraints inside of ProximalProjection update_equilibrium method to skip linear constraint creation in it