MeshIngestor.py

import networkx as nx
from scipy.spatial import cKDTree
import random
import numpy as np
import matplotlib.pyplot as plt

class MeshIngestor:
    # the idea of this class is that, given a file path, it will generate the python dictionary object
    # that represents the mesh generated by CGAL

    # this class will also have a a kd tree representation of the mesh points so that queries can be quick

    def __init__(self, adj_file, locations_file):
        self.adj_file = adj_file
        self.locations_file = locations_file
        self.is_mesh = False
        self.kd_tree = None
        self.graph = None
        self.vertex_coordinates = None

    def generate_mesh(self):
        # read the adj file

        # Initialize an undirected graph
        self.graph = nx.Graph()

        with open(self.adj_file, 'r') as file:
            lines = file.read().split('\n\n')  # Split the file into sections separated by two newlines

            for curr_vertex in lines[:-1]: # last line is newline
                adj_vertices = curr_vertex.strip().split('\n')

                # Extract the main vertex (first line) and add it to the graph
                main_vertex = int(adj_vertices[0][1:])
                self.graph.add_node(main_vertex)

                # Add edges with the main vertex
                for adjacent_vertex in adj_vertices[1:]:
                    adj_vertex = int(adjacent_vertex[1:])
                    self.graph.add_edge(main_vertex, adj_vertex)

        # store point info in kd tree
                    
        self.vertex_coordinates = []

        with open(self.locations_file, 'r') as file:
            lines = file.readlines()

            for line in lines:
                parts = line.strip().split()
                vertex_label = parts[0]
                coordinates = list(map(float, parts[1:]))
                self.vertex_coordinates.append(coordinates)

        # Create a KD tree
        self.kd_tree = cKDTree(self.vertex_coordinates)

        self.vertex_coordinates = np.array(self.vertex_coordinates)

        # indicate that mesh made
        self.is_mesh = True

    def get_point_location(self, idx):
        if not self.is_mesh:
            raise AttributeError("Please initialize the mesh before analysis")
        if idx >= len(self.vertex_coordinates):
            raise IndexError("Index for vertex out of bounds")
        return self.vertex_coordinates[idx]

    def get_nearest_k_points(self, coord, k):
        if not self.is_mesh:
            raise AttributeError("Please initialize the mesh before analysis")
        return self.kd_tree.query(coord, k)

    def get_nearest_point(self, coord):
        if not self.is_mesh:
            raise AttributeError("Please initialize the mesh before analysis")
        return self.kd_tree.query(coord, 1)
    
    def get_a_star_path(self, start_pt, end_pt):
        """
        start_pt: (x,y,z) coordinate of start point
        end_pt: (x,y,z) coordinate of end point

        Returns the a star path from one vertex to the other. If the points requested are 
        not vertices, we start and and at vertices closest to the indicated points. The 
        format is a list of the vertex indices
        """

        # get_nearest_point returns distance, point, hence we access pt[1]
        start_vtx, end_vtx = self.get_nearest_point(start_pt)[1], self.get_nearest_point(end_pt)[1]

        # run networkx a*
        return nx.astar_path(self.graph, start_vtx, end_vtx)
    
    def plot_mesh(self, path=[]):
        """
        spline: list([x,y,z]) coordinates of points on spline
        plot the mesh, and optionally a set of spline points on top of it
        """
        x, y, z = zip(*self.vertex_coordinates)
        fig = plt.figure()
        ax = plt.axes(projection='3d')
        # Create a 3D scatter plot
        ax.scatter(x, y, z, color='blue', alpha=0.2)
        if path is not None:
            path_x = path[:, 0]
            path_y = path[:, 1]
            path_z = path[:, 2]
            ax.scatter(path_x, path_y, path_z, color='red')
        
        # Add labels and a title for clarity
        ax.set_title('Mesh of chicken')
        plt.show()

# if __name__ == '__main__':

#     # Specify the path to your text file
#     adj_path = 'adjacency_matrix.txt'
#     loc_path = 'vertex_lookup.txt'


#     mesh = MeshIngestor(adj_path, loc_path)

#     # Create the graph
#     mesh.generate_mesh()

#     # Print the nodes and edges of the graph

#     # Example: Find the nearest neighbor to a specific point (e.g., the first vertex)
#     nearest_neighbor_distance, nearest_neighbor_index = mesh.kd_tree.query([0,0,0])

#     # Print the results
#     print("Nearest Neighbor:", nearest_neighbor_index)
#     print("Nearest Neighbor Location:", mesh.vertex_coordinates[nearest_neighbor_index])
#     print("Distance:", nearest_neighbor_distance)

#     # get a random point
#     rand_idx = random.randrange(0, len(mesh.vertex_coordinates))
#     rand_pt = mesh.get_point_location(rand_idx)
#     print(rand_idx, rand_pt)
#     nn_dists, nn_pts = mesh.get_nearest_k_points(rand_pt, 10)
#     print(nn_dists, nn_pts)

#     # test astar
#     rand_start_pt = mesh.get_point_location(random.randrange(0, len(mesh.vertex_coordinates)))
#     rand_end_pt = mesh.get_point_location(random.randrange(0, len(mesh.vertex_coordinates)))

#     astar_path = mesh.get_a_star_path(rand_start_pt, rand_end_pt)

#     print(astar_path[:10])