diff --git a/README.md b/README.md index 89a98da..8182920 100644 --- a/README.md +++ b/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: diff --git a/src/aiida_pyscf/calculations/base.py b/src/aiida_pyscf/calculations/base.py index 462c385..d63d13c 100644 --- a/src/aiida_pyscf/calculations/base.py +++ b/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') diff --git a/src/aiida_pyscf/calculations/templates/mean_field.py.j2 b/src/aiida_pyscf/calculations/templates/mean_field.py.j2 index 4e46a16..77f7aef 100644 --- a/src/aiida_pyscf/calculations/templates/mean_field.py.j2 +++ b/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 diff --git a/tests/calculations/test_base.py b/tests/calculations/test_base.py index 81838f1..14fbfe3 100644 --- a/tests/calculations/test_base.py +++ b/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 = { diff --git a/tests/calculations/test_base/test_parameters_mean_field_localize_orbitals.pyr b/tests/calculations/test_base/test_parameters_mean_field_localize_orbitals.pyr new file mode 100644 index 0000000..060da34 --- /dev/null +++ b/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()