From 0cc6ee9873385fc6c3bc6c51460ff012c8f84618 Mon Sep 17 00:00:00 2001
From: "Lori A. Burns" <lori.burns@gmail.com>
Date: Fri, 11 Nov 2022 15:09:48 -0500
Subject: [PATCH] noncontig passthrough

---
 qcelemental/models/molecule.py      | 10 ++++++++-
 qcelemental/molparse/from_schema.py |  8 ++++++--
 qcelemental/tests/test_molecule.py  | 32 +++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/qcelemental/models/molecule.py b/qcelemental/models/molecule.py
index 5944001f..e3690bdd 100644
--- a/qcelemental/models/molecule.py
+++ b/qcelemental/models/molecule.py
@@ -341,14 +341,22 @@ def __init__(self, orient: bool = False, validate: Optional[bool] = None, **kwar
             # original_keys = set(kwargs.keys())  # revive when ready to revisit sparsity
 
             nonphysical = kwargs.pop("nonphysical", False)
+            throw_reorder = not (kwargs.pop("allow_noncontiguous_but_reorder", False))  # experimental
             schema = to_schema(
-                from_schema(kwargs, nonphysical=nonphysical), dtype=kwargs["schema_version"], copy=False, np_out=True
+                from_schema(kwargs, nonphysical=nonphysical, throw_reorder=throw_reorder),
+                dtype=kwargs["schema_version"],
+                copy=False,
+                np_out=True,
             )
             schema = _filter_defaults(schema)
 
             kwargs["validated"] = True
             kwargs = {**kwargs, **schema}  # Allow any extra fields
             validate = True
+        else:
+            # these kwargs only relevant for validation
+            kwargs.pop("nonphysical", None)
+            kwargs.pop("allow_noncontiguous_but_reorder", None)
 
         super().__init__(**kwargs)
 
diff --git a/qcelemental/molparse/from_schema.py b/qcelemental/molparse/from_schema.py
index d1775b4d..61268ac8 100644
--- a/qcelemental/molparse/from_schema.py
+++ b/qcelemental/molparse/from_schema.py
@@ -7,7 +7,7 @@
 from .from_arrays import from_arrays
 
 
-def from_schema(molschema: Dict, *, nonphysical: bool = False, verbose: int = 1) -> Dict:
+def from_schema(molschema: Dict, *, nonphysical: bool = False, throw_reorder: bool = True, verbose: int = 1) -> Dict:
     r"""Construct molecule dictionary representation from non-Psi4 schema.
 
     Parameters
@@ -16,6 +16,10 @@ def from_schema(molschema: Dict, *, nonphysical: bool = False, verbose: int = 1)
         Dictionary form of Molecule following known schema.
     nonphysical
         Do allow masses outside an element's natural range to pass validation?
+    throw_reorder
+        Do, when non-contiguous fragments detected, raise
+        ValidationError (``True``) or to proceed to reorder atoms to
+        contiguize fragments (``False``).
     verbose
         Amount of printing.
 
@@ -54,7 +58,7 @@ def from_schema(molschema: Dict, *, nonphysical: bool = False, verbose: int = 1)
         mass=ms.get("masses", None),
         real=ms.get("real", None),
         elbl=ms.get("atom_labels", None),
-        throw_reorder=True,
+        throw_reorder=throw_reorder,
     )
 
     molrec = from_arrays(
diff --git a/qcelemental/tests/test_molecule.py b/qcelemental/tests/test_molecule.py
index 05b2180a..914b478d 100644
--- a/qcelemental/tests/test_molecule.py
+++ b/qcelemental/tests/test_molecule.py
@@ -436,6 +436,38 @@ def test_charged_fragment():
 #    })
 
 
+@pytest.mark.parametrize(
+    "validate, allow_noncontiguous_but_reorder", [(True, True), (True, False), (False, True), (False, False)]
+)
+def test_noncontig(validate, allow_noncontiguous_but_reorder):
+    mol = {
+        "fragments": [[2, 0], [1]],
+        "geometry": np.array([[0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 0.0]]),
+        "symbols": np.array(["Li", "H", "He"]),
+        "fragment_multiplicities": [4, 2],
+    }
+
+    if validate:
+        if allow_noncontiguous_but_reorder:
+            mol = Molecule(**mol, validate=validate, allow_noncontiguous_but_reorder=allow_noncontiguous_but_reorder)
+            assert compare(["He", "Li", "H"], mol.symbols)
+            assert compare([2, 3, 1], mol.atomic_numbers)
+            assert compare_values([[0, 0, 0], [0, 0, 1], [0, 0, 2]], mol.geometry)
+            assert compare([4, 2], mol.fragment_multiplicities)
+
+        else:
+            with pytest.raises(qcel.ValidationError) as e:
+                Molecule(**mol, validate=validate, allow_noncontiguous_but_reorder=allow_noncontiguous_but_reorder)
+                assert "QCElemental would need to reorder atoms to accommodate non-contiguous fragments" in str(e.value)
+
+    else:
+        mol = Molecule(**mol, validate=validate, allow_noncontiguous_but_reorder=allow_noncontiguous_but_reorder)
+        assert compare(["Li", "H", "He"], mol.symbols)
+        assert compare([3, 1, 2], mol.atomic_numbers)
+        assert compare_values([[0, 0, 1], [0, 0, 2], [0, 0, 0]], mol.geometry)
+        assert compare([4, 2], mol.fragment_multiplicities)
+
+
 @pytest.mark.parametrize("group_fragments, orient", [(True, True), (False, False)])  # original  # Psi4-like
 def test_get_fragment(group_fragments, orient):
     mol = Molecule(