almost lol

src/chembalancer/chembalancer.py

@@ -17,6 +17,7 @@
 from rdkit import Chem
 import requests
 import base64
+import math
 
 
 def get_smiles_from_name(name):
@@ -93,6 +94,7 @@ def get_molecular_formula(smiles):
         return "Invalid SMILES string"
 
 
+
 def balance_chemical_equation(reactant_smiles, product_smiles):
     reactant_counts = [count_atoms(smiles) for smiles in reactant_smiles]
     product_counts = [count_atoms(smiles) for smiles in product_smiles]
@@ -110,9 +112,16 @@ def balance_chemical_equation(reactant_smiles, product_smiles):
             error_message += f"Elements {missing_in_reactants} are in products but not in reactants."
         raise ValueError(error_message)
 
+    # Use the union of elements from reactants and products
     elements = sorted(reactant_elements.union(product_elements))
+
     A_reactants = setup_matrix(elements, reactant_counts)
     A_products = setup_matrix(elements, product_counts)
+
+    # Debugging: Ensure matrices have the same number of rows
+    if A_reactants.shape[0] != A_products.shape[0]:
+        raise ValueError("Internal error: Reactants and products matrices do not match in row dimension.")
+
     A = np.concatenate([A_reactants, -A_products], axis=1)
 
     integer_coefficients = solve_ilp(A)
@@ -127,20 +136,15 @@ def balance_chemical_equation(reactant_smiles, product_smiles):
 
     return reactant_data, product_data
 
-
-
 def setup_matrix(elements, counts):
-    # Assuming counts is a list of dictionaries where each dictionary is the atomic count for a molecule
+    # Setup the matrix based on the element order and counts provided
     matrix = []
     for count in counts:
         row = [count.get(element, 0) for element in elements]
         matrix.append(row)
-
-    # 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
-
+        matrix = matrix.reshape(1, -1)
     return matrix