From f87f7ef4910affe34fb3544aa1039512265947f0 Mon Sep 17 00:00:00 2001
From: Andreas Copan <avcopan@gmail.com>
Date: Tue, 9 Jul 2024 17:09:25 -0400
Subject: [PATCH] New: Guess spin from ChI string

---
 automol/amchi/__init__.py |  2 ++
 automol/amchi/_conv.py    | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/automol/amchi/__init__.py b/automol/amchi/__init__.py
index d2586eb1..6eed5178 100644
--- a/automol/amchi/__init__.py
+++ b/automol/amchi/__init__.py
@@ -96,6 +96,7 @@
 # # derived properties
 from automol.amchi._conv import is_complete
 from automol.amchi._conv import is_valid_multiplicity
+from automol.amchi._conv import guess_spin
 # # derived transformations
 from automol.amchi._conv import add_stereo
 from automol.amchi._conv import expand_stereo
@@ -200,6 +201,7 @@
     # # derived properties
     'is_complete',
     'is_valid_multiplicity',
+    'guess_spin',
     # # derived transformations
     'add_stereo',
     'expand_stereo',
diff --git a/automol/amchi/_conv.py b/automol/amchi/_conv.py
index 7699ead2..f15294ba 100644
--- a/automol/amchi/_conv.py
+++ b/automol/amchi/_conv.py
@@ -4,7 +4,7 @@
 import itertools
 import numbers
 
-from automol import geom
+from automol import form, geom
 from automol import graph as graph_
 from automol.amchi.base import (
     atom_stereo_parities,
@@ -13,12 +13,15 @@
     breaking_bond_keys,
     equivalent,
     forming_bond_keys,
+    formula,
+    formula_layer,
     has_stereo,
     hydrogen_valences,
     is_inchi,
     is_inverted_enantiomer,
     is_reversed_ts,
     isotope_layers,
+    main_layers,
     split,
     standard_form,
     symbols,
@@ -286,6 +289,32 @@ def is_valid_multiplicity(chi, mul):
     return mul in graph_.possible_spin_multiplicities(graph(chi, stereo=False))
 
 
+def guess_spin(chi: str) -> int:
+    """Guess the spin of a ChI string.
+
+    Apart from a special list of common molecules, simply guesses 0 or 1 depending on
+    whether the eletron count is even or odd, respectively.
+
+    Could be made smarter by converting to a graph, but that would be more expensive.
+
+    :param chi: A ChI string
+    :return: The ground-state spin guess
+    """
+    # First, check if this is a case with a hardcoded value
+    lookup_dct = {
+        ("O",): 2,
+        ('O2', ('c', '1-2')): 2,
+    }
+    key = (formula_layer(chi), *sorted(main_layers(chi).items()))
+    val = lookup_dct.get(key, None)
+    if val is not None:
+        return val
+
+    # Otherwise, guess based on the formula
+    fml = formula(chi)
+    return form.electron_count(fml) % 2
+
+
 # # derived transformations
 def add_stereo(chi):
     """Add stereochemistry to a ChI string converting to/from geometry.