From 94c3969eb287857e48bf6316363c9d366261877f Mon Sep 17 00:00:00 2001
From: Sean Kavanagh <s.kavanagh19@imperial.ac.uk>
Date: Mon, 16 Dec 2024 13:25:15 -0500
Subject: [PATCH] Allow structure files or ASE atoms as input to
 `DefectsGenerator` for user convenience

---
 docs/Dev_ToDo.md    |  1 +
 doped/generation.py | 15 ++++++++++++---
 pyproject.toml      |  3 ++-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/docs/Dev_ToDo.md b/docs/Dev_ToDo.md
index 74178c21..189c976e 100644
--- a/docs/Dev_ToDo.md
+++ b/docs/Dev_ToDo.md
@@ -89,6 +89,7 @@
 
 ## SK To-Do for next update:
 - `doped` repo/docs cleanup `TODO`s above
+- Update generation notebook to use structure file string input
 - Quick run through tutorial notebooks to check code all updated and running.
 - Clean up repo, removing old unnecessary git blobs
 - Note in chempots tutorial that LaTeX table generator website can also be used with the `to_csv()` function to generate LaTeX tables for the competing phases.
diff --git a/doped/generation.py b/doped/generation.py
index 556b1474..be1a7691 100644
--- a/doped/generation.py
+++ b/doped/generation.py
@@ -12,7 +12,7 @@
 from functools import partial, reduce
 from itertools import chain
 from multiprocessing import Pool, cpu_count
-from typing import Any, Optional, Union, cast
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
 from unittest.mock import MagicMock
 
 import numpy as np
@@ -53,6 +53,9 @@
 from doped.utils.parsing import reorder_s1_like_s2
 from doped.utils.plotting import format_defect_name
 
+if TYPE_CHECKING:
+    from ase.atoms import Atoms
+
 _dummy_species = DummySpecies("X")  # Dummy species used to keep track of defect coords in the supercell
 
 core._logger.setLevel(logging.CRITICAL)  # avoid unnecessary pymatgen-analysis-defects warnings about
@@ -1198,7 +1201,7 @@ class DefectsGenerator(MSONable):
 
     def __init__(
         self,
-        structure: Structure,
+        structure: Union[Structure, "Atoms", PathLike],
         extrinsic: Optional[Union[str, list, dict]] = None,
         interstitial_coords: Optional[list] = None,
         generate_supercell: bool = True,
@@ -1267,7 +1270,8 @@ def __init__(
 
         Args:
             structure (Structure):
-                Structure of the host material (as a pymatgen Structure object).
+                Structure of the host material, either as a ``pymatgen`` ``Structure``,
+                ``ASE`` ``Atoms`` or path to a structure file (e.g. ``CONTCAR``).
                 If this is not the primitive unit cell, it will be reduced to the
                 primitive cell for defect generation, before supercell generation.
             extrinsic (Union[str, list, dict]):
@@ -1355,6 +1359,11 @@ def __init__(
         """
         self.defects: dict[str, list[Defect]] = {}  # {defect_type: [Defect, ...]}
         self.defect_entries: dict[str, DefectEntry] = {}  # {defect_species: DefectEntry}
+        if isinstance(structure, (str, PathLike)):
+            structure = Structure.from_file(structure)
+        elif not isinstance(structure, Structure):
+            structure = Structure.from_ase_atoms(structure)
+
         self.structure = structure
         self.extrinsic = extrinsic if extrinsic is not None else []
         self.kwargs = kwargs
diff --git a/pyproject.toml b/pyproject.toml
index 26c71cd4..a845900d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -24,6 +24,7 @@ dependencies = [
     "pymatgen>=2024.7.18",  # requirement for spglib>=2.5.0 for SpglibDataset attribute interface
     "pymatgen-analysis-defects>=2023.8.22",
     "shakenbreak>=3.3.4",  # oxi-state functions renaming
+    "ase",  # required by ShakeNBreak anyway
     "pandas",
     "pydefect>=0.8.1",  # for Kumagai (eFNV) correction & PHS analysis
     "filelock",
@@ -128,4 +129,4 @@ wrap-summaries = 79
 wrap-descriptions = 79
 
 [tool.codespell]
-ignore-words-list = "titel,te,bu,aack,unx,nd,efect,vise,mater,sise"  # lowercase!
+ignore-words-list = "titel,te,bu,aack,unx,nd,efect,vise,mater,sise,dscribe"  # lowercase!