Automatic thresholds for projectability disentanglement (dis_proj_auto)

[elinscott]

Motivation

Projectability disentanglement (Qiao et al. 2023) selects and freezes states based on the per-band projectabilities $p_{n\mathbf{k}} = \sum_j |A_{nj}(\mathbf{k})|^2$, controlled by dis_proj_min and dis_proj_max. Good values for these thresholds are far less material-dependent than those for energy-based disentanglement, permitting high-throughput Wannierisation, but there is still some dependence that can't be ignored, especially when the PAO basis is not used and consequently projectabilities are lower than they would be otherwise (see pw2wannier90's atom_proj_ext). Default thresholds can give rise to concrete failures (see e.g. Supplementary Figure 20 of the original paper; replicated below).

This PR adds an option to determine both thresholds automatically at runtime from the projectability distribution itself.

What this PR does

Setting

dis_proj_auto = .true.

pools the projectabilities of all bands and k-points, histograms them into 64 equal-width bins, and splits the distribution into dis_proj_auto_classes classes (default 5, minimum 3) by multi-level Otsu thresholding, i.e. the class boundaries that maximise the between-class variance. dis_proj_min is set to the lowest threshold and dis_proj_max to the highest. More classes widen the disentanglement window between the two thresholds, so dis_proj_auto_classes becomes the one remaining tunable parameter (larger = fewer bands frozen and discarded; smaller = more).

Design decisions

• If the requested class count exceeds the number of populated bins, the count is reduced to what the data supports and a note is printed. Physically, each class needs its own cluster of values, and thresholds can only fall in the gaps between clusters; small systems can have few clusters (the graphene test populates only 12 of 64 bins). If fewer than three clusters exist - for example a manifold with no entanglement, where all projectabilities are nearly equal - the thresholds are genuinely undefined and the run stops with an error telling the user to set the thresholds manually.
• dis_proj_auto requires dis_froz_proj = .true. and is an input error in combination with explicit dis_proj_min/dis_proj_max.
• thresholding is performed over $[0, 1]$ rather than $[p_{min}, p_{max}]$ so that we loudly fail upon edge cases where the range of projectabilities is very narrow

Examples

Graphene

OpenMX s+p+s atomic projectors on C, comparing....

• the default dis_proj_min and _max ("default")
• values I found by trial-and-error ("manual")
• values generated by the new algorithm with dis_proj_auto = True ("auto")

	p_min	p_max	η(E_F+2) (meV)	disentangle steps	minimisation steps	t_total (min)
default	0.01	0.95	23.8	1186	3253	12.2
manual	0.05	0.70	3.9	61	1042	2.7
auto	0.039	0.742	4.0	152	1163	3.2

Bulk Si

OpenMX s+p+d atomic projectors on Si

With dis_froz_max = CBM + 2 eV

	p_min	p_max	η(E_F+2) (meV)	disentangle steps	minimisation steps	t_total (min)
default	0.01	0.95	13.9	398	1052	40.9
auto	0.086	0.742	6.5	415	771	17.5
auto with dis_proj_auto_classes=7	0.055	0.805	5.1	751	824	26.7

Without energy disentanglement

N.B. the left-hand-panel is the fail-case given in the Supp. Mat. of Qiao et al. 2023. In that case the authors painfully created a new pseudopotential in order to get a reliable Wannierization. Here instead we get a reliable Wannierization automatically.

	p_min	p_max	η(E_F+2) (meV)	disentangle steps	minimisation steps	t_total (min)
default	0.01	0.95	3754.8	1165	807	78.2
auto	0.086	0.742	6.6	415	1155	25.8
auto with dis_proj_auto_classes=7	0.055	0.805	160.2	785	693	25.0

Takeaways

• the thresholds generated by Otsu thresholding match hand-tuned quality on graphene and beat the defaults everywhere
• Without an energy-window safety net the defaults can fail outright with even the valence manifold wrong — while dis_proj_auto is robust
• Poorly chosen thresholds waste compute: the default runs consistently need the most disentanglement/minimisation steps and the longest wall time

Validation

• The Fortran Otsu implementation was compared against and matched an independent reference implementation of the same specification on synthetic data (clustered, uniform, degenerate, etc) across class counts 3-6 and 32/64/128 bins
• For the existing testw90_graphene_pdwf, the automatic thresholds come out as [0.039, 0.942] (the hand-tuned values in the test are [0.01, 0.85]), and the total spread improves from 15.80 to 14.00 Ų.
• The full serial test suite passes (95/95).
• All error paths (missing dis_froz_proj, invalid class count, conflicting explicit thresholds, degenerate distribution) exit with clear messages.

Tests and documentation

• New test testw90_graphene_dis_proj_auto (inputs symlinked from testw90_graphene_pdwf)
• dis_proj_auto and dis_proj_auto_classes documented in the user guide, parameter tables and parameters.xml.
• CHANGELOG entry added

Decision for the maintainers

Do we set dis_proj_auto to True by default? As we have seen, the algorithm outperforms the default thresholds.

AI Declaration

I used Claude to help me write and test this code. I have reviewed the code myself.

[elinscott]

@qiaojunfeng

[qiaojunfeng]

Thanks @elinscott, very nice results!

Do we set dis_proj_auto to True by default? As we have seen, the algorithm outperforms the default thresholds.

Yes I'd like to see it set as default, making things more automated is better.

[qiaojunfeng]

I think here it should be 1.0 not 0.85? Because all states with projectability between 0.01 and 1.0 participate into disentanglement.

[elinscott]

Good catch -- fixed

[elinscott]

Thanks @elinscott, very nice results!

Do we set dis_proj_auto to True by default? As we have seen, the algorithm outperforms the default thresholds.

Yes I'd like to see it set as default, making things more automated is better.

Thanks @qiaojunfeng !

OK, if we're going to enable this by default then I have modified the keyword precedence so that this change gives minimal surprising behaviour:

• dis_proj_auto is enabled by default.
• If dis_proj_min/dis_proj_max are explicitly provided, they take precedence and disable the automatic thresholding. This means that...
  • existing PDWF input files with specified thresholds behave as before
  • existing PDWF input files without thresholds specified now use auto rather than defaults (changed behaviour!)
• We give errors if...
  • both dis_proj_auto = .true. and the thresholds are set explicitly
  • with dis_froz_proj = .true., the user sets dis_proj_min xor dis_proj_max
  • with dis_froz_proj = .true., the user sets dis_proj_auto = .false. without providing the thresholds
• dis_proj_auto is silently ignored if dis_froz_proj = .false. (previously a hard error)
• dis_proj_min/dis_proj_max no longer have default values (previously 0.01/0.95): with dis_froz_proj = .true. the thresholds are either determined automatically or set explicitly.

Other changes:

• renamed dis_proj_auto_classes to dis_proj_auto_num_classes to be more explicit
• tweaked the hint in the tutorial to reflect the fact dis_proj_auto runs if the thresholds are not provided

[qiaojunfeng]

Thanks @elinscott, very nice results!

Do we set dis_proj_auto to True by default? As we have seen, the algorithm outperforms the default thresholds.

Yes I'd like to see it set as default, making things more automated is better.

Thanks @qiaojunfeng !

OK, if we're going to enable this by default then I have modified the keyword precedence so that this change gives minimal surprising behaviour:

• dis_proj_auto is enabled by default.

• If dis_proj_min/dis_proj_max are explicitly provided, they take precedence and disable the automatic thresholding. This means that...

  • existing PDWF input files with specified thresholds behave as before
  • existing PDWF input files without thresholds specified now use auto rather than defaults (changed behaviour!)

• We give errors if...

  • both dis_proj_auto = .true. and the thresholds are set explicitly
  • with dis_froz_proj = .true., the user sets dis_proj_min xor dis_proj_max
  • with dis_froz_proj = .true., the user sets dis_proj_auto = .false. without providing the thresholds

• dis_proj_auto is silently ignored if dis_froz_proj = .false. (previously a hard error)

• dis_proj_min/dis_proj_max no longer have default values (previously 0.01/0.95): with dis_froz_proj = .true. the thresholds are either determined automatically or set explicitly.

Other changes:

• renamed dis_proj_auto_classes to dis_proj_auto_num_classes to be more explicit
• tweaked the hint in the tutorial to reflect the fact dis_proj_auto runs if the thresholds are not provided

Thanks for the explanation, all looks good!