[docs]defgenerate_meshgrid_sample(decay:ThreeBodyDecay,resolution:int,x_mandelstam:Literal[1,2,3]=1,y_mandelstam:Literal[1,2,3]=2,)->DataSample:"""Generate a `numpy.meshgrid` sample for plotting with `matplotlib.pyplot`."""boundaries=compute_dalitz_boundaries(decay)returngenerate_sub_meshgrid_sample(decay,resolution,x_range=boundaries[x_mandelstam-1],y_range=boundaries[y_mandelstam-1],x_mandelstam=x_mandelstam,y_mandelstam=y_mandelstam,)
[docs]defgenerate_phasespace_sample(decay:ThreeBodyDecay,n_events:int,seed:int|None=None)->DataSample:r"""Generate a uniform distribution over Dalitz variables :math:`\sigma_{1,2,3}`."""boundaries=compute_dalitz_boundaries(decay)domain_generator=NumpyDomainGenerator(boundaries={"sigma1":boundaries[0],"sigma2":boundaries[1],})phsp_filter=create_phase_space_filter(decay,outside_value=0)phsp_generator=IntensityDistributionGenerator(domain_generator,phsp_filter)rng=NumpyUniformRNG(seed)phsp=phsp_generator.generate(n_events,rng)compute_sigma_z=__create_compute_compute_sigma_z(decay)phsp["sigma3"]=compute_sigma_z(phsp)returnphsp
[docs]defcreate_mass_symbol_mapping(decay:ThreeBodyDecay)->dict[sp.Symbol,float]:return{sp.Symbol(f"m{i}"):decay.states[i].massforiinsorted(decay.states)# ensure that dict keys are sorted by state ID}
def__get_third_mandelstam_index(x_mandelstam:Literal[1,2,3],y_mandelstam:Literal[1,2,3]):ifx_mandelstam==y_mandelstam:msg="x_mandelstam and y_mandelstam must be different"raiseValueError(msg)returnnext(iter({1,2,3}-{x_mandelstam,y_mandelstam}))