From eb24f99fb63a2b7cd44e2d40d235d5e882882961 Mon Sep 17 00:00:00 2001
From: faneshala <fane.shala@epfl.ch>
Date: Sun, 26 May 2024 21:38:52 +0200
Subject: [PATCH] DJ KHALEEEED

---
 src/chembalancer/chembalancer.py        | 30 ++++++++++++++------
 tests/test_balance_chemical_equation.py | 37 +++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_balance_chemical_equation.py

diff --git a/src/chembalancer/chembalancer.py b/src/chembalancer/chembalancer.py
index cbd582c..f330088 100644
--- a/src/chembalancer/chembalancer.py
+++ b/src/chembalancer/chembalancer.py
@@ -92,11 +92,11 @@ def get_molecular_formula(smiles):
     else:
         return "Invalid SMILES string"
 
+
 def balance_chemical_equation(reactant_smiles, product_smiles):
-    """ Balance a chemical equation given reactants and products as SMILES strings. """
     reactant_counts = [count_atoms(smiles) for smiles in reactant_smiles]
     product_counts = [count_atoms(smiles) for smiles in product_smiles]
-    
+
     reactant_elements = set(sum([list(counts.keys()) for counts in reactant_counts], []))
     product_elements = set(sum([list(counts.keys()) for counts in product_counts], []))
 
@@ -114,23 +114,35 @@ def balance_chemical_equation(reactant_smiles, product_smiles):
     A_reactants = setup_matrix(elements, reactant_counts)
     A_products = setup_matrix(elements, product_counts)
     A = np.concatenate([A_reactants, -A_products], axis=1)
-    
+
     integer_coefficients = solve_ilp(A)
+    if integer_coefficients is None or not integer_coefficients:
+        raise ValueError("Failed to solve the balance equation. The system may be underdetermined or inconsistent.")
+
     reactant_coeffs = integer_coefficients[:len(reactant_smiles)]
     product_coeffs = integer_coefficients[len(reactant_smiles):]
-    
+
     reactant_data = [(coeff, get_molecular_formula(smiles)) for coeff, smiles in zip(reactant_coeffs, reactant_smiles)]
     product_data = [(coeff, get_molecular_formula(smiles)) for coeff, smiles in zip(product_coeffs, product_smiles)]
 
     return reactant_data, product_data
+
+
     
-def setup_matrix(elements, compounds):
-    """ Create a stoichiometry matrix for the elements and compounds. """
+def setup_matrix(elements, counts):
+    # Assuming counts is a list of dictionaries where each dictionary is the atomic count for a molecule
     matrix = []
-    for element in elements:
-        row = [compound.get(element, 0) for compound in compounds]
+    for count in counts:
+        row = [count.get(element, 0) for element in elements]
         matrix.append(row)
-    return np.array(matrix, dtype=int)
+    
+    # Ensure matrix is 2D
+    matrix = np.array(matrix)
+    if matrix.ndim == 1:
+        matrix = matrix.reshape(1, -1)  # Reshape to 2D if it's inadvertently 1D
+
+    return matrix
+
 
 def display_reaction(reactants, products):
     """Format and display the chemical reaction."""
diff --git a/tests/test_balance_chemical_equation.py b/tests/test_balance_chemical_equation.py
new file mode 100644
index 0000000..8ff0128
--- /dev/null
+++ b/tests/test_balance_chemical_equation.py
@@ -0,0 +1,37 @@
+import unittest
+import numpy as np
+import sys
+import os
+
+try:
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+except NameError:
+    current_dir = os.path.dirname(os.path.abspath(sys.argv[0]))
+src_path = os.path.join(current_dir, '..', 'src')
+sys.path.insert(0, src_path)
+
+from chembalancer.chembalancer import balance_chemical_equation 
+
+class TestBalanceChemicalEquation:
+    @pytest.fixture(autouse=True)
+    def setup_method(self):
+        pass
+
+    def test_combustion_of_methane(self):
+        # Reactants and products
+        reactant_smiles = ['CH4', 'O2']
+        product_smiles = ['CO2', 'H2O']
+
+        # Call the function to test
+        try:
+            reactant_data, product_data = balance_chemical_equation(reactant_smiles, product_smiles)
+        except ValueError as e:
+            assert "Failed to solve the balance equation." in str(e)
+            return
+
+        # Assertions
+        expected_reactant_data = [(1, 'CH4'), (2, 'O2')]
+        expected_product_data = [(1, 'CO2'), (2, 'H2O')]
+
+        assert reactant_data == expected_reactant_data
+        assert product_data == expected_product_data