<reponame>samrose30/PyPopStar<filename>popstar/tests/test_synthetic.py<gh_stars>0
import time
import pylab as plt
import numpy as np
from popstar import synthetic, reddening, evolution, atmospheres
import pysynphot
import os
import pdb
from scipy.spatial import cKDTree as KDTree

def test_isochrone(plot=False):
    from popstar import synthetic as syn

    logAge = 6.7
    AKs = 2.7
    distance = 4000

    startTime = time.time()
    iso = syn.Isochrone(logAge, AKs, distance)
    print('Test completed in: %d seconds' % (time.time() - startTime))
    # Typically takes 104 - 120 seconds.
    # Limited by pysynphot.Icat call in atmospheres.py

    assert iso.points.meta['LOGAGE'] == logAge
    assert iso.points.meta['AKS'] == AKs
    assert iso.points.meta['DISTANCE'] == distance
    assert len(iso.points) > 100
    
    if plot:
        plt.figure(1) 
        iso.plot_HR_diagram()
        
        plt.figure(2)
        iso.plot_mass_luminosity()

    return iso

def test_iso_wave():
    """
    Test to make sure isochrones generated have spectra with the proper 
    wavelength range, and that the user has control over that wavelength
    range (propagated through IsochronePhot)
    """
    # Define isochrone parameters
    logAge = np.log10(5*10**6.) # Age in log(years)
    AKs = 0.8 # extinction in mags
    dist = 4000 # distance in parsec

    # Define evolution/atmosphere models and extinction law (optional)
    evo_model = evolution.MergedBaraffePisaEkstromParsec() 
    atm_func = atmospheres.get_merged_atmosphere
    red_law = reddening.RedLawHosek18b()

    # Also specify filters for synthetic photometry (optional). Here we use 
    # the HST WFC3-IR F127M, F139M, and F153M filters
    filt_list = ['wfc3,ir,f127m']

    # First, let's make sure the vega spectrum has the proper limits
    vega = synthetic.Vega()

    assert np.min(vega.wave) == 995
    assert np.max(vega.wave) == 100200

    # Make Isochrone object. Will use wave_range = [3000,52000].
    # Make sure range matches to resolution of atmosphere.
    wave_range1 = [3000, 52000]
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10, wave_range=wave_range1,
                            recomp=True)

    test = my_iso.spec_list[0]

    assert np.min(test.wave) == 3010
    assert np.max(test.wave) == 51900

    # Now let's try changing the wave range. Is it carried through
    # properly?
    wave_range2 = [1200, 90000]
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10, wave_range=wave_range2,
                            recomp=True)

    test2 = my_iso.spec_list[0]

    assert np.min(test2.wave) == 1205
    assert np.max(test2.wave) == 89800

    # Does the error exception catch the bad wave_range?
    wave_range3 = [1200, 1000000]
    try:
        my_iso = synthetic.IsochronePhot(logAge, AKs, dist,
                                evo_model=evo_model, atm_func=atm_func,
                                red_law=red_law, filters=filt_list,
                                mass_sampling=10, wave_range=wave_range3,
                                recomp=True)
        print('WAVE TEST FAILED!!! Should have crashed here, wavelength range out of bounds')
        pdb.set_trace()
    except:
        print('Wavelength out of bound condition passed. Test is good')
        pass
    return

def test_IsochronePhot(plot=False):
    from popstar import synthetic as syn
    from popstar import evolution, atmospheres, reddening

    logAge = 6.7
    AKs = 2.7
    distance = 4000
    filt_list = ['wfc3,ir,f127m', 'nirc2,J']
    mass_sampling=1
    iso_dir = 'iso/'

    evo_model = evolution.MISTv1()
    atm_func = atmospheres.get_merged_atmosphere
    redlaw = reddening.RedLawNishiyama09()

    startTime = time.time()
    iso = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo_model,
                                atm_func=atm_func, red_law=redlaw,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)
    endTime = time.time()
    print('IsochronePhot generated in: %d seconds' % (endTime - startTime))
    # Typically takes 120 seconds if file is regenerated.
    # Limited by pysynphot.Icat call in atmospheres.py

    assert iso.points.meta['LOGAGE'] == logAge
    assert iso.points.meta['AKS'] == AKs
    assert iso.points.meta['DISTANCE'] == distance
    assert len(iso.points) > 100

    assert 'm_nirc2_J' in iso.points.colnames

    if plot:
        plt.figure(1) 
        iso.plot_CMD('mag814w', 'mag160w')
        
        plt.figure(2)
        iso.plot_mass_magnitude('mag160w')

    # Finally, let's test the isochronePhot file generation
    assert os.path.exists('{0}/iso_{1:.2f}_{2:4.2f}_{3:4s}_p00.fits'.format(iso_dir, logAge,
                                                                                AKs, str(distance).zfill(5)))
    
    # Check 1: If we try to remake the isochrone, does it read the file rather than
    # making a new one
    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo_model,
                                atm_func=atm_func, red_law=redlaw,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == False

    # Check 2: If we change evo model, atmo model, or redlaw,
    # does IsochronePhot regenerate the isochrone and overwrite the existing one?
    evo2 = evolution.MergedBaraffePisaEkstromParsec()
    mass_sampling=20

    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo2,
                                atm_func=atm_func, red_law=redlaw,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == True

    redlaw2 = reddening.RedLawHosek18b()
    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo2,
                                atm_func=atm_func, red_law=redlaw2,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == True

    atm2 = atmospheres.get_castelli_atmosphere
    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo2,
                                atm_func=atm2, red_law=redlaw2,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == True


    return

def test_ResolvedCluster():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar.imf import imf
    from popstar.imf import multiplicity

    # Define cluster parameters
    logAge = 6.7
    AKs = 2.4
    distance = 4000
    cluster_mass = 10**5.
    mass_sampling=5

    # Test filters
    filt_list = ['nirc2,J', 'nirc2,Kp']

    startTime = time.time()
    
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere

    red_law = reddening.RedLawNishiyama09()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    # Now to create the cluster.
    imf_mass_limits = np.array([0.07, 0.5, 1, np.inf])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    ##########
    # Start without multiplicity
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime))
    
    cluster1 = syn.ResolvedCluster(iso, my_imf1, cluster_mass)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    # Check that stars are returned
    assert len(clust1) > 0

    # Check that the total mass in stars is less than requested (no compact objects).
    cluster_mass_out = clust1['systemMass'].sum()
    assert cluster_mass_out < cluster_mass

    plt.figure(3)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c.')
    plt.gca().invert_yaxis()

    # *** Visual Inspections: ***
    #  - check that points (red) fall between isochrone points (blue)

    ##########
    # Test with multiplicity
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedCluster(iso, my_imf2, cluster_mass)
    clust2 = cluster2.star_systems
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    assert len(clust2) > 0
    assert len(cluster2.companions) > 0
    assert np.sum(clust2['N_companions']) == len(cluster2.companions)

    ##########
    # Plots 
    ##########
    # Plot an IR CMD and compare cluster members to isochrone.
    plt.figure(1)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(clust2['m_nirc2_J'] - clust2['m_nirc2_Kp'], clust2['m_nirc2_J'], 'b.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c-')
    plt.gca().invert_yaxis()
    plt.xlabel('J - Kp (mag)')
    plt.ylabel('J (mag')

    # Plot a mass-magnitude relationship.
    plt.figure(2)
    plt.clf()
    plt.semilogx(clust1['mass'], clust1['m_nirc2_J'], 'r.')
    plt.semilogx(clust2['mass'], clust2['m_nirc2_J'], 'r.')
    plt.gca().invert_yaxis()
    plt.xlabel('Mass (Msun)')
    plt.ylabel('J (mag)')
    
    # # Plot the spectrum of the most massive star
    # idx = cluster.mass.argmax()
    # plt.clf()
    # plt.plot(cluster.stars[idx].wave, cluster.stars[idx].flux, 'k.')

    # # Plot an integrated spectrum of the whole cluster.
    # wave, flux = cluster.get_integrated_spectrum()
    # plt.clf()
    # plt.plot(wave, flux, 'k.')

    return

def test_ResolvedClusterDiffRedden():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar.imf import imf
    from popstar.imf import multiplicity

    logAge = 6.7
    AKs = 2.4
    distance = 4000
    cluster_mass = 10**5.
    deltaAKs = 0.05
    mass_sampling=5

    # Test filters
    filt_list = ['nirc2,J', 'nirc2,Kp']
    
    startTime = time.time()
    
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere

    red_law = reddening.RedLawNishiyama09()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                                mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    imf_mass_limits = np.array([0.07, 0.5, 1, np.inf])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    ##########
    # Start without multiplicity
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime))
    
    cluster1 = syn.ResolvedClusterDiffRedden(iso, my_imf1, cluster_mass, deltaAKs)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    assert len(clust1) > 0
    
    plt.figure(3)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c.')
    plt.gca().invert_yaxis()

    # *** Visual Inspections: ***
    #  - check that points (red) fall between isochrone points (blue)

    ##########
    # Test with multiplicity
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedClusterDiffRedden(iso, my_imf2, cluster_mass, deltaAKs)
    clust2 = cluster2.star_systems
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    assert len(clust2) > 0
    assert len(cluster2.companions) > 0
    assert np.sum(clust2['N_companions']) == len(cluster2.companions)

    ##########
    # Plots 
    ##########
    # Plot an IR CMD and compare cluster members to isochrone.
    plt.figure(1)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(clust2['m_nirc2_J'] - clust2['m_nirc2_Kp'], clust2['m_nirc2_J'], 'b.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c-')
    plt.gca().invert_yaxis()
    plt.xlabel('J - Kp (mag)')
    plt.ylabel('J (mag')

    # Plot a mass-magnitude relationship.
    plt.figure(2)
    plt.clf()
    plt.semilogx(clust1['mass'], clust1['m_nirc2_J'], 'r.')
    plt.semilogx(clust2['mass'], clust2['m_nirc2_J'], 'r.')
    plt.gca().invert_yaxis()
    plt.xlabel('Mass (Msun)')
    plt.ylabel('J (mag)')

    return
    
def test_UnresolvedCluster():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar.imf import imf
    from popstar.imf import multiplicity
    
    log_age = 6.7
    AKs = 0.0
    distance = 4000
    metallicity=0
    cluster_mass = 10**4.

    startTime = time.time()    
    multi = multiplicity.MultiplicityUnresolved()
    imf_in = imf.Kroupa_2001(multiplicity=multi)
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere
    iso = syn.Isochrone(log_age, AKs, distance, metallicity=metallicity,
                            evo_model=evo, atm_func=atm_func, mass_sampling=10)
    print('Made Isochrone: %d seconds' % (time.time() - startTime))

    cluster = syn.UnresolvedCluster(iso, imf_in, cluster_mass)
    print('Constructed unresolved cluster: %d seconds' % (time.time() - startTime))

    # Plot an integrated spectrum of the whole cluster.
    wave = cluster.wave_trim
    flux = cluster.spec_trim
    plt.clf()
    plt.plot(wave, flux, 'k.')

    return

def test_ifmr_multiplicity():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar import ifmr
    from popstar.imf import imf
    from popstar.imf import multiplicity

    # Define cluster parameters
    logAge = 9.7
    AKs = 0.0
    distance = 1000
    cluster_mass = 1e6
    mass_sampling = 5

    # Test all filters
    filt_list = ['nirc2,Kp', 'nirc2,H', 'nirc2,J']

    startTime = time.time()
    
    evo = evolution.MISTv1()
    atm_func = atm.get_merged_atmosphere
    ifmr_obj = ifmr.IFMR()

    red_law = reddening.RedLawNishiyama09()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    # Now to create the cluster.
    imf_mass_limits = np.array([0.07, 0.5, 1, np.inf])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    ##########
    # Start without multiplicity and IFMR
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime)) 
    
    cluster1 = syn.ResolvedCluster(iso, my_imf1, cluster_mass, ifmr=ifmr_obj)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

   
    ##########
    # Test with multiplicity and IFMR
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedCluster(iso, my_imf2, cluster_mass, ifmr=ifmr_obj)
    clust2 = cluster2.star_systems
    comps2 = cluster2.companions
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    ##########
    # Tests
    ##########

    # Check that we have black holes, neutron stars, and white dwarfs in both.
    assert len(np.where(clust1['phase'] == 101)) > 0   # WD
    assert len(np.where(clust2['phase'] == 101)) > 0
    assert len(np.where(clust1['phase'] == 102)) > 0   # NS
    assert len(np.where(clust2['phase'] == 102)) > 0
    assert len(np.where(clust1['phase'] == 103)) > 0   # BH
    assert len(np.where(clust2['phase'] == 103)) > 0

    # Now check that we have companions that are WDs, NSs, and BHs
    assert len(np.where(comps2['phase'] == 101)) > 0
    assert len(np.where(comps2['phase'] == 102)) > 0
    assert len(np.where(comps2['phase'] == 103)) > 0

    # Make sure no funky phase designations (due to interpolation effects)
    # slipped through
    idx = np.where( (clust1['phase'] > 5) & (clust1['phase'] < 101) & (clust1['phase'] != 9) )
    idx2 = np.where( (comps2['phase'] > 5) & (comps2['phase'] < 101) & (comps2['phase'] != 9) )
    assert len(idx[0]) == 0

    return

def test_metallicity():
    """
    Test isochrone generation at different metallicities
    """
    # Define isochrone parameters
    logAge = np.log10(5*10**6.) 
    AKs = 0.8 
    dist = 4000 
    evo_model = evolution.MISTv1()
    atm_func = atmospheres.get_phoenixv16_atmosphere
    red_law = reddening.RedLawHosek18b()
    filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m']

    # Start with a solar metallicity isochrone    
    metallicity= 0.0

    # Make Isochrone object, with high mass_sampling to decrease compute time
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=metallicity,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10)

    # Test isochrone properties
    assert my_iso.points.meta['METAL_IN'] == 0.0
    assert os.path.exists('iso_6.70_0.80_04000_p00.fits')

    # Now for non-solar metallicity
    metallicity= -1.5

    # Make Isochrone object, with high mass_sampling to decrease compute time
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=metallicity,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10)

    metal_act = np.log10(0.00047 / 0.0142) # For Mist isochrones

    # Test isochrone properties
    assert my_iso.points.meta['METAL_IN'] == -1.5
    assert my_iso.points.meta['METAL_ACT'] == metal_act
    assert os.path.exists('iso_6.70_0.80_04000_m15.fits')
    
    return

def test_cluster_mass():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar import ifmr
    from popstar.imf import imf
    from popstar.imf import multiplicity

    # Define cluster parameters
    logAge = 6.7
    AKs = 2.4
    distance = 4000
    cluster_mass = 10**5.
    mass_sampling = 5

    # Test filters
    filt_list = ['nirc2,J', 'nirc2,Kp']

    startTime = time.time()
    
    # Define evolution/atmosphere models and extinction law
    evo = evolution.MISTv1() 
    atm_func = atmospheres.get_merged_atmosphere
    red_law = reddening.RedLawHosek18b()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    # Now to create the cluster.
    imf_mass_limits = np.array([0.2, 0.5, 1, 120.0])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    # IFMR
    my_ifmr = ifmr.IFMR()
    

    ##########
    # Start without multiplicity
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime))
    
    cluster1 = syn.ResolvedCluster(iso, my_imf1, cluster_mass, ifmr=my_ifmr)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    # Check that the total mass is within tolerance of input mass
    cluster_mass_out = clust1['systemMass'].sum()
    assert np.abs(cluster_mass_out - cluster_mass) < 200.0   # within 200 Msun of desired mass.
    print('Cluster Mass: IN = ', cluster_mass, " OUT = ", cluster_mass_out)

    ##########
    # Test with multiplicity
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedCluster(iso, my_imf2, cluster_mass, ifmr=my_ifmr)
    clust2 = cluster2.star_systems
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    # Check that the total mass is within tolerance of input mass
    cluster_mass_out = clust2['systemMass'].sum()
    assert np.abs(cluster_mass_out - cluster_mass) < 200.0   # within 200 Msun of desired mass.
    print('Cluster Mass: IN = ', cluster_mass, " OUT = ", cluster_mass_out)

    return

#=================================#
# Additional timing functions
#=================================#

def time_test_cluster():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar.imf import imf
    from popstar.imf import multiplicity

    logAge = 6.7
    AKs = 2.7
    distance = 4000
    cluster_mass = 10**4

    startTime = time.time()
    
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere
    red_law = reddening.RedLawNishiyama09()
    filt_list = ['nirc2,J', 'nirc2,Kp']
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list)
    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    imf_limits = np.array([0.07, 0.5, 150])
    imf_powers = np.array([-1.3, -2.35])
    multi = multiplicity.MultiplicityUnresolved()
    my_imf = imf.IMF_broken_powerlaw(imf_limits, imf_powers, multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster = syn.ResolvedCluster(iso, my_imf, cluster_mass)
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    return
    
def model_young_cluster_object(resolved=False):
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar.imf import imf
    from popstar.imf import multiplicity

    log_age = 6.5
    AKs = 0.1
    distance = 8000.0
    cluster_mass = 10000.
    
    multi = multiplicity.MultiplicityUnresolved()
    imf_in = imf.Kroupa_2001(multiplicity=multi)
    evo = evolution.MergedPisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere
    iso = syn.Isochrone(log_age, AKs, distance, evo, mass_sampling=10)

    if resolved:
        cluster = syn.ResolvedCluster(iso, imf_in, cluster_mass)
    else:
        cluster = syn.UnresolvedCluster(iso, imf_in, cluster_mass, wave_range=[19000,24000])

    # Plot the spectrum of the most massive star
    idx = cluster.mass_all.argmax()
    print('Most massive star is {0:f} M_sun.'.format(cluster.mass_all[idx]))
    #bigstar = cluster.spec_list_trim[idx]
    plt.figure(1)
    plt.clf()
    plt.plot(cluster.spec_list_trim[idx]._wavetable, cluster.spec_list_trim[idx]._fluxtable, 'k.')

    # Plot an integrated spectrum of the whole cluster.
    wave, flux = cluster.spec_list_trim[idx]._wavetable, cluster.spec_trim
    plt.figure(2)
    plt.clf()
    plt.plot(wave, flux, 'k.')

    return
    
def time_test_mass_match():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar.imf import imf
    from popstar.imf import multiplicity

    log_age = 6.7
    AKs = 2.7
    distance = 4000
    cluster_mass = 5e3
    
    imf_in = imf.Kroupa_2001(multiplicity=None)

    start_time = time.time()
    iso = syn.IsochronePhot(log_age, AKs, distance)
    iso_masses = iso.points['mass']
    print('Generated iso masses in {0:.0f} s'.format(time.time() - start_time))

    start_time = time.time()
    star_masses, isMulti, compMass, sysMass = imf_in.generate_cluster(cluster_mass)
    print('Generated cluster masses in {0:.0f} s'.format(time.time() - start_time))
    
    def match_model_masses1(isoMasses, starMasses):
        indices = np.empty(len(starMasses), dtype=int)
        
        for ii in range(len(starMasses)):
            theMass = starMasses[ii]
            
            dm = np.abs(isoMasses - theMass)
            mdx = dm.argmin()

            # Model mass has to be within 10% of the desired mass
            if (dm[mdx] / theMass) > 0.1:
                indices[ii] = -1
            else:
                indices[ii] = mdx

        return indices
            

    def match_model_masses2(isoMasses, starMasses):
        isoMasses_tmp = isoMasses.reshape((len(isoMasses), 1))
        kdt = KDTree(isoMasses_tmp)
 
        starMasses_tmp = starMasses.reshape((len(starMasses), 1))
        q_results = kdt.query(starMasses_tmp, k=1)
        indices = q_results[1]

        dm_frac = np.abs(starMasses - isoMasses[indices]) / starMasses

        idx = np.where(dm_frac > 0.1)[0]
        indices[idx] = -1
        
        return indices

    print('Test #1 START')
    start_time = time.time()
    idx1 = match_model_masses1(iso_masses, star_masses)
    stop_time = time.time()
    print('Test #1 STOPPED after {0:.0f} seconds'.format(stop_time - start_time))

    print('Test #2 START')
    start_time = time.time()
    idx2 = match_model_masses2(iso_masses, star_masses)
    stop_time = time.time()
    print('Test #2 STOPPED after {0:.0f} seconds'.format(stop_time - start_time))

    return


def test_phot_consistency(filt='all'):
    """
    Test photometric consistency of generated isochrone (IsochronePhot)
    against pre-generated isochrone with native filter sampling. Requires
    consistency to within 0.005 mag.

    Base isochrone is at 5 Myr, AKs = 0, 1000 pc, mass_sampling=10 

    Paramters:
    ----------
    filt: 'all', 'hst', 'vista', 'decam', 'ps1', 'jwst'
        Specify what filter set you want to test

    """
    from astropy.table import Table
    import os
    
    # Load pre-generated isochrone, located in popstar tests directory
    direct = os.path.dirname(__file__)
    orig = Table.read(direct+'/iso_6.70_0.00_01000.fits', format='fits')

    # Generate new isochrone with popstar code
    if filt == 'all':
        filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m',
                         'acs,wfc1,f814w', 'wfc3,ir,f125w', 'wfc3,ir,f160w',
                         'decam,y', 'decam,i', 'decam,z',
                         'decam,u', 'decam,g', 'decam,r',
                         'vista,Y', 'vista,Z', 'vista,J',
                         'vista,H',  'vista,Ks',
                         'ps1,z', 'ps1,g', 'ps1,r',
                         'ps1,i', 'ps1,y',
                         'jwst,F070W', 'jwst,F090W', 'jwst,F115W', 'jwst,F140M',
                         'jwst,F150W', 'jwst,F150W2', 'jwst,F162M', 'jwst,F164N',
                         'jwst,F182M', 'jwst,F187N', 'jwst,F200W', 'jwst,F212N',
                         'jwst,F210M','jwst,F250M', 'jwst,F277W', 'jwst,F300M',
                         'jwst,F322W2', 'jwst,F323N', 'jwst,F335M', 'jwst,F356W',
                         'jwst,F360M', 'jwst,F405N', 'jwst,F410M', 'jwst,F430M',
                         'jwst,F444W', 'jwst,F460M', 'jwst,F466N', 'jwst,F470N',
                         'jwst,F480M', 'nirc2,J', 'nirc2,H', 'nirc2,Kp', 'nirc2,K',
                         'nirc2,Lp', 'nirc2,Hcont',
                         'nirc2,FeII', 'nirc2,Brgamma',
                         'jg,J', 'jg,H', 'jg,K',
                         'nirc1,K', 'nirc1_H', 'ctio_osiris,K', 'ctio_osiris,H',
                         'naco,H', 'naco,Ks', 'ztf,g', 'ztf,r', 'ztf,i']

    elif filt == 'decam':
        filt_list = ['decam,y', 'decam,i', 'decam,z',
                         'decam,u', 'decam,g', 'decam,r']
    elif filt == 'vista':
        filt_list = ['vista,Y', 'vista,Z', 'vista,J',
                         'vista,H', 'vista,Ks']
    elif filt == 'ps1':
        filt_list = ['ps1,z', 'ps1,g', 'ps1,r',  'ps1,i',
                         'ps1,y']
    elif filt == 'jwst':
        filt_list = ['jwst,F070W', 'jwst,F090W', 'jwst,F115W', 'jwst,F140M',
                         'jwst,F150W', 'jwst,F150W2', 'jwst,F162M', 'jwst,F164N',
                         'jwst,F182M', 'jwst,F187N', 'jwst,F200W', 'jwst,F212N',
                         'jwst,F210M','jwst,F250M', 'jwst,F277W', 'jwst,F300M',
                         'jwst,F322W2', 'jwst,F323N', 'jwst,F335M', 'jwst,F356W',
                         'jwst,F360M', 'jwst,F405N', 'jwst,F410M', 'jwst,F430M',
                         'jwst,F444W', 'jwst,F460M', 'jwst,F466N', 'jwst,F470N',
                         'jwst,F480M']
    elif filt == 'hst':
        filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m',
                         'acs,wfc1,f814w', 'wfc3,ir,f125w', 'wfc3,ir,f160w']
    elif filt == 'nirc2':
        filt_list = ['nirc2,J', 'nirc2,H','nirc2,Kp', 'nirc2,K',
                         'nirc2,Lp', 'nirc2,Ms', 'nirc2,Hcont',
                         'nirc2,FeII', 'nirc2,Brgamma']
    elif filt == 'jg':
        filt_list = ['jg,J', 'jg,H', 'jg,K']
    elif filt == 'ztf':
        filt_list = ['ztf,g', 'ztf,r', 'ztf,i']
    elif filt == 'misc':
        filt_list=['nirc1,K', 'nirc1,H', 'ctio_osiris,K', 'ctio_osiris,H',
                       'naco,H', 'naco,Ks']
        
            
    print('Making isochrone')
    iso = synthetic.IsochronePhot(6.7, 0, 1000, mass_sampling=10, filters=filt_list, rebin=True)
    iso = iso.points

    # First find masses that are the same
    foo = []
    for ii in iso['mass']:
        tmp = np.where( orig['mass'] == ii)[0][0]
        foo.append(tmp)
        
    assert len(foo) == len(iso)
    orig = orig[foo]

    # Identify the photometry columns
    cols = list(iso.keys())
    idx = []
    for ii in range(len(cols)):
        if cols[ii].startswith('mag'):
            idx.append(ii)
    mag_cols = np.array(cols)[idx]

    # Test the consistency of each column with the original isochrone
    for ii in mag_cols:
        orig_mag = orig[ii]
        new_mag = iso[ii]

        np.testing.assert_allclose(orig_mag, new_mag, rtol=0.01, err_msg="{0} failed".format(ii))

        # Also report median abs difference
        diff = abs(orig_mag - new_mag)
        print(('{0} median abs diff: {1}'.format(ii, np.median(diff))))


    print(('Phot consistency test successful for {0}'.format(filt)))

    # Remove iso file we just wrote, since it was only a test
    os.remove('iso_6.70_0.00_01000.fits')
    return


# Siess not supported anymore.
# 
# def test_isochrone_siess_mass_current_bug():
#     """
#     Bug found by students in grad stars class. 
#     """
#     # Define isochrone parameters
#     logAges = [6.6, 7.6] # Age in log(years)
#     AKs = 0 # extinction in mags
#     dist = 1000 # distance in parsec
#     metallicity = 0 # Metallicity in [M/H]

#     # Define evolution/atmosphere models and extinction law
#     evo_model = evolution.MergedSiessGenevaPadova()
#     evo_model_name = ['MIST', 'Padova']
#     atm_func = atmospheres.get_merged_atmosphere
#     red_law = reddening.RedLawHosek18b()

#     # Also specify filters for synthetic photometry (optional). Here we use
#     # the HST WFC3-IR F127M, F139M, and F153M filters
#     filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m']

#     # Make Isochrone object. Note that is calculation will take a few minutes, unless the
#     # isochrone has been generated previously.

#     for logAge in logAges:
#         my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=0, evo_model=evo_model, atm_func=atm_func,
#                                          red_law=red_law, filters=filt_list)
#         print(my_iso.save_file)

#     return

import time
import pylab as plt
import numpy as np
from popstar import synthetic, reddening, evolution, atmospheres
import pysynphot
import os
import pdb
from scipy.spatial import cKDTree as KDTree

def test_isochrone(plot=False):
    from popstar import synthetic as syn

    logAge = 6.7
    AKs = 2.7
    distance = 4000

    startTime = time.time()
    iso = syn.Isochrone(logAge, AKs, distance)
    print('Test completed in: %d seconds' % (time.time() - startTime))
    # Typically takes 104 - 120 seconds.
    # Limited by pysynphot.Icat call in atmospheres.py

    assert iso.points.meta['LOGAGE'] == logAge
    assert iso.points.meta['AKS'] == AKs
    assert iso.points.meta['DISTANCE'] == distance
    assert len(iso.points) > 100
    
    if plot:
        plt.figure(1) 
        iso.plot_HR_diagram()
        
        plt.figure(2)
        iso.plot_mass_luminosity()

    return iso

def test_iso_wave():
    """
    Test to make sure isochrones generated have spectra with the proper 
    wavelength range, and that the user has control over that wavelength
    range (propagated through IsochronePhot)
    """
    # Define isochrone parameters
    logAge = np.log10(5*10**6.) # Age in log(years)
    AKs = 0.8 # extinction in mags
    dist = 4000 # distance in parsec

    # Define evolution/atmosphere models and extinction law (optional)
    evo_model = evolution.MergedBaraffePisaEkstromParsec() 
    atm_func = atmospheres.get_merged_atmosphere
    red_law = reddening.RedLawHosek18b()

    # Also specify filters for synthetic photometry (optional). Here we use 
    # the HST WFC3-IR F127M, F139M, and F153M filters
    filt_list = ['wfc3,ir,f127m']

    # First, let's make sure the vega spectrum has the proper limits
    vega = synthetic.Vega()

    assert np.min(vega.wave) == 995
    assert np.max(vega.wave) == 100200

    # Make Isochrone object. Will use wave_range = [3000,52000].
    # Make sure range matches to resolution of atmosphere.
    wave_range1 = [3000, 52000]
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10, wave_range=wave_range1,
                            recomp=True)

    test = my_iso.spec_list[0]

    assert np.min(test.wave) == 3010
    assert np.max(test.wave) == 51900

    # Now let's try changing the wave range. Is it carried through
    # properly?
    wave_range2 = [1200, 90000]
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10, wave_range=wave_range2,
                            recomp=True)

    test2 = my_iso.spec_list[0]

    assert np.min(test2.wave) == 1205
    assert np.max(test2.wave) == 89800

    # Does the error exception catch the bad wave_range?
    wave_range3 = [1200, 1000000]
    try:
        my_iso = synthetic.IsochronePhot(logAge, AKs, dist,
                                evo_model=evo_model, atm_func=atm_func,
                                red_law=red_law, filters=filt_list,
                                mass_sampling=10, wave_range=wave_range3,
                                recomp=True)
        print('WAVE TEST FAILED!!! Should have crashed here, wavelength range out of bounds')
        pdb.set_trace()
    except:
        print('Wavelength out of bound condition passed. Test is good')
        pass
    return

def test_IsochronePhot(plot=False):
    from popstar import synthetic as syn
    from popstar import evolution, atmospheres, reddening

    logAge = 6.7
    AKs = 2.7
    distance = 4000
    filt_list = ['wfc3,ir,f127m', 'nirc2,J']
    mass_sampling=1
    iso_dir = 'iso/'

    evo_model = evolution.MISTv1()
    atm_func = atmospheres.get_merged_atmosphere
    redlaw = reddening.RedLawNishiyama09()

    startTime = time.time()
    iso = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo_model,
                                atm_func=atm_func, red_law=redlaw,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)
    endTime = time.time()
    print('IsochronePhot generated in: %d seconds' % (endTime - startTime))
    # Typically takes 120 seconds if file is regenerated.
    # Limited by pysynphot.Icat call in atmospheres.py

    assert iso.points.meta['LOGAGE'] == logAge
    assert iso.points.meta['AKS'] == AKs
    assert iso.points.meta['DISTANCE'] == distance
    assert len(iso.points) > 100

    assert 'm_nirc2_J' in iso.points.colnames

    if plot:
        plt.figure(1) 
        iso.plot_CMD('mag814w', 'mag160w')
        
        plt.figure(2)
        iso.plot_mass_magnitude('mag160w')

    # Finally, let's test the isochronePhot file generation
    assert os.path.exists('{0}/iso_{1:.2f}_{2:4.2f}_{3:4s}_p00.fits'.format(iso_dir, logAge,
                                                                                AKs, str(distance).zfill(5)))
    
    # Check 1: If we try to remake the isochrone, does it read the file rather than
    # making a new one
    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo_model,
                                atm_func=atm_func, red_law=redlaw,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == False

    # Check 2: If we change evo model, atmo model, or redlaw,
    # does IsochronePhot regenerate the isochrone and overwrite the existing one?
    evo2 = evolution.MergedBaraffePisaEkstromParsec()
    mass_sampling=20

    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo2,
                                atm_func=atm_func, red_law=redlaw,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == True

    redlaw2 = reddening.RedLawHosek18b()
    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo2,
                                atm_func=atm_func, red_law=redlaw2,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == True

    atm2 = atmospheres.get_castelli_atmosphere
    iso_new = syn.IsochronePhot(logAge, AKs, distance, evo_model=evo2,
                                atm_func=atm2, red_law=redlaw2,
                                filters=filt_list,
                                mass_sampling=mass_sampling, iso_dir=iso_dir)

    assert iso_new.recalc == True


    return

def test_ResolvedCluster():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar.imf import imf
    from popstar.imf import multiplicity

    # Define cluster parameters
    logAge = 6.7
    AKs = 2.4
    distance = 4000
    cluster_mass = 10**5.
    mass_sampling=5

    # Test filters
    filt_list = ['nirc2,J', 'nirc2,Kp']

    startTime = time.time()
    
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere

    red_law = reddening.RedLawNishiyama09()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    # Now to create the cluster.
    imf_mass_limits = np.array([0.07, 0.5, 1, np.inf])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    ##########
    # Start without multiplicity
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime))
    
    cluster1 = syn.ResolvedCluster(iso, my_imf1, cluster_mass)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    # Check that stars are returned
    assert len(clust1) > 0

    # Check that the total mass in stars is less than requested (no compact objects).
    cluster_mass_out = clust1['systemMass'].sum()
    assert cluster_mass_out < cluster_mass

    plt.figure(3)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c.')
    plt.gca().invert_yaxis()

    # *** Visual Inspections: ***
    #  - check that points (red) fall between isochrone points (blue)

    ##########
    # Test with multiplicity
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedCluster(iso, my_imf2, cluster_mass)
    clust2 = cluster2.star_systems
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    assert len(clust2) > 0
    assert len(cluster2.companions) > 0
    assert np.sum(clust2['N_companions']) == len(cluster2.companions)

    ##########
    # Plots 
    ##########
    # Plot an IR CMD and compare cluster members to isochrone.
    plt.figure(1)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(clust2['m_nirc2_J'] - clust2['m_nirc2_Kp'], clust2['m_nirc2_J'], 'b.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c-')
    plt.gca().invert_yaxis()
    plt.xlabel('J - Kp (mag)')
    plt.ylabel('J (mag')

    # Plot a mass-magnitude relationship.
    plt.figure(2)
    plt.clf()
    plt.semilogx(clust1['mass'], clust1['m_nirc2_J'], 'r.')
    plt.semilogx(clust2['mass'], clust2['m_nirc2_J'], 'r.')
    plt.gca().invert_yaxis()
    plt.xlabel('Mass (Msun)')
    plt.ylabel('J (mag)')
    
    # # Plot the spectrum of the most massive star
    # idx = cluster.mass.argmax()
    # plt.clf()
    # plt.plot(cluster.stars[idx].wave, cluster.stars[idx].flux, 'k.')

    # # Plot an integrated spectrum of the whole cluster.
    # wave, flux = cluster.get_integrated_spectrum()
    # plt.clf()
    # plt.plot(wave, flux, 'k.')

    return

def test_ResolvedClusterDiffRedden():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar.imf import imf
    from popstar.imf import multiplicity

    logAge = 6.7
    AKs = 2.4
    distance = 4000
    cluster_mass = 10**5.
    deltaAKs = 0.05
    mass_sampling=5

    # Test filters
    filt_list = ['nirc2,J', 'nirc2,Kp']
    
    startTime = time.time()
    
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere

    red_law = reddening.RedLawNishiyama09()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                                mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    imf_mass_limits = np.array([0.07, 0.5, 1, np.inf])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    ##########
    # Start without multiplicity
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime))
    
    cluster1 = syn.ResolvedClusterDiffRedden(iso, my_imf1, cluster_mass, deltaAKs)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    assert len(clust1) > 0
    
    plt.figure(3)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c.')
    plt.gca().invert_yaxis()

    # *** Visual Inspections: ***
    #  - check that points (red) fall between isochrone points (blue)

    ##########
    # Test with multiplicity
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedClusterDiffRedden(iso, my_imf2, cluster_mass, deltaAKs)
    clust2 = cluster2.star_systems
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    assert len(clust2) > 0
    assert len(cluster2.companions) > 0
    assert np.sum(clust2['N_companions']) == len(cluster2.companions)

    ##########
    # Plots 
    ##########
    # Plot an IR CMD and compare cluster members to isochrone.
    plt.figure(1)
    plt.clf()
    plt.plot(clust1['m_nirc2_J'] - clust1['m_nirc2_Kp'], clust1['m_nirc2_J'], 'r.')
    plt.plot(clust2['m_nirc2_J'] - clust2['m_nirc2_Kp'], clust2['m_nirc2_J'], 'b.')
    plt.plot(iso.points['m_nirc2_J'] - iso.points['m_nirc2_Kp'], iso.points['m_nirc2_J'], 'c-')
    plt.gca().invert_yaxis()
    plt.xlabel('J - Kp (mag)')
    plt.ylabel('J (mag')

    # Plot a mass-magnitude relationship.
    plt.figure(2)
    plt.clf()
    plt.semilogx(clust1['mass'], clust1['m_nirc2_J'], 'r.')
    plt.semilogx(clust2['mass'], clust2['m_nirc2_J'], 'r.')
    plt.gca().invert_yaxis()
    plt.xlabel('Mass (Msun)')
    plt.ylabel('J (mag)')

    return
    
def test_UnresolvedCluster():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar.imf import imf
    from popstar.imf import multiplicity
    
    log_age = 6.7
    AKs = 0.0
    distance = 4000
    metallicity=0
    cluster_mass = 10**4.

    startTime = time.time()    
    multi = multiplicity.MultiplicityUnresolved()
    imf_in = imf.Kroupa_2001(multiplicity=multi)
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere
    iso = syn.Isochrone(log_age, AKs, distance, metallicity=metallicity,
                            evo_model=evo, atm_func=atm_func, mass_sampling=10)
    print('Made Isochrone: %d seconds' % (time.time() - startTime))

    cluster = syn.UnresolvedCluster(iso, imf_in, cluster_mass)
    print('Constructed unresolved cluster: %d seconds' % (time.time() - startTime))

    # Plot an integrated spectrum of the whole cluster.
    wave = cluster.wave_trim
    flux = cluster.spec_trim
    plt.clf()
    plt.plot(wave, flux, 'k.')

    return

def test_ifmr_multiplicity():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar import ifmr
    from popstar.imf import imf
    from popstar.imf import multiplicity

    # Define cluster parameters
    logAge = 9.7
    AKs = 0.0
    distance = 1000
    cluster_mass = 1e6
    mass_sampling = 5

    # Test all filters
    filt_list = ['nirc2,Kp', 'nirc2,H', 'nirc2,J']

    startTime = time.time()
    
    evo = evolution.MISTv1()
    atm_func = atm.get_merged_atmosphere
    ifmr_obj = ifmr.IFMR()

    red_law = reddening.RedLawNishiyama09()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    # Now to create the cluster.
    imf_mass_limits = np.array([0.07, 0.5, 1, np.inf])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    ##########
    # Start without multiplicity and IFMR
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime)) 
    
    cluster1 = syn.ResolvedCluster(iso, my_imf1, cluster_mass, ifmr=ifmr_obj)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

   
    ##########
    # Test with multiplicity and IFMR
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedCluster(iso, my_imf2, cluster_mass, ifmr=ifmr_obj)
    clust2 = cluster2.star_systems
    comps2 = cluster2.companions
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    ##########
    # Tests
    ##########

    # Check that we have black holes, neutron stars, and white dwarfs in both.
    assert len(np.where(clust1['phase'] == 101)) > 0   # WD
    assert len(np.where(clust2['phase'] == 101)) > 0
    assert len(np.where(clust1['phase'] == 102)) > 0   # NS
    assert len(np.where(clust2['phase'] == 102)) > 0
    assert len(np.where(clust1['phase'] == 103)) > 0   # BH
    assert len(np.where(clust2['phase'] == 103)) > 0

    # Now check that we have companions that are WDs, NSs, and BHs
    assert len(np.where(comps2['phase'] == 101)) > 0
    assert len(np.where(comps2['phase'] == 102)) > 0
    assert len(np.where(comps2['phase'] == 103)) > 0

    # Make sure no funky phase designations (due to interpolation effects)
    # slipped through
    idx = np.where( (clust1['phase'] > 5) & (clust1['phase'] < 101) & (clust1['phase'] != 9) )
    idx2 = np.where( (comps2['phase'] > 5) & (comps2['phase'] < 101) & (comps2['phase'] != 9) )
    assert len(idx[0]) == 0

    return

def test_metallicity():
    """
    Test isochrone generation at different metallicities
    """
    # Define isochrone parameters
    logAge = np.log10(5*10**6.) 
    AKs = 0.8 
    dist = 4000 
    evo_model = evolution.MISTv1()
    atm_func = atmospheres.get_phoenixv16_atmosphere
    red_law = reddening.RedLawHosek18b()
    filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m']

    # Start with a solar metallicity isochrone    
    metallicity= 0.0

    # Make Isochrone object, with high mass_sampling to decrease compute time
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=metallicity,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10)

    # Test isochrone properties
    assert my_iso.points.meta['METAL_IN'] == 0.0
    assert os.path.exists('iso_6.70_0.80_04000_p00.fits')

    # Now for non-solar metallicity
    metallicity= -1.5

    # Make Isochrone object, with high mass_sampling to decrease compute time
    my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=metallicity,
                            evo_model=evo_model, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=10)

    metal_act = np.log10(0.00047 / 0.0142) # For Mist isochrones

    # Test isochrone properties
    assert my_iso.points.meta['METAL_IN'] == -1.5
    assert my_iso.points.meta['METAL_ACT'] == metal_act
    assert os.path.exists('iso_6.70_0.80_04000_m15.fits')
    
    return

def test_cluster_mass():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar import ifmr
    from popstar.imf import imf
    from popstar.imf import multiplicity

    # Define cluster parameters
    logAge = 6.7
    AKs = 2.4
    distance = 4000
    cluster_mass = 10**5.
    mass_sampling = 5

    # Test filters
    filt_list = ['nirc2,J', 'nirc2,Kp']

    startTime = time.time()
    
    # Define evolution/atmosphere models and extinction law
    evo = evolution.MISTv1() 
    atm_func = atmospheres.get_merged_atmosphere
    red_law = reddening.RedLawHosek18b()
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list,
                            mass_sampling=mass_sampling)

    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    # Now to create the cluster.
    imf_mass_limits = np.array([0.2, 0.5, 1, 120.0])
    imf_powers = np.array([-1.3, -2.3, -2.3])

    # IFMR
    my_ifmr = ifmr.IFMR()
    

    ##########
    # Start without multiplicity
    ##########
    my_imf1 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=None)
    print('Constructed IMF: %d seconds' % (time.time() - startTime))
    
    cluster1 = syn.ResolvedCluster(iso, my_imf1, cluster_mass, ifmr=my_ifmr)
    clust1 = cluster1.star_systems
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    # Check that the total mass is within tolerance of input mass
    cluster_mass_out = clust1['systemMass'].sum()
    assert np.abs(cluster_mass_out - cluster_mass) < 200.0   # within 200 Msun of desired mass.
    print('Cluster Mass: IN = ', cluster_mass, " OUT = ", cluster_mass_out)

    ##########
    # Test with multiplicity
    ##########
    multi = multiplicity.MultiplicityUnresolved()
    my_imf2 = imf.IMF_broken_powerlaw(imf_mass_limits, imf_powers,
                                      multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster2 = syn.ResolvedCluster(iso, my_imf2, cluster_mass, ifmr=my_ifmr)
    clust2 = cluster2.star_systems
    print('Constructed cluster with multiples: %d seconds' % (time.time() - startTime))

    # Check that the total mass is within tolerance of input mass
    cluster_mass_out = clust2['systemMass'].sum()
    assert np.abs(cluster_mass_out - cluster_mass) < 200.0   # within 200 Msun of desired mass.
    print('Cluster Mass: IN = ', cluster_mass, " OUT = ", cluster_mass_out)

    return

#=================================#
# Additional timing functions
#=================================#

def time_test_cluster():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar import reddening
    from popstar.imf import imf
    from popstar.imf import multiplicity

    logAge = 6.7
    AKs = 2.7
    distance = 4000
    cluster_mass = 10**4

    startTime = time.time()
    
    evo = evolution.MergedBaraffePisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere
    red_law = reddening.RedLawNishiyama09()
    filt_list = ['nirc2,J', 'nirc2,Kp']
    
    iso = syn.IsochronePhot(logAge, AKs, distance,
                            evo_model=evo, atm_func=atm_func,
                            red_law=red_law, filters=filt_list)
    print('Constructed isochrone: %d seconds' % (time.time() - startTime))

    imf_limits = np.array([0.07, 0.5, 150])
    imf_powers = np.array([-1.3, -2.35])
    multi = multiplicity.MultiplicityUnresolved()
    my_imf = imf.IMF_broken_powerlaw(imf_limits, imf_powers, multiplicity=multi)
    print('Constructed IMF with multiples: %d seconds' % (time.time() - startTime))
    
    cluster = syn.ResolvedCluster(iso, my_imf, cluster_mass)
    print('Constructed cluster: %d seconds' % (time.time() - startTime))

    return
    
def model_young_cluster_object(resolved=False):
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar.imf import imf
    from popstar.imf import multiplicity

    log_age = 6.5
    AKs = 0.1
    distance = 8000.0
    cluster_mass = 10000.
    
    multi = multiplicity.MultiplicityUnresolved()
    imf_in = imf.Kroupa_2001(multiplicity=multi)
    evo = evolution.MergedPisaEkstromParsec()
    atm_func = atm.get_merged_atmosphere
    iso = syn.Isochrone(log_age, AKs, distance, evo, mass_sampling=10)

    if resolved:
        cluster = syn.ResolvedCluster(iso, imf_in, cluster_mass)
    else:
        cluster = syn.UnresolvedCluster(iso, imf_in, cluster_mass, wave_range=[19000,24000])

    # Plot the spectrum of the most massive star
    idx = cluster.mass_all.argmax()
    print('Most massive star is {0:f} M_sun.'.format(cluster.mass_all[idx]))
    #bigstar = cluster.spec_list_trim[idx]
    plt.figure(1)
    plt.clf()
    plt.plot(cluster.spec_list_trim[idx]._wavetable, cluster.spec_list_trim[idx]._fluxtable, 'k.')

    # Plot an integrated spectrum of the whole cluster.
    wave, flux = cluster.spec_list_trim[idx]._wavetable, cluster.spec_trim
    plt.figure(2)
    plt.clf()
    plt.plot(wave, flux, 'k.')

    return
    
def time_test_mass_match():
    from popstar import synthetic as syn
    from popstar import atmospheres as atm
    from popstar import evolution
    from popstar.imf import imf
    from popstar.imf import multiplicity

    log_age = 6.7
    AKs = 2.7
    distance = 4000
    cluster_mass = 5e3
    
    imf_in = imf.Kroupa_2001(multiplicity=None)

    start_time = time.time()
    iso = syn.IsochronePhot(log_age, AKs, distance)
    iso_masses = iso.points['mass']
    print('Generated iso masses in {0:.0f} s'.format(time.time() - start_time))

    start_time = time.time()
    star_masses, isMulti, compMass, sysMass = imf_in.generate_cluster(cluster_mass)
    print('Generated cluster masses in {0:.0f} s'.format(time.time() - start_time))
    
    def match_model_masses1(isoMasses, starMasses):
        indices = np.empty(len(starMasses), dtype=int)
        
        for ii in range(len(starMasses)):
            theMass = starMasses[ii]
            
            dm = np.abs(isoMasses - theMass)
            mdx = dm.argmin()

            # Model mass has to be within 10% of the desired mass
            if (dm[mdx] / theMass) > 0.1:
                indices[ii] = -1
            else:
                indices[ii] = mdx

        return indices
            

    def match_model_masses2(isoMasses, starMasses):
        isoMasses_tmp = isoMasses.reshape((len(isoMasses), 1))
        kdt = KDTree(isoMasses_tmp)
 
        starMasses_tmp = starMasses.reshape((len(starMasses), 1))
        q_results = kdt.query(starMasses_tmp, k=1)
        indices = q_results[1]

        dm_frac = np.abs(starMasses - isoMasses[indices]) / starMasses

        idx = np.where(dm_frac > 0.1)[0]
        indices[idx] = -1
        
        return indices

    print('Test #1 START')
    start_time = time.time()
    idx1 = match_model_masses1(iso_masses, star_masses)
    stop_time = time.time()
    print('Test #1 STOPPED after {0:.0f} seconds'.format(stop_time - start_time))

    print('Test #2 START')
    start_time = time.time()
    idx2 = match_model_masses2(iso_masses, star_masses)
    stop_time = time.time()
    print('Test #2 STOPPED after {0:.0f} seconds'.format(stop_time - start_time))

    return


def test_phot_consistency(filt='all'):
    """
    Test photometric consistency of generated isochrone (IsochronePhot)
    against pre-generated isochrone with native filter sampling. Requires
    consistency to within 0.005 mag.

    Base isochrone is at 5 Myr, AKs = 0, 1000 pc, mass_sampling=10 

    Paramters:
    ----------
    filt: 'all', 'hst', 'vista', 'decam', 'ps1', 'jwst'
        Specify what filter set you want to test

    """
    from astropy.table import Table
    import os
    
    # Load pre-generated isochrone, located in popstar tests directory
    direct = os.path.dirname(__file__)
    orig = Table.read(direct+'/iso_6.70_0.00_01000.fits', format='fits')

    # Generate new isochrone with popstar code
    if filt == 'all':
        filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m',
                         'acs,wfc1,f814w', 'wfc3,ir,f125w', 'wfc3,ir,f160w',
                         'decam,y', 'decam,i', 'decam,z',
                         'decam,u', 'decam,g', 'decam,r',
                         'vista,Y', 'vista,Z', 'vista,J',
                         'vista,H',  'vista,Ks',
                         'ps1,z', 'ps1,g', 'ps1,r',
                         'ps1,i', 'ps1,y',
                         'jwst,F070W', 'jwst,F090W', 'jwst,F115W', 'jwst,F140M',
                         'jwst,F150W', 'jwst,F150W2', 'jwst,F162M', 'jwst,F164N',
                         'jwst,F182M', 'jwst,F187N', 'jwst,F200W', 'jwst,F212N',
                         'jwst,F210M','jwst,F250M', 'jwst,F277W', 'jwst,F300M',
                         'jwst,F322W2', 'jwst,F323N', 'jwst,F335M', 'jwst,F356W',
                         'jwst,F360M', 'jwst,F405N', 'jwst,F410M', 'jwst,F430M',
                         'jwst,F444W', 'jwst,F460M', 'jwst,F466N', 'jwst,F470N',
                         'jwst,F480M', 'nirc2,J', 'nirc2,H', 'nirc2,Kp', 'nirc2,K',
                         'nirc2,Lp', 'nirc2,Hcont',
                         'nirc2,FeII', 'nirc2,Brgamma',
                         'jg,J', 'jg,H', 'jg,K',
                         'nirc1,K', 'nirc1_H', 'ctio_osiris,K', 'ctio_osiris,H',
                         'naco,H', 'naco,Ks', 'ztf,g', 'ztf,r', 'ztf,i']

    elif filt == 'decam':
        filt_list = ['decam,y', 'decam,i', 'decam,z',
                         'decam,u', 'decam,g', 'decam,r']
    elif filt == 'vista':
        filt_list = ['vista,Y', 'vista,Z', 'vista,J',
                         'vista,H', 'vista,Ks']
    elif filt == 'ps1':
        filt_list = ['ps1,z', 'ps1,g', 'ps1,r',  'ps1,i',
                         'ps1,y']
    elif filt == 'jwst':
        filt_list = ['jwst,F070W', 'jwst,F090W', 'jwst,F115W', 'jwst,F140M',
                         'jwst,F150W', 'jwst,F150W2', 'jwst,F162M', 'jwst,F164N',
                         'jwst,F182M', 'jwst,F187N', 'jwst,F200W', 'jwst,F212N',
                         'jwst,F210M','jwst,F250M', 'jwst,F277W', 'jwst,F300M',
                         'jwst,F322W2', 'jwst,F323N', 'jwst,F335M', 'jwst,F356W',
                         'jwst,F360M', 'jwst,F405N', 'jwst,F410M', 'jwst,F430M',
                         'jwst,F444W', 'jwst,F460M', 'jwst,F466N', 'jwst,F470N',
                         'jwst,F480M']
    elif filt == 'hst':
        filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m',
                         'acs,wfc1,f814w', 'wfc3,ir,f125w', 'wfc3,ir,f160w']
    elif filt == 'nirc2':
        filt_list = ['nirc2,J', 'nirc2,H','nirc2,Kp', 'nirc2,K',
                         'nirc2,Lp', 'nirc2,Ms', 'nirc2,Hcont',
                         'nirc2,FeII', 'nirc2,Brgamma']
    elif filt == 'jg':
        filt_list = ['jg,J', 'jg,H', 'jg,K']
    elif filt == 'ztf':
        filt_list = ['ztf,g', 'ztf,r', 'ztf,i']
    elif filt == 'misc':
        filt_list=['nirc1,K', 'nirc1,H', 'ctio_osiris,K', 'ctio_osiris,H',
                       'naco,H', 'naco,Ks']
        
            
    print('Making isochrone')
    iso = synthetic.IsochronePhot(6.7, 0, 1000, mass_sampling=10, filters=filt_list, rebin=True)
    iso = iso.points

    # First find masses that are the same
    foo = []
    for ii in iso['mass']:
        tmp = np.where( orig['mass'] == ii)[0][0]
        foo.append(tmp)
        
    assert len(foo) == len(iso)
    orig = orig[foo]

    # Identify the photometry columns
    cols = list(iso.keys())
    idx = []
    for ii in range(len(cols)):
        if cols[ii].startswith('mag'):
            idx.append(ii)
    mag_cols = np.array(cols)[idx]

    # Test the consistency of each column with the original isochrone
    for ii in mag_cols:
        orig_mag = orig[ii]
        new_mag = iso[ii]

        np.testing.assert_allclose(orig_mag, new_mag, rtol=0.01, err_msg="{0} failed".format(ii))

        # Also report median abs difference
        diff = abs(orig_mag - new_mag)
        print(('{0} median abs diff: {1}'.format(ii, np.median(diff))))


    print(('Phot consistency test successful for {0}'.format(filt)))

    # Remove iso file we just wrote, since it was only a test
    os.remove('iso_6.70_0.00_01000.fits')
    return


# Siess not supported anymore.
# 
# def test_isochrone_siess_mass_current_bug():
#     """
#     Bug found by students in grad stars class. 
#     """
#     # Define isochrone parameters
#     logAges = [6.6, 7.6] # Age in log(years)
#     AKs = 0 # extinction in mags
#     dist = 1000 # distance in parsec
#     metallicity = 0 # Metallicity in [M/H]

#     # Define evolution/atmosphere models and extinction law
#     evo_model = evolution.MergedSiessGenevaPadova()
#     evo_model_name = ['MIST', 'Padova']
#     atm_func = atmospheres.get_merged_atmosphere
#     red_law = reddening.RedLawHosek18b()

#     # Also specify filters for synthetic photometry (optional). Here we use
#     # the HST WFC3-IR F127M, F139M, and F153M filters
#     filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m']

#     # Make Isochrone object. Note that is calculation will take a few minutes, unless the
#     # isochrone has been generated previously.

#     for logAge in logAges:
#         my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=0, evo_model=evo_model, atm_func=atm_func,
#                                          red_law=red_law, filters=filt_list)
#         print(my_iso.save_file)

#     return
# Typically takes 104 - 120 seconds. # Limited by pysynphot.Icat call in atmospheres.py Test to make sure isochrones generated have spectra with the proper      wavelength range, and that the user has control over that wavelength     range (propagated through IsochronePhot) # Define isochrone parameters # Age in log(years) # extinction in mags # distance in parsec # Define evolution/atmosphere models and extinction law (optional) # Also specify filters for synthetic photometry (optional). Here we use # the HST WFC3-IR F127M, F139M, and F153M filters # First, let's make sure the vega spectrum has the proper limits # Make Isochrone object. Will use wave_range = [3000,52000]. # Make sure range matches to resolution of atmosphere. # Now let's try changing the wave range. Is it carried through # properly? # Does the error exception catch the bad wave_range? # Typically takes 120 seconds if file is regenerated. # Limited by pysynphot.Icat call in atmospheres.py # Finally, let's test the isochronePhot file generation # Check 1: If we try to remake the isochrone, does it read the file rather than # making a new one # Check 2: If we change evo model, atmo model, or redlaw, # does IsochronePhot regenerate the isochrone and overwrite the existing one? # Define cluster parameters # Test filters # Now to create the cluster. ########## # Start without multiplicity ########## # Check that stars are returned # Check that the total mass in stars is less than requested (no compact objects). # *** Visual Inspections: *** #  - check that points (red) fall between isochrone points (blue) ########## # Test with multiplicity ########## ########## # Plots ########## # Plot an IR CMD and compare cluster members to isochrone. # Plot a mass-magnitude relationship. # # Plot the spectrum of the most massive star # idx = cluster.mass.argmax() # plt.clf() # plt.plot(cluster.stars[idx].wave, cluster.stars[idx].flux, 'k.') # # Plot an integrated spectrum of the whole cluster. # wave, flux = cluster.get_integrated_spectrum() # plt.clf() # plt.plot(wave, flux, 'k.') # Test filters ########## # Start without multiplicity ########## # *** Visual Inspections: *** #  - check that points (red) fall between isochrone points (blue) ########## # Test with multiplicity ########## ########## # Plots ########## # Plot an IR CMD and compare cluster members to isochrone. # Plot a mass-magnitude relationship. # Plot an integrated spectrum of the whole cluster. # Define cluster parameters # Test all filters # Now to create the cluster. ########## # Start without multiplicity and IFMR ########## ########## # Test with multiplicity and IFMR ########## ########## # Tests ########## # Check that we have black holes, neutron stars, and white dwarfs in both. # WD # NS # BH # Now check that we have companions that are WDs, NSs, and BHs # Make sure no funky phase designations (due to interpolation effects) # slipped through Test isochrone generation at different metallicities # Define isochrone parameters # Start with a solar metallicity isochrone # Make Isochrone object, with high mass_sampling to decrease compute time # Test isochrone properties # Now for non-solar metallicity # Make Isochrone object, with high mass_sampling to decrease compute time # For Mist isochrones # Test isochrone properties # Define cluster parameters # Test filters # Define evolution/atmosphere models and extinction law # Now to create the cluster. # IFMR ########## # Start without multiplicity ########## # Check that the total mass is within tolerance of input mass # within 200 Msun of desired mass. ########## # Test with multiplicity ########## # Check that the total mass is within tolerance of input mass # within 200 Msun of desired mass. #=================================# # Additional timing functions #=================================# # Plot the spectrum of the most massive star #bigstar = cluster.spec_list_trim[idx] # Plot an integrated spectrum of the whole cluster. # Model mass has to be within 10% of the desired mass #1 START') #1 STOPPED after {0:.0f} seconds'.format(stop_time - start_time)) #2 START') #2 STOPPED after {0:.0f} seconds'.format(stop_time - start_time)) Test photometric consistency of generated isochrone (IsochronePhot)     against pre-generated isochrone with native filter sampling. Requires     consistency to within 0.005 mag.      Base isochrone is at 5 Myr, AKs = 0, 1000 pc, mass_sampling=10       Paramters:     ----------     filt: 'all', 'hst', 'vista', 'decam', 'ps1', 'jwst'         Specify what filter set you want to test # Load pre-generated isochrone, located in popstar tests directory # Generate new isochrone with popstar code # First find masses that are the same # Identify the photometry columns # Test the consistency of each column with the original isochrone # Also report median abs difference # Remove iso file we just wrote, since it was only a test # Siess not supported anymore. # # def test_isochrone_siess_mass_current_bug(): #     """ #     Bug found by students in grad stars class. #     """ #     # Define isochrone parameters #     logAges = [6.6, 7.6] # Age in log(years) #     AKs = 0 # extinction in mags #     dist = 1000 # distance in parsec #     metallicity = 0 # Metallicity in [M/H] #     # Define evolution/atmosphere models and extinction law #     evo_model = evolution.MergedSiessGenevaPadova() #     evo_model_name = ['MIST', 'Padova'] #     atm_func = atmospheres.get_merged_atmosphere #     red_law = reddening.RedLawHosek18b() #     # Also specify filters for synthetic photometry (optional). Here we use #     # the HST WFC3-IR F127M, F139M, and F153M filters #     filt_list = ['wfc3,ir,f127m', 'wfc3,ir,f139m', 'wfc3,ir,f153m'] #     # Make Isochrone object. Note that is calculation will take a few minutes, unless the #     # isochrone has been generated previously. #     for logAge in logAges: #         my_iso = synthetic.IsochronePhot(logAge, AKs, dist, metallicity=0, evo_model=evo_model, atm_func=atm_func, #                                          red_law=red_law, filters=filt_list) #         print(my_iso.save_file) #     return