Merge pull request #140 from Exabyte-io/update/SOF-7381

Update/SOF-7381 Import and reexport all 3rd party from one place

src/py/mat3ra/made/tools/analyze.py

@@ -1,24 +1,21 @@
 from typing import List, Optional
 
 import numpy as np
-from ase import Atoms
-from pymatgen.core import IStructure as PymatgenIStructure
 
 from ..material import Material
 from .convert import decorator_convert_material_args_kwargs_to_atoms, to_pymatgen
-
-PymatgenIStructure = PymatgenIStructure
+from .third_party import ASEAtoms, PymatgenIStructure
 
 
 @decorator_convert_material_args_kwargs_to_atoms
 def get_average_interlayer_distance(
-    interface_atoms: Atoms, tag_substrate: str, tag_film: str, threshold: float = 0.5
+    interface_atoms: ASEAtoms, tag_substrate: str, tag_film: str, threshold: float = 0.5
 ) -> float:
     """
     Calculate the average distance between the top layer of substrate atoms and the bottom layer of film atoms.
 
     Args:
-        interface_atoms (ase.Atoms): The ASE Atoms object containing both sets of atoms.
+        interface_atoms (ase.ASEAtoms): The ASE ASEAtoms object containing both sets of atoms.
         tag_substrate (int): The tag representing the substrate atoms.
         tag_film (int): The tag representing the film atoms.
         threshold (float): The threshold for identifying the top and bottom layers of atoms.
@@ -48,12 +45,12 @@ def get_average_interlayer_distance(
 
 
 @decorator_convert_material_args_kwargs_to_atoms
-def get_surface_area(atoms: Atoms):
+def get_surface_area(atoms: ASEAtoms):
     """
     Calculate the area of the surface perpendicular to the z-axis of the atoms structure.
 
     Args:
-        atoms (ase.Atoms): The Atoms object to calculate the surface area of.
+        atoms (ase.ASEAtoms): The ASEAtoms object to calculate the surface area of.
 
     Returns:
         float: The surface area of the atoms.
@@ -64,12 +61,12 @@ def get_surface_area(atoms: Atoms):
 
 
 @decorator_convert_material_args_kwargs_to_atoms
-def get_chemical_formula(atoms: Atoms):
+def get_chemical_formula(atoms: ASEAtoms):
     """
     Calculate the formula of the atoms structure.
 
     Args:
-        atoms (ase.Atoms): The Atoms object to calculate the formula of.
+        atoms (ase.ASEAtoms): The ASEAtoms object to calculate the formula of.
 
     Returns:
         str: The formula of the atoms.

src/py/mat3ra/made/tools/build/defect/builders.py

@@ -2,15 +2,16 @@
 
 from mat3ra.made.material import Material
 from pydantic import BaseModel
-from pymatgen.analysis.defects.core import (
-    Substitution as PymatgenSubstitution,
-    Vacancy as PymatgenVacancy,
-    Interstitial as PymatgenInterstitial,
-)
-from pymatgen.core import PeriodicSite as PymatgenPeriodicSite
 
+from ...third_party import (
+    PymatgenStructure,
+    PymatgenPeriodicSite,
+    PymatgenVacancy,
+    PymatgenSubstitution,
+    PymatgenInterstitial,
+)
 from ...build import BaseBuilder
-from ...convert import PymatgenStructure, to_pymatgen
+from ...convert import to_pymatgen
 from ..mixins import ConvertGeneratedItemsPymatgenStructureMixin
 from .configuration import PointDefectConfiguration
 

src/py/mat3ra/made/tools/build/mixins.py

@@ -1,4 +1,5 @@
-from ..convert import from_ase, from_pymatgen, ASEAtoms, PymatgenStructure
+from ..convert import from_ase, from_pymatgen
+from ..third_party import ASEAtoms, PymatgenStructure
 
 
 class ConvertGeneratedItemsASEAtomsMixin:

src/py/mat3ra/made/tools/build/slab/builders.py

@@ -1,18 +1,16 @@
-from pymatgen.core.surface import SlabGenerator as PymatgenSlabGenerator
-from ...convert import label_pymatgen_slab_termination
 from typing import List
 from pydantic import BaseModel
 
-
 from mat3ra.made.material import Material
 
-from .termination import Termination
+from ...third_party import PymatgenSlab, PymatgenSlabGenerator, label_pymatgen_slab_termination
 from ...analyze import get_chemical_formula
-from ...convert import to_pymatgen, PymatgenSlab
+from ...convert import to_pymatgen
 from ...build import BaseBuilder
 from ...build.mixins import ConvertGeneratedItemsPymatgenStructureMixin
 from ..supercell import create_supercell
 from .configuration import SlabConfiguration
+from .termination import Termination
 
 
 class SlabSelectorParameters(BaseModel):

src/py/mat3ra/made/tools/build/slab/configuration.py

@@ -1,10 +1,10 @@
 from typing import List, Tuple, Any
+from pydantic import BaseModel
 
 from mat3ra.code.entity import InMemoryEntity
-from pymatgen.symmetry.analyzer import SpacegroupAnalyzer as PymatgenSpacegroupAnalyzer
-from pydantic import BaseModel
 
 from mat3ra.made.material import Material
+from ...third_party import PymatgenSpacegroupAnalyzer
 from ...convert import to_pymatgen, from_pymatgen
 
 

src/py/mat3ra/made/tools/build/supercell.py

@@ -1,15 +1,14 @@
 from typing import List
-from ase import Atoms
-from ase.build.supercells import make_supercell
 
 from mat3ra.made.material import Material
+from ..third_party import ASEAtoms, ase_make_supercell
 from ..utils import decorator_convert_2x2_to_3x3
 from ..convert import from_ase, decorator_convert_material_args_kwargs_to_atoms
 
 
 @decorator_convert_2x2_to_3x3
 @decorator_convert_material_args_kwargs_to_atoms
-def create_supercell(atoms: Atoms, supercell_matrix: List[List[int]]) -> Material:
+def create_supercell(atoms: ASEAtoms, supercell_matrix: List[List[int]]) -> Material:
     """
     Create a supercell of the atoms.
 
@@ -21,5 +20,5 @@ def create_supercell(atoms: Atoms, supercell_matrix: List[List[int]]) -> Materia
         Material: The supercell of the atoms.
     """
 
-    supercell_atoms = make_supercell(atoms, supercell_matrix)
+    supercell_atoms = ase_make_supercell(atoms, supercell_matrix)
     return Material(from_ase(supercell_atoms))

src/py/mat3ra/made/tools/calculate.py

@@ -1,22 +1,19 @@
 from typing import Optional
 
-from ase import Atoms
-from ase.calculators.calculator import Calculator
-from ase.calculators.emt import EMT
-
 from ..material import Material
 from .analyze import get_surface_area
 from .build.interface.utils import get_slab
 from .convert import decorator_convert_material_args_kwargs_to_atoms
+from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
 
 
 @decorator_convert_material_args_kwargs_to_atoms
-def calculate_total_energy(atoms: Atoms, calculator: Calculator):
+def calculate_total_energy(atoms: ASEAtoms, calculator: ASECalculator):
     """
     Set calculator for ASE Atoms and calculate the total energy.
 
     Args:
-        atoms (ase.Atoms): The Atoms object to calculate the energy of.
+        atoms (ASEAtoms): The Atoms object to calculate the energy of.
         calculator (ase.calculators.calculator.Calculator): The calculator to use for the energy calculation.
 
     Returns:
@@ -27,12 +24,12 @@ def calculate_total_energy(atoms: Atoms, calculator: Calculator):
 
 
 @decorator_convert_material_args_kwargs_to_atoms
-def calculate_total_energy_per_atom(atoms: Atoms, calculator: Calculator):
+def calculate_total_energy_per_atom(atoms: ASEAtoms, calculator: ASECalculator):
     """
     Set calculator for ASE Atoms and calculate the total energy per atom.
 
     Args:
-            atoms (ase.Atoms): The Atoms object to calculate the energy of.
+            atoms (ASEAtoms): The Atoms object to calculate the energy of.
             calculator (ase.calculators.calculator.Calculator): The calculator to use for the energy calculation.
 
     Returns:
@@ -42,13 +39,13 @@ def calculate_total_energy_per_atom(atoms: Atoms, calculator: Calculator):
 
 
 @decorator_convert_material_args_kwargs_to_atoms
-def calculate_surface_energy(slab: Atoms, bulk: Atoms, calculator: Calculator):
+def calculate_surface_energy(slab: ASEAtoms, bulk: ASEAtoms, calculator: ASECalculator):
     """
     Calculate the surface energy by subtracting the weighted bulk energy from the slab energy.
 
     Args:
-        slab (ase.Atoms): The slab Atoms object to calculate the surface energy of.
-        bulk (ase.Atoms): The bulk Atoms object to calculate the surface energy of.
+        slab (ASEAtoms): The slab Atoms object to calculate the surface energy of.
+        bulk (ASEAtoms): The bulk Atoms object to calculate the surface energy of.
         calculator (ase.calculators.calculator.Calculator): The calculator to use for the energy calculation.
 
     Returns:
@@ -62,17 +59,19 @@ def calculate_surface_energy(slab: Atoms, bulk: Atoms, calculator: Calculator):
 
 
 @decorator_convert_material_args_kwargs_to_atoms
-def calculate_adhesion_energy(interface: Atoms, substrate_slab: Atoms, film_slab: Atoms, calculator: Calculator):
+def calculate_adhesion_energy(
+    interface: ASEAtoms, substrate_slab: ASEAtoms, film_slab: ASEAtoms, calculator: ASECalculator
+):
     """
     Calculate the adhesion energy.
     The adhesion energy is the difference between the energy of the interface and
     the sum of the energies of the substrate and film.
     According to: 10.1088/0953-8984/27/30/305004
 
     Args:
-        interface (ase.Atoms): The interface Atoms object to calculate the adhesion energy of.
-        substrate_slab (ase.Atoms): The substrate slab Atoms object to calculate the adhesion energy of.
-        film_slab (ase.Atoms): The film slab Atoms object to calculate the adhesion energy of.
+        interface (ASEAtoms): The interface ASEAtoms object to calculate the adhesion energy of.
+        substrate_slab (ASEAtoms): The substrate slab ASEAtoms object to calculate the adhesion energy of.
+        film_slab (ASEAtoms): The film slab ASEAtoms object to calculate the adhesion energy of.
         calculator (ase.calculators.calculator.Calculator): The calculator to use for the energy calculation.
 
     Returns:
@@ -91,7 +90,7 @@ def calculate_interfacial_energy(
     substrate_bulk: Optional[Material] = None,
     film_slab: Optional[Material] = None,
     film_bulk: Optional[Material] = None,
-    calculator: Calculator = EMT(),
+    calculator: ASECalculator = ASECalculatorEMT(),
 ):
     """
     Calculate the interfacial energy.

src/py/mat3ra/made/tools/convert/__init__.py

@@ -5,20 +5,22 @@
 from mat3ra.made.material import Material
 from mat3ra.made.utils import map_array_with_id_value_to_array
 from mat3ra.utils.mixins import RoundNumericValuesMixin
-from pymatgen.io.ase import AseAtomsAdaptor
-from pymatgen.io.vasp.inputs import Poscar
 
-from .utils import (
-    INTERFACE_LABELS_MAP,
+from ..third_party import (
     ASEAtoms,
+    PymatgenAseAtomsAdaptor,
     PymatgenInterface,
     PymatgenLattice,
+    PymatgenPoscar,
     PymatgenSlab,
     PymatgenStructure,
+    label_pymatgen_slab_termination,
+)
+from .utils import (
+    INTERFACE_LABELS_MAP,
     extract_labels_from_pymatgen_structure,
     extract_metadata_from_pymatgen_structure,
     extract_tags_from_ase_atoms,
-    label_pymatgen_slab_termination,
     map_array_to_array_with_id_value,
 )
 
@@ -138,7 +140,7 @@ def to_poscar(material_or_material_data: Union[Material, Dict[str, Any]]) -> str
         str: A POSCAR string.
     """
     structure = to_pymatgen(material_or_material_data)
-    poscar = Poscar(structure)
+    poscar = PymatgenPoscar(structure)
     # For pymatgen `2023.6.23` supporting py3.8 the method name is "get_string"
     # TODO: cleanup the if statement when dropping support for py3.8
     if hasattr(poscar, "get_string"):
@@ -175,7 +177,7 @@ def to_ase(material_or_material_data: Union[Material, Dict[str, Any]]) -> ASEAto
     else:
         material_config = material_or_material_data
     structure = to_pymatgen(material_config)
-    atoms = AseAtomsAdaptor.get_atoms(structure)
+    atoms = PymatgenAseAtomsAdaptor.get_atoms(structure)
 
     atomic_labels = material_config["basis"].get("labels", [])
     if atomic_labels:
@@ -196,7 +198,7 @@ def from_ase(ase_atoms: ASEAtoms) -> Dict[str, Any]:
         dict: A dictionary containing the material information in ESSE format.
     """
     # TODO: check that atomic labels/tags are properly handled
-    structure = AseAtomsAdaptor.get_structure(ase_atoms)
+    structure = PymatgenAseAtomsAdaptor.get_structure(ase_atoms)
     material = from_pymatgen(structure)
     ase_tags = extract_tags_from_ase_atoms(ase_atoms)
     material["basis"]["labels"] = ase_tags

src/py/mat3ra/made/tools/convert/utils.py

@@ -1,22 +1,10 @@
 import json
 from typing import Any, Dict, List, Union
 
-from ase import Atoms as ASEAtoms
 from mat3ra.made.utils import map_array_to_array_with_id_value
 from mat3ra.utils.object import NumpyNDArrayRoundEncoder
-from pymatgen.core.interface import Interface as PymatgenInterface
-from pymatgen.core.interface import label_termination
-from pymatgen.core.structure import Lattice as PymatgenLattice
-from pymatgen.core.structure import Structure as PymatgenStructure
-from pymatgen.core.surface import Slab as PymatgenSlab
-
-# Re-exported imports to allow for both use in type hints and instantiation
-PymatgenLattice = PymatgenLattice
-PymatgenStructure = PymatgenStructure
-PymatgenSlab = PymatgenSlab
-PymatgenInterface = PymatgenInterface
-ASEAtoms = ASEAtoms
-label_pymatgen_slab_termination = label_termination
+
+from ..third_party import ASEAtoms, PymatgenInterface, PymatgenStructure
 
 INTERFACE_LABELS_MAP = {"substrate": 0, "film": 1}
 

src/py/mat3ra/made/tools/modify.py

@@ -1,11 +1,10 @@
 from typing import List, Union
 
 from mat3ra.made.material import Material
-from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer
-from pymatgen.core.structure import Structure
 
 from .analyze import get_atom_indices_within_layer_by_atom_index, get_atom_indices_within_radius_pbc
 from .convert import decorator_convert_material_args_kwargs_to_structure
+from .third_party import PymatgenSpacegroupAnalyzer, PymatgenStructure
 from .utils import translate_to_bottom_pymatgen_structure
 
 
@@ -30,7 +29,7 @@ def filter_by_label(material: Material, label: Union[int, str]) -> Material:
 
 
 @decorator_convert_material_args_kwargs_to_structure
-def translate_to_bottom(structure: Structure, use_conventional_cell: bool = True):
+def translate_to_bottom(structure: PymatgenStructure, use_conventional_cell: bool = True):
     """
     Translate atoms to the bottom of the cell (vacuum on top) to allow for the correct consecutive interface generation.
     If use_conventional_cell is passed, conventional cell is used.
@@ -42,20 +41,20 @@ def translate_to_bottom(structure: Structure, use_conventional_cell: bool = True
         Structure: The normalized pymatgen Structure object.
     """
     if use_conventional_cell:
-        structure = SpacegroupAnalyzer(structure).get_conventional_standard_structure()
+        structure = PymatgenSpacegroupAnalyzer(structure).get_conventional_standard_structure()
     structure = translate_to_bottom_pymatgen_structure(structure)
     return structure
 
 
 @decorator_convert_material_args_kwargs_to_structure
-def wrap_to_unit_cell(structure: Structure):
+def wrap_to_unit_cell(structure: PymatgenStructure):
     """
     Wrap atoms to the cell
 
     Args:
-        structure (Structure): The pymatgen Structure object to normalize.
+        structure (PymatgenStructure): The pymatgen PymatgenStructure object to normalize.
     Returns:
-        Structure: The wrapped pymatgen Structure object.
+        PymatgenStructure: The wrapped pymatgen PymatgenStructure object.
     """
     structure.make_supercell((1, 1, 1), to_unit_cell=True)
     return structure