Skip to content

Instantly share code, notes, and snippets.

@rthompsonj

rthompsonj/dump_to_point_cloud.py

Created September 25, 2017 18:55
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save rthompsonj/fde8b444f233588935d4562346b563fa to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save rthompsonj/fde8b444f233588935d4562346b563fa to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. rthompsonj created this gist Sep 25, 2017.
    83 changes: 83 additions & 0 deletions dump_to_point_cloud.py
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,83 @@
    import struct
    import matplotlib.cm as cm
    import numpy as np
    import os
    import sys
    import yt
    import yt.analysis_modules.sphgr.api as sphgr
    import yt.analysis_modules.sphgr.group_funcs as gf
    from readgadget import *

    SCALE = 5.
    Npoints = 100000

    snap = sys.argv[1]
    output = sys.argv[2]

    basedir = os.path.dirname(os.path.abspath(snap))
    snapname = os.path.basename(snap)

    # load data
    ds = yt.load(snap)
    obj = sphgr.load_sphgr_data('%s/sphgr_%s' % (basedir,snapname), ds)

    #glist = obj.global_particle_lists.halo_glist[::100]
    glist = obj.global_particle_lists.halo_glist
    indexes = np.where(glist > -1)[0]

    #if len(indexes) > Npoints:
    # indexes = indexes[::int(float(len(indexes))/float(Npoints))]

    # center and rotate galaxy
    pos = obj.particle_data['gpos']/obj.boxsize - 0.5 #- (obj.boxsize/2.)
    hsml= obj.particle_data['ghsml']/obj.boxsize
    rho = np.log10(obj.particle_data['grho'].to('g/cm**3')/1.67e-24)

    pos *= SCALE
    hsml*= SCALE

    pos = pos[indexes]
    hsml= hsml[indexes]
    rho = rho[indexes]
    print len(indexes)


    #indexes = np.where(rho > -4)[0]
    #pos = pos[indexes]
    #hsml = hsml[indexes]
    #rho = rho[indexes]

    # map densities onto a color map (rgb)
    rgb = cm.ScalarMappable(cmap=cm.viridis).to_rgba(rho)
    rgb = rgb[:,0:3]

    starpos = obj.particle_data['spos']/obj.boxsize - 0.5
    sage = obj.particle_data['sage']

    starrgb = cm.ScalarMappable(cmap=cm.autumn).to_rgba(sage)
    starrgb = starrgb[:,0:3]
    starhsml = np.full(len(sage),np.median(hsml))

    starpos *= SCALE
    starhsml*= SCALE

    pos = np.concatenate((pos,starpos))
    hsml = np.concatenate((hsml, starhsml))
    rgb = np.concatenate((rgb,starrgb))

    print len(hsml)

    #from pylab import *
    #plot(pos[:,0],pos[:,1],'o',ms=1)
    #plot(rho,t,'o',ms=1)
    #show()

    #"""
    # write binary data
    f = open(output,'wb')
    f.write(struct.pack('<I',len(pos)))
    for data in [pos,rgb,hsml]:
    f.write(data.astype('f').tostring())
    f.write(struct.pack('<I',len(pos)))
    f.close()
    #"""