cleanup (kwargs for SpacegroupAnalyzer, supercell)

fekad · Jan 28, 2021 · 7951435 · 7951435
1 parent bcef468
commit 7951435
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 59 deletions.
diff --git a/examples/4 - Interface for pymatgen.ipynb b/examples/4 - Interface for pymatgen.ipynb
@@ -38,7 +38,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
     "quick_view(struct)"
@@ -47,7 +49,9 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": false
+   },
    "outputs": [],
    "source": [
     "quick_view(struct, conventional=True)"
@@ -75,32 +79,14 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Using extra paramaters"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "All positional arguments will be passed to jsmol's load command. You can use jsmol's load command to build supercell."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
-   "outputs": [],
-   "source": [
-    "quick_view(struct, \"{2 2 2}\", conventional=True)"
+    "## Using extra parameters"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Any extra kw arguments will be passed to th pymatgens CifWriter constructor."
+    "All positional arguments will be passed to jsmol's load command."
    ]
   },
   {
@@ -111,27 +97,7 @@
    },
    "outputs": [],
    "source": [
-    "quick_view(struct, \"{2 2 2}\", conventional=True, symprec=0.01)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
-   "outputs": [],
-   "source": [
-    "quick_view(struct, \"packed\", conventional=True, symprec=0.01)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "quick_view(struct, \"{2 2 2}\", \"packed\", conventional=True, symprec=0.01)"
+    "quick_view(struct, \"packed\", conventional=True)"
    ]
   },
   {
@@ -142,13 +108,6 @@
    "source": [
     "?quick_view"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

diff --git a/jupyter_jsmol/pymatgen.py b/jupyter_jsmol/pymatgen.py
@@ -8,30 +8,46 @@
 
     pymatgen_loaded = True
 except ImportError:
-    raise RuntimeError('To use quick_view, you need to have pymatgen installed.')
+    raise RuntimeError("To use quick_view, you need to have pymatgen installed.")
 
 from jupyter_jsmol import JsmolView
 
-def quick_view(structure: Structure, *args, conventional: bool = False, supercell: list = None, **kwargs) -> JsmolView:
+
+def quick_view(
+    structure: Structure,
+    *args,
+    conventional: bool = False,
+    supercell: list = [1, 1, 1],
+    symprec: float = 0.01,
+    angle_tolerance: float = 5.0
+) -> JsmolView:
     """A function to visualize pymatgen Structure objects in jupyter notebook using jupyter_jsmol package.
 
     Args:
         structure: pymatgen Structure object.
+        *args: Extra arguments for JSmol's load command. Eg. "{2 2 2}", "packed"
         conventional: use conventional cell. Defaults to False.
         supercell: can be used to make supercells with pymatgen.Structure.make_supercell method.
-        *args: Extra arguments for JSmol's load command. Eg. "{2 2 2}", "packed"
-        **kwargs: Kwargs passthru to CifWriter methods. E.g., This allows
-                the passing of parameters like symprec=0.01 for generation of symmetric cifs.
+        symprec: If not none, finds the symmetry of the structure
+            and writes the cif with symmetry information. Passes symprec
+            to the SpacegroupAnalyzer.
+        angle_tolerance: Angle tolerance for symmetry finding. Passes
+            angle_tolerance to the SpacegroupAnalyzer. Used only if symprec
+            is not None.
 
     Returns:
         A jupyter widget object.
     """
 
     s = structure.copy()
     if conventional:
-        s = SpacegroupAnalyzer(s).get_conventional_standard_structure()
+        spga = SpacegroupAnalyzer(s, symprec=symprec, angle_tolerance=angle_tolerance)
+        s = spga.get_conventional_standard_structure()
+
+    cif = CifWriter(
+        s, symprec=symprec, angle_tolerance=angle_tolerance, refine_struct=False
+    )
 
-    if supercell:
-        s.make_supercell(supercell)
+    supercell_str = "{" + " ".join(map(str, supercell)) + "}"
 
-    return JsmolView.from_str(str(CifWriter(s, **kwargs)), *args)
+    return JsmolView.from_str(str(cif), supercell_str, *args)