Callable varying rate function in skysurvey.SNeIa.from_draw() gives negative probabilities and uncorrect size

SkySurvey version

0.14.6

Minimal example

def varying_snia_rate(z):
    ''' Computes the varying SNIa rate according to:

        1. Perley et al. 2020
        2. Barris and Tonry 2006

    Notes:
    ------
    Rates units are approximated to Gpc^-3.yr^-1 (Gpc^-3.yr^-1.h70^3 in litterature)
    '''
    if z <= 0.2:
        rate = 2.35e4
    elif (z > 0.2) & (z <= 0.3):
        rate = 1.7e4
    elif (z > 0.3) & (z <= 0.4):
        rate = 5.3e4
    elif (z > 0.4) & (z <= 0.5):
        rate = 7.3e4
    elif (z > 0.5) & (z <= 0.6):
        rate = 2.04e5
    elif (z > 0.6) & (z <= 0.7):
        rate = 1.49e5
    else:
        rate = 1.78e5
    volume = Planck18.comoving_volume(z).to("Gpc**3").value
    z_rate = volume * rate
    return z_rate

varying_snia_rate = np.vectorize(varying_snia_rate)
snia = skysurvey.SNeIa.from_draw(zmin=0.002, zmax=1.0, nyears=0.001, rate=varying_snia_rate)

Error obtained

ValueError                                Traceback (most recent call last)
Cell In[39], line 30
     27     return z_rate
     29 varying_snia_rate = np.vectorize(varying_snia_rate)
---> 30 snia = skysurvey.SNeIa.from_draw(zmin=0.002, zmax=1.0, nyears=0.001, rate=varying_snia_rate)

File ~/edrisenv/lib/python3.8/site-packages/skysurvey/target/core.py:193, in Target.from_draw(cls, size, model, template, zmax, tstart, tstop, zmin, nyears, skyarea, rate, **kwargs)
    190 if kwargs:
    191     this.update_model_parameter(**kwargs)
--> 193 _ = this.draw( size=size,
    194                zmin=zmin, zmax=zmax,
    195                tstart=tstart, tstop=tstop,
    196                nyears=nyears,
    197                skyarea=skyarea,
    198                inplace=True, # creates self.data
    199                )
    200 return this

File ~/edrisenv/lib/python3.8/site-packages/skysurvey/target/core.py:866, in Target.draw(self, size, zmax, zmin, tstart, tstop, nyears, skyarea, inplace, model, **kwargs)
    863     size = int((self.get_rate(zmax, skyarea=skyarea)-rate_min) * nyears)
    865 # actually draw the data
--> 866 data = drawn_model.draw(size=size, **kwargs)
    867 if inplace:
    868     # lower precision
    869     data = data.astype( {k: str(v).replace("64","32") for k, v in data.dtypes.to_dict().items()})

File ~/edrisenv/lib/python3.8/site-packages/modeldag/modeldag.py:329, in ModelDAG.draw(self, size, limit_to_entries, data, **kwargs)
    326     prior_inputs = None
    328 model = self.get_model(prior_inputs=prior_inputs, **kwargs)
--> 329 return self._draw(model, size=size, limit_to_entries=limit_to_entries,
    330                       data=data)

File ~/edrisenv/lib/python3.8/site-packages/modeldag/modeldag.py:450, in ModelDAG._draw(self, model, size, limit_to_entries, data)
    447 prop = {**params, **inprop}
    448 # 
    449 # Draw it
--> 450 samples = np.asarray(self.draw_param(param_name, **prop))
    452 # and feed
    453 output_name = param_model.get("as", param_name)

File ~/edrisenv/lib/python3.8/site-packages/modeldag/modeldag.py:388, in ModelDAG.draw_param(self, name, func, size, xx, **kwargs)
    385     prop["xx"] = xx
    387 # Draw it.
--> 388 draw_ = func(**{**prop, **kwargs})
    389 if "xx" in func_arguments: # draw_ was actually a pdf
    390     xx_, pdf_ = draw_

File ~/edrisenv/lib/python3.8/site-packages/skysurvey/target/core.py:1087, in Transient.draw_redshift(self, zmax, zmin, zstep, size)
   1085 xx = np.arange(zmin, zmax, zstep)
   1086 pdf = self.getpdf_redshift(xx)
-> 1087 return np.random.choice(np.mean([xx[1:],xx[:-1]], axis=0), 
   1088               size=size, p=pdf/pdf.sum())

File mtrand.pyx:956, in numpy.random.mtrand.RandomState.choice()

ValueError: probabilities are not non-negative

Description of the problem

There are two issues.

The first problem occurs in skysurvey.Transient.getpdf_redshift(). By calling np.diff(), the function may return negative pdf values, which is not compatible with its use as probabilities in np.random.choice() in skysurvey.Transient.draw_redshift().

A second problem occurs in skysurvey.Target.draw(). When the rate is not constant (e.g. 2.35e4 as in Perley 2020), the function only determines the size parameter according to zmin and zmax and does not take into account potential variations in between.

MickaelRigault / skysurvey