better "extras" passing in Molecule for EFP

MolSSI · Sep 3, 2019 · 4b18220 · 4b18220
1 parent 7cb0fc8
commit 4b18220
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 4 deletions.
diff --git a/qcelemental/models/molecule.py b/qcelemental/models/molecule.py
@@ -724,9 +724,9 @@ def from_data(cls,
  raise TypeError("Input type not understood, please supply the 'dtype' kwarg.")
 
  if dtype in ["string", "psi4", "psi4+", "xyz", "xyz+"]:
- mol_dict = from_string(data)
- assert isinstance(mol_dict, dict)
- input_dict = to_schema(mol_dict["qm"], dtype=2)
+ molrec = from_string(data, enable_qm=True, missing_enabled_return_qm='minimal')
+ assert isinstance(molrec, dict)
+ input_dict = to_schema(molrec["qm"], dtype=2)
  validate = True
  elif dtype == "numpy":
  data = np.asarray(data)

diff --git a/qcelemental/molparse/to_schema.py b/qcelemental/molparse/to_schema.py
@@ -94,6 +94,14 @@ def to_schema(molrec: Dict[str, Any],
  if 'connectivity' in molrec:
  molecule['connectivity'] = deepcopy(molrec['connectivity'])
 
+ # EFP extras
+ if 'fragment_files' in molrec:
+ molecule['extras'] = {
+ 'fragment_files': molrec['fragment_files'],
+ 'hint_types': molrec['hint_types'],
+ 'geom_hints': molrec['geom_hints'],
+ }
+
  if dtype == 1:
  qcschema = {'schema_name': 'qcschema_input', 'schema_version': 1, 'molecule': molecule}
  elif dtype == 2:

diff --git a/qcelemental/tests/test_utils.py b/qcelemental/tests/test_utils.py
@@ -55,7 +55,8 @@ def test_unique_everseen(inp, expected):
 @pytest.mark.parametrize("inp,expected", [
  (({1:{"a":"A"},2:{"b":"B"}}, {2:{"c":"C"},3:{"d":"D"}}), {1:{"a":"A"},2:{"b":"B","c":"C"},3:{"d":"D"}}),
  (({1:{"a":"A"},2:{"b":"B","c":None}}, {2:{"c":"C"},3:{"d":"D"}}), {1:{"a":"A"},2:{"b":"B","c":"C"},3:{"d":"D"}}),
- (({1: [None, 1]}, {1: [2, 1],3:{"d":"D"}}), {1:[2, 1], 3:{"d":"D"}})
+ (({1: [None, 1]}, {1: [2, 1],3:{"d":"D"}}), {1:[2, 1], 3:{"d":"D"}}),
+ (({1: True, 'extras':{'aa': [True]}}, {2: None, 'extras':{'bb': 5}}), {1: True, 2: None, 'extras':{'aa': [True], 'bb':5}}),
 ]) # yapf: disable
 def test_updatewitherror(inp, expected):
  assert compare_recursive(expected, qcel.util.update_with_error(inp[0], inp[1]))