From 35b444a66d7db3039218b08ad92a9da644572493 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Beno=C3=AEt=20Gschwind?= <gschwind@gnu-log.net>
Date: Tue, 20 Jun 2023 13:22:17 +0200
Subject: [PATCH] Use user function in LCA computation

---
 lca_algebraic/base_utils.py |  3 ++-
 lca_algebraic/lca.py        | 17 ++++++++++++++---
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/lca_algebraic/base_utils.py b/lca_algebraic/base_utils.py
index 9da11c8..f0456c7 100644
--- a/lca_algebraic/base_utils.py
+++ b/lca_algebraic/base_utils.py
@@ -99,7 +99,8 @@ def _getAmountOrFormula(ex: ExchangeDataset) -> Union[Basic, float]:
     """ Return either a fixed float value or an expression for the amount of this exchange"""
     if 'formula' in ex:
         try:
-            return parse_expr(ex['formula'])
+            local_dict = {x[0].name: x[0] for x in _user_functions.values()}
+            return parse_expr(ex['formula'], local_dict=local_dict)
         except:
             error("Error while parsing formula '%s' : backing to amount" % ex['formula'])
 
diff --git a/lca_algebraic/lca.py b/lca_algebraic/lca.py
index c99730a..478b6d0 100644
--- a/lca_algebraic/lca.py
+++ b/lca_algebraic/lca.py
@@ -4,6 +4,7 @@
 
 from pandas import DataFrame
 from sympy import lambdify, simplify, Function
+from sympy.printing.numpy import NumPyPrinter
 
 from .base_utils import _actName, _getDb, _method_unit
 from .base_utils import _getAmountOrFormula
@@ -137,6 +138,15 @@ def __init__(self, exprOrDict, expanded_params=None, params=None, sobols=None):
         you can provide either the list pf expanded parameters (full vars for enums) for the 'user' param names
         """
 
+        printer = NumPyPrinter({
+            'fully_qualified_modules': False,
+            'inline': True,
+            'allow_unknown_functions': True,
+             'user_functions': { }
+        })
+
+        modules = [{x[0].name : x[1] for x in _user_functions.values()}, 'numpy']
+
         if isinstance(exprOrDict, dict) :
             # Come from JSON serialization
             obj = exprOrDict
@@ -145,8 +155,9 @@ def __init__(self, exprOrDict, expanded_params=None, params=None, sobols=None):
 
             # Full names
             self.expanded_params = _expand_param_names(self.params)
-            self.expr = parse_expr(obj["expr"])
-            self.lambd = lambdify(self.expanded_params, self.expr, 'numpy')
+            local_dict = {x[0].name: x[0] for x in _user_functions.values()}
+            self.expr = parse_expr(obj["expr"], local_dict=local_dict)
+            self.lambd = lambdify(self.expanded_params, self.expr, modules, printer=printer)
             self.sobols = obj["sobols"]
 
         else :
@@ -158,7 +169,7 @@ def __init__(self, exprOrDict, expanded_params=None, params=None, sobols=None):
             if self.params is None :
                 self.params = _expanded_names_to_names(expanded_params)
 
-            self.lambd = lambdify(expanded_params, exprOrDict, 'numpy')
+            self.lambd = lambdify(expanded_params, self.expr, modules, printer=printer)
             self.expanded_params = expanded_params
             self.sobols = sobols