<commit_before>import binary_group_io
import snapshot
import itertools
import numpy as np

class FoFGroup:
    """Store information about each FoF group"""
    def __init__(self, snapshot_num):
        self.snapshot_num = snapshot_num
        self.npart = 0
        self.ids = None
    

class SubfindGroup:
    """Store information about each subfind group"""
    def __init__(self, snapshot_num):
        self.snapshot_num = snapshot_num
        self.npart = 0
        self.pot_min = None
        self.most_bound_particle = None
        self.mass = 0
        self.half_mass_radius = 0
        self.v_max = 0
        self.ids = None
        
def load_FoF_groups(directory, snapshot_num, loadIDs=False, endianness='native'):
    if loadIDs:
        filebase = "%s/groups_%03d/group_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase, endianness)
    else:
        ids = None
        
    #Load first file
    filename = "%s/groups_%03d/group_tab_%03d.0" % (directory, snapshot_num, snapshot_num)
    props = binary_group_io.read_group_file(filename, None, endianness)[0]
    
    groups = []
    
    for i in range(props['num_files']):
        filename = "%s/groups_%03d/group_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
        groups_i = binary_group_io.read_group_file(filename, ids, endianness)[1]
        
        groups.extend(convert_to_fof_groups(groups_i, snapshot_num))
        
    return groups
            
def load_subfind_groups(directory, snapshot_num, loadIDs=False, endianness='native'):
    if loadIDs:
        filebase = "%s/groups_%03d/subhalo_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase, endianness)
    else:
        ids = None
        
    #Load first file
    filename = "%s/groups_%03d/subhalo_tab_%03d.0" % (directory, snapshot_num, snapshot_num)
    props = binary_group_io.read_subfind_file(filename, None, endianness)[0]
    
    groups = []
    subgroups = []
    
    for i in range(props['num_files']):
        filename = "%s/groups_%03d/subhalo_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
        groups_i, subgroups_i = binary_group_io.read_subfind_file(filename, ids, endianness)[1:3]
        
        groups.extend(convert_to_fof_groups(groups_i, snapshot_num))
        subgroups.extend(convert_to_subfind_groups(subgroups_i, snapshot_num))
        
    return groups, subgroups

def convert_to_fof_groups(groups_props, snapshot_num):
    """Convert a dictionary of group properties into a list of group objects"""
    groups = []
    
    for i in range(len(groups_props['npart'])):
        group = FoFGroup(snapshot_num)
        group.npart = groups_props['npart'][i]
        
        if 'ids' in groups_props:
            group.ids = groups_props['ids'][i]
        
        groups.append(group)
    
    return groups

def convert_to_subfind_groups(subgroups_props, snapshot_num):
    subgroups = []
    
    for i in range(len(subgroups_props['npart'])):
        subgroup = SubfindGroup(snapshot_num)
        subgroup.npart = subgroups_props['npart'][i]
        subgroup.pot_min = subgroups_props['pot_min'][i]
        subgroup.most_bound_particle = subgroups_props['most_bound_particle'][i]
        subgroup.mass = subgroups_props['mass'][i]
        subgroup.v_max = subgroups_props['vel_max'][i]
        subgroup.half_mass_radius = subgroups_props['half_mass_radius'][i]
        subgroup.vel = subgroups_props['vel'][i]
        
        if 'ids' in subgroups_props:
            subgroup.ids = subgroups_props['ids'][i]
        
        subgroups.append(subgroup)
    
    return subgroups

def get_subgroup_idx(id, directory, snapshot_num, ids=None):
    if ids is None:
        filebase = "%s/groups_%03d/subhalo_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase)
        
        idx = np.flatnonzero(ids==id)[0]
        del ids
    else:
        idx = np.flatnonzero(ids==id)[0]
    
    #Open file until we find the right one
    subgroup_counter = 0
    for i in itertools.count():
        filename = "%s/groups_%03d/subhalo_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
        props, groups, subgroups = binary_group_io.read_subfind_file(filename)
        
        for j in range(len(subgroups['npart'])):
            if subgroups['offsets'][j]<=idx and subgroups['offsets'][j]+subgroups['npart'][j]>idx:
                del groups
                del subgroups
                
                return subgroup_counter + j
        
        subgroup_counter += props['num_subgroups']

        del groups
        del subgroups

def get_subgroup_ids(subgroup_num, directory, snapshot_num, ids=None):
    if ids is None:
        filebase = "%s/groups_%03d/subhalo_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase)
        
    try:
        iterator = iter(subgroup_num)
        selectedIDs = set()
        for subgroup_num in iterator:
            selectedIDs.update(get_subgroup_ids(subgroup_num, directory, snapshot_num, ids))
        return selectedIDs
    except TypeError: 
        #Open file until we find the right one
        for i in itertools.count():
            filename = "%s/groups_%03d/subhalo_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
            props, groups, subgroups = binary_group_io.read_subfind_file(filename)<commit_msg>Add ability to read difference endianess group files<commit_after>import binary_group_io
import snapshot
import itertools
import numpy as np

class FoFGroup:
    """Store information about each FoF group"""
    def __init__(self, snapshot_num):
        self.snapshot_num = snapshot_num
        self.npart = 0
        self.ids = None
    

class SubfindGroup:
    """Store information about each subfind group"""
    def __init__(self, snapshot_num):
        self.snapshot_num = snapshot_num
        self.npart = 0
        self.pot_min = None
        self.most_bound_particle = None
        self.mass = 0
        self.half_mass_radius = 0
        self.v_max = 0
        self.ids = None
        
def load_FoF_groups(directory, snapshot_num, loadIDs=False, endianness='native'):
    if loadIDs:
        filebase = "%s/groups_%03d/group_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase, endianness)
    else:
        ids = None
        
    #Load first file
    filename = "%s/groups_%03d/group_tab_%03d.0" % (directory, snapshot_num, snapshot_num)
    props = binary_group_io.read_group_file(filename, None, endianness)[0]
    
    groups = []
    
    for i in range(props['num_files']):
        filename = "%s/groups_%03d/group_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
        groups_i = binary_group_io.read_group_file(filename, ids, endianness)[1]
        
        groups.extend(convert_to_fof_groups(groups_i, snapshot_num))
        
    return groups
            
def load_subfind_groups(directory, snapshot_num, loadIDs=False, endianness='native'):
    if loadIDs:
        filebase = "%s/groups_%03d/subhalo_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase, endianness)
    else:
        ids = None
        
    #Load first file
    filename = "%s/groups_%03d/subhalo_tab_%03d.0" % (directory, snapshot_num, snapshot_num)
    props = binary_group_io.read_subfind_file(filename, None, endianness)[0]
    
    groups = []
    subgroups = []
    
    for i in range(props['num_files']):
        filename = "%s/groups_%03d/subhalo_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
        groups_i, subgroups_i = binary_group_io.read_subfind_file(filename, ids, endianness)[1:3]
        
        groups.extend(convert_to_fof_groups(groups_i, snapshot_num))
        subgroups.extend(convert_to_subfind_groups(subgroups_i, snapshot_num))
        
    return groups, subgroups

def convert_to_fof_groups(groups_props, snapshot_num):
    """Convert a dictionary of group properties into a list of group objects"""
    groups = []
    
    for i in range(len(groups_props['npart'])):
        group = FoFGroup(snapshot_num)
        group.npart = groups_props['npart'][i]
        
        if 'ids' in groups_props:
            group.ids = groups_props['ids'][i]
        
        groups.append(group)
    
    return groups

def convert_to_subfind_groups(subgroups_props, snapshot_num):
    subgroups = []
    
    for i in range(len(subgroups_props['npart'])):
        subgroup = SubfindGroup(snapshot_num)
        subgroup.npart = subgroups_props['npart'][i]
        subgroup.pot_min = subgroups_props['pot_min'][i]
        subgroup.most_bound_particle = subgroups_props['most_bound_particle'][i]
        subgroup.mass = subgroups_props['mass'][i]
        subgroup.v_max = subgroups_props['vel_max'][i]
        subgroup.half_mass_radius = subgroups_props['half_mass_radius'][i]
        subgroup.vel = subgroups_props['vel'][i]
        
        if 'ids' in subgroups_props:
            subgroup.ids = subgroups_props['ids'][i]
        
        subgroups.append(subgroup)
    
    return subgroups

def get_subgroup_idx(id, directory, snapshot_num, ids=None, endianness='native'):
    if ids is None:
        filebase = "%s/groups_%03d/subhalo_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase, endianness)
        
        idx = np.flatnonzero(ids==id)[0]
        del ids
    else:
        idx = np.flatnonzero(ids==id)[0]
    
    #Open file until we find the right one
    subgroup_counter = 0
    for i in itertools.count():
        filename = "%s/groups_%03d/subhalo_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
        props, groups, subgroups = binary_group_io.read_subfind_file(filename, endianness)
        
        for j in range(len(subgroups['npart'])):
            if subgroups['offsets'][j]<=idx and subgroups['offsets'][j]+subgroups['npart'][j]>idx:
                del groups
                del subgroups
                
                return subgroup_counter + j
        
        subgroup_counter += props['num_subgroups']

        del groups
        del subgroups

def get_subgroup_ids(subgroup_num, directory, snapshot_num, ids=None, endianness='native'):
    if ids is None:
        filebase = "%s/groups_%03d/subhalo_ids_%03d" % (directory, snapshot_num, snapshot_num)
        ids = binary_group_io.read_IDs(filebase, endianness)
        
    try:
        iterator = iter(subgroup_num)
        selectedIDs = set()
        for subgroup_num in iterator:
            selectedIDs.update(get_subgroup_ids(subgroup_num, directory, snapshot_num, ids))
        return selectedIDs
    except TypeError: 
        #Open file until we find the right one
        for i in itertools.count():
            filename = "%s/groups_%03d/subhalo_tab_%03d.%d" % (directory, snapshot_num, snapshot_num, i)
            props, groups, subgroups = binary_group_io.read_subfind_file(filename, endianness)
{'fraction_non_alphanumeric': 0.1027453112258766, 'fraction_numerical': 0.009966476397571804, 'mean_word_length': 2.8486750348675036, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4113255', 'n_tokens_mistral': 3598, 'n_tokens_neox': 3516, 'n_words': 704}