PyscfCalculation: Add support to localize orbitals (#56)

The `parameters` input now accepts the `localize_orbitals` key. It
accepts a dictionary where the `method` key is required. It should
define one of the orbital localization schemes as implemented in PySCF.
If defined, the plugin will localize the orbitals and override the
molecular orbital coefficients matrix `mean_field_run.mo_coeff`.

README.md

@@ -193,6 +193,27 @@ and energy of each frame in the geometry optimization trajectory, are stored in
 results['trajectory'].get_array('energies')
 ```
 
+### Localizing orbitals
+
+To compute localized orbitals, specify the desired method in the `parameters.localize_orbitals.method` input:
+
+```python
+from ase.build import molecule
+from aiida.engine import run
+from aiida.orm import Dict, StructureData, load_code
+
+builder = load_code('pyscf').get_builder()
+builder.structure = StructureData(ase=molecule('H2O'))
+builder.parameters = Dict({
+    'mean_field': {'method': 'RHF'},
+    'localize_orbitals': {'method': 'ibo'}
+})
+results, node = run.get_node(builder)
+```
+
+The following methods are supported: `boys`, `cholesky`, `edmiston`, `iao`, `ibo`, `lowdin`, `nao`, `orth`, `pipek`,
+`vvo`. For more information, please refer to the [PySCF documentation](https://pyscf.org/user/lo.html).
+
 ### Computing the Hessian
 
 In order to compute the Hessian, specify an empty dictionary for the `hessian` key in the `parameters` input:

src/aiida_pyscf/calculations/base.py

@@ -153,6 +153,15 @@ def validate_parameters(cls, value: Dict | None, _) -> str | None:  # noqa: PLR0
                 'plugin if the `checkpoint` input is provided.'
             )
 
+        if (localize_orbitals := parameters.pop('localize_orbitals', None)) is not None:
+            valid_lo = ('boys', 'cholesky', 'edmiston', 'iao', 'ibo', 'lowdin', 'nao', 'orth', 'pipek', 'vvo')
+            method = localize_orbitals.get('method')
+            if method is None:
+                return f'No method specified in `localize_orbitals` parameters. Choose from: {valid_lo}'
+
+            if method.lower() not in valid_lo:
+                return f'Invalid method `{method}` specified in `localize_orbitals` parameters. Choose from: {valid_lo}'
+
         if (optimizer := parameters.pop('optimizer', None)) is not None:
             valid_solvers = ('geometric', 'berny')
             solver = optimizer.get('solver')

src/aiida_pyscf/calculations/templates/mean_field.py.j2

@@ -16,6 +16,11 @@ density_matrix = None
 mean_field_start = time.perf_counter()
 mean_field_run = mean_field.run(density_matrix)
 
+{% if localize_orbitals %}
+from pyscf import lo
+mean_field_run.mo_coeff = lo.orth_ao(mean_field_run, {{ localize_orbitals.method|render_python }})
+{% endif -%}
+
 results['timings']['mean_field'] = time.perf_counter() - mean_field_start
 results['mean_field'] = {}
 results['mean_field']['is_converged'] = mean_field_run.converged

tests/calculations/test_base.py

@@ -87,6 +87,28 @@ def test_parameters_mean_field(generate_calc_job, generate_inputs_pyscf, file_re
     file_regression.check(content_input_file, encoding='utf-8', extension='.pyr')
 
 
+def test_parameters_mean_field_localize_orbitals(generate_calc_job, generate_inputs_pyscf, file_regression):
+    """Test the ``mean_field.localize_orbitals`` key of the ``parameters`` input."""
+    parameters = {'localize_orbitals': {}}
+    inputs = generate_inputs_pyscf(parameters=parameters)
+
+    with pytest.raises(ValueError, match=r'No method specified in `localize_orbitals` parameters.*'):
+        generate_calc_job(PyscfCalculation, inputs=inputs)
+
+    parameters = {'localize_orbitals': {'method': 'invalid'}}
+    inputs = generate_inputs_pyscf(parameters=parameters)
+
+    with pytest.raises(ValueError, match=r'Invalid method `invalid` specified in `localize_orbitals` parameters.*'):
+        generate_calc_job(PyscfCalculation, inputs=inputs)
+
+    parameters = {'localize_orbitals': {'method': 'ibo'}}
+    inputs = generate_inputs_pyscf(parameters=parameters)
+    tmp_path, _ = generate_calc_job(PyscfCalculation, inputs=inputs)
+
+    content_input_file = (tmp_path / PyscfCalculation.FILENAME_SCRIPT).read_text()
+    file_regression.check(content_input_file, encoding='utf-8', extension='.pyr')
+
+
 def test_parameters_optimizer(generate_calc_job, generate_inputs_pyscf, file_regression):
     """Test the ``optimizer`` key of the ``parameters`` input."""
     parameters = {

tests/calculations/test_base/test_parameters_mean_field_localize_orbitals.pyr

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+"""Script automatically generated by `pyscf.base` plugin of `aiida-pyscf`."""
+
+
+def main():
+    import time
+
+    results = {
+        'timings': {}
+    }
+
+    time_start = time.perf_counter()
+
+
+    # Section: Results
+    def write_results_and_exit(results):
+        import dill
+        import json
+        import sys
+
+        results['timings']['total'] = time.perf_counter() - time_start
+
+        with open('results.json', 'w') as handle:
+            json.dump(results, handle)
+
+        with open('model.pickle', 'wb') as handle:
+            # Need to unset the ``_chkfile`` attribute as it contains an open file handle which cannot be unpickled.
+            mean_field_run._chkfile = None
+            # Need to unset the ``opt`` attribute as it contains ctypes objects containing pointers which cannot be pickled.
+            mean_field_run.opt = None
+            dill.dump(mean_field_run, handle)
+
+        sys.exit(0)
+
+    # Section: Structure definition
+    from pyscf import gto
+    structure = gto.Mole()
+    structure.unit = 'Ang'
+    structure.atom = """
+    O       0.000000000000000      0.000000000000000      0.119262000000000
+    H       0.000000000000000      0.763239000000000     -0.477047000000000
+    H       0.000000000000000     -0.763239000000000     -0.477047000000000
+    """
+    structure.build()
+
+    # Section: Mean field
+    from pyscf import scf
+    mean_field = scf.RHF(structure)
+    mean_field.chkfile = 'checkpoint.chk'
+    density_matrix = None
+    mean_field_start = time.perf_counter()
+    mean_field_run = mean_field.run(density_matrix)
+
+    from pyscf import lo
+    mean_field_run.mo_coeff = lo.orth_ao(mean_field_run, 'ibo')
+    results['timings']['mean_field'] = time.perf_counter() - mean_field_start
+    results['mean_field'] = {}
+    results['mean_field']['is_converged'] = mean_field_run.converged
+    results['mean_field']['total_energy'] = mean_field_run.e_tot
+    results['mean_field']['forces'] = (- mean_field_run.nuc_grad_method().kernel()).tolist()
+
+    if mean_field_run.converged:
+        results['mean_field']['molecular_orbitals'] = {}
+        results['mean_field']['molecular_orbitals']['energies'] = mean_field_run.mo_energy.tolist()
+        results['mean_field']['molecular_orbitals']['labels'] = structure.ao_labels()
+        results['mean_field']['molecular_orbitals']['occupations'] = mean_field_run.mo_occ.tolist()
+    else:
+        write_results_and_exit(results)
+
+
+
+
+
+    write_results_and_exit(results)
+
+
+if __name__ == '__main__':
+    main()