Rename 'opitmizer' parameter to 'geometry_optimizer'

Rename use of 'optimizer' to 'geometry optimizer' and update the tests.
The motivation for this breaking change is to more closely align the
names of the parameters with those used by PySCF. This change also
removes ambiguity around which optimizer - electronic or geometry - and
makes for a more intuitive namespace for outputs.

Previously:
{
"optimizer": {
    "is_converged": true,
    "optimized_coordinates": [
      [
        0,
        0,
        0.6729757879646958
      ],
      [
        0,
        0,
        -0.6729757879646958
      ]
    ]
  }
}

Now, changing to "geometry_optimizer" makes it clear that this stanza in
the results JSON refers to the optimized structure.

This change is a breaking to the API, but will make it easier in the
future to distinguish between settings for the geometry optimizer and
setting for the SCF optimizer.

README.md

@@ -157,9 +157,9 @@ automatically by the plugin based on the `StructureData` input.
 
 ### Optimizing geometry
 
-The geometry can be optimized by specifying the `optimizer` dictionary in the input `parameters`. The `solver` has to be
-specified, and currently the solvers `geometric` and `berny` are supported. The `convergence_parameters` accepts the
-parameters for the selected solver (see
+The geometry can be optimized by specifying the `geometry_optimizer` dictionary in the input `parameters`. The `solver`
+has to be specified, and currently the solvers `geometric` and `berny` are supported. The `convergence_parameters`
+accepts the parameters for the selected solver (see
 [PySCF documentation](https://pyscf.org/user/geomopt.html?highlight=geometry+optimization) for details):
 
 ```python
@@ -171,7 +171,7 @@ builder = load_code('pyscf').get_builder()
 builder.structure = StructureData(ase=molecule('H2O'))
 builder.parameters = Dict({
     'mean_field': {'method': 'RHF'},
-    'optimizer': {
+    'geometry_optimizer': {
         'solver': 'geometric',
         'convergence_parameters': {
             'convergence_energy': 1e-6,  # Eh

src/aiida_pyscf/calculations/base.py

@@ -131,7 +131,17 @@ def define(cls, spec: CalcJobProcessSpec):  # type: ignore[override]
     @classmethod
     def get_valid_parameter_keys(cls) -> tuple[str, ...]:
         """Return list of valid keys for the ``parameters`` input."""
-        return ('mean_field', 'localize_orbitals', 'optimizer', 'cubegen', 'fcidump', 'hessian', 'results', 'structure')
+        valid_keys = (
+            'mean_field',
+            'localize_orbitals',
+            'geometry_optimizer',
+            'cubegen',
+            'fcidump',
+            'hessian',
+            'results',
+            'structure',
+        )
+        return valid_keys
 
     @classmethod
     def validate_parameters(cls, value: Dict | None, _) -> str | None:  # noqa: PLR0911, PLR0912
@@ -171,9 +181,9 @@ def validate_parameters(cls, value: Dict | None, _) -> str | None:  # noqa: PLR0
             if method.lower() not in valid_lo:
                 return f'Invalid method `{method}` specified in `localize_orbitals` parameters. Choose from: {valid_lo}'
 
-        if (optimizer := parameters.get('optimizer')) is not None:
+        if (geometry_optimizer := parameters.get('geometry_optimizer')) is not None:
             valid_solvers = ('geometric', 'berny')
-            solver = optimizer.get('solver')
+            solver = geometry_optimizer.get('solver')
 
             if solver is None:
                 return f'No solver specified in `optimizer` parameters. Choose from: {valid_solvers}'

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

@@ -0,0 +1,24 @@
+# Section: Geometry Optimizer
+convergence_parameters = {}
+{% if geometry_optimizer.convergence_parameters %}
+    {% for key, value in geometry_optimizer.convergence_parameters.items() %}
+convergence_parameters['{{ key }}'] = {{ value|render_python }}
+    {% endfor %}
+{% endif %}
+
+geometry_optimizer_start = time.perf_counter()
+geometry_optimizer = mean_field.Gradients().optimizer(solver='{{ geometry_optimizer.solver }}')
+
+try:
+    geometry_optimizer_run = geometry_optimizer.kernel(convergence_parameters)
+except RuntimeError:
+    results['geometry_optimizer'] = {
+        'is_converged': False
+    }
+else:
+    results['geometry_optimizer'] = {
+        'is_converged': True,
+        'optimized_coordinates': geometry_optimizer_run.atom_coords().tolist(),
+    }
+
+results['timings']['geometry_optimizer'] = time.perf_counter() - geometry_optimizer_start

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

@@ -18,8 +18,8 @@ def main():
 
     {% include 'pyscf/mean_field.py.j2' %}
 
-    {% if optimizer %}
-    {% include 'pyscf/optimizer.py.j2' %}
+    {% if geometry_optimizer %}
+    {% include 'pyscf/geometry_optimizer.py.j2' %}
     {% endif %}
 
     {% if hessian is defined %}

src/aiida_pyscf/parsers/base.py

@@ -65,9 +65,9 @@ def parse(self, retrieved_temporary_folder: str | None = None, **kwargs):  # noq
 
         results_mean_field = parsed_json.setdefault('mean_field', {})
 
-        if 'optimized_coordinates' in parsed_json.get('optimizer', {}):
+        if 'optimized_coordinates' in parsed_json.get('geometry_optimizer', {}):
             structure = self.node.inputs.structure.clone()
-            optimized_coordinates = parsed_json['optimizer'].pop('optimized_coordinates') * ureg.bohr
+            optimized_coordinates = parsed_json['geometry_optimizer'].pop('optimized_coordinates') * ureg.bohr
             structure.reset_sites_positions(optimized_coordinates.to(ureg.angstrom).magnitude.tolist())
             self.out('structure', structure)
 
@@ -118,7 +118,7 @@ def parse(self, retrieved_temporary_folder: str | None = None, **kwargs):  # noq
         if results_mean_field['is_converged'] is False:
             return self.handle_failure('ERROR_ELECTRONIC_CONVERGENCE_NOT_REACHED', override_scheduler=True)
 
-        if 'optimizer' in parsed_json and parsed_json['optimizer']['is_converged'] is False:
+        if 'geometry_optimizer' in parsed_json and parsed_json['geometry_optimizer']['is_converged'] is False:
             return self.handle_failure('ERROR_IONIC_CONVERGENCE_NOT_REACHED', override_scheduler=True)
 
         return ExitCode(0)

tests/calculations/test_base.py

@@ -109,10 +109,10 @@ def test_parameters_mean_field_localize_orbitals(generate_calc_job, generate_inp
     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."""
+def test_parameters_geometry_optimizer(generate_calc_job, generate_inputs_pyscf, file_regression):
+    """Test the ``geometry_optimizer`` key of the ``parameters`` input."""
     parameters = {
-        'optimizer': {
+        'geometry_optimizer': {
             'solver': 'geomeTRIC',
             'convergence_parameters': {
                 'convergence_energy': 2.0,
@@ -216,11 +216,11 @@ def test_invalid_parameters_unknown_arguments(generate_calc_job, generate_inputs
         }, r'Invalid solver `solve-this` specified in `optimizer` parameters'),
     )
 )
-def test_invalid_parameters_optimizer(generate_calc_job, generate_inputs_pyscf, parameters, expected):
-    """Test validation of ``parameters.optimizer``."""
+def test_invalid_parameters_geometry_optimizer(generate_calc_job, generate_inputs_pyscf, parameters, expected):
+    """Test validation of ``parameters.geometry_optimizer``."""
     with pytest.raises(ValueError, match=expected):
-        generate_calc_job(PyscfCalculation, inputs=generate_inputs_pyscf(parameters=Dict({'optimizer': parameters})))
-
+        generate_calc_job(
+            PyscfCalculation, inputs=generate_inputs_pyscf(parameters=Dict({'geometry_optimizer': parameters})))
 
 @pytest.mark.parametrize(
     'parameters, expected', (

tests/calculations/test_base/test_parameters_geometry_optimizer.pyr

@@ -0,0 +1,94 @@
+#!/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 json
+        import pickle
+        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:
+            pickle.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)
+
+    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)
+
+    # Section: Geometry Optimizer
+    convergence_parameters = {}
+    convergence_parameters['string'] = 'value'
+    convergence_parameters['convergence_energy'] = 2.0
+
+    geometry_optimizer_start = time.perf_counter()
+    geometry_optimizer = mean_field.Gradients().optimizer(solver='geomeTRIC')
+
+    try:
+        geometry_optimizer_run = geometry_optimizer.kernel(convergence_parameters)
+    except RuntimeError:
+        results['geometry_optimizer'] = {
+            'is_converged': False
+        }
+    else:
+        results['geometry_optimizer'] = {
+            'is_converged': True,
+            'optimized_coordinates': geometry_optimizer_run.atom_coords().tolist(),
+        }
+
+    results['timings']['geometry_optimizer'] = time.perf_counter() - geometry_optimizer_start
+
+
+
+
+    write_results_and_exit(results)
+
+
+if __name__ == '__main__':
+    main()

tests/parsers/fixtures/base/relax/results.json

@@ -1 +1 @@
-{"timings": {"mean_field": 0.07242159199813614, "optimizer": 0.37188204899575794, "total": 0.9836381720015197}, "mean_field": {"is_converged": true, "total_energy": -1.11690055771897, "forces": [[-0.0, -0.0, -0.025062730160341773], [-0.0, -0.0, 0.025062730160341773]], "molecular_orbitals": {"energies": [-0.5797286564236249, 0.6740804576465153], "labels": ["0 H 1s    ", "1 H 1s    "], "occupations": [2.0, 0.0]}}, "optimizer": {"is_converged": true, "optimized_coordinates": [[0.0, 0.0, 0.6729757879646958], [0.0, 0.0, -0.6729757879646958]]}}
+{"timings": {"mean_field": 0.07242159199813614, "geometry_optimizer": 0.37188204899575794, "total": 0.9836381720015197}, "mean_field": {"is_converged": true, "total_energy": -1.11690055771897, "forces": [[-0.0, -0.0, -0.025062730160341773], [-0.0, -0.0, 0.025062730160341773]], "molecular_orbitals": {"energies": [-0.5797286564236249, 0.6740804576465153], "labels": ["0 H 1s    ", "1 H 1s    "], "occupations": [2.0, 0.0]}}, "geometry_optimizer": {"is_converged": true, "optimized_coordinates": [[0.0, 0.0, 0.6729757879646958], [0.0, 0.0, -0.6729757879646958]]}}

tests/parsers/test_base/test_relax.yml

@@ -4,6 +4,8 @@ energies:
 - -30.408872159833294
 - -30.40888413284324
 parameters:
+  geometry_optimizer:
+    is_converged: true
   mean_field:
     forces:
     - - -0.0
@@ -26,11 +28,9 @@ parameters:
       - 0.0
     total_energy: -30.39241247445083
     total_energy_units: eV
-  optimizer:
-    is_converged: true
   timings:
+    geometry_optimizer: 0.37188204899575794
     mean_field: 0.07242159199813614
-    optimizer: 0.37188204899575794
     total: 0.9836381720015197
 positions:
 - 0.0

tests/test_integration.py

@@ -47,7 +47,7 @@ def test_pyscf_base_geometry_optimization(
     builder.parameters = orm.Dict(
         {
             'mean_field': {'method': 'RHF'},
-            'optimizer': {
+            'geometry_optimizer': {
                 'solver': 'geomeTRIC',
             },
         }
@@ -59,7 +59,7 @@ def test_pyscf_base_geometry_optimization(
 
     # Timings differ from run to run, so we just check the expected keys exist and are floats.
     timings = parameters.pop('timings')
-    for key in ('total', 'mean_field', 'optimizer'):
+    for key in ('total', 'mean_field', 'geometry_optimizer'):
         assert key in timings
         assert isinstance(timings[key], float)
 

tests/test_integration/test_pyscf_base_geometry_optimization.csv

@@ -1,10 +1,10 @@
 ,cell,positions,forces,total_energy,mo_energies
-0,0,4.0584834845214003e-15,-5.9556344379907998e-15,-2039.8853743663999,-550.86280025027997
-1,0,6.7404881421057998e-14,1.4450896956753999e-14,,-34.375426862456003
-2,0,0.14564171835587,2.2269701395491999,,-16.629598134599
-3,0,-6.6395809507747998e-15,1.5404180728208e-15,,-12.323304634735999
-4,0,0.75805955516664003,0.64803113397734002,,-10.637428057751
-5,0,-0.4902368591654,-1.1134850697747001,,16.200273277781999
-6,0,3.1805437496249999e-15,4.4152163651700002e-15,,19.796075801491
-7,0,-0.75805955516652002,-0.64803113397735002,,
-8,0,-0.49023685916538001,-1.1134850697747001,,
+0,0,3.4503165637247998e-16,2.8146553209427001e-15,-2039.8853743663999,-550.86280025027997
+1,0,-6.5778452218681002e-14,2.1765548502765e-14,,-34.375426862456003
+2,0,0.14564171835590001,2.2269701395491999,,-16.629598134599
+3,0,-4.5013598495410002e-16,1.6299521569328999e-15,,-12.323304634735999
+4,0,0.75805955516654999,0.64803113397735002,,-10.637428057751
+5,0,-0.49023685916532,-1.1134850697747001,,16.200273277781999
+6,0,-6.2197045838828996e-16,-4.4446074778754997e-15,,19.796075801491
+7,0,-0.75805955516662005,-0.64803113397732004,,
+8,0,-0.49023685916539,-1.1134850697747001,,

tests/test_integration/test_pyscf_base_geometry_optimization.yml

@@ -1,3 +1,5 @@
+geometry_optimizer:
+  is_converged: true
 mean_field:
   forces_units: eV/Å
   is_converged: true
@@ -19,5 +21,3 @@ mean_field:
     - 0.0
     - 0.0
   total_energy_units: eV
-optimizer:
-  is_converged: true

tests/test_integration/test_pyscf_base_mean_field.csv

@@ -1,10 +1,10 @@
 ,forces,total_energy,mo_energies
-0,-3.6983894632066003e-15,-2039.8853743663999,-550.86280025027997
-1,4.2817472464455997e-15,,-34.375426862456003
-2,2.2269701395494002,,-16.629598134599
-3,-1.2896218568954999e-15,,-12.323304634735999
-4,0.64803113397731005,,-10.637428057751
-5,-1.1134850697747001,,16.200273277781999
-6,4.9880113201019998e-15,,19.796075801491
-7,-0.64803113397735002,,
+0,-1.5033186023723e-16,-2039.8853743663999,-550.86280025027997
+1,1.3023647874605e-14,,-34.375426862455001
+2,2.2269701395491999,,-16.629598134599
+3,2.5515673915859e-16,,-12.323304634735999
+4,0.64803113397727996,,-10.637428057751
+5,-1.1134850697747001,,16.200273277783001
+6,-1.0482487892139e-16,,19.796075801491
+7,-0.64803113397727996,,
 8,-1.1134850697747001,,