<commit_before>package gtet_util

import (
	"fmt"
	"os"
	"path"
	
	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/io"
)

const (
	rockstarMemoDir = "rockstar"
	rockstarMemoFile = "halo_%d.dat"
	rockstarShortMemoFile = "halo_short_%d.dat"

	RockstarShortMemoNum = 10 * 1000
)

// This function does fairly large heap allocations even when it doesn't need
// to. Consider passing it a buffer.
func ReadRockstar(
	snap int, ids []int, valFlags ...halo.Val,
) ([][]float64, error) {
	// Find binFile.
	memoDir, err := MemoDir()
	if err != nil { return nil, err }
	dir := path.Join(memoDir, rockstarMemoDir)
	if !PathExists(dir) { os.Mkdir(dir, 0777) }

	binFile := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
	shortBinFile := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))

	// This wastes a read the first time it's called. You need to decide if you
	// care. (Answer: probably.)
	vals, err := readRockstar(
		shortBinFile, RockstarShortMemoNum, snap, ids, valFlags,
	)
	if err == nil { return vals, err }
	return readRockstar(binFile, -1, snap, ids, valFlags)
}

func readRockstar(
	binFile string, n, snap int, ids []int, valFlags []halo.Val,
) ([][]float64, error) {
	// If binFile doesn't exist, create it.
	if !PathExists(binFile) {
		rockstarDir, err := RockstarDir()
		if err != nil { return nil, err }
		hlists, err := DirContents(rockstarDir)
		if err != nil { return nil, err }
		if n == -1 {
			err = halo.RockstarConvert(hlists[snap - 1], binFile)
			if err != nil { return nil, err }
		} else {
			err = halo.RockstarConvertTopN(hlists[snap - 1], binFile, n)
			if err != nil { return nil, err}
		}
	}

	// Get cosmo header.
	gtetFmt, err := GtetFmt()
	if err != nil { return nil, err }
	files, err := DirContents(fmt.Sprintf(gtetFmt, snap))
	hd := &io.SheetHeader{}
	err = io.ReadSheetHeaderAt(files[0], hd)

	rids, rvals, err := halo.ReadBinaryRockstarVals(
		binFile, &hd.Cosmo, valFlags...,
	)
	if err != nil { return nil, err }
	
	// Select out only the IDs we want.
	vals := make([][]float64, len(rvals))
	found := make([]bool, len(ids))
	for i := range vals { vals[i] = make([]float64, len(ids)) }
	// I think this looping strategy has better cache properties. But it will
	// End up doing more match checks than it should.
	for i, rid := range rids {
		for j, id := range ids {
			// IDs match. Copy over the values.
			if id == rid {
				for vi := range vals { vals[vi][j] = rvals[vi][i] }
				found[j] = true
				continue
			}
		}
	}

	for i := range found {
		if !found[i] {
			return nil, fmt.Errorf("Could not find ID %d.", ids[i])
		}
	}

	return vals, nil
}<commit_msg>Wrote ReadSortedRockstarIDs.<commit_after>package gtet_util

import (
	"fmt"
	"os"
	"path"
	"sort"
	
	"github.com/phil-mansfield/gotetra/render/halo"
	"github.com/phil-mansfield/gotetra/render/io"
)

const (
	rockstarMemoDir = "rockstar"
	rockstarMemoFile = "halo_%d.dat"
	rockstarShortMemoFile = "halo_short_%d.dat"

	RockstarShortMemoNum = 10 * 1000
)

func ReadSortedRockstarIDs(snap, maxID int, flag halo.Val) ([]int, error) {
	memoDir, err := MemoDir()
	if err != nil { return nil, err }
	dir := path.Join(memoDir, rockstarMemoDir)
	if !PathExists(dir) { os.Mkdir(dir, 0777) }

	var vals [][]float64
	if maxID >= RockstarShortMemoNum {
		file := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
		vals, err = readRockstar(
			file, RockstarShortMemoNum, snap, nil, halo.ID, flag,
		)
		if err != nil { return nil, err }
	} else {
		file := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))
		vals, err = readRockstar(file, -1, snap, nil, halo.ID, flag)
		if err != nil { return nil, err }
	}
	
	fids, ms := vals[0], vals[1]
	ids := make([]int, len(fids))
	for i := range ids { ids[i] = int(fids[i]) }

	if len(ids) < maxID {
		return nil, fmt.Errorf(
			"ID %d too large for snapshot %d", maxID, snap,
		)
	}

	sortRockstar(ids, ms)
	return ids[:maxID+1], nil
}

type massSet struct {
	ids []int
	ms []float64
}

func (set massSet) Len() int { return len(set.ids) }
// We're reverse sorting.
func (set massSet) Less(i, j int) bool { return set.ms[i] > set.ms[j] }
func (set massSet) Swap(i, j int) {
	set.ms[i], set.ms[j] = set.ms[j], set.ms[i]
	set.ids[i], set.ids[j] = set.ids[j], set.ids[i]
}

func sortRockstar(ids []int, ms []float64) {
	set := massSet{ ids, ms }
	sort.Sort(set)
}

// This function does fairly large heap allocations even when it doesn't need
// to. Consider passing it a buffer.
func ReadRockstar(
	snap int, ids []int, valFlags ...halo.Val,
) ([][]float64, error) {
	// Find binFile.
	memoDir, err := MemoDir()
	if err != nil { return nil, err }
	dir := path.Join(memoDir, rockstarMemoDir)
	if !PathExists(dir) { os.Mkdir(dir, 0777) }

	binFile := path.Join(dir, fmt.Sprintf(rockstarMemoFile, snap))
	shortBinFile := path.Join(dir, fmt.Sprintf(rockstarShortMemoFile, snap))

	// This wastes a read the first time it's called. You need to decide if you
	// care. (Answer: probably.)
	vals, err := readRockstar(
		shortBinFile, RockstarShortMemoNum, snap, ids, valFlags...,
	)
	if err == nil { return vals, err }
	return readRockstar(binFile, -1, snap, ids, valFlags...)
}

func readRockstar(
	binFile string, n, snap int, ids []int, valFlags ...halo.Val,
) ([][]float64, error) {
	// If binFile doesn't exist, create it.
	if !PathExists(binFile) {
		rockstarDir, err := RockstarDir()
		if err != nil { return nil, err }
		hlists, err := DirContents(rockstarDir)
		if err != nil { return nil, err }
		if n == -1 {
			err = halo.RockstarConvert(hlists[snap - 1], binFile)
			if err != nil { return nil, err }
		} else {
			err = halo.RockstarConvertTopN(hlists[snap - 1], binFile, n)
			if err != nil { return nil, err}
		}
	}

	// Get cosmo header.
	gtetFmt, err := GtetFmt()
	if err != nil { return nil, err }
	files, err := DirContents(fmt.Sprintf(gtetFmt, snap))
	hd := &io.SheetHeader{}
	err = io.ReadSheetHeaderAt(files[0], hd)

	rids, rvals, err := halo.ReadBinaryRockstarVals(
		binFile, &hd.Cosmo, valFlags...,
	)
	if err != nil { return nil, err }
	
	// Select out only the IDs we want.
	vals := make([][]float64, len(rvals))
	if ids == nil { return vals, nil }

	found := make([]bool, len(ids))
	for i := range vals { vals[i] = make([]float64, len(ids)) }
	// I think this looping strategy has better cache properties. But it will
	// End up doing more match checks than it should.
	for i, rid := range rids {
		for j, id := range ids {
			// IDs match. Copy over the values.
			if id == rid {
				for vi := range vals { vals[vi][j] = rvals[vi][i] }
				found[j] = true
				continue
			}
		}
	}

	for i := range found {
		if !found[i] {
			return nil, fmt.Errorf("Could not find ID %d.", ids[i])
		}
	}

	return vals, nil
}
{'fraction_non_alphanumeric': 0.15356027947078935, 'fraction_numerical': 0.008919280511372083, 'mean_word_length': 3.4973262032085564, 'pattern_counts': {'":': 0, '<': 4, '<?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': '3283565', 'n_tokens_mistral': 2754, 'n_tokens_neox': 2542, 'n_words': 748}