[BUG] Ensemble.split() with Julia leaves forces_qspace = 0: the Fourier gradient ignores the forces (wrong gradient, minimization diverges) #436

Summary

With the Julia extension active, an ensemble returned by Ensemble.split() has forces_qspace equal to zero. The Fourier (Julia) gradient then ignores the real forces, and a minimization on a split ensemble produces a wrong gradient and diverges.

Where (master 70c3060d, same code in v1.6 / 1.6.1)

When fourier_gradient is on, the ensemble keeps q-space copies of displacements and forces. The Julia gradient uses only those copies:

delta_forces = self.forces_qspace - self.sscha_forces_qspace      # Modules/Ensemble.py L2561, get_fourier_gradient

These copies are built from the real-space forces in Ensemble.init() (L434 onwards). load_bin and compute_ensemble call self.init() after the forces are set (e.g. L1242), so those paths are consistent.

Ensemble.split() (L4194-4235) does not do this:

1. ens = Ensemble(self.dyn_0, ...) and ens.init_from_structures(structs) (L4219). At this point the forces are not known yet, so init_from_structures sets self.forces_qspace = np.zeros_like(self.u_disps_qspace) (L1306).
2. The forces are copied in place into the real-space array: ens.forces[:, :, :] = self.forces[split_mask, :, :] (L4225). Nothing updates forces_qspace.
3. It calls the real-space ens.update_weights(...) (L4231), not update_weights_fourier, and never calls ens.init().

So in the split ensemble forces_qspace == 0, and the Julia gradient becomes <u (0 - f_sscha)> instead of <u (f - f_sscha)>.

Observed (cubic perovskite, 5 atoms, 4x4x4, 10032 configurations, Julia 1.12.5)

Minimizations on split subsets of 1000, 2000, 5000 and 8000 configurations all ran away to |gc| ~ 1.38e5 and stopped on the Kong-Liu threshold (KL/N ~ 0.46). The final value was the same whatever N was, i.e. independent of the data.

Minimal reproduction (Julia available)

ens = sscha.Ensemble.Ensemble(dyn, T, supercell=dyn.GetSupercell()); ens.load_bin(data_dir, 0)
sub = ens.split(np.ones(ens.N, dtype=bool))
print(np.abs(ens.forces_qspace).max(), np.abs(sub.forces_qspace).max())   # nonzero vs 0.0

Suggested fix

Call ens.init() at the end of split(), after forces and stresses are copied. Alternatively, assign the arrays (ens.forces = ...) so that the q-space copies are rebuilt.

Related

• Open PR Ensemble: opt-in linear-memory q-space mode, plus four fixes on the Fourier path #428 fixes units and sign of the Fourier quantities in this same block of init_from_structures(), but it does not touch forces_qspace after split(), so this issue remains after Ensemble: opt-in linear-memory q-space mode, plus four fixes on the Fourier path #428.
• get_preconditioned_gradient_parallel itself calls split() at every step. It is only used without Julia, so this bug does not combine with the constant-error issue ([BUG] Without Julia, the FC-gradient error is a constant matrix of ones (get_preconditioned_gradient_parallel), so the FC stop criterion is meaningless #435), but any future use of split() with Julia is affected.

Workaround we use

Slice the *_pop0.npy arrays and reload with load_bin, which calls init(). On the full set this reproduces the original minimization exactly (same 22 steps, |gc| = 0.9615).

Environment

python-sscha 1.6.0 (installed), checked against tag v1.6.1 and master; cellconstructor 1.6.0; Julia 1.12.5.