def test_custom_emitters(plot=True):
    from jetset.jet_model import Jet

    from jetset.jet_emitters import EmittersDistribution
    import numpy as np

    def distr_func_bkn(gamma_break, gamma, s1, s2):
        return np.power(gamma, -s1) * (1. +
                                       (gamma / gamma_break))**(-(s2 - s1))

    n_e = EmittersDistribution('custom_bkn', spectral_type='bkn')
    n_e.add_par('gamma_break',
                par_type='turn-over-energy',
                val=1E3,
                vmin=1.,
                vmax=None,
                unit='lorentz-factor')
    n_e.add_par('s1',
                par_type='LE_spectral_slope',
                val=2.5,
                vmin=-10.,
                vmax=10,
                unit='')
    n_e.add_par('s2',
                par_type='LE_spectral_slope',
                val=3.2,
                vmin=-10.,
                vmax=10,
                unit='')
    n_e.set_distr_func(distr_func_bkn)
    n_e.parameters.show_pars()
    n_e.parameters.s1.val = 2.0
    n_e.parameters.s2.val = 3.5
    if plot is True:
        n_e.plot()

    my_jet = Jet(emitters_distribution=n_e)
    my_jet.Norm_distr = True
    my_jet.parameters.N.val = 5E4
    my_jet.eval()
    np.testing.assert_allclose(my_jet.emitters_distribution.eval_N(),
                               my_jet.parameters.N.val,
                               rtol=1E-5)
    print(my_jet.emitters_distribution.eval_N(), my_jet.parameters.N.val)
    print(n_e.eval_N(), my_jet.parameters.N.val)
    assert (my_jet.emitters_distribution.emitters_type == 'electrons')
    my_jet.save_model('test_jet_custom_emitters.pkl')
    my_jet = Jet.load_model('test_jet_custom_emitters.pkl')
    my_jet.eval()
def custom_emitters(plot=True):
    from jetset.jet_model import Jet

    from jetset.jet_emitters import EmittersDistribution
    import numpy as np

    def distr_func_bkn(gamma_break, gamma, s1, s2):
        return np.power(gamma, -s1) * (1. +
                                       (gamma / gamma_break))**(-(s2 - s1))

    n_e = EmittersDistribution('bkn')
    n_e.add_par('gamma_break',
                par_type='turn-over-energy',
                val=1E3,
                vmin=1.,
                vmax=None,
                unit='lorentz-factor')
    n_e.add_par('s1',
                par_type='LE_spectral_slope',
                val=2.5,
                vmin=-10.,
                vmax=10,
                unit='')
    n_e.add_par('s2',
                par_type='LE_spectral_slope',
                val=3.2,
                vmin=-10.,
                vmax=10,
                unit='')
    n_e.set_distr_func(distr_func_bkn)
    n_e.parameters.show_pars()
    n_e.parameters.s1.val = 2.0
    n_e.parameters.s2.val = 3.5
    if plot is True:
        n_e.plot()

    my_jet = Jet(electron_distribution=n_e)
    my_jet.Norm_distr = True
    my_jet.parameters.N.val = 5E4
    my_jet.eval()
    diff = np.fabs(
        np.trapz(n_e.n_gamma_e, n_e.gamma_e) - my_jet.parameters.N.val)
    print('diff', diff)
    assert (diff < 1E-3)