Created
August 16, 2022 19:45
-
-
Save trbritt/9c9976f3338422519321b1453e87539d to your computer and use it in GitHub Desktop.
Generate a 3D surface from 2D array (with 1D X and Y arrays for location and scaling) and create corresponding object in Blender
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/python | |
| """\ | |
| 3D plot of all spectra in a single 2D scan, using Blender to render surfaces. | |
| Must have Blender installed. Uses bundled version of Python. | |
| Requires Blender | |
| Usage: | |
| Unless you build Blender yourself as a stand-along package, your scripts | |
| must be run within a Blender process: | |
| blender plot_scene.blend --background | |
| --python [your_script.py] | |
| Also, the Python interpreter in Blender will not be able to find this | |
| module. The easiest thing to do is copy the contents of blenderplot.py | |
| to your script, and write your own main(). | |
| """ | |
| import bpy | |
| import bmesh | |
| import bpy_extras | |
| import mathutils | |
| import addon_utils | |
| import numpy as np | |
| from mathutils import Vector | |
| import struct | |
| from bpy.props import PointerProperty | |
| from matplotlib.cm import get_cmap | |
| from matplotlib.colors import rgb2hex | |
| # ============================================================= | |
| # ADD GRADIENTS PLUGIN MANUALLY | |
| # ============================================================= | |
| bpy.context.scene.render.engine = "CYCLES" | |
| bl_info = { | |
| "name": "gradients", | |
| "author": "trbritt", | |
| "version": (0, 0, 4), | |
| "blender": (2, 7, 6), | |
| "category": "Node", | |
| "wiki_url": "", | |
| "tracker_url": "", | |
| } | |
| def getviz(name): | |
| cmap = get_cmap(name) | |
| colors = list() | |
| for i in range(16): | |
| colors.append(rgb2hex(cmap(i / 16))[1:]) | |
| return [name, " ".join(colors)] | |
| hexviz = { | |
| 0: [ | |
| "lines", | |
| "0072bd d95319 edb120 7e2f8e 77ac30 4dbeee a2142f 0072bd d95319 edb120 7e2f8e 77ac30 4dbeee a2142f 0072bd d95319", | |
| ], | |
| 1: [ | |
| "pink", | |
| "3c0000 653636 814c4c 985d5d ac6c6c be7878 c69184 cda68e d4b898 dac9a1 e1d9aa e7e7b2 ededc8 f3f3dc f9f9ee ffffff", | |
| ], | |
| 2: [ | |
| "copper", | |
| "000000 150d08 2b1b11 402819 553522 6a422a 805033 955d3b aa6a44 bf784c d48555 ea925d ff9f65 ffad6e ffba76 ffc77f", | |
| ], | |
| 3: [ | |
| "bone", | |
| "000005 0f0f1a 1e1e2e 2d2d42 3c3c56 4a4a6a 595f79 687388 778797 869ba6 95afb5 a4c3c3 bad2d2 d1e1e1 e8f0f0 ffffff", | |
| ], | |
| 4: [ | |
| "gray", | |
| "000000 111111 222222 333333 444444 555555 666666 777777 888888 999999 aaaaaa bbbbbb cccccc dddddd eeeeee ffffff", | |
| ], | |
| 5: [ | |
| "winter", | |
| "0000ff 0011f7 0022ee 0033e6 0044dd 0055d5 0066cc 0077c3 0088bb 0099b3 00aaaa 00bba2 00cc99 00dd91 00ee88 00ff80", | |
| ], | |
| 6: [ | |
| "autumn", | |
| "ff0000 ff1100 ff2200 ff3300 ff4400 ff5500 ff6600 ff7700 ff8800 ff9900 ffaa00 ffbb00 ffcc00 ffdd00 ffee00 ffff00", | |
| ], | |
| 7: [ | |
| "summer", | |
| "008066 118866 229166 339966 44a266 55aa66 66b366 77bb66 88c366 99cc66 aad466 bbdd66 cce666 ddee66 eef766 ffff66", | |
| ], | |
| 8: [ | |
| "spring", | |
| "ff00ff ff11ee ff22dd ff33cc ff44bb ff55aa ff6699 ff7788 ff8877 ff9966 ffaa55 ffbb44 ffcc33 ffdd22 ffee11 ffff00", | |
| ], | |
| 9: [ | |
| "cool", | |
| "00ffff 11eeff 22ddff 33ccff 44bbff 55aaff 6699ff 7788ff 8877ff 9966ff aa55ff bb44ff bf40ff dd22ff ee11ff ff00ff", | |
| ], | |
| 10: [ | |
| "hot", | |
| "2b0000 550000 800000 aa0000 d50000 ff0000 ff0000 ff5500 ff8000 ffaa00 ffd500 ffff00 ffff40 ffff80 ffffbf ffffff", | |
| ], | |
| 11: [ | |
| "hsv", | |
| "ff0000 ff6000 ffbf00 dfff00 80ff00 20ff00 00ff40 00ff9f 00ffff 009fff 0040ff 2000ff 8000ff df00ff ff00bf ff0060", | |
| ], | |
| 12: [ | |
| "jet", | |
| "0000bf 0000ff 0040ff 0080ff 00bfff 00ffff 40ffbf 80ff80 bfff40 ffff00 ffbf00 ff8000 ff4000 ff0000 bf0000 800000", | |
| ], | |
| 13: [ | |
| "parula", | |
| "352a87 3145bc 0265e1 0f77db 1388d3 079ccf 07aac1 20b4ad 49bc94 7abf7c a5be6b cabb5c ecb94c fec634 f6dd22 f9fb0e", | |
| ], | |
| } | |
| hexviz[14] = getviz("magma") | |
| hexviz[15] = getviz("terrain") | |
| def hex_to_rgb(rgb_str): | |
| int_tuple = struct.unpack("BBB", bytes.fromhex(rgb_str)) | |
| return tuple([val / 255 for val in int_tuple]) | |
| def generate_pie_nodes(context, nodes, origin): | |
| node_tree = context.space_data.node_tree | |
| x, y = context.space_data.cursor_location | |
| ColorRamp = nodes.new(type="ShaderNodeValToRGB") | |
| location = (353.9348 + x, 2.7287 + y) if origin == "search" else (0, 0) | |
| ColorRamp.location = Vector(location) | |
| def generate_pcts_from_hexviz(mode): | |
| key = int(mode) | |
| hexes = hexviz[key][1].split(" ") | |
| pcts = [] | |
| for i, j in enumerate(hexes): | |
| pcts.append([round(1 / 16 * 100, 4), list(hex_to_rgb(j)) + [1.0]]) | |
| return pcts | |
| def slice_mutuals(elements, mode): | |
| procents_and_colors = generate_pcts_from_hexviz(mode) | |
| # this sections adds or removes elements if you update your | |
| # procents_and_colors list with more / fewer elements | |
| diff = len(elements) - len(procents_and_colors) | |
| if diff > 0: | |
| for i in range(abs(diff)): | |
| elements.remove(elements[-1]) | |
| elif diff < 0: | |
| for i in range(abs(diff)): | |
| elements.new(position=0.0) | |
| position = 0 | |
| for idx, section in enumerate(procents_and_colors): | |
| elements[idx].color = section[1] | |
| elements[idx].position = position | |
| position += section[0] / 100.0 | |
| def make_slices(ctx, nodes, origin): | |
| cRamp = nodes.get("ColorRamp") | |
| if not cRamp: | |
| generate_pie_nodes(ctx, nodes, origin) | |
| elements = nodes["ColorRamp"].color_ramp.elements | |
| else: | |
| elements = cRamp.color_ramp.elements | |
| mode = ctx.scene.gradients_props.selected_mode | |
| slice_mutuals(elements, mode) | |
| def external_make_slices(cRamp, mode=0): | |
| mode = str(mode) | |
| elements = cRamp.color_ramp.elements | |
| slice_mutuals(elements, mode) | |
| class GradientsOps(bpy.types.Operator): | |
| bl_label = "Gradient Operator" | |
| bl_idname = "scene.gradient_pusher" | |
| origin = bpy.props.StringProperty(default="search") | |
| scripted = bpy.props.BoolProperty(default=False) | |
| def execute(self, context): | |
| try: | |
| space = context.space_data | |
| node_tree = space.node_tree | |
| nodes = node_tree.nodes | |
| make_slices(context, nodes, self.origin) | |
| except: | |
| print("external mode") | |
| bpy.app.driver_namespace["external_gradients"] = external_make_slices | |
| return {"FINISHED"} | |
| class GradientsPanel(bpy.types.Panel): | |
| """Creates a Panel in the scene context of the properties editor""" | |
| bl_label = "Gradient Demo" | |
| bl_idname = "GRADIENT_PT_loader" | |
| bl_space_type = "NODE_EDITOR" | |
| bl_region_type = "UI" | |
| @classmethod | |
| def poll(self, context): | |
| try: | |
| return context.space_data.edit_tree.name == "Shader Nodetree" | |
| except: | |
| return False | |
| def draw(self, context): | |
| layout = self.layout | |
| scn = context.scene | |
| layout.label(text="Pick a gradient") | |
| r = layout.row(align=True) | |
| for i in range(16): | |
| r.prop(scn.gradients_props, "color_" + str(i), text="") | |
| r = layout.row() | |
| r.prop(scn.gradients_props, "selected_mode", text="") | |
| r = layout.row() | |
| r.operator("scene.gradient_pusher", text="set gradient").origin = "button" | |
| class GradientsProperties(bpy.types.PropertyGroup): | |
| @classmethod | |
| def register(cls): | |
| Scn = bpy.types.Scene | |
| Scn.gradients_props = PointerProperty( | |
| name="internal global properties", | |
| description="shared properties between operators", | |
| type=cls, | |
| ) | |
| def updateColors(self, context): | |
| props = context.scene.gradients_props | |
| key = int(props.selected_mode) | |
| hexes = hexviz[key][1].split(" ") | |
| for i, j in enumerate(hexes): | |
| exec("props.color_" + str(i) + " = hex_to_rgb(j)") | |
| for i in range(16): | |
| k = " = bpy.props.FloatVectorProperty(subtype='COLOR', min=0.0, max=1.0)" | |
| exec("cls.color_" + str(i) + k) | |
| mode_options = [(str(i), hexviz[i][0], "", i) for i in range(14)] | |
| cls.selected_mode = bpy.props.EnumProperty( | |
| items=mode_options, | |
| description="offers....", | |
| default="0", | |
| update=updateColors, | |
| ) | |
| @classmethod | |
| def unregister(cls): | |
| del bpy.types.Scene.gradients_props | |
| def register(): | |
| bpy.types.Scene.GRADIENTS_CHOICE = bpy.props.StringProperty() | |
| bpy.utils.register_class(GradientsProperties) | |
| bpy.utils.register_class(GradientsPanel) | |
| bpy.utils.register_class(GradientsOps) | |
| def unregister(): | |
| bpy.utils.unregister_class(GradientsOps) | |
| bpy.utils.unregister_class(GradientsPanel) | |
| bpy.utils.unregister_class(GradientsProperties) | |
| del bpy.types.Scene.OCTAVE_GRADIENTS_CHOICE | |
| def include_only_one_collection( | |
| view_layer: bpy.types.ViewLayer, collection_include: bpy.types.Collection | |
| ): | |
| for layer_collection in view_layer.layer_collection.children: | |
| if layer_collection.collection != collection_include: | |
| layer_collection.exclude = True | |
| else: | |
| layer_collection.exclude = False | |
| def do_colorRamp(material_name, mode=12): | |
| # make it was added already, this checks before running the alternative code | |
| driver_namespace = bpy.app.driver_namespace | |
| if "external_gradients" not in driver_namespace: | |
| bpy.ops.scene.gradient_pusher() | |
| # the 'external_gradients' function is now in drive, this will get a reference to it | |
| external_gradients = bpy.app.driver_namespace["external_gradients"] | |
| # add material, or reuse existing one | |
| mymat = bpy.data.materials.get(material_name) | |
| if not mymat: | |
| print("didnt fnid mat") | |
| mymat = bpy.data.materials.new(material_name) | |
| mymat.use_nodes = True | |
| nodes = mymat.node_tree.nodes | |
| if "ColorRamp" not in nodes: | |
| nodes.new(type="ShaderNodeValToRGB") | |
| ColorRamp = nodes["ColorRamp"] | |
| external_gradients(ColorRamp, mode) # jet=12 | |
| # hook the ColorRamp up to the DiffuseBSDF node. | |
| nodes.remove(nodes.get("Principled BSDF")) | |
| matout = nodes.get("Material Output") | |
| DiffuseBSDF = nodes.new(type="ShaderNodeBsdfDiffuse") | |
| SeparateXYZ = nodes.new(type="ShaderNodeSeparateXYZ") | |
| TextureCoordinate = nodes.new(type="ShaderNodeTexCoord") | |
| VectorDisp = nodes.new(type="ShaderNodeVectorDisplacement") | |
| VectorDisp.space = "WORLD" | |
| Multiply = nodes.new(type="ShaderNodeMath") | |
| Multiply.operation = "MULTIPLY" | |
| mymat.node_tree.links.new(TextureCoordinate.outputs[3], VectorDisp.inputs[0]) | |
| mymat.node_tree.links.new(VectorDisp.outputs[0], SeparateXYZ.inputs[0]) | |
| mymat.node_tree.links.new(SeparateXYZ.outputs[2], Multiply.inputs[0]) | |
| mymat.node_tree.links.new(Multiply.outputs[0], ColorRamp.inputs[0]) | |
| mymat.node_tree.links.new(ColorRamp.outputs[0], DiffuseBSDF.inputs[0]) | |
| mymat.node_tree.links.new(DiffuseBSDF.outputs[0], matout.inputs[0]) | |
| # positioning for clarity | |
| TextureCoordinate.location = (-650, 0) | |
| VectorDisp.location = (-500, 0) | |
| SeparateXYZ.location = (-350, 0) | |
| Multiply.location = (-150, 0) | |
| ColorRamp.location = (50, 10) | |
| DiffuseBSDF.location = (350, 0) | |
| matout.location = (600, 10) | |
| # ============================================================= | |
| # RENDER AXES AND 2D DATA | |
| # ============================================================= | |
| class BlenderAxis: | |
| """Class for storing, manipulating, and rendering a 3D surface plot and | |
| axis in Blender. Output uses Blender internal rendering for | |
| plotting Output to both | |
| pdf (raster plot, vector axes and text) or png (all raster). All | |
| raster output has a resolution of self.dpi | |
| """ | |
| dpi = 1200 # note, 128 is macbook screen resolution | |
| x_size = 5 # inches | |
| y_size = 5 | |
| dx = 0.1 # regrid spacing | |
| Zscale = 20 | |
| DX = 1 # wireframe spacing | |
| DY = 1 | |
| zticks = [0.0, 0.2, 0.4] | |
| nslices = 25.0 # for contour plot | |
| xticks = [-3, 0, 3, 6] # Data units | |
| yticks = [1, 5, 9, 13] | |
| ticklen = 0.5 # Blender units | |
| def __init__(self, prefix, X, Y, Z): | |
| """Take X (1D), Y (1D) and Z (2D) {numpy.array}s and return a new | |
| instance of BlenderAxis. Add data, wireframe, and axes to current | |
| Blender scene. | |
| """ | |
| self.prefix = prefix | |
| # Get current Blender scene | |
| self.scn = bpy.context.scene | |
| self.collection = bpy.context.collection | |
| # Regrid data | |
| self.X, self.Y, self.Z = self.regrid(X, Y, Z) | |
| # Save axis of rotation, and center camera vertically on plot | |
| self.rot_axis = mathutils.Vector( | |
| ( | |
| np.mean(self.X[[0, -1]]), | |
| np.mean(self.Y[[0, -1]]), | |
| np.mean(np.array([0, self.Z.max() * self.Zscale])), | |
| ) | |
| ) | |
| bpy.data.objects["Camera.001"].location[2] = self.rot_axis[2] | |
| self.scn.cursor.location = self.rot_axis | |
| # Get data objects | |
| self.data_raw = self.get_raw(X, Y, Z) | |
| self.data = self.get_mesh(self.X, self.Y, self.Z) | |
| self.data_wireframe = self.get_wireframe(self.X, self.Y, self.Z) | |
| self.data_contours = self.get_contours(self.X, self.Y, self.Z) | |
| self.axX, self.axY, self.axZ = self.get_axis(X, Y, Z) | |
| # Add objects to scene | |
| self.collection.objects.link(self.data_raw) | |
| self.collection.objects.link(self.data) | |
| self.collection.objects.link(self.data_wireframe) | |
| self.collection.objects.link(self.data_contours) | |
| self.collection.objects.link(self.axX) | |
| self.collection.objects.link(self.axY) | |
| self.collection.objects.link(self.axZ) | |
| # Set smooth flag for shading | |
| self.data_raw.select_set(state=True) | |
| self.data.select_set(state=True) | |
| self.data_wireframe.select_set(state=True) | |
| bpy.ops.object.shade_smooth() | |
| self.set_layers(2, 2, 2, 1) | |
| # Set material to Matlab-style color map (stored in .blend file) | |
| self.data_raw.data.materials.append(bpy.data.materials["Material.004"]) | |
| self.data.data.materials.append(bpy.data.materials["Material.003"]) | |
| self.data_wireframe.data.materials.append(bpy.data.materials["Material.004"]) | |
| self.data_contours.data.materials.append(bpy.data.materials["Material.003"]) | |
| do_colorRamp("Material.004", mode=15) | |
| do_colorRamp("Material.003", mode=15) | |
| # Set image resolution | |
| self.scn.render.resolution_x = self.x_size * self.dpi | |
| self.scn.render.resolution_y = self.y_size * self.dpi | |
| return | |
| def set_layers(self, raw_layer, surface_layer, wireframe_layer, contours_layer): | |
| """Set layers of objects (raw, surface, wireframe). 0 and 1 are active, | |
| only 1 will cast shadows. | |
| """ | |
| coll_data_raw = bpy.context.blend_data.collections.new("Data raw") | |
| self.scn.collection.children.link(coll_data_raw) | |
| coll_data_raw.objects.link(self.data_raw) | |
| coll_data = bpy.context.blend_data.collections.new("Data") | |
| self.scn.collection.children.link(coll_data) | |
| coll_data.objects.link(self.data) | |
| coll_data_wireframe = bpy.context.blend_data.collections.new("Data wireframe") | |
| self.scn.collection.children.link(coll_data_wireframe) | |
| coll_data_wireframe.objects.link(self.data_wireframe) | |
| coll_data_contours = bpy.context.blend_data.collections.new("Data contours") | |
| self.scn.collection.children.link(coll_data_contours) | |
| coll_data_contours.objects.link(self.data_contours) | |
| return | |
| def adjust_camera(self, cam_el, foc, cam_type): | |
| # Elevate camera | |
| cam = bpy.data.objects["Camera.001"] | |
| cam_loc_save = cam.location.copy() | |
| cam_rot_save = cam.rotation_euler.copy() | |
| cam_rot_matrix = mathutils.Matrix.Rotation(-1 * cam_el * np.pi / 180, 3, "X") | |
| cam.location = cam_rot_matrix @ (cam.location - self.rot_axis) + self.rot_axis | |
| cam.rotation_euler.rotate_axis("X", -1 * cam_el * np.pi / 180) | |
| # Deal with lens-focal-length--scale factor | |
| foc_rot_axis = self.rot_axis.copy() | |
| foc_rot_axis[0] = cam.location[0] | |
| cam.location = foc * (cam.location - foc_rot_axis) + foc_rot_axis | |
| lens_save = bpy.data.cameras["Camera.001"].lens | |
| bpy.data.cameras["Camera.001"].lens *= foc | |
| # Switch camera to orthographic projection if requested, | |
| # keeping the same field of view at the center of the plot | |
| # self.scn.update() | |
| # API changes | |
| view_layer = bpy.context.view_layer | |
| view_layer.update() | |
| if cam_type == "O": | |
| # Calculate ortho scale | |
| x0 = bpy_extras.object_utils.world_to_camera_view( | |
| self.scn, | |
| cam, | |
| mathutils.Vector([self.X[0], self.rot_axis[1], self.rot_axis[2]]), | |
| )[0] | |
| x1 = bpy_extras.object_utils.world_to_camera_view( | |
| self.scn, | |
| cam, | |
| mathutils.Vector([self.X[-1], self.rot_axis[1], self.rot_axis[2]]), | |
| )[0] | |
| bpy.data.cameras["Camera.001"].type = "ORTHO" | |
| bpy.data.cameras["Camera.001"].ortho_scale = (self.X[-1] - self.X[0]) / ( | |
| x1 - x0 | |
| ) | |
| def rotate_view(self, rot_angle): | |
| # Rotate objects | |
| rot_axis = (np.mean(self.X[[0, -1]]), np.mean(self.Y[[0, -1]]), 0) | |
| self.rotate_z(self.data_raw, rot_axis, rot_angle) | |
| self.rotate_z(self.data, rot_axis, rot_angle) | |
| self.rotate_z(self.data_wireframe, rot_axis, rot_angle) | |
| self.rotate_z(self.data_contours, rot_axis, rot_angle) | |
| self.rotate_z(self.axX, rot_axis, rot_angle) | |
| self.rotate_z(self.axY, rot_axis, rot_angle) | |
| self.rotate_z(self.axZ, rot_axis, rot_angle) | |
| def write_image( | |
| self, filename, rot_angle, cam_el, foc, cam_type, surf_alpha, wire_alpha | |
| ): | |
| """Rotate (copy of) axis by rot_angle, elevate camera by cam_el, output | |
| .png and .tex files into temporary directory, process with pdfLaTeX, | |
| crop with pdfcrop, and convert back to .png with ghostscript. | |
| Arguments: | |
| filename -- (string) filename | |
| rot_angle -- (float) angle in degrees to rotate the plot around | |
| local z | |
| cam_el -- (float) angle to elevate the camera, still pointing at the | |
| center of the plot | |
| foc -- (float) scale the focal length of the lens, maintaining the | |
| field of view at the plot distance | |
| cam_type -- (string) 'O' for orthographic, 'P' (actually anything | |
| but 'O' for now) for perspective | |
| surf_alpha -- (float) alpha for surface-like things | |
| wire_alpha -- (float) alpha for wireframe-like things | |
| (wireframes of data, raw spectra if this is a series | |
| of spectra) | |
| """ | |
| # Save current axis | |
| data_raw_save = bmesh.new() | |
| data_raw_save.from_mesh(self.data_raw.data) | |
| data_save = bmesh.new() | |
| data_save.from_mesh(self.data.data) | |
| data_wireframe_save = bmesh.new() | |
| data_wireframe_save.from_mesh(self.data_wireframe.data) | |
| data_contours_save = bmesh.new() | |
| data_contours_save.from_mesh(self.data_contours.data) | |
| axX_save = bmesh.new() | |
| axX_save.from_mesh(self.axX.data) | |
| axY_save = bmesh.new() | |
| axY_save.from_mesh(self.axY.data) | |
| axZ_save = bmesh.new() | |
| axZ_save.from_mesh(self.axZ.data) | |
| # Elevate camera | |
| cam = bpy.data.objects["Camera.001"] | |
| cam_loc_save = cam.location.copy() | |
| cam_rot_save = cam.rotation_euler.copy() | |
| cam_rot_matrix = mathutils.Matrix.Rotation(-1 * cam_el * np.pi / 180, 3, "X") | |
| cam.location = cam_rot_matrix @ (cam.location - self.rot_axis) + self.rot_axis | |
| cam.rotation_euler.rotate_axis("X", -1 * cam_el * np.pi / 180) | |
| # Deal with lens-focal-length--scale factor | |
| foc_rot_axis = self.rot_axis.copy() | |
| foc_rot_axis[0] = cam.location[0] | |
| cam.location = foc * (cam.location - foc_rot_axis) + foc_rot_axis | |
| lens_save = bpy.data.cameras["Camera.001"].lens | |
| bpy.data.cameras["Camera.001"].lens *= foc | |
| # Switch camera to orthographic projection if requested, | |
| # keeping the same field of view at the center of the plot | |
| # self.scn.update() | |
| # API changes | |
| view_layer = bpy.context.view_layer | |
| view_layer.update() | |
| if cam_type == "O": | |
| # Calculate ortho scale | |
| x0 = bpy_extras.object_utils.world_to_camera_view( | |
| self.scn, | |
| cam, | |
| mathutils.Vector([self.X[0], self.rot_axis[1], self.rot_axis[2]]), | |
| )[0] | |
| x1 = bpy_extras.object_utils.world_to_camera_view( | |
| self.scn, | |
| cam, | |
| mathutils.Vector([self.X[-1], self.rot_axis[1], self.rot_axis[2]]), | |
| )[0] | |
| bpy.data.cameras["Camera.001"].type = "ORTHO" | |
| bpy.data.cameras["Camera.001"].ortho_scale = (self.X[-1] - self.X[0]) / ( | |
| x1 - x0 | |
| ) | |
| # Rotate objects | |
| rot_axis = (np.mean(self.X[[0, -1]]), np.mean(self.Y[[0, -1]]), 0) | |
| self.rotate_z(self.data_raw, rot_axis, rot_angle) | |
| self.rotate_z(self.data, rot_axis, rot_angle) | |
| self.rotate_z(self.data_wireframe, rot_axis, rot_angle) | |
| self.rotate_z(self.data_contours, rot_axis, rot_angle) | |
| self.rotate_z(self.axX, rot_axis, rot_angle) | |
| self.rotate_z(self.axY, rot_axis, rot_angle) | |
| self.rotate_z(self.axZ, rot_axis, rot_angle) | |
| # Save .png | |
| self.scn.render.filepath = filename | |
| bpy.ops.render.render(write_still=True) | |
| # Restore unrotated plot and axes | |
| data_raw_save.to_mesh(self.data_raw.data) | |
| data_raw_save.free() | |
| data_save.to_mesh(self.data.data) | |
| data_save.free() | |
| data_wireframe_save.to_mesh(self.data_wireframe.data) | |
| data_wireframe_save.free() | |
| data_contours_save.to_mesh(self.data_contours.data) | |
| data_contours_save.free() | |
| axX_save.to_mesh(self.axX.data) | |
| axX_save.free() | |
| axY_save.to_mesh(self.axY.data) | |
| axY_save.free() | |
| axZ_save.to_mesh(self.axZ.data) | |
| axZ_save.free() | |
| # Restore camera | |
| cam.location = cam_loc_save | |
| cam.rotation_euler = cam_rot_save | |
| bpy.data.cameras["Camera.001"].type = "PERSP" | |
| bpy.data.cameras["Camera.001"].lens = lens_save | |
| return | |
| def rotate_z(self, ob, center, angle): | |
| """Rotate Blender object through 'angle' degrees about a vector passing | |
| through center parallel to the z axis. | |
| """ | |
| # Make a copy of the mesh from the object | |
| mesh = bmesh.new() | |
| mesh.from_mesh(ob.data) | |
| # Rotate the mesh | |
| bmesh.ops.rotate( | |
| mesh, | |
| cent=center, | |
| matrix=mathutils.Matrix.Rotation(angle * np.pi / 180, 3, "Z"), | |
| verts=mesh.verts, | |
| ) | |
| # Write the mesh back to the object | |
| mesh.to_mesh(ob.data) | |
| mesh.free() | |
| return | |
| def regrid(self, X, Y, Z): | |
| """Interpolate onto an evenly-spaced grid in x and y, spacing dx. Return a | |
| tuple of new (X, Y, Z). For both input and output, X and Y are | |
| 1D, and Z is 2D. | |
| """ | |
| newX = np.ogrid[X[0] : X[-1] : self.dx] | |
| newY = np.ogrid[Y[0] : Y[-1] : self.dx] | |
| # Iterate through rows of Z and interpolate using newX | |
| tempZ = np.zeros((len(Y), len(newX))) | |
| for y in range(len(Y)): | |
| tempZ[y, :] = np.interp(newX, X, Z[y, :]) | |
| # Iterate through columns of Z and interpolate using newY | |
| newZ = np.zeros((len(newY), len(newX))) | |
| for x in range(len(newX)): | |
| newZ[:, x] = np.interp(newY, Y, tempZ[:, x]) | |
| return newX, newY, newZ | |
| def get_axis(self, X, Y, Z): | |
| """Return three objects with the positions of all the X, Y, and Z tick | |
| marks (and end points, and positions of axis labels) as | |
| vertices. This is useful so that Blender methods for object | |
| rotation and determining 2D coordinates can be used on the axis | |
| points. | |
| """ | |
| self.zticks = [self.Zscale * z for z in self.zticks] | |
| Z = self.Zscale * Z | |
| axX = np.concatenate((np.array([X[0]]), self.xticks, np.array([X[-1]]))) | |
| axY = np.concatenate((np.array([Y[0]]), self.yticks, np.array([Y[-1]]))) | |
| axZ = np.concatenate((np.array([0]), self.zticks, np.array([Z.max()]))) | |
| # X axis | |
| verts = [] | |
| meshX = bpy.data.meshes.new(f"{self.prefix}axisMeshX") | |
| obX = bpy.data.objects.new(f"{self.prefix}axisX", meshX) | |
| obX.location = (0, 0, 0) | |
| obX.show_name = False | |
| for x in axX: | |
| verts.append([x, axY[0], axZ[0]]) | |
| verts.append([x, axY[0] - self.ticklen, axZ[0]]) | |
| verts.append([axX[0], axY[0] - 8 * self.ticklen, axZ[0]]) | |
| verts.append([(axX[0] + axX[-1]) / 2, axY[0] - 8 * self.ticklen, axZ[0]]) | |
| verts.append([axX[-1], axY[0] - 8 * self.ticklen, axZ[0]]) | |
| meshX.from_pydata(verts, [], []) | |
| meshX.update() | |
| # Y axis | |
| verts = [] | |
| meshY = bpy.data.meshes.new(f"{self.prefix}axisMeshY") | |
| obY = bpy.data.objects.new(f"{self.prefix}axisY", meshY) | |
| obY.location = (0, 0, 0) | |
| obY.show_name = False | |
| for y in axY: | |
| verts.append([axX[0], y, axZ[0]]) | |
| verts.append([axX[0] - self.ticklen, y, axZ[0]]) | |
| verts.append([axX[0] - 8 * self.ticklen, axY[0], axZ[0]]) | |
| verts.append([axX[0] - 8 * self.ticklen, (axY[0] + axY[-1]) / 2, axZ[0]]) | |
| verts.append([axX[0] - 8 * self.ticklen, axY[-1], axZ[0]]) | |
| meshY.from_pydata(verts, [], []) | |
| meshY.update() | |
| # Z axis | |
| verts = [] | |
| meshZ = bpy.data.meshes.new(f"{self.prefix}axisMeshZ") | |
| obZ = bpy.data.objects.new(f"{self.prefix}axisZ", meshZ) | |
| obZ.location = (0, 0, 0) | |
| obZ.show_name = False | |
| for z in axZ: | |
| verts.append([axX[-1], axY[0], z]) | |
| verts.append([axX[-1], axY[0] - self.ticklen, z]) | |
| verts.append([axX[-1], axY[0] - 12 * self.ticklen, axZ[0]]) | |
| verts.append([axX[-1], axY[0] - 12 * self.ticklen, (axZ[0] + axZ[-1]) / 2]) | |
| verts.append([axX[-1], axY[0] - 12 * self.ticklen, axZ[-1]]) | |
| meshZ.from_pydata(verts, [], []) | |
| meshZ.update() | |
| return obX, obY, obZ | |
| def get_wireframe(self, X, Y, Z): | |
| """Generate wireframe as a sparse version of X, Y, Z mesh (but connected | |
| along contour of surface) | |
| """ | |
| first = np.ceil(X[0] / self.dx) * self.dx | |
| last = np.floor(X[-1] / self.dx) * self.dx | |
| temp = np.ogrid[first : last : self.dx] | |
| newX = np.concatenate((np.array([X[0]]), temp, np.array([X[-1]]))) | |
| first = np.ceil(Y[0] / self.dx) * self.dx | |
| last = np.floor(Y[-1] / self.dx) * self.dx | |
| temp = np.ogrid[first : last : self.dx] | |
| newY = np.concatenate((np.array([Y[0]]), temp, np.array([Y[-1]]))) | |
| # Iterate through rows of Z and interpolate using newX | |
| tempZ = np.zeros((len(Y), len(newX))) | |
| for y in range(len(Y)): | |
| tempZ[y, :] = np.interp(newX, X, Z[y, :]) | |
| # Iterate through columns of Z and interpolate using newY | |
| newZ = np.zeros((len(newY), len(newX))) | |
| for x in range(len(newX)): | |
| newZ[:, x] = np.interp(newY, Y, tempZ[:, x]) | |
| # Get indices of multiples of DX in new X and Y arrays | |
| xI = [] | |
| for n in range(newX.shape[0]): | |
| # For DX not an integer, the following is more likely to work | |
| # than just np.mod() | |
| if np.mod(np.round(np.mod(newX[n], self.DX), decimals=6), self.DX) == 0: | |
| xI.append(n) | |
| yI = [] | |
| for n in range(newY.shape[0]): | |
| if np.mod(np.round(np.mod(newY[n], self.DY), decimals=6), self.DY) == 0: | |
| yI.append(n) | |
| # For new arrays, construct a list of vertices and edges | |
| verts = [] | |
| nverts = -1 | |
| edges = [] | |
| for n in yI: | |
| for m in range(1, newX.shape[0]): | |
| if m == 1: | |
| verts.append([newX[m - 1], newY[n], newZ[n, m - 1]]) | |
| nverts += 1 | |
| verts.append([newX[m], newY[n], newZ[n, m]]) | |
| nverts += 1 | |
| edges.append([nverts - 1, nverts]) | |
| for m in xI: | |
| for n in range(1, newY.shape[0]): | |
| if n == 1: | |
| verts.append([newX[m], newY[n - 1], newZ[n - 1, m]]) | |
| nverts += 1 | |
| verts.append([newX[m], newY[n], newZ[n, m]]) | |
| nverts += 1 | |
| edges.append([nverts - 1, nverts]) | |
| # Make a new Blender mesh | |
| mesh = bpy.data.meshes.new(f"{self.prefix}wireframeMesh") | |
| ob = bpy.data.objects.new(f"{self.prefix}wireframe", mesh) | |
| ob.location = (0, 0, 0) | |
| ob.show_name = False | |
| mesh.from_pydata(verts, edges, []) | |
| mesh.update() | |
| # Scale z coordinates | |
| for v in mesh.vertices: | |
| v.co[2] *= self.Zscale | |
| # Add skin modifier | |
| mod = ob.modifiers.new("wireframeskin", "SKIN") | |
| mod.use_smooth_shade = True | |
| for v in ob.data.skin_vertices[0].data: | |
| v.radius = 0.03, 0.03 | |
| v.use_root = True | |
| return ob | |
| def get_mesh(self, X, Y, Z): | |
| """Construct filled Blender mesh from Z array, using X and Y vectors for X | |
| and Y positions of points. Return a Blender object. | |
| """ | |
| # Get an empty Blender mesh | |
| mesh = bpy.data.meshes.new(f"{self.prefix}mesh") | |
| ob = bpy.data.objects.new(f"{self.prefix}", mesh) | |
| ob.location = (0, 0, 0) | |
| ob.show_name = False | |
| # Iterate rows, then columns of Z, to add vertices to mesh | |
| verts = [] | |
| for n in range(Z.shape[0]): | |
| for m in range(Z.shape[1]): | |
| verts.append([X[m], Y[n], Z[n, m]]) | |
| # Iterate rows, then columns of Z, to add faces to mesh. For | |
| # each point, use three nearest points (right, bottom, and | |
| # diagonal bottom right) to make face. Therefore, stop one | |
| # from the last column and one from last row | |
| faces = [] | |
| for n in range(Z.shape[0] - 1): | |
| for m in range(Z.shape[1] - 1): | |
| curr = m + n * Z.shape[1] | |
| faces.append([curr, curr + 1, curr + Z.shape[1] + 1, curr + Z.shape[1]]) | |
| mesh.from_pydata(verts, [], faces) | |
| mesh.update() | |
| for v in mesh.vertices: | |
| v.co[2] *= self.Zscale | |
| return ob | |
| def get_raw(self, X, Y, Z): | |
| """Construct a wireframe consisting just the rows of data, not connected to | |
| each other (useful for plotting a series of spectra, for example). | |
| """ | |
| mesh = bpy.data.meshes.new(f"{self.prefix}rawMesh") | |
| ob = bpy.data.objects.new(f"{self.prefix}raw", mesh) | |
| ob.location = (0, 0, 0) | |
| ob.show_name = False | |
| # Iterate rows, then columns of Z, to add vertices to mesh | |
| verts = [] | |
| for n in range(Z.shape[0]): | |
| for m in range(Z.shape[1]): | |
| verts.append([X[m], Y[n], Z[n, m]]) | |
| # Iterate rows then columns of Z, adding edges to mesh | |
| edges = [] | |
| for n in range(Z.shape[0]): | |
| for m in range(Z.shape[1] - 1): | |
| curr = m + n * Z.shape[1] | |
| edges.append([curr, curr + 1]) | |
| mesh.from_pydata(verts, edges, []) | |
| mesh.update() | |
| # Scale z coordinates | |
| for v in mesh.vertices: | |
| v.co[2] *= self.Zscale | |
| # Add skin modifier | |
| mod = ob.modifiers.new("rawdataskin", "SKIN") | |
| mod.use_smooth_shade = True | |
| for v in ob.data.skin_vertices[0].data: | |
| v.radius = 0.05, 0.05 | |
| v.use_root = True | |
| return ob | |
| def update_raw(self, X, Y, Z): | |
| """Replace mesh in '2Draw' with a new one from given data. For | |
| convenience, put '2Draw' in an active layer before returning. | |
| """ | |
| bpy.ops.object.select_all(action="DESELECT") | |
| self.data_raw.select_set(state=True) | |
| bpy.ops.object.delete() | |
| self.data_raw = self.get_raw(X, Y, Z) | |
| self.collection.objects.link(self.data_raw) | |
| self.data_raw.select_set(state=True) | |
| bpy.ops.object.shade_smooth() | |
| self.data_raw.data.materials.append(bpy.data.materials["Material.004"]) | |
| self.set_layers(1, 2, 2, 2) | |
| return | |
| def get_contours(self, X, Y, Z): | |
| """Contour plot. Build a solid of the plot, and use the bisect tool to | |
| make 'nslices' (instance variable) equi-Z-value filled contours | |
| of the data. An orthographic projection of this thing viewed | |
| from directly above should look exactly like a filled contour | |
| plot. | |
| """ | |
| print("Calculating contours ... ") | |
| # Make a solid mesh by connecting the data surface to a | |
| # projection of itself onto the z=0 plane | |
| solid_mesh = bpy.data.meshes.new(f"{self.prefix}solidMesh") | |
| solid_ob = bpy.data.objects.new(f"{self.prefix}solid", solid_mesh) | |
| solid_ob.location = (0, 0, 0) | |
| solid_ob.show_name = False | |
| # Add verts | |
| verts = [] | |
| for n in range(Z.shape[0]): | |
| for m in range(Z.shape[1]): | |
| verts.append([X[m], Y[n], Z[n, m] * self.Zscale]) | |
| # Add same verts at z=0 | |
| for n in range(Z.shape[0]): | |
| for m in range(Z.shape[1]): | |
| verts.append([X[m], Y[n], -0.5]) | |
| # Add data surface faces | |
| faces = [] | |
| for n in range(Z.shape[0] - 1): | |
| for m in range(Z.shape[1] - 1): | |
| curr = m + n * Z.shape[1] | |
| faces.append([curr, curr + 1, curr + Z.shape[1] + 1, curr + Z.shape[1]]) | |
| # Add z=0 faces | |
| offset = Z.size | |
| for n in range(Z.shape[0] - 1): | |
| for m in range(Z.shape[1] - 1): | |
| curr = m + n * Z.shape[1] + offset | |
| faces.append([curr, curr + 1, curr + Z.shape[1] + 1, curr + Z.shape[1]]) | |
| # Add front and back faces | |
| c1 = 0 | |
| c2 = offset - Z.shape[1] | |
| for n in range(Z.shape[1] - 1): | |
| faces.append([c1, c1 + 1, c1 + offset + 1, c1 + offset]) | |
| faces.append([c2, c2 + 1, c2 + offset + 1, c2 + offset]) | |
| c1 += 1 | |
| c2 += 1 | |
| # Add side faces | |
| c1 = 0 | |
| c2 = Z.shape[1] - 1 | |
| for n in range(Z.shape[0] - 1): | |
| faces.append([c1, c1 + Z.shape[1], c1 + offset + Z.shape[1], c1 + offset]) | |
| faces.append([c2, c2 + Z.shape[1], c2 + offset + Z.shape[1], c2 + offset]) | |
| c1 += Z.shape[1] | |
| c2 += Z.shape[1] | |
| solid_mesh.from_pydata(verts, [], faces) | |
| solid_mesh.update() | |
| # self.scn.objects.link(solid_ob) | |
| # API changes | |
| self.collection.objects.link(solid_ob) | |
| # Get slices | |
| slicesZ = self.Zscale * np.ogrid[0 : Z.max() : ((Z.max() - 0.0) / self.nslices)] | |
| verts = [] | |
| edges = [] | |
| faces = [] | |
| index_offset = 0 | |
| for n in range(1, slicesZ.size): | |
| solid_bmesh = bmesh.new() | |
| solid_bmesh.from_mesh(solid_mesh) | |
| bmesh.ops.bisect_plane( | |
| solid_bmesh, | |
| geom=solid_bmesh.verts[:] + solid_bmesh.edges[:] + solid_bmesh.faces[:], | |
| plane_co=mathutils.Vector([0, 0, slicesZ[n]]), | |
| plane_no=mathutils.Vector([0, 0, 1]), | |
| clear_outer=True, | |
| clear_inner=True, | |
| ) | |
| bmesh.ops.triangle_fill( | |
| solid_bmesh, use_beauty=True, edges=solid_bmesh.edges[:] | |
| ) | |
| # Add to verts, edges, and faces lists | |
| solid_bmesh.verts.index_update() | |
| for v in solid_bmesh.verts: | |
| verts.append([v.co.x, v.co.y, v.co.z]) | |
| for e in solid_bmesh.edges: | |
| edges.append( | |
| [e.verts[0].index + index_offset, e.verts[1].index + index_offset] | |
| ) | |
| for f in solid_bmesh.faces: | |
| face = [] | |
| for v in f.verts: | |
| face.append(v.index + index_offset) | |
| faces.append(face) | |
| index_offset += len(solid_bmesh.verts) | |
| solid_bmesh.free() | |
| # New mesh and object for contours | |
| contours_mesh = bpy.data.meshes.new(f"{self.prefix}contourMesh") | |
| contours_ob = bpy.data.objects.new(f"{self.prefix}contour", contours_mesh) | |
| contours_ob.location = (0, 0, 0) | |
| contours_ob.show_name = False | |
| contours_mesh.from_pydata(verts, edges, faces) | |
| contours_mesh.update() | |
| bpy.ops.object.select_all(action="DESELECT") | |
| solid_ob.select_set(state=True) | |
| bpy.ops.object.delete() | |
| bpy.data.meshes.remove(solid_mesh) | |
| return contours_ob | |
| def main(): | |
| """main""" | |
| import numpy as np | |
| data = np.load("energy.npy") | |
| X = np.linspace(-15, 15, 1000) | |
| Y = np.linspace(-15, 15, 1000) | |
| # Get new BlenderAxis using the loaded data | |
| STRIDE = 1 | |
| data = (data - data.min()) / (data.max() - data.min()) | |
| PLOTTABLE = data[::STRIDE, ::STRIDE] | |
| baxis = BlenderAxis('2D', X[::STRIDE], Y[::STRIDE], PLOTTABLE) | |
| baxis.set_layers(2, 2, 2, 1) # Make the contour plot visible; hide all others | |
| el = 20 | |
| az = 0 | |
| foc = 1 | |
| for key in bpy.data.objects.keys()[2:8]: | |
| bpy.data.objects.get(key).hide_viewport = True | |
| bpy.data.objects.get(key).hide_render = True | |
| baxis.adjust_camera(el, foc, "P") | |
| baxis.rotate_view(az) | |
| # # Orthographic plot, top view | |
| # foc = 1 | |
| # surf_alpha = 1.0 | |
| # wire_alpha = 1.0 | |
| # el = 90 | |
| # az = 0 | |
| # baxis.write_image(r'C:\Users\trist\Desktop\Personal_projects\cover_art\example00.png', | |
| # az, el, foc, 'O', surf_alpha, wire_alpha) | |
| # # Perspective plot, view from bottom left | |
| # el = 20 | |
| # az = 30 | |
| # baxis.write_image(r'C:\Users\trist\Desktop\Personal_projects\cover_art\example01.png', | |
| # az, el, foc, 'P', surf_alpha, wire_alpha) | |
| return | |
| if __name__ == "__main__": | |
| register() | |
| f = addon_utils.enable("gradients") | |
| if not f: | |
| print("Blender can't find the addon, this means it isn't installed properly") | |
| main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment