code: #  Copyright (c) 2021 zfit
  
  import tensorflow as tf
  
  import zfit
  from zfit import z
  
  
  class CustomPDF2D(zfit.pdf.BasePDF):
      """My custom, 2 dimensional pdf. The axes are: Energy, Momentum.
  
      """
  
      def __init__(self, param1, param2, param3, obs, name="CustomPDF", ):
          # we can now do complicated stuff here if needed
          # only thing: we have to specify explicitly here what is which parameter
          params = {'super_param': param1,  # we can change/compose etc parameters
                    'param2': param2, 'param3': param3}
          super().__init__(obs, params, name=name)
  
      def _unnormalized_pdf(self, x):
          energy, momentum = x.unstack_x()
          param1 = self.params['super_param']
          param2 = self.params['param2']
          param3 = self.params['param3']
  
          # just a fantasy function
          probs = param1 * tf.cos(energy ** 2) + tf.math.log(param2 * momentum ** 2) + param3
          return probs
  
  
  # add an analytic integral
  
  # define the integral function
  def integral_full(limits, norm_range, params, model):
      lower, upper = limits.rect_limits    # for a more detailed guide, see the space.py example
      param1 = params['super_param']
      param2 = params['param2']
      param3 = params['param3']
  
      lower = z.convert_to_tensor(lower)
      upper = z.convert_to_tensor(upper)
  
      # calculate the integral here, dummy integral, wrong!
      integral = param1 * param2 * param3 + z.reduce_sum([lower, upper])
      return integral
  
  
  # define the space over which it is defined. Here, we use the axes
  lower_full = (-10, zfit.Space.ANY_LOWER)
  upper_full = (10, zfit.Space.ANY_UPPER)
  integral_full_limits = zfit.Space(axes=(0, 1),
                                    limits=(lower_full, upper_full))
  
  CustomPDF2D.register_analytic_integral(func=integral_full,
                                         limits=integral_full_limits)
  
  
  # define the partial integral function
  def integral_axis1(x, limits, norm_range, params, model):
      data_0 = x.unstack_x()  # data from axis 0
  
      param1 = params['super_param']
      param2 = params['param2']
      param3 = params['param3']
  
      lower, upper = limits.limit1d  # for a more detailed guide, see the space.py example
      lower = z.convert_to_tensor(lower)  # the limits are now 1-D, for axis 1
      upper = z.convert_to_tensor(upper)
  
      # calculate the integral here, dummy integral
      integral = data_0 ** 2 * param1 * param2 * param3 + z.reduce_sum([lower, upper])
      # notice that the returned shape will be in the same as data_0, e.g. the number of events given in x
      return integral
  
  
  # define the space over which it is defined. Here, we use the axes
  lower_axis1 = ((zfit.Space.ANY_LOWER,),)
  upper_axis1 = ((zfit.Space.ANY_UPPER,),)
  integral_axis1_limits = zfit.Space(axes=(1,),  # axes one corresponds to the second obs, here obs2
                                     limits=(lower_axis1, upper_axis1))
  
  CustomPDF2D.register_analytic_integral(func=integral_axis1,
                                         limits=integral_axis1_limits)
  
  if __name__ == '__main__':
      import numpy as np
  
      obs = zfit.Space('obs1', (-10, 10)) * zfit.Space('obs2', (-3, 5))
      pdf = CustomPDF2D(1, 2, 3, obs=obs)
      sample = pdf.sample(n=1000)
      pdf.pdf([[2., 2.5], [5.4, 3.2]])
      x_part = zfit.Data.from_numpy(array=np.array([2.1, 2.2, 3.2]), obs='obs1')
  
      # integrate over obs2 with limits 1, 2 for the `x_part`. This will use the analytic integral above
      pdf.partial_integrate(x=x_part, limits=zfit.Space('obs2', (1, 2)))
      # we can explicitly call the analytic integral. Without registering it (e.g. comment the line with the `register`
      # and run again), it will raise an error
      pdf.partial_analytic_integrate(x=x_part, limits=zfit.Space('obs2', (1, 2)))
  
apis: numpy.array tensorflow.cos tensorflow.math.log