treeReaderWithSplit.py

Go to the documentation of this file.

import numpy as np
import matplotlib.pyplot as plt
from scipy import io
import os
import math
import argparse
import matplotlib.cm as cm
 
import re
from collections import defaultdict
 
def round_sig(x, sig):
    return float(np.format_float_positional(x, precision=sig, unique=False, trim='k'))
 
def read_patches(file_path, size_group):
    with open(file_path, "r") as f:
        content = f.read()
    
    # Regular expression to extract offset and size
    patch_pattern = re.compile(
        r'begin patch.*?offset ([\d.eE+\-]+ [\d.eE+\-]+ [\d.eE+\-]+).*?'
        r'size ([\d.eE+\-]+ [\d.eE+\-]+ [\d.eE+\-]+)', 
        re.DOTALL
    )
 
    # Find all matches
    for match in patch_pattern.findall(content):
        offset, size = match
        size_group[size].append(offset)  # Group offsets by size
 
    return size_group
 
 
#main start here
#input parameters
parser = argparse.ArgumentParser(description='Tree Reader')
parser.add_argument("-file-path", "--file-path", dest="file", type=str, required="True",  help="file path" )
 
args = parser.parse_args()
 
#folder_path = 'case1'
#input parameters, case5 is bottomup(one big tree), case4 is topdown(bad pattern), [not in the online version]
folder_path=args.file
prefix = 'grid-'
# List all files in the folder with the specified prefix
datafiles = [f for f in os.listdir(folder_path) if f.startswith(prefix)]
 
dimension=3
 
#extract grid infomation
size_groups = []   # Dictionary to store {size: [offsets]}, notice we create dictionaries for every single subtree (i.e., every single file)
for file_path in datafiles: 
    size_group = defaultdict(list)
    read_patches(folder_path+'/'+file_path, size_group)
    size_groups.append(size_group)
 
 
#not all trees contain every level of the whole tree so we do a full search here.
total_level=-1
for size_group in size_groups:
    total_level=max(total_level, len(size_group.keys()))
 
 
patch_number_for_each_level = defaultdict(list)
for size_group in size_groups:
    for size, offsets in size_group.items():
        print(f"Size: {size} -> Count: {len(offsets)}")
        patch_number_for_each_level[size].append(len(offsets))
    #print("Offsets:")
    #for offset in offsets:
    #    print(f"  {offset}")
    print("-" * 40)
#print(total_patch_count)
print("Total Cell Count:")
for size, count in patch_number_for_each_level.items():
    print(f"Size: {size} -> Count: {sum(count)}")
 
 
#find the domain info directly from grid plot
max_size=0
 
for size_group in size_groups:
    for size in size_group.keys():
        size_tuple = tuple(map(float, size.split()))
        if size_tuple[0]>max_size:
            max_size=size_tuple[0]
            max_size_key=size
#print(max_size)
 
leftmost_offset=[999,999,999]
for size_group in size_groups:
    for offset in size_group[max_size_key]:
        offset_tuple = tuple(map(float, offset.split()))
        for m in range(dimension):
            if offset_tuple[m]<leftmost_offset[m]:
                leftmost_offset[m]=offset_tuple[m]
                leftmost_offset_key=offset
#print(leftmost_offset, leftmost_offset_key)
 
domain_offset = tuple(leftmost_offset)#tuple(map(float, leftmost_offset_key.split()))
domain_size = tuple(map(float, max_size_key.split()))
domain_size = tuple( round_sig(x * 3,5) for x in domain_size)
 
print("Domain Offset:", domain_offset, "Domain Size:", domain_size)
 
 
#preparing tree matrices
#total_level=len(size_group.keys())
#for i in range(total_level):
#   print("level "+str(i)+" Matrix total entries:", (3*(3**i))**dimension,"size: "+str(domain_size[0]/3/3**i))
tree_matrices = []
tree_patches  = []
 
#print(len(tree_matrix))
 
#calculate the tree matrices
index=0
for size_group in size_groups:
    tree_matrix   = [None]*total_level
    tree_patch    = [None]*total_level
    for size, offsets in size_group.items():
        size_tuple = tuple(map(float, size.split()))
        level=int( math.log( int(round_sig(domain_size[0]/size_tuple[0],1)), 3) )
        matrix_size=(3**level)
        #print(level, matrix_size)
 
        matrix=np.zeros((matrix_size, matrix_size, matrix_size), dtype=bool)
        for offset in offsets:
            offset_tuple = tuple(map(float, offset.split()))
            index=[-1,-1,-1]
            for i in range(len(offset_tuple)):
                index[i]=round((offset_tuple[i]-domain_offset[i])/size_tuple[i])
            matrix[tuple(index)]=1
 
        #print(len(tree_matrix))
        tree_matrix[level-1]=matrix
        tree_patch[level-1]=len(offsets)
    tree_matrices.append(tree_matrix)
    tree_patches.append(tree_patch)
 
#for i in range(len(datafiles)):
#    print(tree_patches[i])
 
#colors=["black","red"]
colors = plt.colormaps['tab10'] 
 
#plotting stuff
fig = plt.figure(figsize=(28, 8))
axes=[]
for k in range(total_level):
    axes.append(fig.add_subplot(1,total_level,k+1, projection='3d'))
 
for j in range(len(datafiles)):
    color=colors(j)
    for i in range(total_level):
        x, y, z = np.where(tree_matrices[j][i])
        axes[i].scatter(x, y, z, color=color, marker='s',s=50*27/3**(1+i), depthshade=False)  # 'c' controls color, 'marker' controls shape
 
        # Set labels
        axes[i].set_xlabel("X"); axes[i].set_ylabel("Y"); axes[i].set_zlabel("Z") 
        axes[i].set_xlim([0,3**(1+i)-1]); axes[i].set_ylim([0,3**(1+i)-1]); axes[i].set_zlim([0,3**(1+i)-1]);
        axes[i].set_title("Level="+str(i))
 
# Show the 3D plot
plt.show()
 
plt.savefig("TreeStructure.pdf")
plt.savefig("TreeStrucutre.png")