print_rollers.py

Go to the documentation of this file.

import sys
import os
from numpy.core.numeric import full
import pyvista as pv
import numpy as np
import argparse
 
from numpy.testing._private.utils import integer_repr
 
MAX_NB_OF_ROLLERS = 10
 
def split_string_to_vector3(txt):
    x = txt.split(",")
    nb_list = []
    for i in x:
        nb_list.append(float(i))
    if len(nb_list) != 3:
        exit("coords and orientation require 3 numbers: x,y,z")
    return nb_list
 
 
def is_not_vtk(file):
    return not file.endswith('vtk')
 
 
def get_euler_angles_from_vectors(vec1, vec2):
 
    a, b = (vec1 / np.linalg.norm(vec1)).reshape(3), (vec2 /
                                                      np.linalg.norm(vec2)).reshape(3)
    v = np.cross(a, b)
    c = np.dot(a, b)
    s = np.linalg.norm(v)
    if s <= 1e-12:
        return np.array([0, 0, 0])
    kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
    R = np.eye(3) + kmat + kmat.dot(kmat) * ((1 - c) / (s ** 2))
 
    sy = np.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    angle0 = np.arctan2(R[2, 1], R[2, 2]) * 180 / np.pi
    angle1 = np.arctan2(-R[2, 0], sy) * 180 / np.pi
    angle2 = np.arctan2(R[1, 0], R[0, 0]) * 180 / np.pi
 
    return np.array([angle0, angle1, angle2])
 
def get_capped_cone(origin, orientation, radius1, angle, height, resolution):
 
    orientation_np = np.array(orientation)
    scale = np.linalg.norm(orientation_np)
    norm_orientation = orientation_np / scale
 
    full_h = radius1 / np.tan(angle * np.pi / 180)
    new_origin = origin + norm_orientation * (full_h - height) / 2.
    body_test = pv.Cone(center=new_origin, direction=orientation, height=full_h, radius=radius1, capping=False, angle=None, resolution=resolution)
    cut_origin = origin + norm_orientation * (height / 2.)
    slice = body_test.slice(normal=orientation,origin=cut_origin, generate_triangles=True)
    slice = pv.PolyData(slice.points).delaunay_2d()
    body_test = body_test.clip(normal=orientation, origin=cut_origin, invert=True)
    body = pv.MultiBlock([body_test, slice]).combine()
 
    return body
 
class RigidBody:
    def __init__(self):
        self.body_type = 'plane'
        self.orientation = [0, 0, 1]
        self.radius1 = 5
        self.inner_radius = 1
        self.height = 1
        self.angle = 45
        self.angle_a = 45
        self.angle_b = 45
        self.resolution = 64
    
    def set_body_data_from_vertex(self, body_type_id, orientation_vec, roller_data):
        types = {0: 'plane',
                 1: 'sphere',
                 2: 'cylinder',
                 3: 'cone',
                 4: 'torus',
                 5: 'roller',
                 6: 'stl',
                 7: 'nurbs',
                 8: 'last'
                 }
 
        self.body_type = types[body_type_id]
        self.orientation = orientation_vec
        self.radius1 = roller_data[0]
        self.inner_radius = roller_data[1]
        # sometimes we use infinite cones or cylinders
        self.height = roller_data[2] if not np.isinf(roller_data[2]) else self.radius1 * 3
        self.angle = roller_data[3]
        self.angle_a = roller_data[4]
        self.angle_b = roller_data[5]
        
    def get_body(self, origin):
        # origin=[0, 0, 0]
        orientation_np = np.array(self.orientation)
        scale = np.linalg.norm(orientation_np)
        normz_orien = orientation_np / scale
        
        body = pv.Box()
 
        if(self.body_type == 'cylinder'):
            body = pv.Cylinder(center=origin, direction=self.orientation, radius=self.radius1,
                            height=self.height, resolution=self.resolution)
        elif(self.body_type == 'sphere'):
            body = pv.Sphere(center=origin, radius=self.radius1, phi_resolution=self.resolution)
        elif (self.body_type == 'torus'):
            pyvista_radius = self.radius1 - self.inner_radius
            body1 = pv.ParametricTorus(ringradius=pyvista_radius, crosssectionradius=self.inner_radius)
            body2 = pv.Cylinder(center=[0,0,0], direction=[0, 0, 1], radius=pyvista_radius,
                            height=self.inner_radius*2, resolution=self.resolution)
            roller = [body1, body2]
            body = pv.MultiBlock(roller).combine()
            
            
            angles = get_euler_angles_from_vectors(np.array([0., 0., 1.]), normz_orien)
            # print(angles[0], angles[1], angles[2])
            body.rotate_x(angles[0])
            body.rotate_y(angles[1])
            body.rotate_z(angles[2])
 
            transform_matrix = np.array(
                [[1, 0, 0, origin[0]], [0, 1, 0, origin[1]], [0, 0, 1, origin[2]], [0, 0, 0, 1]])
            body.transform(transform_matrix, inplace=True)
 
        elif(self.body_type == 'plane'):
            body = pv.Plane(center=origin, direction=self.orientation, i_resolution=self.resolution,
                            j_resolution=self.resolution, i_size=scale, j_size=scale)
        
        elif(self.body_type == 'cone'):
            body = get_capped_cone(origin, self.orientation, self.radius1, self.angle, self.height, self.resolution)
            # body.plot(color="tan", smooth_shading=True)
            # body.save("test_mesh_0.vtu")
        elif(self.body_type == 'roller'):
            # torus
            pyvista_radius = self.radius1 - self.inner_radius
            body1 = pv.ParametricTorus(ringradius=pyvista_radius, crosssectionradius=self.inner_radius)
            angles = get_euler_angles_from_vectors(np.array([0., 0., 1.]), normz_orien)
            body1.rotate_x(angles[0])
            body1.rotate_y(angles[1])
            body1.rotate_z(angles[2])
            transform_matrix = np.array(
                [[1, 0, 0, origin[0]], [0, 1, 0, origin[1]], [0, 0, 1, origin[2]], [0, 0, 0, 1]])
            body1.transform(transform_matrix, inplace=True)
 
            # cone1
            slope_a = np.tan(self.angle_a * np.pi / 180)
            offset_a = (self.inner_radius * slope_a / np.sqrt(1 + slope_a * slope_a)) + self.height / 2
            new_origin = origin + normz_orien * offset_a
            cone_rad = pyvista_radius + self.inner_radius / np.sqrt(1 + slope_a * slope_a)
            body2 = get_capped_cone(new_origin, self.orientation, cone_rad, self.angle_a, self.height, self.resolution)
 
            # cone2
            slope_b = np.tan(self.angle_b * np.pi / 180)
            offset_b = (self.inner_radius * slope_b / np.sqrt(1 + slope_b * slope_b)) + self.height / 2
            new_origin = origin - normz_orien * offset_b
            cone_rad = pyvista_radius + self.inner_radius / np.sqrt(1 + slope_b * slope_b)
            orientation_flip = normz_orien * -1
            body3 = get_capped_cone(new_origin, orientation_flip, cone_rad, self.angle_b, self.height, self.resolution)
            body = pv.MultiBlock([body1, body2, body3]).combine()
        else:
            exit("this body type is not supported")
 
        # body.plot()
        return body
 
 
 
if __name__ == '__main__':
 
    # print("Number of arguments:", len(sys.argv), "arguments.")
    # print("Argument List:", str(sys.argv))
    resolution = 64
    rigid_bodies_dict = {}
 
    if "-resolution" in sys.argv:
        resolution = int(sys.argv[sys.argv.index("-resolution") + 1])
 
    # for x in range(MAX_NB_OF_ROLLERS):
 
    #     body_arg = "-body" + str(x)
    #     if body_arg in sys.argv:
    #         rigid_body = RigidBody()
    #         rigid_body.resolution = resolution
    #         rigid_body.body_type = sys.argv[sys.argv.index(body_arg) + 1]
 
    #         orientation_str = '-direction' + str(x)
    #         if orientation_str in sys.argv:
    #             orient_text = sys.argv[sys.argv.index(orientation_str) + 1]
    #             rigid_body.orientation = split_string_to_vector3(orient_text)
 
    #         if rigid_body.body_type == 'torus' or rigid_body.body_type == 'cylinder' or rigid_body.body_type == 'sphere' or rigid_body.body_type == 'cone' or rigid_body.body_type == 'roller':
    #             if "-radius" + str(x) in sys.argv:
    #                 rigid_body.radius1 = float(
    #                     sys.argv[sys.argv.index("-radius" + str(x)) + 1])
    #             else:
    #                 exit("-radius x must be specified for cylinder, torus or sphere")
    #         if rigid_body.body_type == 'cone' or rigid_body.body_type == 'cylinder':
    #             if "-height" + str(x) in sys.argv:
    #                 rigid_body.height = float(
    #                     sys.argv[sys.argv.index("-height" + str(x)) + 1])
    #         if rigid_body.body_type == 'torus':
    #             if "-inner_radius" + str(x) in sys.argv:
    #                 rigid_body.inner_radius = float(
    #                     sys.argv[sys.argv.index("-inner_radius" + str(x)) + 1])
    #             else:
    #                 exit(
    #                     "-inner_radius x must be specified for cylinder, torus or sphere")
    #         if rigid_body.body_type == 'cone':
    #             if "-angle" + str(x) in sys.argv:
    #                 rigid_body.angle = float(
    #                     sys.argv[sys.argv.index("-angle" + str(x)) + 1])
    #             else:
    #                 exit(
    #                     "-angle x must be specified for cone")
    #         if rigid_body.body_type == 'roller':
    #             if "-angle_a" + str(x) in sys.argv:
    #                 rigid_body.angle_a = float(
    #                     sys.argv[sys.argv.index("-angle_a" + str(x)) + 1])
    #             else:
    #                 exit(
    #                     "-angle_a must be specified for roller")
    #             if "-angle_b" + str(x) in sys.argv:
    #                 rigid_body.angle_b = float(
    #                     sys.argv[sys.argv.index("-angle_b" + str(x)) + 1])
    #             else:
    #                 exit(
    #                     "-angle_b must be specified for roller")
    #             if "-fillet" + str(x) in sys.argv:
    #                 rigid_body.inner_radius = float(
    #                     sys.argv[sys.argv.index("-fillet" + str(x)) + 1])
    #                 rigid_body.height = 3 * rigid_body.inner_radius
    #             else:
    #                 exit(
    #                     "-fillet must be specified for roller")
    #             if "-height" + str(x) in sys.argv:
    #                 rigid_body.height = float(
    #                     sys.argv[sys.argv.index("-height" + str(x)) + 1])
    #         rigid_bodies_dict[x] = rigid_body
 
    # print("Number of bodies defined: " + str(len(rigid_bodies_dict)) + ".")
    parser = argparse.ArgumentParser(
        description="Convert multiple vtk roller files to vtk with rollers geometry", add_help=False)
    parser.add_argument(
        "-file_name", help="list of vtk files or a regexp mask", nargs='+')
    args, unknown = parser.parse_known_args()
 
    args.file_name.sort(key=os.path.getmtime)
 
    # my_files = ["/Users/karollewandowski/mofem_install/users_modules_debug/multifield_plasticity/out_roller_2.vtk"]
    # args.file_name = my_files
 
    for file in args.file_name:
        file_vtu = os.path.splitext(file)[0] + ".vtu"
        if os.path.exists(file_vtu):
            if os.path.getmtime(file) < os.path.getmtime(file_vtu):
                print("skipping {} file".format(file))
                continue
        print("Processing " + file + " file.")
        if not file.endswith('vtk'):
            exit("Provide vtk file or files (-file_name)")
        mesh = pv.read(file)
 
        # if mesh.n_points != len(rigid_bodies_dict):
        #     print("vtk file does not have enough data points")
        #     continue
 
        index = 0
        all_bodies = []
        for point in mesh.points:
            
            id_array = mesh['_ROLLER_ID'][index]
            body_type_id = mesh['_BODY_TYPE'][index]
            orientation_vec = mesh['_ORIENTATION'][index]
            roller_data = mesh['_ROLLER_DATA'][index]
 
            rigid_body = RigidBody()
            rigid_body.resolution = resolution
            rigid_body.set_body_data_from_vertex(body_type_id, orientation_vec, roller_data);
 
            # if id_array not in rigid_bodies_dict:
            #     exit("vtk files and input have inconsistent ids")
            all_bodies.append(rigid_body.get_body(point))
            index = index + 1
 
        mesh = pv.MultiBlock(all_bodies).combine()
        mesh.save(file[:-4] + ".vtu")
    exit()