Generate a 3D model based on a 2D equation
# stl-surface.py | |
# Generate a 3D model based on a 2D equation | |
# The model will be rectangular with a flat base. The top surface is based on | |
# a provided equation in "surface_function". The file name can be set with the | |
# output_filename variable. The x and y width of the model and the grid spacing | |
# is defined by the following parameters. | |
# x_spacing | |
# y_spacing | |
# x_number_of_points | |
# y_number_of_points | |
# There may be a warning that model isn't closed. Most likely this is an error | |
# caused by the way numpy-stl tests the model. This usually only happens on | |
# models with a large number of faces. Ignore this warning but check the model | |
# to be sure. | |
import numpy as np | |
from stl import mesh | |
import math | |
from typing import NamedTuple, List | |
# Simple coordinate storage class | |
class Vertex(NamedTuple): | |
x: float | |
y: float | |
z: float | |
# The function that describes the upper surface | |
def surface_function(x, y): | |
return 10 - 2*(1 - math.cos(2 * x * math.pi / 20)) *\ | |
(1 - math.cos(2 * y * math.pi / 20)) | |
# Define dimensions and display parameters | |
output_filename = "model.stl" | |
x_spacing = 1 | |
x_number_of_points = 101 | |
y_spacing = 1 | |
y_number_of_points = 101 | |
print("x size = " + str(x_spacing * (x_number_of_points - 1))) | |
print("y size = " + str(y_spacing * (y_number_of_points - 1))) | |
number_of_top_coordinates = x_number_of_points * y_number_of_points | |
number_of_top_faces = (x_number_of_points - 1) * \ | |
(y_number_of_points - 1) * 2 | |
print("number of coordinates in the upper surface = " + | |
str(number_of_top_coordinates)) | |
print("number of faces in the upper surface = " + str(number_of_top_faces)) | |
total_number_of_faces = number_of_top_faces * 2 +\ | |
4 * (y_number_of_points + x_number_of_points - 2) | |
print("total number of faces in the model = " + str(total_number_of_faces)) | |
# Storage for vertex coordinates using the x and y index of the coordinates | |
top_vertices = dict() | |
bottom_vertices = dict() | |
# Create the vertices for the top and bottom surfaces | |
for y_index in range(y_number_of_points): | |
for x_index in range(x_number_of_points): | |
x_coord = x_index * x_spacing | |
y_coord = y_index * y_spacing | |
# Create the vertices for the top surface. These are defined by | |
# surface_function | |
top_vertices[(x_index, y_index)] =\ | |
Vertex(x_coord, y_coord, surface_function(x_coord, y_coord)) | |
# Create the vertices for the bottom flat surface at z = 0 | |
bottom_vertices[(x_index, y_index)] =\ | |
Vertex(x_coord, y_coord, 0) | |
# Preallocate storage for the triangles that make up the upper and lower faces. | |
# I've chosen to preallocate storage for the face data instead of constantly | |
# growing the list. It shouldn't make a difference for models with a small | |
# number of faces, but it seems to improve speed for larger models. | |
top_faces: List[tuple or None] = [None] * \ | |
((x_number_of_points - 1) * | |
(y_number_of_points - 1) * 2) | |
bottom_faces: List[tuple or None] = [None] * \ | |
((x_number_of_points - 1) * | |
(y_number_of_points - 1) * 2) | |
# Every vertex in the grid (apart from ones on the top and right sides)\ | |
# correspond to 2 triangular faces. For example, in a grid with spacing of 1, | |
# the coordinate of (x, y) would correspond to two triangles. The order of the | |
# coordinates in each face is important as it defines the outward facing | |
# direction of the face based on the right hand rule. | |
# (x, y), (x + 1, y), (x + 1, y + 1) | |
# (x, y), (x + 1, y + 1), (x, y + 1) | |
# *---*---* | |
# | / | / | | |
# *---*--- | |
# | / | / | | |
# *---*---* | |
counter = 0 | |
for y_index in range(y_number_of_points - 1): | |
for x_index in range(x_number_of_points - 1): | |
# Add faces for the top surface by adding the coordinates of three | |
# vertices to a tuple | |
top_faces[counter * 2] = ((top_vertices[x_index, y_index], | |
top_vertices[x_index + 1, y_index + 1], | |
top_vertices[x_index, y_index + 1])) | |
top_faces[counter * 2 + 1] = ((top_vertices[x_index, y_index], | |
top_vertices[x_index + 1, y_index], | |
top_vertices[x_index + 1, y_index + 1])) | |
# Add faces for the bottom surface | |
bottom_faces[counter * 2] =\ | |
((bottom_vertices[x_index, y_index], | |
bottom_vertices[x_index, y_index + 1], | |
bottom_vertices[x_index + 1, y_index + 1])) | |
bottom_faces[counter * 2 + 1] =\ | |
((bottom_vertices[x_index, y_index], | |
bottom_vertices[x_index + 1, y_index + 1], | |
bottom_vertices[x_index + 1, y_index])) | |
counter += 1 | |
# Add faces along the edge of the model to close it. These faces are parallel | |
# to the x-axis | |
x_faces_1: List[tuple or None] = [None] * ((x_number_of_points - 1) * 2) | |
x_faces_2: List[tuple or None] = [None] * ((x_number_of_points - 1) * 2) | |
counter = 0 | |
for x_index in range(x_number_of_points - 1): | |
x_faces_1[counter * 2] = ((top_vertices[x_index, 0], | |
bottom_vertices[x_index, 0], | |
bottom_vertices[x_index + 1, 0])) | |
x_faces_2[counter * 2] =\ | |
((top_vertices[x_index, y_number_of_points - 1], | |
bottom_vertices[x_index + 1, y_number_of_points - 1], | |
bottom_vertices[x_index, y_number_of_points - 1])) | |
x_faces_1[counter * 2 + 1] =\ | |
((top_vertices[x_index + 1, 0], | |
top_vertices[x_index, 0], | |
bottom_vertices[x_index + 1, 0])) | |
x_faces_2[counter * 2 + 1] =\ | |
((top_vertices[x_index, y_number_of_points - 1], | |
top_vertices[x_index + 1, y_number_of_points - 1], | |
bottom_vertices[x_index + 1, y_number_of_points - 1])) | |
counter += 1 | |
# Add faces along the edge of the model to close it. These faces are parallel | |
# to the y-axis | |
y_faces_1: List[tuple or None] = [None] * ((y_number_of_points - 1) * 2) | |
y_faces_2: List[tuple or None] = [None] * ((y_number_of_points - 1) * 2) | |
counter = 0 | |
for y_index in range(y_number_of_points - 1): | |
y_faces_1[counter * 2] = ((top_vertices[0, y_index], | |
bottom_vertices[0, y_index + 1], | |
bottom_vertices[0, y_index])) | |
y_faces_2[counter * 2] =\ | |
((top_vertices[x_number_of_points - 1, y_index], | |
bottom_vertices[x_number_of_points - 1, y_index], | |
bottom_vertices[x_number_of_points - 1, y_index + 1])) | |
y_faces_1[counter * 2 + 1] = ((top_vertices[0, y_index], | |
top_vertices[0, y_index + 1], | |
bottom_vertices[0, y_index + 1])) | |
y_faces_2[counter * 2 + 1] =\ | |
((top_vertices[x_number_of_points - 1, y_index + 1], | |
top_vertices[x_number_of_points - 1, y_index], | |
bottom_vertices[x_number_of_points - 1, y_index + 1])) | |
counter += 1 | |
# Combine all the faces | |
all_faces = top_faces + bottom_faces +\ | |
x_faces_1 + x_faces_2 +\ | |
y_faces_1 + y_faces_2 | |
model = mesh.Mesh(np.zeros(total_number_of_faces*4, dtype=mesh.Mesh.dtype)) | |
# Create the model | |
for index, face in enumerate(all_faces): | |
for vertex_index in range(3): | |
model.vectors[index][vertex_index] = np.array([face[vertex_index].x, | |
face[vertex_index].y, | |
face[vertex_index].z]) | |
# Print model properties | |
volume, cog, inertia = model.get_mass_properties() | |
print("Volume of the model = " + str(volume)) | |
print("Centre of gravity of the model [x, y, z] = " + str(cog)) | |
print("Mass moment of inertia matrix at centre of gravity =\n" + str(inertia)) | |
# Save the model | |
model.save(output_filename) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment