From c9132ca4f3af40eb6248e8ee7c666450f7090a04 Mon Sep 17 00:00:00 2001
From: Jonas Eschle <mayou36@jonas.eschle.com>
Date: Sun, 8 Oct 2023 13:56:15 +0100
Subject: [PATCH] fix: minor conv fixes, removed class level caching

---
 zfit_physics/models/pdf_conv.py | 73 +++++++++++++--------------------
 1 file changed, 29 insertions(+), 44 deletions(-)

diff --git a/zfit_physics/models/pdf_conv.py b/zfit_physics/models/pdf_conv.py
index 985b26a..2554508 100644
--- a/zfit_physics/models/pdf_conv.py
+++ b/zfit_physics/models/pdf_conv.py
@@ -51,8 +51,6 @@ def __init__(
         # func = lambda x: func.unnormalized_pdf(x=x)
         # kernel = lambda x: kernel.unnormalized_pdf(x=x)
 
-        self._grid_points = None  # true vars  # callable func of reco - true vars
-        self._func_values = None
         self.conv_limits = limits
         self._ndraws = ndraws
         self._experimental_pdf_normalized = experimental_pdf_normalized
@@ -60,36 +58,29 @@ def __init__(
     @z.function
     def _unnormalized_pdf(self, x):
         limits = self.conv_limits
-        # area = limits.area()  # new spaces
         area = limits.rect_area()[0]  # new spaces
 
-        samples = self._grid_points
-        func_values = self._func_values
-        # if func_values is None:
-        if True:
-            # create sample for numerical integral
-            lower, upper = limits.rect_limits
-            lower = z.convert_to_tensor(lower, dtype=self.dtype)
-            upper = z.convert_to_tensor(upper, dtype=self.dtype)
-            samples_normed = tfp.mcmc.sample_halton_sequence(
-                dim=limits.n_obs,
-                num_results=self._ndraws,
-                dtype=self.dtype,
-                randomized=False,
-            )
-            samples = samples_normed * (upper - lower) + lower  # samples is [0, 1], stretch it
-            samples = zfit.Data.from_tensor(obs=limits, tensor=samples)
-            self._grid_points = samples
-
-            func_values = self.pdfs[0].pdf(samples, norm=False)  # func of true vars
-            self._func_values = func_values
+        # create sample for numerical integral
+        lower, upper = limits.rect_limits
+        lower = z.convert_to_tensor(lower, dtype=self.dtype)
+        upper = z.convert_to_tensor(upper, dtype=self.dtype)
+        samples_normed = tfp.mcmc.sample_halton_sequence(
+            dim=limits.n_obs,
+            num_results=self._ndraws,
+            dtype=self.dtype,
+            randomized=False,
+        )
+        samples = samples_normed * (upper - lower) + lower  # samples is [0, 1], stretch it
+        samples = zfit.Data.from_tensor(obs=limits, tensor=samples)
+
+        func_values = self.pdfs[0].pdf(samples, norm=False)  # func of true vars
 
         return tf.map_fn(
             lambda xi: area * tf.reduce_mean(func_values * self.pdfs[1].pdf(xi - samples.value(), norm=False)),
             x.value(),
         )
-        # func of reco vars
 
+    @zfit.supports(norm=True)
     @z.function
     def _pdf(self, x, norm_range):
         if not self._experimental_pdf_normalized:
@@ -99,26 +90,20 @@ def _pdf(self, x, norm_range):
         # area = limits.area()  # new spaces
         area = limits.rect_area()[0]  # new spaces
 
-        samples = self._grid_points
-        func_values = self._func_values
-        # if func_values is None:
-        if True:
-            # create sample for numerical integral
-            lower, upper = limits.rect_limits
-            lower = z.convert_to_tensor(lower, dtype=self.dtype)
-            upper = z.convert_to_tensor(upper, dtype=self.dtype)
-            samples_normed = tfp.mcmc.sample_halton_sequence(
-                dim=limits.n_obs,
-                num_results=self._ndraws,
-                dtype=self.dtype,
-                randomized=False,
-            )
-            samples = samples_normed * (upper - lower) + lower  # samples is [0, 1], stretch it
-            samples = zfit.Data.from_tensor(obs=limits, tensor=samples)
-            self._grid_points = samples
-
-            func_values = self.pdfs[0].pdf(samples)  # func of true vars
-            self._func_values = func_values
+        # create sample for numerical integral
+        lower, upper = limits.rect_limits
+        lower = z.convert_to_tensor(lower, dtype=self.dtype)
+        upper = z.convert_to_tensor(upper, dtype=self.dtype)
+        samples_normed = tfp.mcmc.sample_halton_sequence(
+            dim=limits.n_obs,
+            num_results=self._ndraws,
+            dtype=self.dtype,
+            randomized=False,
+        )
+        samples = samples_normed * (upper - lower) + lower  # samples is [0, 1], stretch it
+        samples = zfit.Data.from_tensor(obs=limits, tensor=samples)
+
+        func_values = self.pdfs[0].pdf(samples)  # func of true vars
 
         return tf.map_fn(
             lambda xi: area * tf.reduce_mean(func_values * self.pdfs[1].pdf(xi - samples.value())),