2dgaussianzfit

Problem when defining a 2d gaussian in zfit

I want to create a custom two dimensional gaussian PDF in Zfit but I thing i am doing it wrong.

class MyG(zfit.pdf.ZPDF):
    _N_OBS = 2  # dimension, can be omitted
    _PARAMS = ['mean_x', 'std_x', 'mean_y', 'std_y']  # the name 
    of the parameters
@zfit.supports()
    
def _unnormalized_pdf(self, x, params):
    x1,y1 = z.unstack_x(x)
    mean_x = params['mean_x']
    std_x  = params['std_x']
    mean_y = params['mean_y']
    std_y  = params['std_y']
            
    A = np.pi * 2 * std_x * std_y
    
    exponent = -((x[0] - mean_x) ** 2 / (2 * std_x ** 2) + (x[1] - mean_y) ** 2 / (2 * std_y ** 2))
    

    return z.exp(exponent) / A

meanx = zfit.Parameter("meanx", 5) 
meany = zfit.Parameter("meany", 5)
stdx = zfit.Parameter("stdx", 1)
stdy = zfit.Parameter("stdy", 1)

xobs = zfit.Space('xobs', limits = (1, 10))
yobs = zfit.Space('yobs', limits = (1, 10))  
obsxy = yobs * xobs 

pdf = MyG(obs=obsxy, mean_x = meanx, mean_y=meany, std_x = stdx, std_y = stdy)
x = np.linspace(1, 10, 100)
y = np.linspace(1, 10, 100)
X, Y = np.meshgrid(x, y)
grid_points = np.column_stack([X.ravel(), Y.ravel()])
pdf_values = custom_pdf.pdf(grid_points).numpy()
Z = pdf_values.reshape(X.shape)
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')  
ax.plot_surface(X, Y, Z, cmap="viridis", alpha=0.7)
plt.show

Somehow this gives me this distribution. enter image description here

Can someone point out where did i go wrong with this?

Solution

  • I am not exactly sure where this is going wrong, as the snippet has a few minor errors, I assume they are just typos.

    (minor improvements: you don't need the unstack. znp.exp or similar instead of z)

    Correcting these and running it I can get the PDF as expected, see the code snippet and the plot. Are you running on a new enough zfit version? If not, I would suggest to upgrade if possible.

    import numpy as np
from matplotlib import pyplot as plt

import zfit
import zfit.z.numpy as znp


class MyG(zfit.pdf.ZPDF):
    _N_OBS = 2  # dimension, can be omitted
    _PARAMS = ['mean_x', 'std_x', 'mean_y', 'std_y']  # the name of the parameters

    @zfit.supports()
    def _unnormalized_pdf(self, x, params):
        mean_x = params['mean_x']
        std_x  = params['std_x']
        mean_y = params['mean_y']
        std_y  = params['std_y']

        A = np.pi * 2 * std_x * std_y

        exponent = -((x[0] - mean_x) ** 2 / (2 * std_x ** 2) + (x[1] - mean_y) ** 2 / (2 * std_y ** 2))


        return znp.exp(exponent) / A

meanx = zfit.Parameter("meanx", 5)
meany = zfit.Parameter("meany", 5)
stdx = zfit.Parameter("stdx", 1)
stdy = zfit.Parameter("stdy", 1)

xobs = zfit.Space('xobs', limits = (1, 10))
yobs = zfit.Space('yobs', limits = (1, 10))
obsxy = yobs * xobs

pdf = MyG(obs=obsxy, mean_x = meanx, mean_y=meany, std_x = stdx, std_y = stdy)
x = np.linspace(1, 10, 100)
y = np.linspace(1, 10, 100)
X, Y = np.meshgrid(x, y)
grid_points = np.column_stack([X.ravel(), Y.ravel()])
pdf_values = pdf.pdf(grid_points).numpy()
Z = pdf_values.reshape(X.shape)
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z, cmap="viridis", alpha=0.7)
plt.show()

    enter image description here