├── LICENSE.md ├── README.md ├── illustris_python ├── __init__.py ├── cartesian.py ├── groupcat.py ├── lhalotree.py ├── snapshot.py ├── sublink.py ├── tests │ ├── __init__.py │ ├── __main__.py │ ├── groupcat_test.py │ ├── snapshot_test.py │ └── sublink_test.py └── util.py ├── requirements.txt └── setup.py /LICENSE.md: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017, illustris & illustris_python developers 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are met: 6 | 7 | 1. Redistributions of source code must retain the above copyright notice, this 8 | list of conditions and the following disclaimer. 9 | 2. Redistributions in binary form must reproduce the above copyright notice, 10 | this list of conditions and the following disclaimer in the documentation 11 | and/or other materials provided with the distribution. 12 | 13 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND 14 | ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 15 | WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 16 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR 17 | ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES 18 | (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; 19 | LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND 20 | ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 21 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS 22 | SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 23 | 24 | The views and conclusions contained in the software and documentation are those 25 | of the authors and should not be interpreted as representing official policies, 26 | either expressed or implied, of the FreeBSD Project. 27 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # README # 2 | 3 | The Illustris Simulation: Public Data Release 4 | 5 | Example code (Python). 6 | 7 | See the [Illustris Website Data Access Page](http://www.illustris-project.org/data/) for details. 8 | 9 | # Install 10 | 11 | 12 | ``` 13 | git clone git@github.com:illustristng/illustris_python.git 14 | cd illustris_python 15 | pip install . 16 | ``` 17 | -------------------------------------------------------------------------------- /illustris_python/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = ["groupcat", "snapshot", "util", "sublink", "lhalotree", "cartesian"] 2 | 3 | from . import groupcat, snapshot, util, sublink, lhalotree, cartesian 4 | -------------------------------------------------------------------------------- /illustris_python/cartesian.py: -------------------------------------------------------------------------------- 1 | """ Illustris Simulation: Public Data Release. 2 | cartesian.py: File I/O related to the cartesian output files (THESAN only). """ 3 | from __future__ import print_function 4 | 5 | import numpy as np 6 | import h5py 7 | import six 8 | from os.path import isfile 9 | 10 | 11 | def cartPath(basePath, cartNum, chunkNum=0): 12 | """ Return absolute path to a cartesian HDF5 file (modify as needed). """ 13 | filePath_list = [ f'{basePath}/cartesian_{cartNum:03d}/cartesian_{cartNum:03d}.{chunkNum}.hdf5', 14 | ] 15 | 16 | for filePath in filePath_list: 17 | if isfile(filePath): 18 | return filePath 19 | 20 | raise ValueError("No cartesian file found!") 21 | 22 | def getNumPixel(header): 23 | """ Calculate number of pixels (per dimension) given a cartesian header. """ 24 | return header['NumPixels'] 25 | 26 | def loadSubset(basePath, cartNum, fields=None, bbox=None, sq=True): 27 | """ Load a subset of fields in the cartesian grids. 28 | If bbox is specified, load only that subset of data. bbox should have the 29 | form [[start_i, start_j, start_k], [end_i, end_j, end_k]], where i,j,k are 30 | the indices for x,y,z dimensions. Notice the last index is *inclusive*. 31 | If sq is True, return a numpy array instead of a dict if len(fields)==1. """ 32 | result = {} 33 | 34 | # make sure fields is not a single element 35 | if isinstance(fields, six.string_types): 36 | fields = [fields] 37 | 38 | # load header from first chunk 39 | with h5py.File(cartPath(basePath, cartNum), 'r') as f: 40 | header = dict(f['Header'].attrs.items()) 41 | nPix = getNumPixel(header) 42 | 43 | # decide global read size, starting file chunk, and starting file chunk offset 44 | if bbox: 45 | load_all = False 46 | start_i, start_j, start_k = bbox[0] 47 | end_i, end_j, end_k = bbox[1] 48 | assert(start_i>=0) 49 | assert(start_j>=0) 50 | assert(start_k>=0) 51 | assert(end_i=bbox[0,0]) & (local_pixels_i<=bbox[1,0]) &\ 109 | (local_pixels_j>=bbox[0,1]) & (local_pixels_j<=bbox[1,1]) &\ 110 | (local_pixels_k>=bbox[0,2]) & (local_pixels_k<=bbox[1,2]) 111 | numToReadLocal = len(np.where(pixToReadLocal)[0]) 112 | 113 | # loop over each requested field for this particle type 114 | for i, field in enumerate(fields): 115 | result[field][wOffset:wOffset+numToReadLocal] = f[field][pixToReadLocal] 116 | 117 | wOffset += numToReadLocal 118 | numToRead -= numToReadLocal 119 | 120 | fileOffset += numPixelsLocal 121 | fileNum += 1 122 | 123 | # verify we read the correct number 124 | if origNumToRead != wOffset: 125 | raise Exception("Read ["+str(wOffset)+"] particles, but was expecting ["+str(origNumToRead)+"]") 126 | 127 | # only a single field? then return the array instead of a single item dict 128 | if sq and len(fields) == 1: 129 | return result[fields[0]] 130 | 131 | return result 132 | 133 | -------------------------------------------------------------------------------- /illustris_python/groupcat.py: -------------------------------------------------------------------------------- 1 | """ Illustris Simulation: Public Data Release. 2 | groupcat.py: File I/O related to the FoF and Subfind group catalogs. """ 3 | from __future__ import print_function 4 | 5 | import six 6 | from os.path import isfile,expanduser,join 7 | import numpy as np 8 | import h5py 9 | from pathlib import Path 10 | 11 | 12 | def gcPath(basePath, snapNum, chunkNum=0): 13 | """ Return absolute path to a group catalog HDF5 file (modify as needed). """ 14 | gcPath = basePath + '/groups_%03d/' % snapNum 15 | filePath1 = gcPath + 'groups_%03d.%d.hdf5' % (snapNum, chunkNum) 16 | filePath2 = gcPath + 'fof_subhalo_tab_%03d.%d.hdf5' % (snapNum, chunkNum) 17 | 18 | if isfile(expanduser(filePath1)): 19 | return filePath1 20 | return filePath2 21 | 22 | 23 | def offsetPath(basePath, snapNum): 24 | """ Return absolute path to a separate offset file (modify as needed). """ 25 | offsetPath = join(Path(basePath).parent, 'postprocessing/offsets/offsets_%03d.hdf5' % snapNum) 26 | 27 | return offsetPath 28 | 29 | 30 | def loadObjects(basePath, snapNum, gName, nName, fields): 31 | """ Load either halo or subhalo information from the group catalog. """ 32 | result = {} 33 | 34 | # make sure fields is not a single element 35 | if isinstance(fields, six.string_types): 36 | fields = [fields] 37 | 38 | # load header from first chunk 39 | with h5py.File(gcPath(basePath, snapNum), 'r') as f: 40 | 41 | header = dict(f['Header'].attrs.items()) 42 | 43 | if 'N'+nName+'_Total' not in header and nName == 'subgroups': 44 | nName = 'subhalos' # alternate convention 45 | 46 | result['count'] = np.int64(f['Header'].attrs['N' + nName + '_Total']) 47 | 48 | if not result['count']: 49 | print('warning: zero groups, empty return (snap=' + str(snapNum) + ').') 50 | return result 51 | 52 | # if fields not specified, load everything 53 | if not fields: 54 | fields = list(f[gName].keys()) 55 | 56 | for field in fields: 57 | # verify existence 58 | if field not in f[gName].keys(): 59 | raise Exception("Group catalog does not have requested field [" + field + "]!") 60 | 61 | # replace local length with global 62 | shape = list(f[gName][field].shape) 63 | shape[0] = result['count'] 64 | 65 | # allocate within return dict 66 | result[field] = np.zeros(shape, dtype=f[gName][field].dtype) 67 | 68 | # loop over chunks 69 | wOffset = 0 70 | 71 | for i in range(header['NumFiles']): 72 | f = h5py.File(gcPath(basePath, snapNum, i), 'r') 73 | 74 | if not f['Header'].attrs['N'+nName+'_ThisFile']: 75 | continue # empty file chunk 76 | 77 | # loop over each requested field 78 | for field in fields: 79 | if field not in f[gName].keys(): 80 | raise Exception("Group catalog does not have requested field [" + field + "]!") 81 | 82 | # shape and type 83 | shape = f[gName][field].shape 84 | 85 | # read data local to the current file 86 | if len(shape) == 1: 87 | result[field][wOffset:wOffset+shape[0]] = f[gName][field][0:shape[0]] 88 | else: 89 | result[field][wOffset:wOffset+shape[0], :] = f[gName][field][0:shape[0], :] 90 | 91 | wOffset += shape[0] 92 | f.close() 93 | 94 | # only a single field? then return the array instead of a single item dict 95 | if len(fields) == 1: 96 | return result[fields[0]] 97 | 98 | return result 99 | 100 | 101 | def loadSubhalos(basePath, snapNum, fields=None): 102 | """ Load all subhalo information from the entire group catalog for one snapshot 103 | (optionally restrict to a subset given by fields). """ 104 | 105 | return loadObjects(basePath, snapNum, "Subhalo", "subgroups", fields) 106 | 107 | 108 | def loadHalos(basePath, snapNum, fields=None): 109 | """ Load all halo information from the entire group catalog for one snapshot 110 | (optionally restrict to a subset given by fields). """ 111 | 112 | return loadObjects(basePath, snapNum, "Group", "groups", fields) 113 | 114 | 115 | def loadHeader(basePath, snapNum): 116 | """ Load the group catalog header. """ 117 | with h5py.File(gcPath(basePath, snapNum), 'r') as f: 118 | header = dict(f['Header'].attrs.items()) 119 | 120 | return header 121 | 122 | 123 | def load(basePath, snapNum): 124 | """ Load complete group catalog all at once. """ 125 | r = {} 126 | r['subhalos'] = loadSubhalos(basePath, snapNum) 127 | r['halos'] = loadHalos(basePath, snapNum) 128 | r['header'] = loadHeader(basePath, snapNum) 129 | return r 130 | 131 | 132 | def loadSingle(basePath, snapNum, haloID=-1, subhaloID=-1): 133 | """ Return complete group catalog information for one halo or subhalo. """ 134 | if (haloID < 0 and subhaloID < 0) or (haloID >= 0 and subhaloID >= 0): 135 | raise Exception("Must specify either haloID or subhaloID (and not both).") 136 | 137 | gName = "Subhalo" if subhaloID >= 0 else "Group" 138 | searchID = subhaloID if subhaloID >= 0 else haloID 139 | 140 | # old or new format 141 | if 'fof_subhalo' in gcPath(basePath, snapNum): 142 | # use separate 'offsets_nnn.hdf5' files 143 | with h5py.File(offsetPath(basePath, snapNum), 'r') as f: 144 | offsets = f['FileOffsets/'+gName][()] 145 | else: 146 | # use header of group catalog 147 | with h5py.File(gcPath(basePath, snapNum), 'r') as f: 148 | offsets = f['Header'].attrs['FileOffsets_'+gName] 149 | 150 | offsets = searchID - offsets 151 | fileNum = np.max(np.where(offsets >= 0)) 152 | groupOffset = offsets[fileNum] 153 | 154 | # load halo/subhalo fields into a dict 155 | result = {} 156 | 157 | with h5py.File(gcPath(basePath, snapNum, fileNum), 'r') as f: 158 | for haloProp in f[gName].keys(): 159 | result[haloProp] = f[gName][haloProp][groupOffset] 160 | 161 | return result 162 | -------------------------------------------------------------------------------- /illustris_python/lhalotree.py: -------------------------------------------------------------------------------- 1 | """ Illustris Simulation: Public Data Release. 2 | lhalotree.py: File I/O related to the LHaloTree merger tree files. """ 3 | 4 | import numpy as np 5 | import h5py 6 | import six 7 | 8 | from .groupcat import gcPath, offsetPath 9 | from os.path import isfile 10 | 11 | 12 | def treePath(basePath, chunkNum=0): 13 | """ Return absolute path to a LHaloTree HDF5 file (modify as needed). """ 14 | 15 | filePath_list = [ basePath + '/trees/treedata/' + 'trees_sf1_135.' + str(chunkNum) + '.hdf5', 16 | basePath + '/../postprocessing/trees/LHaloTree/trees_sf1_099.' + str(chunkNum) + '.hdf5', #new path scheme for TNG 17 | basePath + '/../postprocessing/trees/LHaloTree/trees_sf1_080.' + str(chunkNum) + '.hdf5', #Thesan 18 | ] 19 | 20 | for filePath in filePath_list: 21 | if isfile(filePath): 22 | return filePath 23 | 24 | raise ValueError("No tree file found!") 25 | 26 | 27 | def treeOffsets(basePath, snapNum, id): 28 | """ Handle offset loading for a LHaloTree merger tree cutout. """ 29 | # load groupcat chunk offsets from header of first file (old or new format) 30 | if 'fof_subhalo' in gcPath(basePath, snapNum): 31 | # load groupcat chunk offsets from separate 'offsets_nnn.hdf5' files 32 | with h5py.File(offsetPath(basePath, snapNum), 'r') as f: 33 | groupFileOffsets = f['FileOffsets/Subhalo'][()] 34 | 35 | offsetFile = offsetPath(basePath, snapNum) 36 | prefix = 'Subhalo/LHaloTree/' 37 | 38 | groupOffset = id 39 | else: 40 | # load groupcat chunk offsets from header of first file 41 | with h5py.File(gcPath(basePath, snapNum), 'r') as f: 42 | groupFileOffsets = f['Header'].attrs['FileOffsets_Subhalo'] 43 | 44 | # calculate target groups file chunk which contains this id 45 | groupFileOffsets = int(id) - groupFileOffsets 46 | fileNum = np.max(np.where(groupFileOffsets >= 0)) 47 | groupOffset = groupFileOffsets[fileNum] 48 | 49 | offsetFile = gcPath(basePath, snapNum, fileNum) 50 | prefix = 'Offsets/Subhalo_LHaloTree' 51 | 52 | with h5py.File(offsetFile, 'r') as f: 53 | # load the merger tree offsets of this subgroup 54 | TreeFile = f[prefix+'File'][groupOffset] 55 | TreeIndex = f[prefix+'Index'][groupOffset] 56 | TreeNum = f[prefix+'Num'][groupOffset] 57 | return TreeFile, TreeIndex, TreeNum 58 | 59 | 60 | def singleNodeFlat(conn, index, data_in, data_out, count, onlyMPB): 61 | """ Recursive helper function: Add a single tree node. """ 62 | data_out[count] = data_in[index] 63 | 64 | count += 1 65 | count = recProgenitorFlat(conn, index, data_in, data_out, count, onlyMPB) 66 | 67 | return count 68 | 69 | 70 | def recProgenitorFlat(conn, start_index, data_in, data_out, count, onlyMPB): 71 | """ Recursive helper function: Flatten out the unordered LHaloTree, one data field at a time. """ 72 | firstProg = conn["FirstProgenitor"][start_index] 73 | 74 | if firstProg < 0: 75 | return count 76 | 77 | # depth-ordered traversal (down mpb) 78 | count = singleNodeFlat(conn, firstProg, data_in, data_out, count, onlyMPB) 79 | 80 | # explore breadth 81 | if not onlyMPB: 82 | nextProg = conn["NextProgenitor"][firstProg] 83 | 84 | while nextProg >= 0: 85 | count = singleNodeFlat(conn, nextProg, data_in, data_out, count, onlyMPB) 86 | 87 | nextProg = conn["NextProgenitor"][nextProg] 88 | firstProg = conn["FirstProgenitor"][firstProg] 89 | 90 | return count 91 | 92 | 93 | def loadTree(basePath, snapNum, id, fields=None, onlyMPB=False): 94 | """ Load portion of LHaloTree, for a given subhalo, re-arranging into a flat format. """ 95 | TreeFile, TreeIndex, TreeNum = treeOffsets(basePath, snapNum, id) 96 | 97 | if TreeNum == -1: 98 | print('Warning, empty return. Subhalo [%d] at snapNum [%d] not in tree.' % (id, snapNum)) 99 | return None 100 | 101 | # config 102 | gName = 'Tree' + str(TreeNum) # group name containing this subhalo 103 | nRows = None # we do not know in advance the size of the tree 104 | 105 | # make sure fields is not a single element 106 | if isinstance(fields, six.string_types): 107 | fields = [fields] 108 | 109 | fTree = h5py.File(treePath(basePath, TreeFile), 'r') 110 | 111 | # if no fields requested, return everything 112 | if not fields: 113 | fields = list(fTree[gName].keys()) 114 | 115 | # verify existence of requested fields 116 | for field in fields: 117 | if field not in fTree[gName].keys(): 118 | raise Exception('Error: Requested field '+field+' not in tree.') 119 | 120 | # load connectivity for this entire TreeX group 121 | connFields = ['FirstProgenitor', 'NextProgenitor'] 122 | conn = {} 123 | 124 | for field in connFields: 125 | conn[field] = fTree[gName][field][:] 126 | 127 | # determine sub-tree size with dummy walk 128 | dummy = np.zeros(conn['FirstProgenitor'].shape, dtype='int32') 129 | nRows = singleNodeFlat(conn, TreeIndex, dummy, dummy, 0, onlyMPB) 130 | 131 | result = {} 132 | result['count'] = nRows 133 | 134 | # walk through connectivity, one data field at a time 135 | for field in fields: 136 | # load field for entire tree? doing so is much faster than randomly accessing the disk 137 | # during walk, assuming that the sub-tree is a large fraction of the full tree, and that 138 | # the sub-tree is large in the absolute sense. the decision is heuristic, and can be 139 | # modified (if you have the tree on a fast SSD, could disable the full load). 140 | if nRows < 1000: # and float(nRows)/len(result['FirstProgenitor']) > 0.1 141 | # do not load, walk with single disk reads 142 | full_data = fTree[gName][field] 143 | else: 144 | # pre-load all, walk in-memory 145 | full_data = fTree[gName][field][:] 146 | 147 | # allocate the data array in the sub-tree 148 | dtype = fTree[gName][field].dtype 149 | shape = list(fTree[gName][field].shape) 150 | shape[0] = nRows 151 | 152 | data = np.zeros(shape, dtype=dtype) 153 | 154 | # walk the tree, depth-first 155 | count = singleNodeFlat(conn, TreeIndex, full_data, data, 0, onlyMPB) 156 | 157 | # save field 158 | result[field] = data 159 | 160 | fTree.close() 161 | 162 | # only a single field? then return the array instead of a single item dict 163 | if len(fields) == 1: 164 | return result[fields[0]] 165 | 166 | return result 167 | -------------------------------------------------------------------------------- /illustris_python/snapshot.py: -------------------------------------------------------------------------------- 1 | """ Illustris Simulation: Public Data Release. 2 | snapshot.py: File I/O related to the snapshot files. """ 3 | from __future__ import print_function 4 | 5 | import numpy as np 6 | import h5py 7 | import six 8 | from os.path import isfile 9 | 10 | from .util import partTypeNum 11 | from .groupcat import gcPath, offsetPath, loadSingle 12 | 13 | 14 | def snapPath(basePath, snapNum, chunkNum=0): 15 | """ Return absolute path to a snapshot HDF5 file (modify as needed). """ 16 | snapPath = basePath + '/snapdir_' + str(snapNum).zfill(3) + '/' 17 | filePath1 = snapPath + 'snap_' + str(snapNum).zfill(3) + '.' + str(chunkNum) + '.hdf5' 18 | filePath2 = filePath1.replace('/snap_', '/snapshot_') 19 | 20 | if isfile(filePath1): 21 | return filePath1 22 | return filePath2 23 | 24 | def getNumPart(header): 25 | """ Calculate number of particles of all types given a snapshot header. """ 26 | if 'NumPart_Total_HighWord' not in header: 27 | return np.int64(header['NumPart_Total']) # new uint64 convention 28 | 29 | nTypes = 6 30 | 31 | nPart = np.zeros(nTypes, dtype=np.int64) 32 | for j in range(nTypes): 33 | nPart[j] = header['NumPart_Total'][j] | (np.int64(header['NumPart_Total_HighWord'][j]) << 32) 34 | 35 | return nPart 36 | 37 | 38 | def loadSubset(basePath, snapNum, partType, fields=None, subset=None, mdi=None, sq=True, float32=False): 39 | """ Load a subset of fields for all particles/cells of a given partType. 40 | If offset and length specified, load only that subset of the partType. 41 | If mdi is specified, must be a list of integers of the same length as fields, 42 | giving for each field the multi-dimensional index (on the second dimension) to load. 43 | For example, fields=['Coordinates', 'Masses'] and mdi=[1, None] returns a 1D array 44 | of y-Coordinates only, together with Masses. 45 | If sq is True, return a numpy array instead of a dict if len(fields)==1. 46 | If float32 is True, load any float64 datatype arrays directly as float32 (save memory). """ 47 | result = {} 48 | 49 | ptNum = partTypeNum(partType) 50 | gName = "PartType" + str(ptNum) 51 | 52 | # make sure fields is not a single element 53 | if isinstance(fields, six.string_types): 54 | fields = [fields] 55 | 56 | # load header from first chunk 57 | with h5py.File(snapPath(basePath, snapNum), 'r') as f: 58 | 59 | header = dict(f['Header'].attrs.items()) 60 | nPart = getNumPart(header) 61 | 62 | # decide global read size, starting file chunk, and starting file chunk offset 63 | if subset: 64 | offsetsThisType = subset['offsetType'][ptNum] - subset['snapOffsets'][ptNum, :] 65 | 66 | fileNum = np.max(np.where(offsetsThisType >= 0)) 67 | fileOff = offsetsThisType[fileNum] 68 | numToRead = subset['lenType'][ptNum] 69 | else: 70 | fileNum = 0 71 | fileOff = 0 72 | numToRead = nPart[ptNum] 73 | 74 | result['count'] = numToRead 75 | 76 | if not numToRead: 77 | # print('warning: no particles of requested type, empty return.') 78 | return result 79 | 80 | # find a chunk with this particle type 81 | i = 1 82 | while gName not in f: 83 | f = h5py.File(snapPath(basePath, snapNum, i), 'r') 84 | i += 1 85 | 86 | # if fields not specified, load everything 87 | if not fields: 88 | fields = list(f[gName].keys()) 89 | 90 | for i, field in enumerate(fields): 91 | # verify existence 92 | if field not in f[gName].keys(): 93 | raise Exception("Particle type ["+str(ptNum)+"] does not have field ["+field+"]") 94 | 95 | # replace local length with global 96 | shape = list(f[gName][field].shape) 97 | shape[0] = numToRead 98 | 99 | # multi-dimensional index slice load 100 | if mdi is not None and mdi[i] is not None: 101 | if len(shape) != 2: 102 | raise Exception("Read error: mdi requested on non-2D field ["+field+"]") 103 | shape = [shape[0]] 104 | 105 | # allocate within return dict 106 | dtype = f[gName][field].dtype 107 | if dtype == np.float64 and float32: dtype = np.float32 108 | result[field] = np.zeros(shape, dtype=dtype) 109 | 110 | # loop over chunks 111 | wOffset = 0 112 | origNumToRead = numToRead 113 | 114 | while numToRead: 115 | f = h5py.File(snapPath(basePath, snapNum, fileNum), 'r') 116 | 117 | # no particles of requested type in this file chunk? 118 | if gName not in f: 119 | f.close() 120 | fileNum += 1 121 | fileOff = 0 122 | continue 123 | 124 | # set local read length for this file chunk, truncate to be within the local size 125 | numTypeLocal = f['Header'].attrs['NumPart_ThisFile'][ptNum] 126 | 127 | numToReadLocal = numToRead 128 | 129 | if fileOff + numToReadLocal > numTypeLocal: 130 | numToReadLocal = int(numTypeLocal - fileOff) 131 | 132 | #print('['+str(fileNum).rjust(3)+'] off='+str(fileOff)+' read ['+str(numToReadLocal)+\ 133 | # '] of ['+str(numTypeLocal)+'] remaining = '+str(numToRead-numToReadLocal)) 134 | 135 | # loop over each requested field for this particle type 136 | for i, field in enumerate(fields): 137 | # read data local to the current file 138 | if mdi is None or mdi[i] is None: 139 | result[field][wOffset:wOffset+numToReadLocal] = f[gName][field][fileOff:fileOff+numToReadLocal] 140 | else: 141 | result[field][wOffset:wOffset+numToReadLocal] = f[gName][field][fileOff:fileOff+numToReadLocal, mdi[i]] 142 | 143 | wOffset += numToReadLocal 144 | numToRead -= numToReadLocal 145 | fileNum += 1 146 | fileOff = 0 # start at beginning of all file chunks other than the first 147 | 148 | f.close() 149 | 150 | # verify we read the correct number 151 | if origNumToRead != wOffset: 152 | raise Exception("Read ["+str(wOffset)+"] particles, but was expecting ["+str(origNumToRead)+"]") 153 | 154 | # only a single field? then return the array instead of a single item dict 155 | if sq and len(fields) == 1: 156 | return result[fields[0]] 157 | 158 | return result 159 | 160 | 161 | def getSnapOffsets(basePath, snapNum, id, type): 162 | """ Compute offsets within snapshot for a particular group/subgroup. """ 163 | r = {} 164 | 165 | # old or new format 166 | if 'fof_subhalo' in gcPath(basePath, snapNum): 167 | # use separate 'offsets_nnn.hdf5' files 168 | with h5py.File(offsetPath(basePath, snapNum), 'r') as f: 169 | groupFileOffsets = f['FileOffsets/'+type][()] 170 | r['snapOffsets'] = np.transpose(f['FileOffsets/SnapByType'][()]) # consistency 171 | else: 172 | # load groupcat chunk offsets from header of first file 173 | with h5py.File(gcPath(basePath, snapNum), 'r') as f: 174 | groupFileOffsets = f['Header'].attrs['FileOffsets_'+type] 175 | r['snapOffsets'] = f['Header'].attrs['FileOffsets_Snap'] 176 | 177 | # calculate target groups file chunk which contains this id 178 | groupFileOffsets = int(id) - groupFileOffsets 179 | fileNum = np.max(np.where(groupFileOffsets >= 0)) 180 | groupOffset = groupFileOffsets[fileNum] 181 | 182 | # load the length (by type) of this group/subgroup from the group catalog 183 | with h5py.File(gcPath(basePath, snapNum, fileNum), 'r') as f: 184 | r['lenType'] = f[type][type+'LenType'][groupOffset, :] 185 | 186 | # old or new format: load the offset (by type) of this group/subgroup within the snapshot 187 | if 'fof_subhalo' in gcPath(basePath, snapNum): 188 | with h5py.File(offsetPath(basePath, snapNum), 'r') as f: 189 | r['offsetType'] = f[type+'/SnapByType'][id, :] 190 | 191 | # add TNG-Cluster specific offsets if present 192 | if 'OriginalZooms' in f: 193 | for key in f['OriginalZooms']: 194 | r[key] = f['OriginalZooms'][key][()] 195 | else: 196 | with h5py.File(gcPath(basePath, snapNum, fileNum), 'r') as f: 197 | r['offsetType'] = f['Offsets'][type+'_SnapByType'][groupOffset, :] 198 | 199 | return r 200 | 201 | 202 | def loadSubhalo(basePath, snapNum, id, partType, fields=None): 203 | """ Load all particles/cells of one type for a specific subhalo 204 | (optionally restricted to a subset fields). """ 205 | # load subhalo length, compute offset, call loadSubset 206 | subset = getSnapOffsets(basePath, snapNum, id, "Subhalo") 207 | return loadSubset(basePath, snapNum, partType, fields, subset=subset) 208 | 209 | 210 | def loadHalo(basePath, snapNum, id, partType, fields=None): 211 | """ Load all particles/cells of one type for a specific halo 212 | (optionally restricted to a subset fields). """ 213 | # load halo length, compute offset, call loadSubset 214 | subset = getSnapOffsets(basePath, snapNum, id, "Group") 215 | return loadSubset(basePath, snapNum, partType, fields, subset=subset) 216 | 217 | 218 | def loadOriginalZoom(basePath, snapNum, id, partType, fields=None): 219 | """ Load all particles/cells of one type corresponding to an 220 | original (entire) zoom simulation. TNG-Cluster specific. 221 | (optionally restricted to a subset fields). """ 222 | # load fuzz length, compute offset, call loadSubset 223 | subset = getSnapOffsets(basePath, snapNum, id, "Group") 224 | 225 | # identify original halo ID and corresponding index 226 | halo = loadSingle(basePath, snapNum, haloID=id) 227 | assert 'GroupOrigHaloID' in halo, 'Error: loadOriginalZoom() only for the TNG-Cluster simulation.' 228 | orig_index = np.where(subset['HaloIDs'] == halo['GroupOrigHaloID'])[0][0] 229 | 230 | # (1) load all FoF particles/cells 231 | subset['lenType'] = subset['GroupsTotalLengthByType'][orig_index, :] 232 | subset['offsetType'] = subset['GroupsSnapOffsetByType'][orig_index, :] 233 | 234 | data1 = loadSubset(basePath, snapNum, partType, fields, subset=subset) 235 | 236 | # (2) load all non-FoF particles/cells 237 | subset['lenType'] = subset['OuterFuzzTotalLengthByType'][orig_index, :] 238 | subset['offsetType'] = subset['OuterFuzzSnapOffsetByType'][orig_index, :] 239 | 240 | data2 = loadSubset(basePath, snapNum, partType, fields, subset=subset) 241 | 242 | # combine and return 243 | if isinstance(data1, np.ndarray): 244 | # protect against empty data 245 | if isinstance(data2, dict): 246 | return data1 247 | return np.concatenate((data1,data2), axis=0) 248 | 249 | # protect against empty data 250 | if data1["count"] == 0: 251 | return data2 252 | elif data2["count"] == 0: 253 | return data1 254 | 255 | data = {'count':data1['count']+data2['count']} 256 | for key in data1.keys(): 257 | if key == 'count': continue 258 | data[key] = np.concatenate((data1[key],data2[key]), axis=0) 259 | return data 260 | 261 | -------------------------------------------------------------------------------- /illustris_python/sublink.py: -------------------------------------------------------------------------------- 1 | """ Illustris Simulation: Public Data Release. 2 | sublink.py: File I/O related to the Sublink merger tree files. """ 3 | 4 | import numpy as np 5 | import h5py 6 | import glob 7 | import six 8 | import os 9 | 10 | from .groupcat import gcPath, offsetPath 11 | from .util import partTypeNum 12 | 13 | 14 | def treePath(basePath, treeName, chunkNum=0): 15 | """ Return absolute path to a SubLink HDF5 file (modify as needed). """ 16 | # tree_path = '/trees/' + treeName + '/' + 'tree_extended.' + str(chunkNum) + '.hdf5' 17 | tree_path = os.path.join('trees', treeName, 'tree_extended.' + str(chunkNum) + '.hdf5') 18 | 19 | _path = os.path.join(basePath, tree_path) 20 | if len(glob.glob(_path)): 21 | return _path 22 | 23 | # new path scheme 24 | _path = os.path.join(basePath, os.path.pardir, 'postprocessing', tree_path) 25 | if len(glob.glob(_path)): 26 | return _path 27 | 28 | # try one or more alternative path schemes before failing 29 | _path = os.path.join(basePath, 'postprocessing', tree_path) 30 | if len(glob.glob(_path)): 31 | return _path 32 | 33 | raise ValueError("Could not construct treePath from basePath = '{}'".format(basePath)) 34 | 35 | 36 | def treeOffsets(basePath, snapNum, id, treeName): 37 | """ Handle offset loading for a SubLink merger tree cutout. """ 38 | # old or new format 39 | if 'fof_subhalo' in gcPath(basePath, snapNum) or treeName == "SubLink_gal": 40 | # load groupcat chunk offsets from separate 'offsets_nnn.hdf5' files 41 | with h5py.File(offsetPath(basePath, snapNum), 'r') as f: 42 | groupFileOffsets = f['FileOffsets/Subhalo'][()] 43 | 44 | offsetFile = offsetPath(basePath, snapNum) 45 | prefix = 'Subhalo/' + treeName + '/' 46 | 47 | groupOffset = id 48 | else: 49 | # load groupcat chunk offsets from header of first file 50 | with h5py.File(gcPath(basePath, snapNum), 'r') as f: 51 | groupFileOffsets = f['Header'].attrs['FileOffsets_Subhalo'] 52 | 53 | # calculate target groups file chunk which contains this id 54 | groupFileOffsets = int(id) - groupFileOffsets 55 | fileNum = np.max(np.where(groupFileOffsets >= 0)) 56 | groupOffset = groupFileOffsets[fileNum] 57 | 58 | offsetFile = gcPath(basePath, snapNum, fileNum) 59 | prefix = 'Offsets/Subhalo_Sublink' 60 | 61 | with h5py.File(offsetFile, 'r') as f: 62 | # load the merger tree offsets of this subgroup 63 | RowNum = f[prefix+'RowNum'][groupOffset] 64 | LastProgID = f[prefix+'LastProgenitorID'][groupOffset] 65 | SubhaloID = f[prefix+'SubhaloID'][groupOffset] 66 | return RowNum, LastProgID, SubhaloID 67 | 68 | offsetCache = dict() 69 | 70 | def subLinkOffsets(basePath, treeName, cache=True): 71 | # create quick offset table for rows in the SubLink files 72 | if cache is True: 73 | cache = offsetCache 74 | 75 | if type(cache) is dict: 76 | path = os.path.join(basePath, treeName) 77 | try: 78 | return cache[path] 79 | except KeyError: 80 | pass 81 | 82 | search_path = treePath(basePath, treeName, '*') 83 | numTreeFiles = len(glob.glob(search_path)) 84 | if numTreeFiles == 0: 85 | raise ValueError("No tree files found! for path '{}'".format(search_path)) 86 | offsets = np.zeros(numTreeFiles, dtype='int64') 87 | 88 | for i in range(numTreeFiles-1): 89 | with h5py.File(treePath(basePath, treeName, i), 'r') as f: 90 | offsets[i+1] = offsets[i] + f['SubhaloID'].shape[0] 91 | 92 | if type(cache) is dict: 93 | cache[path] = offsets 94 | 95 | return offsets 96 | 97 | def loadTree(basePath, snapNum, id, fields=None, onlyMPB=False, onlyMDB=False, treeName="SubLink", cache=True): 98 | """ Load portion of Sublink tree, for a given subhalo, in its existing flat format. 99 | (optionally restricted to a subset fields).""" 100 | # the tree is all subhalos between SubhaloID and LastProgenitorID 101 | RowNum, LastProgID, SubhaloID = treeOffsets(basePath, snapNum, id, treeName) 102 | 103 | if RowNum == -1: 104 | print('Warning, empty return. Subhalo [%d] at snapNum [%d] not in tree.' % (id, snapNum)) 105 | return None 106 | 107 | rowStart = RowNum 108 | rowEnd = RowNum + (LastProgID - SubhaloID) 109 | 110 | # make sure fields is not a single element 111 | if isinstance(fields, six.string_types): 112 | fields = [fields] 113 | 114 | offsets = subLinkOffsets(basePath, treeName, cache) 115 | 116 | # find the tree file chunk containing this row 117 | rowOffsets = rowStart - offsets 118 | 119 | try: 120 | fileNum = np.max(np.where(rowOffsets >= 0)) 121 | except ValueError as err: 122 | print("ERROR: ", err) 123 | print("rowStart = {}, offsets = {}, rowOffsets = {}".format(rowStart, offsets, rowOffsets)) 124 | print(np.where(rowOffsets >= 0)) 125 | raise 126 | fileOff = rowOffsets[fileNum] 127 | 128 | # load only main progenitor branch? in this case, get MainLeafProgenitorID now 129 | if onlyMPB: 130 | with h5py.File(treePath(basePath, treeName, fileNum), 'r') as f: 131 | MainLeafProgenitorID = f['MainLeafProgenitorID'][fileOff] 132 | 133 | # re-calculate rowEnd 134 | rowEnd = RowNum + (MainLeafProgenitorID - SubhaloID) 135 | 136 | # load only main descendant branch (e.g. from z=0 descendant to current subhalo) 137 | if onlyMDB: 138 | with h5py.File(treePath(basePath, treeName, fileNum),'r') as f: 139 | RootDescendantID = f['RootDescendantID'][fileOff] 140 | 141 | # re-calculate tree subset (rowStart), either single branch to root descendant, or 142 | # subset of tree ending at this subhalo if this subhalo is not on the MPB of that 143 | # root descendant 144 | rowStart = RowNum - (SubhaloID - RootDescendantID) + 1 145 | rowEnd = RowNum + 1 146 | fileOff -= (rowEnd - rowStart) 147 | 148 | # calculate number of rows to load 149 | nRows = rowEnd - rowStart + 1 150 | 151 | # read 152 | result = {'count': nRows} 153 | 154 | with h5py.File(treePath(basePath, treeName, fileNum), 'r') as f: 155 | # if no fields requested, return all fields 156 | if not fields: 157 | fields = list(f.keys()) 158 | 159 | if fileOff + nRows > f['SubfindID'].shape[0]: 160 | raise Exception('Should not occur. Each tree is contained within a single file.') 161 | 162 | # loop over each requested field 163 | for field in fields: 164 | if field not in f.keys(): 165 | raise Exception("SubLink tree does not have field ["+field+"]") 166 | 167 | # read 168 | result[field] = f[field][fileOff:fileOff+nRows] 169 | 170 | # only a single field? then return the array instead of a single item dict 171 | if len(fields) == 1: 172 | return result[fields[0]] 173 | 174 | return result 175 | 176 | 177 | def maxPastMass(tree, index, partType='stars'): 178 | """ Get maximum past mass (of the given partType) along the main branch of a subhalo 179 | specified by index within this tree. """ 180 | ptNum = partTypeNum(partType) 181 | 182 | branchSize = tree['MainLeafProgenitorID'][index] - tree['SubhaloID'][index] + 1 183 | masses = tree['SubhaloMassType'][index: index + branchSize, ptNum] 184 | return np.max(masses) 185 | 186 | 187 | def numMergers(tree, minMassRatio=1e-10, massPartType='stars', index=0, alongFullTree=False): 188 | """ Calculate the number of mergers, along the main progenitor branch, in this sub-tree 189 | (optionally above some mass ratio threshold). If alongFullTree, count across the full 190 | sub-tree and not only along the MPB. """ 191 | # verify the input sub-tree has the required fields 192 | reqFields = ['SubhaloID', 'NextProgenitorID', 'MainLeafProgenitorID', 193 | 'FirstProgenitorID', 'SubhaloMassType'] 194 | 195 | if not set(reqFields).issubset(tree.keys()): 196 | raise Exception('Error: Input tree needs to have loaded fields: '+', '.join(reqFields)) 197 | 198 | num = 0 199 | invMassRatio = 1.0 / minMassRatio 200 | 201 | # walk back main progenitor branch 202 | rootID = tree['SubhaloID'][index] 203 | fpID = tree['FirstProgenitorID'][index] 204 | 205 | while fpID != -1: 206 | fpIndex = index + (fpID - rootID) 207 | fpMass = maxPastMass(tree, fpIndex, massPartType) 208 | 209 | # explore breadth 210 | npID = tree['NextProgenitorID'][fpIndex] 211 | 212 | while npID != -1: 213 | npIndex = index + (npID - rootID) 214 | npMass = maxPastMass(tree, npIndex, massPartType) 215 | 216 | # count if both masses are non-zero, and ratio exceeds threshold 217 | if fpMass > 0.0 and npMass > 0.0: 218 | ratio = npMass / fpMass 219 | 220 | if ratio >= minMassRatio and ratio <= invMassRatio: 221 | num += 1 222 | 223 | npID = tree['NextProgenitorID'][npIndex] 224 | 225 | # count along full tree instead of just along the MPB? (non-standard) 226 | if alongFullTree: 227 | if tree['FirstProgenitorID'][npIndex] != -1: 228 | numSubtree = numMergers(tree, minMassRatio=minMassRatio, massPartType=massPartType, index=npIndex) 229 | num += numSubtree 230 | 231 | fpID = tree['FirstProgenitorID'][fpIndex] 232 | 233 | return num 234 | -------------------------------------------------------------------------------- /illustris_python/tests/__init__.py: -------------------------------------------------------------------------------- 1 | """Tests. 2 | """ 3 | 4 | import os 5 | import sys 6 | 7 | BASE_PATH_ILLUSTRIS_1 = "/virgotng/universe/Illustris/L75n1820FP/output/" 8 | 9 | # Add path to directory containing 'illustris_python' module 10 | # e.g. if this file is in '/n/home00/lkelley/illustris/illustris_python/tests/' 11 | this_path = os.path.realpath(__file__) 12 | ill_py_path = os.path.abspath(os.path.join(this_path, os.path.pardir, os.path.pardir, os.path.pardir)) 13 | sys.path.append(ill_py_path) 14 | import illustris_python as ill 15 | -------------------------------------------------------------------------------- /illustris_python/tests/__main__.py: -------------------------------------------------------------------------------- 1 | """ 2 | """ 3 | 4 | import os 5 | import sys 6 | import nose 7 | 8 | 9 | if __name__ == "__main__": 10 | module_name = sys.modules[__name__].__file__ 11 | # print("sys.argv = ", sys.argv) 12 | 13 | # Run the full directory 14 | module_name = os.path.split(module_name)[0] 15 | nose_args = [sys.argv[0], module_name] 16 | if len(sys.argv) > 1: 17 | nose_args.extend(sys.argv[1:]) 18 | # print("nose_args = ", nose_args) 19 | result = nose.run(argv=nose_args) 20 | -------------------------------------------------------------------------------- /illustris_python/tests/groupcat_test.py: -------------------------------------------------------------------------------- 1 | """Tests for the `illustris_python.groupcat` submodule. 2 | 3 | Running Tests 4 | ------------- 5 | To run all tests, this script can be executed as: 6 | `$ python tests/groupcat_test.py [-v] [--nocapture]` 7 | from the root directory. 8 | 9 | Alternatively, `nosetests` can be run and it will find the tests: 10 | `$ nosetests [-v] [--nocapture]` 11 | 12 | To run particular tests (for example), 13 | `$ nosetests tests/groupcat_test.py:test_groupcat_loadSubhalos` 14 | 15 | To include coverage information, 16 | `$ nosetests --with-coverage --cover-package=.` 17 | 18 | """ 19 | 20 | import os 21 | 22 | from nose.tools import assert_true, assert_equal, assert_raises 23 | import numpy as np 24 | 25 | # `illustris_python` is imported as `ill` in local `__init__.py` 26 | from . import ill, BASE_PATH_ILLUSTRIS_1 27 | 28 | 29 | # ========================= 30 | # ==== loadHalos ==== 31 | # ========================= 32 | 33 | 34 | def test_groupcat_loadHalos_field(): 35 | fields = ['GroupFirstSub'] 36 | snap = 135 37 | group_first_sub = ill.groupcat.loadHalos(BASE_PATH_ILLUSTRIS_1, snap, fields=fields) 38 | print("group_first_sub.shape = ", group_first_sub.shape) 39 | assert_equal(group_first_sub.shape, (7713601,)) 40 | print("group_first_sub = ", group_first_sub) 41 | assert_true(np.all(group_first_sub[:3] == [0, 16937, 30430])) 42 | return 43 | 44 | 45 | def test_groupcat_loadHalos_all_fields(): 46 | snap = 135 47 | num_fields = 23 48 | # Illustris-1, snap 135 49 | cat_shape = (7713601,) 50 | first_key = 'GroupBHMass' 51 | all_fields = ill.groupcat.loadHalos(BASE_PATH_ILLUSTRIS_1, snap) 52 | print("len(all_fields.keys()) = {} (should be {})".format(len(all_fields.keys()), num_fields)) 53 | assert_equal(len(all_fields.keys()), num_fields) 54 | key = sorted(all_fields.keys())[0] 55 | print("all_fields.keys()[0] = '{}' (should be '{}')".format(key, first_key)) 56 | assert_equal(key, first_key) 57 | shape = np.shape(all_fields[key]) 58 | print("np.shape(all_fields[{}]) = {} (should be {})".format( 59 | key, shape, cat_shape)) 60 | assert_equal(shape, cat_shape) 61 | return 62 | 63 | 64 | def test_groupcat_loadHalos_1(): 65 | fields = ['GroupFirstSub'] 66 | snap = 135 67 | # Construct a path that should not be found: fail 68 | fail_path = os.path.join(BASE_PATH_ILLUSTRIS_1, 'failure') 69 | print("path '{}' should not be found".format(fail_path)) 70 | # `OSError` is raised in python3 (but in py3 OSError == IOError), `IOError` in python2 71 | assert_raises(IOError, ill.groupcat.loadHalos, fail_path, snap, fields=fields) 72 | return 73 | 74 | 75 | def test_groupcat_loadHalos_2(): 76 | fields = ['GroupFirstSub'] 77 | snap = 136 78 | # Construct a path that should not be found: fail 79 | print("snap '{}' should not be found".format(snap)) 80 | # `OSError` is raised in python3 (but in py3 OSError == IOError), `IOError` in python2 81 | assert_raises(IOError, ill.groupcat.loadHalos, BASE_PATH_ILLUSTRIS_1, snap, fields=fields) 82 | return 83 | 84 | 85 | def test_groupcat_loadHalos_3(): 86 | # This field should not be found 87 | fields = ['GroupFailSub'] 88 | snap = 100 89 | # Construct a path that should not be found: fail 90 | print("fields '{}' should not be found".format(fields)) 91 | assert_raises(Exception, ill.groupcat.loadHalos, BASE_PATH_ILLUSTRIS_1, snap, fields=fields) 92 | return 93 | 94 | 95 | # ========================== 96 | # ==== loadSingle ==== 97 | # ========================== 98 | 99 | 100 | def test_groupcat_loadSingle(): 101 | # Gas fractions for the first 5 subhalos 102 | gas_frac = [0.0344649, 0.00273708, 0.0223776, 0.0256707, 0.0134044] 103 | 104 | ptNumGas = ill.snapshot.partTypeNum('gas') # 0 105 | ptNumStars = ill.snapshot.partTypeNum('stars') # 4 106 | for i in range(5): 107 | # all_fields = ill.groupcat.loadSingle(BASE_PATH_ILLUSTRIS_1, 135, subhaloID=group_first_sub[i]) 108 | all_fields = ill.groupcat.loadSingle(BASE_PATH_ILLUSTRIS_1, 135, subhaloID=i) 109 | gas_mass = all_fields['SubhaloMassInHalfRadType'][ptNumGas] 110 | stars_mass = all_fields['SubhaloMassInHalfRadType'][ptNumStars] 111 | frac = gas_mass / (gas_mass + stars_mass) 112 | # print(i, group_first_sub[i], frac) 113 | print("subhalo {} with gas frac '{}' (should be '{}')".format(i, frac, gas_frac[i])) 114 | assert_true(np.isclose(frac, gas_frac[i])) 115 | 116 | return 117 | 118 | 119 | # ============================ 120 | # ==== loadSubhalos ==== 121 | # ============================ 122 | 123 | 124 | def test_groupcat_loadSubhalos(): 125 | fields = ['SubhaloMass', 'SubhaloSFRinRad'] 126 | snap = 135 127 | subhalos = ill.groupcat.loadSubhalos(BASE_PATH_ILLUSTRIS_1, snap, fields=fields) 128 | print("subhalos['SubhaloMass'] = ", subhalos['SubhaloMass'].shape) 129 | assert_true(subhalos['SubhaloMass'].shape == (4366546,)) 130 | print("subhalos['SubhaloMass'] = ", subhalos['SubhaloMass']) 131 | assert_true( 132 | np.allclose(subhalos['SubhaloMass'][:3], [2.21748203e+04, 2.21866333e+03, 5.73408325e+02])) 133 | 134 | return 135 | -------------------------------------------------------------------------------- /illustris_python/tests/snapshot_test.py: -------------------------------------------------------------------------------- 1 | """Tests for the `illustris_python.snapshot` submodule. 2 | 3 | Running Tests 4 | ------------- 5 | To run all tests, this script can be executed as: 6 | `$ python tests/snapshot_test.py [-v] [--nocapture]` 7 | from the root directory. 8 | 9 | Alternatively, `nosetests` can be run and it will find the tests: 10 | `$ nosetests [-v] [--nocapture]` 11 | 12 | To run particular tests (for example), 13 | `$ nosetests tests/snapshot_test.py:test_snapshot_partTypeNum_1` 14 | 15 | To include coverage information, 16 | `$ nosetests --with-coverage --cover-package=.` 17 | 18 | """ 19 | import numpy as np 20 | from nose.tools import assert_equal, assert_raises, assert_true 21 | 22 | # `illustris_python` is imported as `ill` in local `__init__.py` 23 | from . import ill, BASE_PATH_ILLUSTRIS_1 24 | 25 | 26 | def test_snapshot_partTypeNum_1(): 27 | names = ['gas', 'dm', 'tracers', 'stars', 'blackhole', 'GaS', 'blackholes'] 28 | nums = [0, 1, 3, 4, 5, 0, 5] 29 | 30 | for name, num in zip(names, nums): 31 | pn = ill.snapshot.partTypeNum(name) 32 | print("\npartTypeNum('{}') = '{}' (should be '{}')".format(name, pn, num)) 33 | assert_equal(pn, num) 34 | 35 | return 36 | 37 | 38 | def test_snapshot_partTypeNum_2(): 39 | # These should fail and raise an exception 40 | names = ['peanuts', 'monkeys'] 41 | nums = [0, 1] 42 | 43 | for name, num in zip(names, nums): 44 | print("\npartTypeNum('{}') should raise `Exception`".format(name)) 45 | assert_raises(Exception, ill.snapshot.partTypeNum, name) 46 | 47 | return 48 | 49 | 50 | ''' 51 | # Too slow 52 | def test_loadSubset(): 53 | from datetime import datetime 54 | snap = 135 55 | fields = ['Masses'] 56 | beg = datetime.now() 57 | gas_mass = ill.snapshot.loadSubset(BASE_PATH_ILLUSTRIS_1, snap, 'gas', fields=fields) 58 | print("Loaded after '{}'".format(datetime.now() - beg)) 59 | print(np.shape(gas_mass)) 60 | print(np.log10(np.mean(gas_mass, dtype='double')*1e10/0.704)) 61 | return 62 | ''' 63 | 64 | 65 | def test_loadHalo(): 66 | snap = 135 67 | halo_num = 100 68 | 69 | # Values for Illustris-1, snap=135, halo 100 70 | coords = [[19484.6576131, 20662.6423522], 71 | [54581.7254122, 55598.2078751], 72 | [60272.0348192, 61453.9991835]] 73 | stars_count = 981545 74 | 75 | stars = ill.snapshot.loadHalo(BASE_PATH_ILLUSTRIS_1, snap, halo_num, 'stars') 76 | assert_equal(stars["count"], stars_count) 77 | for i in range(3): 78 | _min = np.min(stars['Coordinates'][:, i]) 79 | _max = np.max(stars['Coordinates'][:, i]) 80 | print("Coords axis '{}' min, max: {}, {} (should be {}, {})".format( 81 | i, _min, _max, coords[i][0], coords[i][1])) 82 | assert_true(np.isclose(_min, coords[i][0])) 83 | assert_true(np.isclose(_max, coords[i][1])) 84 | 85 | return 86 | -------------------------------------------------------------------------------- /illustris_python/tests/sublink_test.py: -------------------------------------------------------------------------------- 1 | """Tests for the `illustris_python.sublink` submodule. 2 | 3 | Running Tests 4 | ------------- 5 | To run all tests, this script can be executed as: 6 | `$ python tests/sublink_test.py [-v] [--nocapture]` 7 | from the root directory. 8 | 9 | Alternatively, `nosetests` can be run and it will find the tests: 10 | `$ nosetests [-v] [--nocapture]` 11 | 12 | To run particular tests (for example), 13 | `$ nosetests tests/sublink_test.py:test_loadTree` 14 | 15 | To include coverage information, 16 | `$ nosetests --with-coverage --cover-package=.` 17 | 18 | """ 19 | import os 20 | import glob 21 | import numpy as np 22 | from nose.tools import assert_equal, assert_raises, assert_true, assert_false 23 | 24 | # `illustris_python` is imported as `ill` in local `__init__.py` 25 | from . import ill, BASE_PATH_ILLUSTRIS_1 26 | 27 | 28 | def test_treePath_1(): 29 | tree_name = "SubLink" 30 | _path = ill.sublink.treePath(BASE_PATH_ILLUSTRIS_1, tree_name, '*') 31 | paths = glob.glob(_path) 32 | assert_false(len(paths) == 0) 33 | assert_true(os.path.exists(paths[0])) 34 | return 35 | 36 | 37 | def test_treePath_2(): 38 | # Construct a path that should fail 39 | tree_name = "SubLinkFail" 40 | assert_raises(ValueError, ill.sublink.treePath, BASE_PATH_ILLUSTRIS_1, tree_name, '*') 41 | return 42 | 43 | 44 | def test_loadTree(): 45 | fields = ['SubhaloMass', 'SubfindID', 'SnapNum'] 46 | snap = 135 47 | start = 100 48 | 49 | # Values for Illustris-1, snap=135, start=100 50 | snap_num_last = [22, 25, 22, 21, 22] 51 | subhalo_mass_last = [0.0104475, 0.0105625, 0.0133463, 0.0138612, 0.0117906] 52 | 53 | group_first_sub = ill.groupcat.loadHalos(BASE_PATH_ILLUSTRIS_1, snap, fields=['GroupFirstSub']) 54 | 55 | for ii, nn, mm in zip(range(start, start+5), snap_num_last, subhalo_mass_last): 56 | tree = ill.sublink.loadTree( 57 | BASE_PATH_ILLUSTRIS_1, 135, group_first_sub[ii], fields=fields, onlyMPB=True) 58 | assert_equal(tree['SnapNum'][-1], nn) 59 | assert_true(np.isclose(tree['SubhaloMass'][-1], mm)) 60 | 61 | return 62 | 63 | 64 | def test_numMergers(): 65 | snap = 135 66 | ratio = 1.0/5.0 67 | start = 100 68 | 69 | # Values for Illustris-1, snap=135, start=100 70 | num_mergers = [2, 2, 3, 4, 3] 71 | 72 | group_first_sub = ill.groupcat.loadHalos(BASE_PATH_ILLUSTRIS_1, snap, fields=['GroupFirstSub']) 73 | 74 | # the following fields are required for the walk and the mass ratio analysis 75 | fields = ['SubhaloID', 'NextProgenitorID', 'MainLeafProgenitorID', 76 | 'FirstProgenitorID', 'SubhaloMassType'] 77 | for i, nm in zip(range(start, start+5), num_mergers): 78 | tree = ill.sublink.loadTree(BASE_PATH_ILLUSTRIS_1, snap, group_first_sub[i], fields=fields) 79 | _num_merg = ill.sublink.numMergers(tree, minMassRatio=ratio) 80 | print("group_first_sub[{}] = {}, num_mergers = {} (should be {})".format( 81 | i, group_first_sub[i], _num_merg, nm)) 82 | assert_equal(_num_merg, nm) 83 | 84 | return 85 | -------------------------------------------------------------------------------- /illustris_python/util.py: -------------------------------------------------------------------------------- 1 | """ Illustris Simulation: Public Data Release. 2 | util.py: Various helper functions. """ 3 | 4 | def partTypeNum(partType): 5 | """ Mapping between common names and numeric particle types. """ 6 | if str(partType).isdigit(): 7 | return int(partType) 8 | 9 | if str(partType).lower() in ['gas','cells']: 10 | return 0 11 | if str(partType).lower() in ['dm','darkmatter']: 12 | return 1 13 | if str(partType).lower() in ['dmlowres']: 14 | return 2 # only zoom simulations, not present in full periodic boxes 15 | if str(partType).lower() in ['tracer','tracers','tracermc','trmc']: 16 | return 3 17 | if str(partType).lower() in ['star','stars','stellar']: 18 | return 4 # only those with GFM_StellarFormationTime>0 19 | if str(partType).lower() in ['wind']: 20 | return 4 # only those with GFM_StellarFormationTime<0 21 | if str(partType).lower() in ['bh','bhs','blackhole','blackholes']: 22 | return 5 23 | 24 | raise Exception("Unknown particle type name.") 25 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # Packages required for basic usage of `illustris_python` package 2 | numpy 3 | h5py 4 | six 5 | 6 | # Packages required for unit tests 7 | nose 8 | coverage 9 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup( 4 | name='illustris_python', 5 | version='1.0.0', 6 | packages=["illustris_python"], 7 | install_requires=["numpy", "h5py", "six"], 8 | tests_require=["nose","coverage"], 9 | ) 10 | --------------------------------------------------------------------------------