Release upper version limit on PySCF (#61)

To enable check-pointing of the state of a converged calculation, Sebastiaan rolled some custom serialization with Dill. This worked until some data structures were changed in PySCF around v2.4, breaking this functionality. This forced an upper version limit on PySCF. Recently, in PySCF v2.7, serialization via Pickle is natively supported. This commit removes the upper version constraints on PySCF by making the needed changes to the templates and parser, and regenerating the reference file for the tests. Other versions were also bumped, where appropriate.
microsoft · Nov 8, 2024 · b461f27 · b461f27
1 parent 6a99d37
commit b461f27
Show file tree

Hide file tree

Showing 21 changed files with 72 additions and 111 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -5,7 +5,8 @@ requires = ['flit_core>=3.4,<4']
 [project]
 authors = [
   {name = 'Sebastiaan P. Huber', email = '[email protected]'},
-  {name = 'Adam Grofe', email = '[email protected]'}
+  {name = 'Adam Grofe', email = '[email protected]'},
+  {name = 'Conrad Johnston', email = '[email protected]'}
 ]
 classifiers = [
   'Development Status :: 3 - Alpha',
@@ -20,12 +21,11 @@ classifiers = [
   'Topic :: Scientific/Engineering'
 ]
 dependencies = [
-  'aiida-core[atomic_tools]~=2.5',
+  'aiida-core[atomic_tools]~=2.6',
   'aiida-shell>=0.5.3',
-  'dill',
   'numpy',
   'pint',
-  'pyscf[geomopt]~=2.2,<2.4'
+  'pyscf[geomopt] ~= 2.7'
 ]
 dynamic = ['description', 'version']
 keywords = ['aiida', 'workflows', 'pyscf']
@@ -45,13 +45,13 @@ requires-python = '>=3.9'
 
 [project.optional-dependencies]
 pre-commit = [
-  'mypy==1.8.0',
+  'mypy==1.13.0',
   'pre-commit~=2.17'
 ]
 tests = [
   'packaging',
   'pgtest~=1.3,>=1.3.1',
-  'pytest~=7.2',
+  'pytest~=8.3',
   'pytest-regressions'
 ]
 
@@ -90,10 +90,10 @@ module = 'aiida_pyscf.*'
 ignore_missing_imports = true
 module = [
   'ase.*',
-  'dill.*',
   'pint.*',
   'plumpy.*',
-  'ruamel.*'
+  'ruamel.*',
+  'aiida_shell.*'
 ]
 
 [tool.pytest.ini_options]

diff --git a/src/aiida_pyscf/calculations/templates/results.py.j2 b/src/aiida_pyscf/calculations/templates/results.py.j2
@@ -1,7 +1,7 @@
 # Section: Results
 def write_results_and_exit(results):
-    import dill
     import json
+    import pickle
     import sys
 
     results['timings']['total'] = time.perf_counter() - time_start
@@ -11,11 +11,8 @@ def write_results_and_exit(results):
 
     {% if results.pickle_model %}
     with open('{{ results.filename_model }}', '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)
+        pickle.dump(mean_field_run, handle)
+
     {% endif %}
 
     sys.exit(0)
diff --git a/src/aiida_pyscf/parsers/base.py b/src/aiida_pyscf/parsers/base.py
@@ -4,8 +4,8 @@
 
 import json
 import pathlib
+import pickle
 
-import dill
 import numpy
 from aiida.engine import ExitCode
 from aiida.orm import ArrayData, Dict, FolderData, SinglefileData, TrajectoryData
@@ -54,11 +54,11 @@ def parse(self, retrieved_temporary_folder: str | None = None, **kwargs):  # noq
 
         try:
             with retrieved.open(PyscfCalculation.FILENAME_MODEL, 'rb') as handle:
-                model = dill.load(handle)
+                model = pickle.load(handle)
         except FileNotFoundError:
             if parameters.get('results', {}).get('pickle_model', True):
                 self.logger.warning(f'The pickled model file `{PyscfCalculation.FILENAME_MODEL}` could not be read.')
-        except dill.UnpicklingError:
+        except pickle.UnpicklingError:
             self.logger.warning(f'The pickled model file `{PyscfCalculation.FILENAME_MODEL}` could not be unpickled.')
         else:
             self.out('model', PickledData(model))
@@ -156,7 +156,7 @@ def batch(iterable, batch_size):
             symbols=[site.kind_name for site in self.node.inputs.structure.sites],
             positions=numpy.array(positions),
         )
-        energies = (numpy.array(energies) * ureg.hartree).to(ureg.electron_volt).magnitude  # type: ignore[attr-defined]
+        energies = (numpy.array(energies) * ureg.hartree).to(ureg.electron_volt).magnitude
         trajectory.set_array('energies', numpy.array(energies))
 
         return trajectory

diff --git a/tests/calculations/test_base/test_checkpoint.pyr b/tests/calculations/test_base/test_checkpoint.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_default.pyr b/tests/calculations/test_base/test_default.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_cubegen_parameters0_.pyr b/tests/calculations/test_base/test_parameters_cubegen_parameters0_.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_cubegen_parameters1_.pyr b/tests/calculations/test_base/test_parameters_cubegen_parameters1_.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_cubegen_parameters2_.pyr b/tests/calculations/test_base/test_parameters_cubegen_parameters2_.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_cubegen_parameters3_.pyr b/tests/calculations/test_base/test_parameters_cubegen_parameters3_.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_fcidump.pyr b/tests/calculations/test_base/test_parameters_fcidump.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_hessian.pyr b/tests/calculations/test_base/test_parameters_hessian.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_mean_field.pyr b/tests/calculations/test_base/test_parameters_mean_field.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

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
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_optimizer.pyr b/tests/calculations/test_base/test_parameters_optimizer.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start
@@ -24,11 +24,8 @@ def main():
             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)
+            pickle.dump(mean_field_run, handle)
+
 
         sys.exit(0)
 

diff --git a/tests/calculations/test_base/test_parameters_pickle_model.pyr b/tests/calculations/test_base/test_parameters_pickle_model.pyr
@@ -14,8 +14,8 @@ def main():
 
     # Section: Results
     def write_results_and_exit(results):
-        import dill
         import json
+        import pickle
         import sys
 
         results['timings']['total'] = time.perf_counter() - time_start