def test_interpsnapshotKeplerPotential_eval_naz():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 1.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,51),
                                            zgrid=(0.,0.2,51),
                                            logR=False,
                                            interpPot=True,
                                            zsym=True,
                                            numcores=1)
    spaz= potential.InterpSnapshotRZPotential(s,
                                              rgrid=(0.01,2.,51),
                                              zgrid=(0.,0.2,51),
                                              logR=False,
                                              interpPot=True,
                                              zsym=True,
                                              numcores=1,nazimuths=12)
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)
    zs= numpy.linspace(-0.2,0.2,41)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp(r,z)-spaz(r,z))/sp(r,z)) < 10.**-10., 'RZPot interpolation w/ InterpSnapShotPotential of KeplerPotential with different nazimuths fails at (R,z) = (%g,%g)' % (r,z)
    #This tests within the grid, with vector evaluation
    rs= numpy.linspace(0.01,2.,10)
    zs= numpy.linspace(-0.2,0.2,20)
    mr,mz= numpy.meshgrid(rs,zs)
    mr= mr.flatten()
    mz= mz.flatten()
    assert numpy.all(numpy.fabs((sp(mr,mz)-spaz(mr,mz))/sp(mr,mz)) < 10.**-5.), 'RZPot interpolation w/ interpRZPotential with different nazimimuths fails for vector input'
    return None
def test_interpsnapshotKeplerpotential_Rzderiv():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 2.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,101),
                                            zgrid=(0.,0.2,101),
                                            logR=False,
                                            interpPot=True,
                                            interpepifreq=True,
                                            interpverticalfreq=True,
                                            zsym=True)
    kp= potential.KeplerPotential(amp=2.) #should be the same
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)[1:]
    zs= numpy.linspace(-0.2,0.2,41)
    zs= zs[zs != 0.]# avoid zero
    # Test, but give small |z| a less constraining 
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)) < 10.**-4.*(1.+19.*(numpy.fabs(z) < 0.05)), 'RZPot interpolation of Rzderiv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g; value is %g' % (r,z,numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)),kp.Rzderiv(r,z))
    #This tests within the grid
    rs= numpy.linspace(0.01,2.,10)[1:]
    zs= numpy.linspace(-0.2,0.2,20)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)) < 10.**-4.*(1.+19.*(numpy.fabs(z) < 0.05)), 'RZPot interpolation of Rzderiv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)))
    return None
def test_interpsnapshotKeplerPotential_z2deriv():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 2.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,101),
                                            zgrid=(0.,0.2,101),
                                            logR=False,
                                            interpPot=True,
                                            interpverticalfreq=True,
                                            zsym=True)
    kp= potential.KeplerPotential(amp=2.) #should be the same
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)[1:]
    zs= numpy.linspace(-0.2,0.2,41)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)) < 10.**-4., 'RZPot interpolation of z2deriv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)))
    #This tests within the grid
    rs= numpy.linspace(0.01,2.,10)[1:]
    zs= numpy.linspace(-0.2,0.2,20)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)) < 2.*10.**-4., 'RZPot interpolation of z2deriv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)))
    return None
def test_interpsnapshotKeplerPotential_normalize_units():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 4.
    s['eps']= 0.
    s['pos'].units= 'kpc'
    s['vel'].units= 'km s**-1'
    sp= potential.InterpSnapshotRZPotential(s,
                                          rgrid=(0.01,3.,201),
                                          zgrid=(0.,0.2,201),
                                          logR=False,
                                          interpPot=True,
                                          zsym=True)
    #Currently unnormalized
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    # Normalize
    sp.normalize(R0=1.)
    assert numpy.fabs(sp.Rforce(1.,0.)+1.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    # De normalize
    sp.denormalize()
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    # Also test when R0 =/= 1
    sp.normalize(R0=2.)
    assert numpy.fabs(sp.Rforce(1.,0.)+1.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    # De normalize
    sp.denormalize()
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    return None
def test_interpsnapshotKeplerPotential_eval():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 1.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,201),
                                            zgrid=(0.,0.2,201),
                                            logR=False,
                                            interpPot=True,
                                            zsym=True,
                                            numcores=1)
    kp= potential.KeplerPotential(amp=1.) #should be the same
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)
    zs= numpy.linspace(-0.2,0.2,41)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp(r,z)-kp(r,z))/kp(r,z)) < 10.**-10., 'RZPot interpolation w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g)' % (r,z)
    #This tests within the grid
    rs= numpy.linspace(0.01,2.,10)
    zs= numpy.linspace(-0.2,0.2,20)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp(r,z)-kp(r,z))/kp(r,z)) < 10.**-5., 'RZPot interpolation w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp(r,z)-kp(r,z))/kp(r,z)))           
    #Test all at the same time to use vector evaluation
    mr,mz= numpy.meshgrid(rs,zs)
    mr= mr.flatten()
    mz= mz.flatten()
    assert numpy.all(numpy.fabs((sp(mr,mz)-kp(mr,mz))/kp(mr,mz)) < 10.**-5.), 'RZPot interpolation w/ interpRZPotential fails for vector input'
    return None
def test_interpsnapshotKeplerPotential_verticalfreq():
    # Set up a snapshot with just one unit mass at the origin
    s = pynbody.new(star=1)
    s['mass'] = 2.
    s['eps'] = 0.
    sp = potential.InterpSnapshotRZPotential(s,
                                             rgrid=(0.01, 2., 101),
                                             zgrid=(0., 0.2, 101),
                                             logR=False,
                                             interpPot=True,
                                             interpverticalfreq=True,
                                             zsym=True)
    kp = potential.KeplerPotential(normalize=2.)  #should be the same
    #This just tests on the grid
    rs = numpy.linspace(0.01, 2., 21)[1:]
    for r in rs:
        assert numpy.fabs(
            (sp.verticalfreq(r) - kp.verticalfreq(r)) / kp.verticalfreq(r)
        ) < 10.**-4., 'RZPot interpolation of verticalfreq w/ InterpSnapShotPotential of KeplerPotential fails at R = %g by %g' % (
            r,
            numpy.fabs((sp.verticalfreq(r) - kp.verticalfreq(r)) /
                       kp.verticalfreq(r)))
    #This tests within the grid
    rs = numpy.linspace(0.01, 2., 10)[1:]
    for r in rs:
        assert numpy.fabs(
            (sp.verticalfreq(r) - kp.verticalfreq(r)) / kp.verticalfreq(r)
        ) < 10.**-4., 'RZPot interpolation of verticalfreq w/ InterpSnapShotPotential of KeplerPotential fails at R = %g by %g' % (
            r,
            numpy.fabs((sp.verticalfreq(r) - kp.verticalfreq(r)) /
                       kp.verticalfreq(r)))
    return None
def test_interpsnapshotKeplerPotential_logR_eval():
    # Set up a snapshot with just one unit mass at the origin
    s = pynbody.new(star=1)
    s['mass'] = 1.
    s['eps'] = 0.
    sp = potential.InterpSnapshotRZPotential(s,
                                             rgrid=(numpy.log(0.01),
                                                    numpy.log(20.), 251),
                                             logR=True,
                                             zgrid=(0., 0.2, 201),
                                             interpPot=True,
                                             zsym=True)
    kp = potential.KeplerPotential(amp=1.)  #should be the same
    rs = numpy.linspace(0.02, 16., 20)
    zs = numpy.linspace(-0.15, 0.15, 40)
    mr, mz = numpy.meshgrid(rs, zs)
    mr = mr.flatten()
    mz = mz.flatten()
    assert numpy.all(
        numpy.fabs((sp(mr, mz) - kp(mr, mz)) / kp(mr, mz)) < 10.**-5.
    ), 'RZPot interpolation w/ interpRZPotential fails for vector input, w/ logR at (R,z) = (%f,%f) by %g' % (
        mr[numpy.argmax(numpy.fabs(
            (sp(mr, mz) - kp(mr, mz)) / kp(mr, mz)))], mz[numpy.argmax(
                numpy.fabs((sp(mr, mz) - kp(mr, mz)) / kp(mr, mz)))],
        numpy.amax(numpy.fabs((sp(mr, mz) - kp(mr, mz)) / kp(mr, mz))))
    return None
def test_interpsnapshotrzpotential_nopynbody():
    # Test that if we cannot load pynbody, we get an ImportError
    from galpy.potential_src import SnapshotRZPotential
    SnapshotRZPotential._PYNBODY_LOADED= False
    try:
        sp= potential.InterpSnapshotRZPotential(1.) #1. doesn't matter
    except ImportError: pass
    else:
        raise AssertionError("InterpSnapshotRZPotential w/o pynbody should have raised an error, but didn't")
    return None
def test_interpsnapshotKeplerPotential_normalize():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 4.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,3.,201),
                                            zgrid=(0.,0.2,201),
                                            logR=False,
                                            interpPot=True,
                                            interpepifreq=True,
                                            interpverticalfreq=True,
                                            zsym=True)
    nkp= potential.KeplerPotential(normalize=1.) # normalized Kepler
    ukp= potential.KeplerPotential(amp=4.) # Un-normalized Kepler
    #Currently unnormalized
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp(1.,0.1)-ukp(1.,0.1)) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.epifreq(1.)-ukp.epifreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.verticalfreq(1.)-ukp.verticalfreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    # Normalize
    sp.normalize(R0=1.)
    assert numpy.fabs(sp.Rforce(1.,0.)+1.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    assert numpy.fabs(sp(1.,0.1)-nkp(1.,0.1)) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.epifreq(1.)-nkp.epifreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.verticalfreq(1.)-nkp.verticalfreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    # De normalize
    sp.denormalize()
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    assert numpy.fabs(sp(1.,0.1)-ukp(1.,0.1)) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.epifreq(1.)-ukp.epifreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.verticalfreq(1.)-ukp.verticalfreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    # Also test when R0 =/= 1
    sp.normalize(R0=2.)
    assert numpy.fabs(sp.Rforce(1.,0.)+1.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    assert numpy.fabs(sp(1.,0.1)-nkp(1.,0.1)) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.epifreq(1.)-nkp.epifreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.verticalfreq(1.)-nkp.verticalfreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    # De normalize
    sp.denormalize()
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    assert numpy.fabs(sp(1.,0.1)-ukp(1.,0.1)) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.epifreq(1.)-ukp.epifreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    assert numpy.fabs(sp.verticalfreq(1.)-ukp.verticalfreq(1.)) < 10.**-4., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    return None
def test_interpsnapshotKeplerPotential_noc_eval():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 1.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                          rgrid=(0.01,2.,201),
                                          zgrid=(0.,0.2,201),
                                          logR=False,
                                          interpPot=True,
                                          zsym=True,
                                          enable_c=False)
    kp= potential.KeplerPotential(amp=1.) #should be the same
    #Test all at the same time to use vector evaluation
    rs= numpy.linspace(0.01,2.,10)
    zs= numpy.linspace(-0.2,0.2,20)
    mr,mz= numpy.meshgrid(rs,zs)
    mr= mr.flatten()
    mz= mz.flatten()
    assert numpy.all(numpy.fabs((sp(mr,mz)-kp(mr,mz))/kp(mr,mz)) < 10.**-5.), 'RZPot interpolation w/ interpRZPotential fails for vector input, without enable_c'
    return None