add algebraic trick for ipu_eigh initialisation #104
Conversation
| """Each call updates density_matrix attempting to minimize energy(density_matrix, ... ). """ | ||
| density_matrix, V_xc, diff_JK, O, H_core, L_inv = vals[:6] # All (N, N) matrices | ||
| E_nuc, occupancy, ERI, grid_weights, grid_AO, diis_history, log = vals[6:] # Varying types/shapes. | ||
| density_matrix, V_xc, diff_JK, O, H_core, L_inv, prev_eigvects = vals[:7] # All (N, N) matrices |
There was a problem hiding this comment.
Assuming initialization prev_eigvects=np.eye(N).
|
|
||
| # Step 2: Solve eigh (L_inv turns generalized eigh into eigh). | ||
| eigvects = L_inv.T @ linalg_eigh(L_inv @ H @ L_inv.T, opts)[1] # (N, N) | ||
| linalg_input = L_inv @ H @ L_inv.T |
There was a problem hiding this comment.
We might consider merging all the eigh tricks together:
L_inv.T @ prev_eigvects @ eigh(prev_eigvects.T @ L_inv @ H @ L_inv.T @ prev_eigvects)
This hints that the L_inv trick (to turn generalized eigh into eigh) is the same as the initial_guess stuff. Unfortunately, this has the downside of removing variable names so less documentation.
| # Log matrices from all DFT iterations (not used by DFT algorithm). | ||
| N = H_core.shape[0] | ||
| log = {"matrices": np.zeros((opts.its, 4, N, N)), "E_xc": np.zeros((opts.its)), "energy": np.zeros((opts.its, 5))} | ||
| log = {"matrices": np.zeros((opts.its, 4, N, N)), "E_xc": np.zeros((opts.its)), "energy": np.zeros((opts.its, 5)), "eigenvectors" : np.zeros((opts.its, N, N))} |
There was a problem hiding this comment.
Assuming this is mainly for plotting. Could you try to profile at.[i].set(..) compared to jax.lax.dynamic_update_slice(..)? I've had issues with the at[i].set(..) being 10x slower due to unfortunate memory layout.
| #assert False | ||
|
|
||
| eigvects, eigvals = ipu_eigh(x, sort_eigenvalues=True, num_iters=12) | ||
| eigvects, eigvals = ipu_eigh(x, sort_eigenvalues=True, num_iters=opts.ipu_eigh_its) |
|
Looks great! Minor comments: |
|
We need a before/after graph or table to show the effect of the change, all other things being fixed |
Do you mean before/after of |
The change:
ipu_eighinitializationipu_eigh_itsparameter so that we can now select the number ofipu_eighiterations from CLI--visualise_eigvects true