From e21192375778af0321ae981512421e522c5c990d Mon Sep 17 00:00:00 2001
From: Johannes Hjorth <hjorth@kth.se>
Date: Fri, 6 Dec 2024 15:08:03 +0100
Subject: [PATCH] Working on units for sbtab import

---
 snudda/utils/sbtab_to_snudda.py | 59 ++++++++++++++++++++++-----------
 1 file changed, 40 insertions(+), 19 deletions(-)

diff --git a/snudda/utils/sbtab_to_snudda.py b/snudda/utils/sbtab_to_snudda.py
index 8a510f551..330aae80f 100644
--- a/snudda/utils/sbtab_to_snudda.py
+++ b/snudda/utils/sbtab_to_snudda.py
@@ -2,6 +2,8 @@
 import numpy as np
 import pandas as pd
 import json
+import quantities as pq
+
 
 # See convert_sbtab_to_json.sh -- example
 
@@ -43,6 +45,12 @@ def __init__(self, sbtab_path, out_filename,
         self.parameters_data = None
         self.constants_data = None
 
+        self.default_concentration_unit = "molar"
+        self.unit_dict = dict()
+
+        self.unit_dict["nmol_unit"] = pq.UnitQuantity('nanomole', pq.nano * pq.mole, symbol='nmol')
+        self.unit_dict["nM_unit"] = pq.UnitQuantity('nanomolar', pq.nano * pq.molar, symbol='nM')
+
         self.parameters = dict()
 
         self.data = dict()
@@ -74,16 +82,14 @@ def _get_optional_value(self, row, value_name, default_value=0):
 
     def parse(self):
 
-        import quantities as pq
-
         self.data["species"] = dict()
         self.data["rates"] = dict()
         self.data["reactions"] = dict()
 
         original_sbtab_units = dict()
 
-        nmol_unit = pq.UnitQuantity('nanomole', pq.nano * pq.mole, symbol='nmol')
-        nM_unit = pq.UnitQuantity('nanomolar', pq.nano * pq.molar, symbol='nM')
+        nmol_unit = self.unit_dict["nmol_unit"]
+        nM_unit = self.unit_dict["nM_unit"]
 
         conc_unit = nM_unit
 
@@ -103,7 +109,7 @@ def parse(self):
                              }
 
             self.data["species"][species_name] = species_data
-            original_sbtab_units[species_name] = species_unit.symbol  # Wilhelm är bombsäker(!)
+            original_sbtab_units[species_name] = species_unit.symbol 
 
         self.get_parameters()  # TODO, not finished yet!
 
@@ -122,9 +128,6 @@ def get_parameters(self):
             # We assume that reaction name kf_R0, kr_R0 are related to R0 reaction.
             reaction_name = row["!Name"].split("_")[-1]
 
-            import pdb
-            pdb.set_trace()
-
             reaction_row = self.reactions_data[self.reactions_data["!ID"] == reaction_name]
 
             if len(reaction_row) != 1:
@@ -146,33 +149,51 @@ def get_parameters(self):
                 raise KeyError(f"Unable to find reaction {parameter_name}, we assume "
                                f"it is before compounds in {self.reactions_filename} column !KineticLaw")
 
-            unit_str = ""
+            original_unit_str = ""
+            target_unit_str = ""
 
             # Next we replace components with their units, and use the quantities library to calculate conversion factor
             for rc in reaction_components:
                 rc_name = rc.split("^")[0]
-                compound_row = self.compounds_data[self.constants_data["!Name"] == rc_name]
+
+                try:
+                    compound_row = self.compounds_data[self.compounds_data["!Name"] == rc_name]
+                except:
+                    import traceback
+                    print(traceback.format_exc())
+                    import pdb
+                    pdb.set_trace()
 
                 if len(compound_row) != 1:
                     raise KeyError(f"The compound {rc_name} does not appear on a unique row in {self.compounds_filename}")
 
-                if unit_str != "":
-                    unit_str += " * "
+                if original_unit_str != "":
+                    original_unit_str += " * "
+                    target_unit_str += " * "
 
-                unit_str += f"({compound_row['!Unit'][0]})"
+                import pdb
+                pdb.set_trace()
 
-                if "^" in rc:
-                    unit_str += f"^{rc.split('^')[1]}"
+                # TODO: 2024-12-06, verify unit scaling for k_f and k_r
+
+                original_unit_str += f"({compound_row['!Unit'].iloc[0]})"
+                target_unit_str += f"({self.default_concentration_unit})"
 
+                if "^" in rc:
+                    original_unit_str += f"^{rc.split('^')[1]}"
+                    target_unit_str += f"^{rc.split('^')[1]}"
 
+            nmol_unit = self.unit_dict["nmol_unit"]
+            nM_unit = self.unit_dict["nM_unit"]
 
+            original_unit = pq.CompoundUnit(original_unit_str)
+            target_unit = pq.CompoundUnit(target_unit_str)
 
-            import pdb
-            pdb.set_trace()
+            scale = original_unit.rescale(target_unit)
 
             # We need to convert to SI units...
 
-            scale = row["!Scale"]
+            # scale = row["!Scale"]
 
             # Vi behöver ta hänsyn till om det är log10 eller annan skala, eller använda linjära värdet -- vilket är bäst?
             # Hämta ut K_forward, K_backwards (om det finns), spara i self.parameters
@@ -187,7 +208,7 @@ def get_parameters(self):
 
             # TODO: Continue here!!
 
-            self.parameters[parameter_name] = parameter_value
+            self.parameters[parameter_name] = parameter_value * scale
 
             import pdb
             pdb.set_trace()