From cc64388f0dc0e56d8a3649e4dc1ca0cf5b759f04 Mon Sep 17 00:00:00 2001
From: Xiaorui Dong <xiaorui@mit.edu>
Date: Fri, 20 Oct 2023 14:38:51 -0400
Subject: [PATCH] add reaction element count

---
 rdmc/reaction.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/rdmc/reaction.py b/rdmc/reaction.py
index 8c832152..4df357a0 100644
--- a/rdmc/reaction.py
+++ b/rdmc/reaction.py
@@ -139,6 +139,20 @@ def is_num_atoms_balanced(self) -> bool:
         """
         return self.reactant_complex.GetNumAtoms() == self.product_complex.GetNumAtoms()
 
+    @property
+    def reactant_element_count(self) -> dict:
+        """
+        The element count in the reactant(s) and product(s).
+        """
+        return dict(Counter(self.reactant_complex.GetElementSymbols()))
+
+    @property
+    def product_element_count(self) -> dict:
+        """
+        The element count in the reactant(s) and product(s).
+        """
+        return dict(Counter(self.product_complex.GetElementSymbols()))
+
     @property
     def is_element_balanced(self) -> bool:
         """
@@ -477,8 +491,7 @@ def has_same_reactants(
         Returns:
             bool: Whether the reaction has the same reactants as the other reaction.
         """
-        return self.is_same_reactants(other.reactant_complex,
-                                      resonance=resonance)
+        return self.is_same_reactants(other.reactant_complex, resonance=resonance)
 
     def is_same_reactants(
         self,
@@ -495,9 +508,7 @@ def is_same_reactants(
         Returns:
             bool: Whether the reaction has the same reactants as the given reactants or reactant complex.
         """
-        return is_same_complex(self.reactant_complex,
-                               reactants,
-                               resonance=resonance)
+        return is_same_complex(self.reactant_complex, reactants, resonance=resonance)
 
     def has_same_products(
         self,
@@ -513,8 +524,7 @@ def has_same_products(
         Returns:
             bool: Whether the reaction has the same products as the other reaction.
         """
-        return self.is_same_products(other.product_complex,
-                                     resonance=resonance)
+        return self.is_same_products(other.product_complex, resonance=resonance)
 
     def is_same_products(
         self,
@@ -531,6 +541,4 @@ def is_same_products(
         Returns:
             bool: Whether the reaction has the same products as the given products or product complex.
         """
-        return is_same_complex(self.product_complex,
-                               products,
-                               resonance=resonance)
+        return is_same_complex(self.product_complex, products, resonance=resonance)