Lodge It

from __future__ import print_function

import matplotlib
matplotlib.use('Agg')
import yt,caesar
import yt.utilities.physical_constants as const
import numpy as np
import sys
from multiprocessing import Pool

from yt import derived_field
#from fast_histogram import histogram1d


from os.path import expanduser
home = expanduser("~")

#this is just to insert the lookup tables in the same path
sys.path.insert(0,"/blue/narayanan/desika.narayanan/despotic_lookup_table_generator/")

from lookup_table_reader_z_units import *
from astropy.io import fits
import matplotlib.pyplot as plt
import pdb,ipdb

from astropy import constants as astropy_constants
from astropy import units as u

from glob2 import glob
import itertools


from datetime import datetime
from filter_galaxy import filter_galaxy

#==============================================================================
#RUN PARAMETERS
#==============================================================================
caesarfile = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m25n512/output/Groups/caesar_0160_z2.000.hdf5'
caesarfile = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m25n512/output/Groups/caesar_0088_z4.957.hdf5'

snapshot = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m25n512/output/snapshot_160.hdf5'
snapshot = '/orange/narayanan/desika.narayanan/gizmo_runs/simba/m25n512/output/snapshot_088.hdf5'

galaxies = [int(sys.argv[1])]

boxsize=500
zoom_boxsize = 50
oref=0
n_ref=64

unit_base = {'UnitLength_in_cm'         : 3.08568e+21,
             'UnitMass_in_g'            :   1.989e+43,
             'UnitVelocity_in_cm_per_s' :      100000,
             'UnitTime_in_s' : 1}

bbox = [[-boxsize,boxsize],
        [-boxsize,boxsize],
        [-boxsize,boxsize]]

lookup_npzfilename = "/blue/narayanan/desika.narayanan/despotic_lookup_table_generator/high_res_abu_z.npz"
#lookup_npzfilename ='/ufrc/narayanan/pg3552/despotic_lookup_table_generator/testintIntensity_final.npz'
bin_width = 10 #km/s
dim=2048
#dim=1024

turb_factor = 0.25
outdir = '/blue/narayanan/desika.narayanan/paper/ngvla_memo_narayanan/ppv_files/simba/m25n512/intIntensity/production_0.25/'#680a6ee_oldtable_intTB/'
outdir = '/blue/narayanan/desika.narayanan/paper/ngvla_memo_narayanan/ppv_files/simba/m25n512/intIntensity/production_0.25_snap88/'#680a6ee_oldtable_intTB/'
#outdir = '/blue/narayanan/desika.narayanan/paper/ngvla_memo_narayanan/ppv_files/simba/m25n512/intTb/production_0.25_snap160/'#680a6ee_oldtable_intTB/'
#outdir = './'
nlevels = 10

PARTICLES = True #flag says if we're going to do thinsg on a
                 #particle-by-particle basis (and then smooth), vs doing smoothing first, and then doing the CO calculation.
#INTENSITY_FLAG = True
units = 'intIntensity'
#units = 'intTb'
nprocesses=1



                #==============================================================================

def _fmol(field,data):
    return ds.arr(h2_mass/(h2_mass+hi_mass)*data[("PartType0","Density")],'code_mass/code_length**3')

def _WCO(field,data):
    return ds.arr(wco_particles,'K*km*s**(-1)')

def _neutralhydrogen(field,data):
    return data[("PartType0","NeutralHydrogenAbundance")]*data[("PartType0","Density")]


#def _WCO10(field,data):
#    return ds.arr(wco_particles[0,:],'K*km*s**(-1)')

def _manual_renorm_density(field,data):
    return ds.arr(data["PartType0","density"],'code_mass/code_length**3')

def _WCO10(field,data):
    return ds.arr(wco_particles[0,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO21(field,data):
    return ds.arr(wco_particles[1,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO32(field,data):
    return ds.arr(wco_particles[2,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO43(field,data):
    return ds.arr(wco_particles[3,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO54(field,data):
    return ds.arr(wco_particles[4,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')


def _WCO65(field,data):
    return ds.arr(wco_particles[5,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO76(field,data):
    return ds.arr(wco_particles[6,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO87(field,data):
    return ds.arr(wco_particles[7,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO98(field,data):
    return ds.arr(wco_particles[8,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCO109(field,data):
    return ds.arr(wco_particles[9,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCI10(field,data):
    return ds.arr(wci_particles[0,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCI21(field,data):
    return ds.arr(wci_particles[1,:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')

def _WCII(field,data):
    return ds.arr(wcii_particles[:]*data[("PartType0","Density")],'K*km*s**(-1)*code_mass/code_length**3')


table = TableReader(lookup_npzfilename)
table.limitsMode = "leave"
#table.limitsMode='clip'
table.copyPoints = True

def find_between( s, first, last ):
    try:
        start = s.index( first ) + len( first )
        end = s.index( last, start )
        return s[start:end]
    except ValueError:
        return ""

def generate_lines(package):
    table = TableReader(lookup_npzfilename)
    table.limitsMode = "leave"

    redshift_particles = package[0]
    column_particles = package[1]
    metal_particles = package[2]
    nh_particles = package[3]
    sfr_particles = package[4]
    
    output_new = table.getValues(
        redshift_particles, column_particles,
        metal_particles, nh_particles, sfr_particles,
        units=units, log_input=False, log_output=False)
    
    wco_particles = output_new['co']
    
    return wco_particles


    
#find the overlapping lists of gadget files and pdfiles

gizmo_file_list = []
pd_file_list = []
pd_snaps = []
gizmo_snaps = []


obj = caesar.load(caesarfile)


for galaxy in galaxies:


    try:
        filter_galaxy(snapshot,caesarfile,galaxy,'temp.hdf5',halos=False)
        snapshot = 'temp.hdf5'
    except:
        print("Error in filter_galaxy: not filtering")
        pass



    ds = yt.load(snapshot)
    if ds.cosmological_simulation == 0:
        ds = yt.load(snapshot,unit_base=unit_base,bounding_box=bbox)

    ds.index
    ad = ds.all_data()



    center = obj.galaxies[galaxy].pos.in_units('code_length')

    #glist = obj.galaxies[galaxy].glist
    glist = np.arange(len(ad[("PartType0","Masses")]))

    nparticles = len(glist)


    mass = ad[('PartType0','Masses')][glist].in_units('Msun')#*ad[('PartType0', 'NeutralHydrogenAbundance')]
    hsml = ad[('PartType0','SmoothingLength')][glist].in_units('pc')
    dens = ad[('PartType0', 'Density')][glist].in_units('g/cm**3')#*ad[('PartType0', 'NeutralHydrogenAbundance')]


    redshift_particles = np.repeat(ds.current_redshift,nparticles)

    metal_particles = ad[("PartType0",'Metallicity_00')][glist]/0.013
    #metal_particles[metal_particles < 0.1] = 0.1 #have to set this floor or else we'll get negative CO fluxes

    #DEBUG 071019
    #metal_particles = np.ones(len(glist))
    
    wco_particles = np.zeros(nparticles)
    wci_particles = np.zeros(nparticles)
    wcii_particles = np.zeros(nparticles)

    sfr_particles = np.repeat(obj.galaxies[galaxy].sfr.value,nparticles)

    nh_particles = (dens/const.mass_hydrogen).in_units('cm**-3').value

    #THIS IS SOMETHING WE CAN LOOK AT IF WE DON'T HAVE SUFFICIENT COLUMN: THE HSML AS THE SIZE
    column_particles = (dens*hsml).in_units('Msun/pc**2').value
    column_particles *= (4./3*np.pi)

    #DEBUG
    for i in range(len(column_particles)):
        column_particles[i] =np.max((75,column_particles[i]))
        
    #volume = mass/dens
    #r = ((3./4*volume/np.pi)**(1./3)).in_units('cm')
    #column_particles = (dens*r).in_units('Msun/pc**2').value
    

    #THIS IS SOMETHING WE CAN LOOK AT IF WE DON'T HAVE SUFFICIENT
    #DENSITY: THE TEMPERATURE IS SET AS CONSTANT BASED ON CO
    #OBSERVATIONS, BUT INSTEAD PERHAPS WE SHOULD TAKE THE MASS
    #WEIGHTED VALUES FROM LI, NARAYANAN, DAVE & KRUMHOLZ FOR ENTIRE
    #CLOUDS

    gamma = 1.4
    cs = np.sqrt(gamma*astropy_constants.k_B/astropy_constants.m_p*10.*u.K) #assuming temp of 10K
    sigma_vir = 2.2 * (mass/1.e5)**1./4
    sigma_vir = (sigma_vir.value)*u.km/u.s
    mach = sigma_vir.cgs/cs.cgs
    sigma_p_sq = np.log(1.+ 3.*mach**2./4)
    turbulent_compression_factor =np.exp(sigma_p_sq/2.)
    
    if turb_factor > 0:
        #THIS IS SOMETHING WE CAN LOOK AT IF WE DON'T GET SUFFICIENT COLUMN: THE COLUMN DENSITY * THE TURBULENT COMPRESSION FACTOR
        
        nh_particles *= turb_factor*turbulent_compression_factor.flatten().value
        column_particles *= turb_factor*turbulent_compression_factor.flatten().value

    p = Pool(processes=nprocesses)
    nchunks = nprocesses
    chunk_start_indices = []
    chunk_start_indices.append(0) #the start index is obviously 0
    #this should just be int(nparticles/nchunks) but in case particles < nchunks, we need to ensure that this is at least  1
    delta_chunk_indices = np.max([int(nparticles / nchunks),1]) 
    print ('delta_chunk_indices = ',delta_chunk_indices)
    for n in range(1,nchunks):
        chunk_start_indices.append(chunk_start_indices[n-1]+delta_chunk_indices)

    redshift_list_of_chunks = []
    column_list_of_chunks = []
    metal_list_of_chunks = []
    nh_list_of_chunks = []
    sfr_list_of_chunks = []
    packages = []

    for n in range(nchunks):
        redshift_list_chunk = redshift_particles[chunk_start_indices[n]:chunk_start_indices[n]+delta_chunk_indices]
        column_list_chunk = column_particles[chunk_start_indices[n]:chunk_start_indices[n]+delta_chunk_indices]
        metal_list_chunk = metal_particles[chunk_start_indices[n]:chunk_start_indices[n]+delta_chunk_indices]
        nh_list_chunk = nh_particles[chunk_start_indices[n]:chunk_start_indices[n]+delta_chunk_indices]
        sfr_list_chunk = sfr_particles[chunk_start_indices[n]:chunk_start_indices[n]+delta_chunk_indices]
        
        #if we're on the last chunk, we might not have the full list included, so need to make sure that we have that here
        if n == nchunks-1: 
            redshift_list_chunk = redshift_particles[chunk_start_indices[n]::]
            column_list_chunk = column_particles[chunk_start_indices[n]::]
            metal_list_chunk = metal_particles[chunk_start_indices[n]::]
            nh_list_chunk = nh_particles[chunk_start_indices[n]::]
            sfr_list_chunk = sfr_particles[chunk_start_indices[n]::]

        redshift_list_of_chunks.append(redshift_list_chunk)
        column_list_of_chunks.append(column_list_chunk)
        metal_list_of_chunks.append(metal_list_chunk)
        nh_list_of_chunks.append(nh_list_chunk)
        sfr_list_of_chunks.append(sfr_list_chunk)
        
        packages.append([redshift_list_chunk,column_list_chunk,metal_list_chunk,nh_list_chunk,sfr_list_chunk])

    '''#this is the code to actually run the pool.map version of the
    table generation.  doesn't really seem necessary but if you end up
    executing it, just need to change the return statements for generate_lines to also return  CI and CII

    t1=datetime.now()
    chunk_sol = p.map(generate_lines,[arg for arg in packages])
    wco_partidcles = np.concatenate((chunk_sol[0::]),axis=1)
    t2=datetime.now()
    print ('Execution time map = '+str(t2-t1))

    '''



    t1=datetime.now()
    output_new = table.getValues(
        redshift_particles, column_particles,
        metal_particles, nh_particles, sfr_particles,
        units=units, log_input=False, log_output=False)
    t2=datetime.now()
    print ('Execution time map = '+str(t2-t1))



    wco_particles = output_new['co']*ad[('PartType0', 'NeutralHydrogenAbundance')]
    wci_particles = output_new['ci']*ad[('PartType0', 'NeutralHydrogenAbundance')]
    wcii_particles = output_new['cii']*ad[('PartType0', 'NeutralHydrogenAbundance')]



    try:
        h2_abu_particles = ad[('PartType0', 'FractionH2')][glist]
        hi_abu_particles = 1.-h2_abu_particles
    except:
        h2_abu_particles =  output_new['h2']
        hi_abu_particles = output_new['hi']

    h2_mass = mass*h2_abu_particles
    hi_mass = mass*hi_abu_particles

    if units=='lumPerH':
    #if we're doing in terms of erg/s/h then we need to multiply by the number of H nucleons
        num_hydrogen = mass.in_units('g')/const.mass_hydrogen.in_units('g')
        wco_particles *= num_hydrogen
        wci_particles *= num_hydrogen
        wcii_particles *= num_hydrogen




    #now for wco
    #wco = np.random.random(len(ad["PartType0","Masses"]))
    wco_particles_dict = {}
    wci_particles_dict = {}
    wcii_particles_dict = {}

    for i in range(nlevels):
        wco_particles_dict['wco_particles'+str(i+1)+str(i)] = wco_particles[i,:]
    for i in range(2):
        wci_particles_dict['wci_particles'+str(i+1)+str(i)] = wci_particles[i,:]
    
    wcii_particles_dict['wcii_particles10'] = wcii_particles[:]
            
    print("adding CO/CI/CII fields to the yt dataset")

    ds.add_field(("PartType0","wco10"),function=_WCO10,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco21"),function=_WCO21,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco32"),function=_WCO32,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco43"),function=_WCO43,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco54"),function=_WCO54,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco65"),function=_WCO65,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco76"),function=_WCO76,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco87"),function=_WCO87,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco98"),function=_WCO98,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wco109"),function=_WCO98,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    
    ds.add_field(("PartType0","wci10"),function=_WCI10,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wci21"),function=_WCI21,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    ds.add_field(("PartType0","wcii"),function=_WCII,units='K*km*s**(-1)*code_mass/code_length**3',particle_type=True)
    
    
    print("adding deposited data fields to the yt dataset")

    #ds.add_field(("PartType0","wco"),function=_WCO,units='K*km*s**(-1)',particle_type=True)
    #ds.add_deposited_particle_field(("PartType0","wco"),'sum')
    for i in range(nlevels):
        ds.add_deposited_particle_field(("PartType0","wco"+str(i+1)+str(i)),'sum')
        ds.add_deposited_particle_field(("PartType0","wco"+str(i+1)+str(i)),'sum')
        

    for i in range(2):
        ds.add_deposited_particle_field(("PartType0","wci"+str(i+1)+str(i)),'sum')
    ds.add_deposited_particle_field(("PartType0","wcii"),'sum')


    #first, compute the renormalization factor that we used to correctly deposit the wco particles
    ds.add_field(("PartType0","manual_renorm_density"),function=_manual_renorm_density,units='code_mass/code_length**3',particle_type=True)
    ds.add_deposited_particle_field(("PartType0","manual_renorm_density"),"sum")
    px_h_renorm = yt.ProjectionPlot(ds,'z',('deposit','PartType0_sum_manual_renorm_density'),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    px_h_renorm.save('renorm.png')
    #proj = ds.proj(('PartType0_sum_manual_renorm_density'),1,center=obj.galaxies[galaxy].pos.in_units('code_length'))
    #h_renorm_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
    #h_renorm_frb = h_renorm_frb[('PartType0_sum_manual_renorm_density')].in_units('g/cm**2')
    h_renorm_frb = px_h_renorm.frb[('deposit','PartType0_sum_manual_renorm_density')]

    #second, compute the wco/ci/cii map and divide out the renormalization factor 
    px_co = yt.ProjectionPlot(ds,'z',('deposit','PartType0_sum_wco10'),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    px_co.save('co_proj.png')
    #proj = ds.proj(('deposit','PartType0_sum_wco10'),1,center=obj.galaxies[galaxy].pos.in_units('code_length'))
    #co_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
    #co_frb = (co_frb[('deposit','PartType0_sum_wco10')]/h_renorm_frb).in_units('K*km/s')
    co_frb = px_co.frb[('deposit', 'PartType0_sum_wco10')]/h_renorm_frb
    print("in the middle, co_frb = ",np.sum(co_frb[~np.isnan(co_frb)]))
    

    px_ci = yt.ProjectionPlot(ds,'z',('deposit','PartType0_sum_wci10'),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    px_ci.save('ci_proj.png')
    ci_frb = px_ci.frb[('deposit', 'PartType0_sum_wci10')]/h_renorm_frb

    px_cii = yt.ProjectionPlot(ds,'z',('deposit','PartType0_sum_wcii'),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    px_cii.save('cii_proj.png')
    cii_frb = px_cii.frb[('deposit', 'PartType0_sum_wcii')]/h_renorm_frb


    #third, get the real column density from the SPH particles
    px_h = yt.ProjectionPlot(ds,'z',('deposit','PartType0_smoothed_density'),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    px_h.set_unit(('deposit','PartType0_smoothed_density'),'g*cm**(-2)')
    px_h.save('density.png')

    #proj = ds.proj(('deposit','PartType0_smoothed_density'),1,center=obj.galaxies[galaxy].pos.in_units('code_length'))
    #h_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
    #h_frb=(h_frb[('deposit','PartType0_smoothed_density')]/const.mass_hydrogen).in_units('cm**-2')
    h_frb = px_h.frb[('deposit','PartType0_smoothed_density')]/const.mass_hydrogen

    #fourth get fmol
    ds.add_field(("PartType0","fmol"),function=_fmol,units='code_mass/code_length**3',particle_type=True)
    ds.add_deposited_particle_field(("PartType0","fmol"),"sum")
    px_fmol = yt.ProjectionPlot(ds,'z',("deposit","PartType0_sum_fmol"),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    px_fmol.save('junk3.png')
    #proj = ds.proj(("deposit","PartType0_sum_fmol"),1,center=obj.galaxies[galaxy].pos.in_units('code_length'))
    #fmol_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
    #fmol_frb = fmol_frb[("deposit","PartType0_sum_fmol")]/h_renorm_frb
    fmol_frb = px_fmol.frb[('deposit','PartType0_sum_fmol')]/h_renorm_frb
    

    #fifth, deposit the neutral hydrogen abundance just so we can include it with the xco. this isn't used for anything else really
    ds.add_field(("PartType0","NeutralGasDensity"),function=_neutralhydrogen,units='code_mass/code_length**3',particle_type=True)
    ds.add_deposited_particle_field(("PartType0","NeutralGasDensity"),"sum")
    px_NHA = yt.ProjectionPlot(ds,'z',("deposit","PartType0_sum_NeutralGasDensity"),center=obj.galaxies[galaxy].pos.in_units('code_length'),width=(zoom_boxsize,'kpc'))
    NHA_frb = px_NHA.frb[("deposit","PartType0_sum_NeutralGasDensity")]/h_renorm_frb
    NHA_frb[np.isnan(NHA_frb)] = 0


    xco_frb = NHA_frb*h_frb/co_frb*fmol_frb




    #Plot the histograms of the CO map


    

    fig = plt.figure()
    ax1 = fig.add_subplot(221)

    #we set up this 1D array where we mask out the nans since we might
    #have a ton after filtering the galaxy and putting a bunch of 0s
    #in the grid
    xco_hist = xco_frb.value.ravel()[~np.isnan(xco_frb.value.ravel())]
    NH_hist = h_frb.value.ravel()[~np.isnan(xco_frb.value.ravel())]
    fmol_hist = fmol_frb.value.ravel()[~np.isnan(xco_frb.value.ravel())]
    co_frb_hist = np.take(co_frb.value.ravel(),~np.isnan(xco_frb.value.ravel()))



#    plotting - leaving out for now

    p,b,h = plt.hist(np.log10(xco_hist[xco_hist>0]),bins=50)#,weights=co_frb_hist[xco_hist>0])
    ax1.set_yscale('log')
    ax1.set_xlabel(r'log$_\mathrm{10}$(cm$^{-2}$ (K-km/s)$^{-1}$)')
    ax1.set_ylabel('N')
    ax1.set_xlim([15,30])
    ax1.set_title('Xco')
    
    ax2 = fig.add_subplot(222)
    p,b,h = plt.hist(np.log10(NH_hist[np.nonzero(NH_hist)]),bins=50,weights=co_frb_hist[np.nonzero(NH_hist)])
    #p,b,h = plt.hist(np.log10(column_particles),bins=50)
    ax2.set_yscale('log')
    ax2.set_xlabel(r'log$_\mathrm{10}$(cm$^{-2}$)')
    ax2.set_ylabel('N')
    ax2.set_title(r'N_\mathrm{H2}$')


    ax3 = fig.add_subplot(223)
    p,b,h = plt.hist(np.log10(fmol_hist[np.nonzero(fmol_hist)]),bins=50,weights=co_frb_hist[np.nonzero(fmol_hist)])
    ax3.set_yscale('log')
    ax3.set_xlabel(r'log fraction')
    ax3.set_ylabel('N')
    ax3.set_title(r'f$_\mathrm{mol}$')
    
    ax4 = fig.add_subplot(224)
    p,b,h = plt.hist(np.log10(co_frb_hist[co_frb_hist>0]),bins=50)#,weights=co_frb_hist[np.nonzero(co_frb_hist)])
    ax4.set_yscale('log')
    ax4.set_ylabel('N')
    ax4.set_xlabel(r'K-km/s')
    ax4.set_title('Wco')

    fig.subplots_adjust(hspace=0.55, wspace=0.55)
    fig.savefig(outdir+'xco_hist.'+str(galaxy)+'.png',dpi=300)




    #Plot the CO map itself
    fig = plt.figure()
    ax = fig.add_subplot(111)

    for jlower in range(9):

        proj = ds.proj(('PartType0_sum_manual_renorm_density'),2,center=obj.galaxies[galaxy].pos.in_units('code_length'))
        h_renorm_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
        h_renorm_frb = h_renorm_frb[('PartType0_sum_manual_renorm_density')]
        proj = ds.proj(('deposit','PartType0_sum_wco'+str(jlower+1)+str(jlower)),2,center=obj.galaxies[galaxy].pos.in_units('code_length'))
        co_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
        co_frb = (co_frb[('deposit','PartType0_sum_wco'+str(jlower+1)+str(jlower))]/h_renorm_frb).in_units('K*km/s')
        
        print("at the end, co_frb = ",np.sum(co_frb[~np.isnan(co_frb)]))

        #imgplot=plt.imshow(np.log10(co_frb),origin='lower',interpolation='nearest',extent=[-zoom_boxsize/2,zoom_boxsize/2,-zoom_boxsize/2,zoom_boxsize/2])
        #imgplot.set_cmap('viridis')
        #cb = plt.colorbar(imgplot)
        #cb.set_label('log($\mathrm{W_{CO (J=1-0)}}$ K-km/s)')
        #plt.savefig('comap.png',dpi=300)
        
        print("Writing CO to disk")
        hdu = fits.PrimaryHDU(co_frb)
        hdu.writeto(outdir+'/galaxy'+str(galaxy)+'.co'+str(jlower+1)+str(jlower)+'.fits',overwrite=True)


    #Plot the CII Map
    fig = plt.figure()
    ax = fig.add_subplot(111)

    proj = ds.proj(('deposit','PartType0_sum_wcii'),2,center=obj.galaxies[galaxy].pos.in_units('code_length'))
    cii_frb = proj.to_frb((zoom_boxsize,'kpc'),(dim,dim),obj.galaxies[galaxy].pos.in_units('code_length'))
    cii_frb = (cii_frb[('deposit','PartType0_sum_wcii')]/h_renorm_frb).in_units('K*km/s')

    imgplot=plt.imshow(np.log10(cii_frb),origin='lower',interpolation='nearest',extent=[-zoom_boxsize/2,zoom_boxsize/2,-zoom_boxsize/2,zoom_boxsize/2])
    imgplot.set_cmap('viridis')
    cb = plt.colorbar(imgplot)
    cb.set_label('log($\mathrm{W_{CII (J=1-0)}}$ K-km/s)')
    plt.savefig('ciimap.png',dpi=300)
    
    print("Writing CII to disk")
    hdu = fits.PrimaryHDU(cii_frb)
    hdu.writeto(outdir+'/galaxy'+str(galaxy)+'.cii.fits',overwrite=True)