Skip to content

Commit

Permalink
fix: Correctly handle translation of group names from old to new format
Browse files Browse the repository at this point in the history
  • Loading branch information
abensonca committed Sep 3, 2024
1 parent 0e3a5ad commit 1f77baa
Showing 1 changed file with 22 additions and 5 deletions.
27 changes: 22 additions & 5 deletions scripts/aux/extractSingleTree.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ def restricted_integer(x):
return x

# Parse command line arguments.
parser = argparse.ArgumentParser(prog='extractSingleTree.py',description='Extract a single tree from a Galacticus merger tree file to its own file')
parser = argparse.ArgumentParser(prog='extractSingleTree.py',description='Extract a single forest from a Galacticus merger tree file to its own file')
parser.add_argument('fromFile' )
parser.add_argument('toFile' )
parser.add_argument('forestIndex',type=restricted_integer)
Expand All @@ -36,7 +36,8 @@ def restricted_integer(x):
forestName = "treeIndex" if fromFormatVersion == 1 else "forestIndex"
halosName = "haloTrees" if fromFormatVersion == 1 else "forestHalos"

# Find the tree to extract.
# Find the forest to extract.
print("Locating forest...")
forestIndex = fromFile[forestName][forestName ][:]
firstNode = fromFile[forestName]['firstNode' ][:]
nodeCount = fromFile[forestName]['numberOfNodes'][:]
Expand All @@ -48,8 +49,10 @@ def restricted_integer(x):
start = firstNode[np.nonzero(selected)]
count = nodeCount[np.nonzero(selected)]
end = start+count
print("...done")

# Read all halo datasets.
print("Extracting datasets...")
toFile.create_group("forestHalos")
for name, h5obj in fromFile[halosName].items():
# Ignore groups.
Expand All @@ -58,8 +61,10 @@ def restricted_integer(x):
# Skip the particleIndex datasets as we will need to recompute them later.
if name == "particleIndexStart" or name == "particleIndexCount":
continue
print(" "+name)
dataset = fromFile[halosName][name][:]
toFile["forestHalos"].create_dataset(name,data=dataset[start[0]:end[0],...])
print("...done")

# Create the forestIndex group.
toFile.create_group("forestIndex")
Expand All @@ -69,6 +74,7 @@ def restricted_integer(x):

# Copy the particle data if present.
if "particleIndexStart" in fromFile[halosName]:
print("Extracting particle data...")
particleIndexStart = fromFile[halosName]['particleIndexStart'][start[0]:end[0]]
particleIndexCount = fromFile[halosName]['particleIndexCount'][start[0]:end[0]]
indices = set()
Expand All @@ -91,15 +97,26 @@ def restricted_integer(x):
particleIndexStartNew[i] = np.nonzero(indices == particleIndexStart[i])[0][0]
toFile[halosName].create_dataset("particleIndexStart",data=particleIndexStartNew)
toFile[halosName].create_dataset("particleIndexCount",data=particleIndexCount )
print("...done")

# Set format version.
toFile.attrs["formatVersion"] = 2

# Copy all attributes.
print("Copying attributes....")
for name, h5obj in fromFile.items():
# Ignore non-groups.
if not isinstance(h5obj,h5py.Group):
continue
if not name in toFile:
toFile.create_group(name)
toFile[name].attrs.update(h5obj.attrs)
# Translate names between versions.
nameTo = name
if name == "treeIndex":
nameTo = "forestIndex"
if name == "haloTrees":
nameTo = "forestHalos"
# Copy the attributes.
print(" "+name+" -> "+nameTo)
if not nameTo in toFile:
toFile.create_group(nameTo)
toFile[nameTo].attrs.update(h5obj.attrs)
print("...done")

0 comments on commit 1f77baa

Please sign in to comment.