Add lambdify and evaluate functions to SymbolicHandler

hyper_surrogate/scripts/generate_deformations.py

@@ -50,4 +50,16 @@
     end_time = time.time()
     print(f"Generated {args.batch_size} pk2 tensors in {end_time - start_time:.5f} seconds.")
     # average per entry
-    print(f"Average time per entry: {(end_time - start_time) / args.batch_size:.5f} seconds.")
+    # print(f"Average time per entry: {(end_time - start_time) / args.batch_size:.5f} seconds.")
+
+    # evaluate lambdify
+    pk2_func = h.lambdify(pk2, *sef_params.keys())
+    start_time = time.time()
+    pk2_func_iterator = h.evaluate_iterator(pk2_func, c, 1, 1)
+    b = np.array([next(pk2_func_iterator) for _ in range(args.batch_size)])
+    end_time = time.time()
+    print(f"Generated {args.batch_size} pk2 lambdified tensors in {end_time - start_time:.5f} seconds.")
+
+    print(a)
+    print("#########")
+    print(b)

hyper_surrogate/symbolic.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Iterable
 
 import numpy as np
 import sympy as sym
@@ -15,6 +15,15 @@ class SymbolicHandler:
     def __init__(self) -> None:
         self.c_tensor = self._c_tensor()
 
+    def c_symbols(self) -> Any:
+        """
+        Return the c_tensor flattened symbols.
+
+        Returns:
+            list: A list of c_tensor symbols.
+        """
+        return [self.c_tensor[i, j] for i in range(3) for j in range(3)]
+
     def _c_tensor(self) -> sym.Matrix:
         """
         Create a 3x3 matrix of symbols.
@@ -97,7 +106,12 @@ def cmat_tensor(self, pk2: sym.Matrix) -> sym.ImmutableDenseNDimArray:
             ]
         )
 
-    def substitute(self, symbolic_tensor: sym.MutableDenseMatrix, numerical_c_tensor: np.ndarray, *args: dict) -> Any:
+    def substitute(
+        self,
+        symbolic_tensor: sym.MutableDenseMatrix,
+        numerical_c_tensor: np.ndarray,
+        *args: dict,
+    ) -> Any:
         """
         Automatically substitute numerical values from a given 3x3 numerical matrix into c_tensor.
 
@@ -122,7 +136,10 @@ def substitute(self, symbolic_tensor: sym.MutableDenseMatrix, numerical_c_tensor
         return symbolic_tensor.subs(substitutions)
 
     def substitute_iterator(
-        self, symbolic_tensor: sym.MutableDenseMatrix, numerical_c_tensors: np.ndarray, *args: dict
+        self,
+        symbolic_tensor: sym.MutableDenseMatrix,
+        numerical_c_tensors: np.ndarray,
+        *args: dict,
     ) -> Any:
         """
         Automatically substitute numerical values from a given 3x3 numerical matrix into c_tensor.
@@ -140,3 +157,43 @@ def substitute_iterator(
         """
         for numerical_c_tensor in numerical_c_tensors:
             yield self.substitute(symbolic_tensor, numerical_c_tensor, *args)
+
+    def lambdify(self, symbolic_tensor: sym.Matrix, *args: Iterable[Any]) -> Any:
+        """
+        Create a lambdified function from a symbolic tensor that can be used for numerical evaluation.
+
+        Args:
+            tensor (sym.Expr or sym.Matrix): The symbolic tensor to be lambdified.
+            args (dict): Additional substitution lists of symbols.
+        Returns:
+            function: A function that can be used to numerically evaluate the tensor with specific values.
+        """
+
+        return sym.lambdify((self.c_symbols(), *args), symbolic_tensor, modules="numpy")
+
+    def evaluate(self, lambdified_tensor: Any, *args: Any) -> Any:
+        """
+        Evaluate a lambdified tensor with specific values.
+
+        Args:
+            lambdified_tensor (function): A lambdified tensor function.
+            args (dict): Additional substitution lists of symbols.
+
+        Returns:
+            Any: The evaluated tensor.
+        """
+        return lambdified_tensor(*args)
+
+    def evaluate_iterator(self, lambdified_tensor: Any, numerical_c_tensors: np.ndarray, *args: Any) -> Any:
+        """
+        Evaluate a lambdified tensor with specific values.
+
+        Args:
+            lambdified_tensor (function): A lambdified tensor function.
+            args (dict): Additional substitution lists of symbols.
+
+        Returns:
+            Any: The evaluated tensor.
+        """
+        for numerical_c_tensor in numerical_c_tensors:
+            yield self.evaluate(lambdified_tensor, numerical_c_tensor.flatten(), *args)

tests/test_symbolic_rules.py

@@ -1,3 +1,5 @@
+import logging
+
 import numpy as np
 import pytest
 import sympy as sym
@@ -50,6 +52,13 @@ def cmat(handler, pk2) -> sym.ImmutableDenseNDimArray:
 # testing
 
 
+def test_c_symbols(handler):
+    # assert c_symbols
+    c_symbols = handler.c_symbols()
+    assert isinstance(c_symbols, list)
+    assert all(isinstance(c, sym.Symbol) for c in c_symbols)
+
+
 def test_symbolic_pk2_cmat(pk2, cmat):
     # derivative of sef in order to c_tensor
 
@@ -98,3 +107,29 @@ def test_symbolic_subs_in_cmat(handler, cmat, right_cauchys, sef_args):
         )
         for subs in handler.substitute_iterator(cmat, right_cauchys, sef_args)
     )
+
+
+def test_lambdify(handler, sef_args, right_cauchys, pk2):
+    c = right_cauchys[0]
+    C10, C01 = sym.symbols("C10 C01")
+    # pk2
+    pk2_func = handler.lambdify(pk2, *sef_args.keys())
+    pk2_value = handler.evaluate(pk2_func, c.flatten(), 1, 1)
+    assert isinstance(pk2_value, np.ndarray)
+    assert pk2_value.shape == (3, 3)
+    assert all(isinstance(pk2_value[i, j], float) for i in range(3) for j in range(3))
+
+
+def test_lambdify_iterator(handler, sef_args, right_cauchys, pk2):
+    logging.info(right_cauchys.shape)
+    # pk2 function
+    pk2_func = handler.lambdify(pk2, *sef_args.keys())
+    pk2_values = list(handler.evaluate_iterator(pk2_func, right_cauchys, 1, 1))
+    assert isinstance(pk2_values, list)
+    assert all(isinstance(pk2_value, np.ndarray) for pk2_value in pk2_values)
+    assert all(pk2_value.shape == (3, 3) for pk2_value in pk2_values)
+    assert all(isinstance(pk2_value[i, j], float) for pk2_value in pk2_values for i in range(3) for j in range(3))
+
+    # for evals in handler.evaluate_iterator(pk2_func, right_cauchys, 1, 1):
+    #     logging.info(evals)
+    # logging.info(list(handler.evaluate_iterator(pk2_func, right_cauchys, 1, 1)))