Release upper version limit on PySCF

ConradJohnston · ConradJohnston · commit 8a0838412368 · 2024-11-07T11:48:35.000-08:00
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.
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 = 'mail@sphuber.net'},
-  {name = 'Adam Grofe', email = 'v-adamgrofe@microsoft.com'}
+  {name = 'Adam Grofe', email = 'adamgrofe@microsoft.com'},
+  {name = 'Conrad Johnston', email = 'conrad.johnston@microsoft.com'}
 ]
 classifiers = [
   'Development Status :: 3 - Alpha',
@@ -20,19 +21,18 @@ 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']
 license = {file = 'LICENSE.txt'}
 name = 'aiida-pyscf'
 readme = 'README.md'
-requires-python = '>=3.9'
+requires-python = '>=3.9,<3.12'
 
 [project.entry-points.'aiida.calculations']
 'pyscf.base' = 'aiida_pyscf.calculations.base:PyscfCalculation'
@@ -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
diff --git a/tests/calculations/test_base/test_parameters_structure.pyr b/tests/calculations/test_base/test_parameters_structure.pyr
diff --git a/tests/parsers/fixtures/base/model_valid/aiida.out b/tests/parsers/fixtures/base/model_valid/aiida.out
diff --git a/tests/parsers/fixtures/base/model_valid/model.pickle b/tests/parsers/fixtures/base/model_valid/model.pickle
diff --git a/tests/parsers/fixtures/base/model_valid/results.json b/tests/parsers/fixtures/base/model_valid/results.json
diff --git a/tests/test_integration/test_pyscf_base_geometry_optimization.csv b/tests/test_integration/test_pyscf_base_geometry_optimization.csv
diff --git a/tests/test_integration/test_pyscf_base_mean_field.csv b/tests/test_integration/test_pyscf_base_mean_field.csv
