From 64963c0d07a36cfe2ae0ec75a01000451be31b1e Mon Sep 17 00:00:00 2001
From: Remco de Boer <29308176+redeboer@users.noreply.github.com>
Date: Tue, 16 May 2023 15:40:51 +0200
Subject: [PATCH] ENH: move data arguments to front for `partial`

This is required for defining a gradient over the parameters
---
 src/tensorwaves/function/sympy/__init__.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/tensorwaves/function/sympy/__init__.py b/src/tensorwaves/function/sympy/__init__.py
index b8c8f327..ffd5c34a 100644
--- a/src/tensorwaves/function/sympy/__init__.py
+++ b/src/tensorwaves/function/sympy/__init__.py
@@ -109,7 +109,10 @@ def create_parametrized_function(
       [0.0, 0.0, 0.0, 0.0, 0.0]
     """
     free_symbols = _get_free_symbols(expression)
-    sorted_symbols = sorted(free_symbols, key=lambda s: s.name)
+    parameter_set = set(parameters)
+    parameter_symbols = sorted(free_symbols & parameter_set, key=lambda s: s.name)
+    data_symbols = sorted(free_symbols - parameter_set, key=lambda s: s.name)
+    sorted_symbols = tuple(data_symbols + parameter_symbols)  # for partial+gradient
     lambdified_function = _lambdify_normal_or_fast(
         expression=expression,
         symbols=sorted_symbols,