pyiron-workflow-atomistics 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (46) hide show
  1. pyiron_workflow_atomistics-0.0.3/LICENSE +29 -0
  2. pyiron_workflow_atomistics-0.0.3/MANIFEST.in +6 -0
  3. pyiron_workflow_atomistics-0.0.3/PKG-INFO +169 -0
  4. pyiron_workflow_atomistics-0.0.3/docs/README.md +107 -0
  5. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/__init__.py +6 -0
  6. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/_version.py +21 -0
  7. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/ase_calculator.py +326 -0
  8. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/ase_engine.py +326 -0
  9. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/bulk.py +308 -0
  10. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/bulk_defect/point.py +40 -0
  11. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/bulk_defect/vacancy.py +33 -0
  12. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/calculator.py +207 -0
  13. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/dataclass_storage.py +163 -0
  14. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/featurisers.py +286 -0
  15. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/__init__.py +0 -0
  16. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/analysis.py +623 -0
  17. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/cleavage.py +779 -0
  18. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/dataclass_storage.py +79 -0
  19. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/gb_code/__init__.py +0 -0
  20. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/gb_code/constructor.py +410 -0
  21. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/gb_code/csl_generator.py +928 -0
  22. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/gb_code/gb_generator.py +583 -0
  23. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/gb_code/searcher.py +513 -0
  24. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/gb_study.py +263 -0
  25. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/hcp_generator.py +302 -0
  26. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/optimiser.py +488 -0
  27. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/segregation.py +204 -0
  28. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/gb/utils.py +38 -0
  29. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/pyiron_module_template/__init__.py +3 -0
  30. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/pyiron_module_template/_version.py +716 -0
  31. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/structure_manipulator/__init__.py +0 -0
  32. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/structure_manipulator/interstitials.py +322 -0
  33. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/structure_manipulator/tools.py +144 -0
  34. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/surface/__init__.py +0 -0
  35. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/surface/builder.py +298 -0
  36. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/surface/surface_study.py +151 -0
  37. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics/utils.py +174 -0
  38. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics.egg-info/PKG-INFO +169 -0
  39. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics.egg-info/SOURCES.txt +44 -0
  40. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics.egg-info/dependency_links.txt +1 -0
  41. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics.egg-info/requires.txt +9 -0
  42. pyiron_workflow_atomistics-0.0.3/pyiron_workflow_atomistics.egg-info/top_level.txt +1 -0
  43. pyiron_workflow_atomistics-0.0.3/pyproject.toml +53 -0
  44. pyiron_workflow_atomistics-0.0.3/setup.cfg +4 -0
  45. pyiron_workflow_atomistics-0.0.3/setup.py +8 -0
  46. pyiron_workflow_atomistics-0.0.3/versioneer.py +2277 -0
@@ -0,0 +1,29 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2024, Max-Planck-Institut für Nachhaltige Materialien GmbH - Computational Materials Design (CM) Department
4
+ All rights reserved.
5
+
6
+ Redistribution and use in source and binary forms, with or without
7
+ modification, are permitted provided that the following conditions are met:
8
+
9
+ * Redistributions of source code must retain the above copyright notice, this
10
+ list of conditions and the following disclaimer.
11
+
12
+ * Redistributions in binary form must reproduce the above copyright notice,
13
+ this list of conditions and the following disclaimer in the documentation
14
+ and/or other materials provided with the distribution.
15
+
16
+ * Neither the name of the copyright holder nor the names of its
17
+ contributors may be used to endorse or promote products derived from
18
+ this software without specific prior written permission.
19
+
20
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,6 @@
1
+ include LICENSE
2
+ include versioneer.py
3
+ include docs/README.md
4
+ include pyproject.toml
5
+ include pyiron_workflow_atomistics/_version.py
6
+ recursive-include pyiron_workflow_atomistics *.py
@@ -0,0 +1,169 @@
1
+ Metadata-Version: 2.4
2
+ Name: pyiron_workflow_atomistics
3
+ Version: 0.0.3
4
+ Summary: A pyiron module for atomistic simulation workflows
5
+ Author-email: pyiron team <pyiron@mpie.de>
6
+ License: BSD 3-Clause License
7
+
8
+ Copyright (c) 2024, Max-Planck-Institut für Nachhaltige Materialien GmbH - Computational Materials Design (CM) Department
9
+ All rights reserved.
10
+
11
+ Redistribution and use in source and binary forms, with or without
12
+ modification, are permitted provided that the following conditions are met:
13
+
14
+ * Redistributions of source code must retain the above copyright notice, this
15
+ list of conditions and the following disclaimer.
16
+
17
+ * Redistributions in binary form must reproduce the above copyright notice,
18
+ this list of conditions and the following disclaimer in the documentation
19
+ and/or other materials provided with the distribution.
20
+
21
+ * Neither the name of the copyright holder nor the names of its
22
+ contributors may be used to endorse or promote products derived from
23
+ this software without specific prior written permission.
24
+
25
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
26
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
27
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
28
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
29
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
30
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
31
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
32
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
33
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
34
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
35
+
36
+ Project-URL: Homepage, https://pyiron.org/
37
+ Project-URL: Documentation, https://pyiron_workflow_atomistics.readthedocs.io
38
+ Project-URL: Repository, https://github.com/pyiron/pyiron_workflow_atomistics
39
+ Keywords: pyiron,atomistics,workflow
40
+ Classifier: Development Status :: 3 - Alpha
41
+ Classifier: Topic :: Scientific/Engineering
42
+ Classifier: License :: OSI Approved :: BSD License
43
+ Classifier: Intended Audience :: Science/Research
44
+ Classifier: Operating System :: OS Independent
45
+ Classifier: Programming Language :: Python :: 3.9
46
+ Classifier: Programming Language :: Python :: 3.10
47
+ Classifier: Programming Language :: Python :: 3.11
48
+ Classifier: Programming Language :: Python :: 3.12
49
+ Requires-Python: <3.14,>=3.9
50
+ Description-Content-Type: text/markdown
51
+ License-File: LICENSE
52
+ Requires-Dist: numpy<1.27.0,>=1.22.0
53
+ Requires-Dist: pandas<2.4.0,>=1.3.0
54
+ Requires-Dist: matplotlib<3.11.0,>=3.10.6
55
+ Requires-Dist: ase<3.27.0,>=3.21.0
56
+ Requires-Dist: scipy<1.17.0,>=1.4.0
57
+ Requires-Dist: pyiron-workflow<=0.15.2,>=0.14.1
58
+ Requires-Dist: pymatgen<2025.7.0,>=2024.8.8
59
+ Requires-Dist: pyiron_snippets<=0.2.0,>=0.1.0
60
+ Requires-Dist: scikit-learn<1.8.0,>=1.0.0
61
+ Dynamic: license-file
62
+
63
+ # pyiron_workflow_atomistics
64
+
65
+ ## Overview
66
+
67
+ This repository contains a pyiron module for atomistic simulation workflows, providing tools and utilities for working with atomic structures, grain boundaries, and various atomistic calculations.
68
+
69
+ ## Features
70
+
71
+ - **Grain Boundary Analysis**: Tools for analyzing and manipulating grain boundaries, including:
72
+ - GB plane detection and analysis
73
+ - Cleavage plane identification
74
+ - Structure manipulation for GB studies
75
+
76
+ - **Structure Manipulation**: Utilities for working with atomic structures:
77
+ - Bulk structure handling
78
+ - Structure featurization
79
+ - Calculator integration
80
+
81
+ - **Workflow Integration**: Seamless integration with pyiron workflow system for:
82
+ - Automated structure calculations
83
+ - Data processing and analysis
84
+ - Results visualization
85
+
86
+ ## Installation
87
+
88
+ The package can be installed via pip:
89
+
90
+ ```bash
91
+ pip install pyiron_workflow_atomistics
92
+ ```
93
+
94
+ Or via conda:
95
+
96
+ ```bash
97
+ conda install -c conda-forge pyiron_workflow_atomistics
98
+ ```
99
+
100
+ ## Dependencies
101
+
102
+ The package requires:
103
+ - Python >= 3.9, < 3.13
104
+ - numpy < 2.0.0
105
+ - pandas >= 1.3.0
106
+ - matplotlib >= 3.4.0
107
+ - ase >= 3.22.0
108
+ - scipy >= 1.7.0
109
+ - pyiron_workflow
110
+ - pymatgen >= 2024.8.8
111
+ - pyiron_snippets
112
+ - scikit-learn >= 1.0.0
113
+
114
+ ## Usage
115
+
116
+ ### Grain Boundary Analysis
117
+
118
+ ```python
119
+ from pyiron_workflow_atomistics.gb.analysis import find_GB_plane
120
+ from pyiron_workflow_atomistics.gb.cleavage import cleave_gb_structure
121
+
122
+ # Find GB plane in a structure
123
+ gb_info = find_GB_plane(atoms, featuriser, axis="c")
124
+
125
+ # Cleave structure at GB
126
+ cleaved_structures, cleavage_planes = cleave_gb_structure(
127
+ base_structure=atoms,
128
+ axis_to_cleave="c",
129
+ target_coord=target_coord
130
+ )
131
+ ```
132
+
133
+ ### Structure Calculations
134
+
135
+ ```python
136
+ from pyiron_workflow_atomistics.calculator import calculate_structure_node
137
+
138
+ # Run structure calculations
139
+ results = calculate_structure_node(
140
+ structure=atoms,
141
+ calc=calculator,
142
+ output_dir="calculations"
143
+ )
144
+ ```
145
+
146
+ ## Documentation
147
+
148
+ For detailed documentation, visit our [ReadTheDocs page](https://pyiron_workflow_atomistics.readthedocs.io).
149
+
150
+ ## Contributing
151
+
152
+ We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.rst) for details.
153
+
154
+ ## License
155
+
156
+ This project is licensed under the BSD License - see the [LICENSE](LICENSE) file for details.
157
+
158
+ ## Citation
159
+
160
+ If you use this package in your research, please cite:
161
+
162
+ ```bibtex
163
+ @software{pyiron_workflow_atomistics,
164
+ author = {pyiron team},
165
+ title = {pyiron_workflow_atomistics},
166
+ year = {2024},
167
+ url = {https://github.com/pyiron/pyiron_workflow_atomistics}
168
+ }
169
+ ```
@@ -0,0 +1,107 @@
1
+ # pyiron_workflow_atomistics
2
+
3
+ ## Overview
4
+
5
+ This repository contains a pyiron module for atomistic simulation workflows, providing tools and utilities for working with atomic structures, grain boundaries, and various atomistic calculations.
6
+
7
+ ## Features
8
+
9
+ - **Grain Boundary Analysis**: Tools for analyzing and manipulating grain boundaries, including:
10
+ - GB plane detection and analysis
11
+ - Cleavage plane identification
12
+ - Structure manipulation for GB studies
13
+
14
+ - **Structure Manipulation**: Utilities for working with atomic structures:
15
+ - Bulk structure handling
16
+ - Structure featurization
17
+ - Calculator integration
18
+
19
+ - **Workflow Integration**: Seamless integration with pyiron workflow system for:
20
+ - Automated structure calculations
21
+ - Data processing and analysis
22
+ - Results visualization
23
+
24
+ ## Installation
25
+
26
+ The package can be installed via pip:
27
+
28
+ ```bash
29
+ pip install pyiron_workflow_atomistics
30
+ ```
31
+
32
+ Or via conda:
33
+
34
+ ```bash
35
+ conda install -c conda-forge pyiron_workflow_atomistics
36
+ ```
37
+
38
+ ## Dependencies
39
+
40
+ The package requires:
41
+ - Python >= 3.9, < 3.13
42
+ - numpy < 2.0.0
43
+ - pandas >= 1.3.0
44
+ - matplotlib >= 3.4.0
45
+ - ase >= 3.22.0
46
+ - scipy >= 1.7.0
47
+ - pyiron_workflow
48
+ - pymatgen >= 2024.8.8
49
+ - pyiron_snippets
50
+ - scikit-learn >= 1.0.0
51
+
52
+ ## Usage
53
+
54
+ ### Grain Boundary Analysis
55
+
56
+ ```python
57
+ from pyiron_workflow_atomistics.gb.analysis import find_GB_plane
58
+ from pyiron_workflow_atomistics.gb.cleavage import cleave_gb_structure
59
+
60
+ # Find GB plane in a structure
61
+ gb_info = find_GB_plane(atoms, featuriser, axis="c")
62
+
63
+ # Cleave structure at GB
64
+ cleaved_structures, cleavage_planes = cleave_gb_structure(
65
+ base_structure=atoms,
66
+ axis_to_cleave="c",
67
+ target_coord=target_coord
68
+ )
69
+ ```
70
+
71
+ ### Structure Calculations
72
+
73
+ ```python
74
+ from pyiron_workflow_atomistics.calculator import calculate_structure_node
75
+
76
+ # Run structure calculations
77
+ results = calculate_structure_node(
78
+ structure=atoms,
79
+ calc=calculator,
80
+ output_dir="calculations"
81
+ )
82
+ ```
83
+
84
+ ## Documentation
85
+
86
+ For detailed documentation, visit our [ReadTheDocs page](https://pyiron_workflow_atomistics.readthedocs.io).
87
+
88
+ ## Contributing
89
+
90
+ We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.rst) for details.
91
+
92
+ ## License
93
+
94
+ This project is licensed under the BSD License - see the [LICENSE](LICENSE) file for details.
95
+
96
+ ## Citation
97
+
98
+ If you use this package in your research, please cite:
99
+
100
+ ```bibtex
101
+ @software{pyiron_workflow_atomistics,
102
+ author = {pyiron team},
103
+ title = {pyiron_workflow_atomistics},
104
+ year = {2024},
105
+ url = {https://github.com/pyiron/pyiron_workflow_atomistics}
106
+ }
107
+ ```
@@ -0,0 +1,6 @@
1
+ from . import _version
2
+
3
+ __version__ = _version.get_versions()["version"]
4
+
5
+ from . import _version
6
+ __version__ = _version.get_versions()['version']
@@ -0,0 +1,21 @@
1
+
2
+ # This file was generated by 'versioneer.py' (0.29) from
3
+ # revision-control system data, or from the parent directory name of an
4
+ # unpacked source archive. Distribution tarballs contain a pre-generated copy
5
+ # of this file.
6
+
7
+ import json
8
+
9
+ version_json = '''
10
+ {
11
+ "date": "2025-11-02T19:31:27+0100",
12
+ "dirty": true,
13
+ "error": null,
14
+ "full-revisionid": "bf2c5ea91d7e3cb7e01fdd83c155d02a8961f49d",
15
+ "version": "0.0.3"
16
+ }
17
+ ''' # END VERSION_JSON
18
+
19
+
20
+ def get_versions():
21
+ return json.loads(version_json)
@@ -0,0 +1,326 @@
1
+ import json
2
+ import os
3
+ from typing import Any, Callable, Dict, List, Optional, Tuple
4
+
5
+ import numpy as np
6
+ import pandas as pd
7
+ import pyiron_workflow as pwf
8
+ from ase import Atoms
9
+ from ase.calculators.calculator import Calculator
10
+ from ase.io import write as ase_write
11
+ from ase.optimize import BFGS
12
+
13
+
14
+ def ase_calc_structure(
15
+ structure: Atoms,
16
+ calc: Calculator,
17
+ optimizer_class=BFGS,
18
+ optimizer_kwargs: Optional[Dict[str, Any]] = None,
19
+ record_interval: int = 1,
20
+ fmax: float = 0.01,
21
+ max_steps: int = 10000,
22
+ properties: Tuple[str, ...] = ("energy", "forces", "stresses", "volume"),
23
+ write_to_disk: bool = False,
24
+ working_directory: str = "calc_output",
25
+ initial_struct_path: Optional[str] = "initial_structure.xyz",
26
+ initial_results_path: Optional[str] = "initial_results.json",
27
+ traj_struct_path: Optional[str] = "trajectory.xyz",
28
+ traj_results_path: Optional[str] = "trajectory_results.json",
29
+ final_struct_path: Optional[str] = "final_structure.xyz",
30
+ final_results_path: Optional[str] = "final_results.json",
31
+ data_pickle: str = "job_data.pkl.gz",
32
+ ) -> Dict[str, Any]:
33
+ """
34
+ Relax an ASE Atoms object with a customizable optimizer and recording interval,
35
+ attach properties to each snapshot and write extended XYZ,
36
+ and store trajectory data as a pickled DataFrame including structures.
37
+
38
+ Returns
39
+ -------
40
+ dict with keys:
41
+ - initial: {{structure, results}}
42
+ - trajectory: list of {{structure, results}}
43
+ - final: {{structure, results}}
44
+ - converged: bool
45
+ """
46
+ # Setup
47
+ props = [p.strip() for p in properties]
48
+ os.makedirs(working_directory, exist_ok=True)
49
+ optimizer_kwargs = optimizer_kwargs or {}
50
+
51
+ def gather(atoms: Atoms) -> Dict[str, Any]:
52
+ all_results: Dict[str, Any] = {
53
+ "energy": atoms.get_potential_energy(),
54
+ "forces": atoms.get_forces().tolist(),
55
+ "cell": atoms.get_cell().tolist(),
56
+ "volume": atoms.get_volume(),
57
+ "positions": atoms.get_positions().tolist(),
58
+ "numbers": atoms.get_atomic_numbers().tolist(),
59
+ "masses": atoms.get_masses().tolist(),
60
+ }
61
+ if "stresses" in props:
62
+ try:
63
+ all_results["stresses"] = atoms.get_stress().tolist()
64
+ except Exception:
65
+ pass
66
+ mapping = {
67
+ "charges": "get_charges",
68
+ "dipole": "get_dipole_moment",
69
+ "magmoms": "get_magnetic_moments",
70
+ "virial": "get_virial",
71
+ "pressure": "get_pressure",
72
+ }
73
+ for key, method in mapping.items():
74
+ if key in props:
75
+ try:
76
+ val = getattr(atoms, method)()
77
+ all_results[key] = val.tolist() if hasattr(val, "tolist") else val
78
+ except Exception:
79
+ pass
80
+ missing = [p for p in props if p not in all_results]
81
+ if missing:
82
+ raise KeyError(f"Requested properties not available: {missing}")
83
+ return {p: all_results[p] for p in props}
84
+
85
+ def attach_props(atoms: Atoms, results: Dict[str, Any]):
86
+ # Attach energy
87
+ if "energy" in results:
88
+ atoms.info["energy"] = results["energy"]
89
+ # Attach forces
90
+ if "forces" in results:
91
+ atoms.set_array("forces", np.array(results["forces"]))
92
+ # Attach stresses
93
+ if "stresses" in results:
94
+ atoms.info["stresses"] = results["stresses"]
95
+ return atoms
96
+
97
+ atoms = structure.copy()
98
+ atoms.calc = calc
99
+
100
+ # Initial snapshot
101
+ initial_res = gather(atoms)
102
+ initial_atoms = attach_props(atoms.copy(), initial_res)
103
+ initial = {"structure": initial_atoms, "results": initial_res}
104
+ if write_to_disk and initial_struct_path:
105
+ ase_write(os.path.join(working_directory, initial_struct_path), initial_atoms)
106
+ if write_to_disk and initial_results_path:
107
+ with open(os.path.join(working_directory, initial_results_path), "w") as f:
108
+ json.dump(initial_res, f, indent=2)
109
+
110
+ # Trajectory recording
111
+ trajectory: List[Dict[str, Any]] = []
112
+
113
+ def record_step():
114
+ snap = atoms.copy()
115
+ snap_res = gather(atoms)
116
+ snap_att = attach_props(snap, snap_res)
117
+ trajectory.append({"structure": snap_att, "results": snap_res})
118
+ if write_to_disk and traj_struct_path:
119
+ ase_write(
120
+ os.path.join(working_directory, traj_struct_path), snap_att, append=True
121
+ )
122
+
123
+ # Optimize
124
+ optimizer = optimizer_class(atoms, **optimizer_kwargs)
125
+ optimizer.attach(record_step, interval=record_interval)
126
+ converged = optimizer.run(fmax=fmax, steps=max_steps)
127
+
128
+ # Write trajectory results JSON
129
+ if write_to_disk and traj_results_path:
130
+ traj_res_list = [step["results"] for step in trajectory]
131
+ with open(os.path.join(working_directory, traj_results_path), "w") as f:
132
+ json.dump(traj_res_list, f, indent=2)
133
+
134
+ # Final snapshot
135
+ final_res = gather(atoms)
136
+ final_atoms = attach_props(atoms.copy(), final_res)
137
+ # print(final_res)
138
+ final = {"structure": final_atoms, "results": final_res}
139
+ if write_to_disk and final_struct_path:
140
+ ase_write(os.path.join(working_directory, final_struct_path), final_atoms)
141
+ if write_to_disk and final_results_path:
142
+ with open(os.path.join(working_directory, final_results_path), "w") as f:
143
+ json.dump(final_res, f, indent=2)
144
+
145
+ # Build DataFrame including structures
146
+ df = pd.DataFrame(
147
+ [{"structure": step["structure"], **step["results"]} for step in trajectory]
148
+ )
149
+ df.to_pickle(os.path.join(working_directory, data_pickle), compression="gzip")
150
+
151
+ return {
152
+ "initial": initial,
153
+ "trajectory": trajectory,
154
+ "final": final,
155
+ "converged": bool(converged),
156
+ }
157
+
158
+
159
+ def ase_calculate_structure_node_interface(
160
+ structure: Atoms,
161
+ calc: Calculator,
162
+ optimizer_class=BFGS,
163
+ optimizer_kwargs: dict[str, Any] | None = None,
164
+ record_interval: int = 1,
165
+ fmax: float = 0.01,
166
+ max_steps: int = 10000,
167
+ properties: Tuple[str, ...] = ("energy", "forces", "stresses"),
168
+ write_to_disk: bool = False,
169
+ working_directory: str = "calc_output",
170
+ initial_struct_path: Optional[str] = "initial_structure.xyz",
171
+ initial_results_path: Optional[str] = "initial_results.json",
172
+ traj_struct_path: Optional[str] = "trajectory.xyz",
173
+ traj_results_path: Optional[str] = "trajectory_results.json",
174
+ final_struct_path: Optional[str] = "final_structure.xyz",
175
+ final_results_path: Optional[str] = "final_results.json",
176
+ data_pickle: str = "job_data.pkl.gz",
177
+ calc_structure_fn: Callable[..., Any] = ase_calc_structure,
178
+ ) -> Tuple[Atoms, Dict[str, Any], bool]:
179
+ """
180
+ Node wrapper to call calc_structure_fn (default: ase_calc_structure)
181
+ with all ASE kwargs forwarded from node inputs.
182
+ """
183
+ out = calc_structure_fn(
184
+ structure=structure,
185
+ calc=calc,
186
+ optimizer_class=optimizer_class,
187
+ optimizer_kwargs=optimizer_kwargs,
188
+ record_interval=record_interval,
189
+ fmax=fmax,
190
+ max_steps=max_steps,
191
+ properties=properties,
192
+ write_to_disk=write_to_disk,
193
+ working_directory=working_directory,
194
+ initial_struct_path=initial_struct_path,
195
+ initial_results_path=initial_results_path,
196
+ traj_struct_path=traj_struct_path,
197
+ traj_results_path=traj_results_path,
198
+ final_struct_path=final_struct_path,
199
+ final_results_path=final_results_path,
200
+ data_pickle=data_pickle,
201
+ )
202
+ atoms = out["final"]["structure"]
203
+ final_results = out["final"]["results"]
204
+ converged = out["converged"]
205
+ converged = bool(converged)
206
+ # print(atoms, final_results, converged)
207
+ return atoms, final_results, converged
208
+
209
+
210
+ @pwf.as_function_node("atoms", "results", "converged")
211
+ def calculate_structure_node(
212
+ structure: Atoms,
213
+ # calc_structure_fn: Callable[..., Any] = ase_calculate_structure_node_interface,
214
+ calc_structure_fn=ase_calculate_structure_node_interface,
215
+ calc_structure_fn_kwargs: dict[str, Any] | None = None,
216
+ ) -> Tuple[Atoms, dict[str, Any], bool]:
217
+ if calc_structure_fn_kwargs is None:
218
+ calc_structure_fn_kwargs = {}
219
+ atoms, final_results, converged = calc_structure_fn(
220
+ structure=structure, **calc_structure_fn_kwargs
221
+ )
222
+ # print(calc_structure_fn_kwargs["working_directory"])
223
+ return atoms, final_results, converged
224
+
225
+
226
+ @pwf.as_function_node("output")
227
+ def extract_values(results_list, key):
228
+ """
229
+ Extract a list of values for a specified key from a list of result dictionaries.
230
+
231
+ Parameters
232
+ ----------
233
+ results_list : list of dict
234
+ Each dict should contain the specified key.
235
+ key : str
236
+ The dictionary key to extract values for (e.g., 'energy', 'volume').
237
+
238
+ Returns
239
+ -------
240
+ values : list
241
+ List of values corresponding to key from each dict.
242
+
243
+ Raises
244
+ ------
245
+ KeyError
246
+ If any entry in results_list is missing the specified key.
247
+ """
248
+ try:
249
+ extracted_values = [entry[key] for entry in results_list]
250
+ except Exception as e:
251
+ # print(results_list, key)
252
+ print(f"Error {e} when trying to parse output")
253
+ extracted_values = np.nan
254
+ return extracted_values
255
+
256
+
257
+ @pwf.as_function_node("full_calc_kwargs2")
258
+ def fillin_default_calckwargs(
259
+ calc_kwargs: dict[str, Any],
260
+ default_values: dict[str, Any] | None | str = None,
261
+ remove_keys: list[str] | None = None,
262
+ ) -> dict[str, Any]:
263
+ # 1) overlay any user-supplied default overrides
264
+ built_in = {}
265
+ if isinstance(default_values, dict):
266
+ built_in.update(default_values)
267
+
268
+ # 2) start with everything user passed in
269
+ full: dict[str, Any] = dict(calc_kwargs)
270
+
271
+ # 3) fill in missing built-ins
272
+ for key, default in built_in.items():
273
+ full.setdefault(key, default)
274
+
275
+ # 4) ensure properties is a tuple
276
+ if "properties" in full:
277
+ full["properties"] = tuple(full["properties"])
278
+
279
+ # 5) remove any keys requested
280
+ if remove_keys:
281
+ for key in remove_keys:
282
+ full.pop(key, None)
283
+
284
+ return full
285
+
286
+
287
+ @pwf.as_function_node("kwargs_variant")
288
+ def generate_kwargs_variant(
289
+ base_kwargs: dict[str, Any],
290
+ key: str,
291
+ value: Any,
292
+ ):
293
+ from copy import deepcopy
294
+
295
+ kwargs_variant = deepcopy(base_kwargs)
296
+ kwargs_variant[key] = value
297
+ return kwargs_variant
298
+
299
+
300
+ @pwf.as_function_node("kwargs_variants")
301
+ def generate_kwargs_variants(
302
+ base_kwargs: dict[str, Any],
303
+ key: str,
304
+ values: list[Any],
305
+ ):
306
+ """
307
+ Given a base kwargs dict, produce one dict per value in `values`,
308
+ each with `key` set to that value (overriding any existing entry).
309
+
310
+ Parameters
311
+ ----------
312
+ base_kwargs
313
+ The original kwargs to copy.
314
+ key
315
+ The dict key whose value you want to vary.
316
+ values
317
+ A list of values to assign to `key`.
318
+
319
+ Returns
320
+ -------
321
+ List of dicts
322
+ Each is a shallow copy of base_kwargs with base_kwargs[key] = value.
323
+ """
324
+ return_kwargs = [{**base_kwargs, key: v} for v in values]
325
+ # print(return_kwargs)
326
+ return return_kwargs