digichem-core 6.0.0rc1__py3-none-any.whl

digichem/data/batoms/batoms-renderer.py

@@ -0,0 +1,351 @@
+# Render something with Beautiful Atoms.
+# To use this script, run something like:
+#
+# # To disabled possibly conflicting packages outside of the conda env.
+# PYTHONNOUSERSITE=1
+# blender -b -P batoms-renderer.py
+#
+# Where 'blender' is the path to the Beautiful Atoms Blender executable.
+
+import sys
+import argparse
+import itertools
+import bpy
+import yaml
+import math
+from pathlib import Path
+import logging
+
+import ase.io
+from batoms import Batoms
+from batoms.utils.butils import object_mode
+
+def add_molecule(
+        cube_file,
+        name,
+        rotations = None,
+        visible = True,
+        isovalues = None,
+        isotype = "both",
+        primary_color = [1, 0.058, 0.0, 0.55],
+        secondary_color = [0.1, 0.1, 0.9, 0.55],
+        style = "default"
+        ):
+    """
+    """
+    rotations = [] if rotations is None else rotations
+    isovalues = [] if isovalues is None else isovalues
+    
+    surface_settings = []
+    for isovalue in isovalues:
+        if isotype in ["both", "positive"]:
+            surface_settings.append({'level': isovalue, 'color': primary_color, 'material_style': style})
+        
+        if isotype in ["both", "negative"]:
+            surface_settings.append({'level': -isovalue, 'color': secondary_color, 'material_style': style})
+    
+    
+    # Load the input data.
+    cube = ase.io.read(cube_file, format="cube", read_data=True, full_output=True)
+    
+    # The centre of the cube is often offset, fix that by shifting the atoms.
+    cube["atoms"].translate(-cube["origin"][0:3])
+    
+    # Get the mol object.
+    mol = Batoms(name, from_ase = cube["atoms"])
+    
+    # Set some look and feel options.    
+    # Change molecule style.
+    mol.model_style = 1
+    
+    # Hide cell boundaries.
+    mol.cell.hide = True
+    
+    # Colour tuning.
+    # Carbon to 'black'.
+    try:
+        mol["C"].color = (0.095, 0.095, 0.095, 1)
+    except AttributeError:
+        pass
+    try:
+        mol["B"].color = (1.0, 0.396, 0.468, 1)
+    except AttributeError:
+        pass
+    
+    if not visible:
+        mol.hide = True
+    
+    # Add volumes.
+    if len(surface_settings) != 0:
+        mol.volumetric_data['surface'] = cube['data']
+        
+        for index, settings in enumerate(surface_settings):
+            mol.isosurface.settings[index+1] = settings
+        
+        mol.isosurface.draw()
+    
+    # Now move the entire molecule (isosurface and all) back to the origin.
+    mol.translate(cube["origin"][0:3])
+    
+    # Fix the origin point so we can still rotate properly.
+    # For some reason, this code moves the bond objects to a new location?
+#     object_mode()
+#     bpy.ops.object.select_all(action='DESELECT')
+#     mol.obj.select_set(True)
+#     bpy.context.view_layer.objects.active = mol.obj
+#     bpy.ops.object.origin_set(type='ORIGIN_CURSOR')
+    
+    # If we have any rotations, apply those.
+    for axis, angle in rotations:
+        # Convert to degree.
+        # Not used.
+        degree = angle * (180/math.pi)
+        
+        if axis == 0:
+            axis = "x"
+        
+        elif axis == 1:
+            axis = "y"
+        
+        elif axis == 2:
+            axis = "z"
+        
+        else:
+            raise ValueError("Unknown rotation axis '{}'".format(axis))
+        
+        #mol.rotate(degree, axis)
+        
+        # Taken from batoms source (collection.py: 51)
+        object_mode()
+        bpy.ops.object.select_all(action='DESELECT')
+        mol.obj.select_set(True)
+        bpy.context.view_layer.objects.active = mol.obj
+        bpy.ops.transform.rotate(value=angle, orient_axis=axis.upper(),
+                                 center_override = (0,0,0))
+    
+    return mol
+
+# Adapted from https://blender.stackexchange.com/questions/5898/how-can-i-create-a-cylinder-linking-two-points-with-python?newreg=f372ba9448694f5b97879a6dab963cee
+def draw_primitive(start, end, radius, mesh_type, color, collection = None):
+    """
+    """
+    dx = end[0] - start[0]
+    dy = end[1] - start[1]
+    dz = end[2] - start[2]  
+    dist = math.sqrt(dx**2 + dy**2 + dz**2)
+    
+    # First draw the shape.
+    if mesh_type == "cylinder":
+        bpy.ops.mesh.primitive_cylinder_add(
+            radius = radius, 
+            depth = dist,
+            vertices = 32,
+            location = (dx/2 + start[0], dy/2 + start[1], dz/2 + start[2]) 
+        )
+    elif mesh_type == "cone":
+        bpy.ops.mesh.primitive_cone_add(
+            radius1 = radius,
+            depth = dist,
+            vertices = 32,
+            location = (dx/2 + start[0], dy/2 + start[1], dz/2 + start[2])
+        )
+    else:
+        raise ValueError("Unknown mesh type '{}'".format(mesh_type))
+        
+    # Get a reference to the object we just made.
+    obj = bpy.context.active_object
+    
+    phi = math.atan2(dy, dx)
+    theta = math.acos(dz/dist) 
+    
+    bpy.context.object.rotation_euler[1] = theta 
+    bpy.context.object.rotation_euler[2] = phi
+    
+    # Get material
+    mat = bpy.data.materials.new(name="Material")
+    
+    # assign to 1st material slot
+    obj.active_material = mat
+        
+    mat.use_nodes = True
+    tree = mat.node_tree
+    nodes = tree.nodes
+    bsdf = nodes["Principled BSDF"]
+    bsdf.inputs["Base Color"].default_value = color
+    bsdf.inputs["Metallic"].default_value = 0.1
+    bsdf.inputs["Specular"].default_value = 0.2
+    bsdf.inputs["Roughness"].default_value = 0.2
+    mat.diffuse_color = color
+    
+    # If we've been asked to, asign our new object to a given collection.
+    # First unlink from any old collections
+    for coll in obj.users_collection:
+        # Unlink the object
+        coll.objects.unlink(obj)
+
+    # Link each object to the target collection
+    collection.objects.link(obj)
+
+    
+def draw_arrow(start, end, radius, color, split = 0.9, collection = None):
+    # Decide what proportion of the total vector to dedicate to the arrow stem and head.
+    dx = end[0] - start[0]
+    dy = end[1] - start[1]
+    dz = end[2] - start[2]
+    dist = math.sqrt(dx**2 + dy**2 + dz**2)
+    
+    join = (dx * split + start[0], dy * split + start[1], dz * split + start[2])
+    draw_primitive(start, join, radius, "cylinder", color, collection = collection)
+    draw_primitive(join, end, radius*2, "cone", color, collection = collection)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog='Beautiful Atoms Renderer',
+        description='Render images with BAtoms')
+    
+    parser.add_argument("cube_file", help = "Path to the cube file to read")
+    parser.add_argument("output", help = "File to write to")
+    parser.add_argument("--second_cube", help = "Optional second cube file to read additional isosurface data from", default = None)
+    parser.add_argument("--isovalues", help = "List of isovalues to render", nargs = "*", type = float, default = [])
+    parser.add_argument("--isotype", help = "Whether to render positive, negative or both isosurfaces for each isovalue", choices = ["positive", "negative", "both"], default = "both")
+    parser.add_argument("--isocolor", help = "The colouring method to use for isosurfaces", choices = ["sign", "cube"], default = "sign")
+    parser.add_argument("--primary-color", help = "RGBA for one of the colors to use for isosurfaces", type = float, nargs = 4, default = [0.1, 0.1, 0.9, 0.7])
+    parser.add_argument("--secondary-color", help = "RGBA for the other color to use for isosurfaces", type = float, nargs = 4, default = [1, 0.058, 0.0, 0.7])
+    parser.add_argument("--style", help = "Material style for isosurfaces", choices = ('default', 'metallic', 'plastic', 'ceramic', 'mirror'), default = "default")
+    parser.add_argument("--cpus", help = "Number of parallel CPUs to use for rendering", type = int, default = 1)
+    parser.add_argument("--use-gpu", help = "Whether to enable GPU rendering", action = "store_true")
+    parser.add_argument("--orientation", help = "The orientation to render from, as x, y, z values", nargs = 3, type = float, default = [0, 0, 1])
+    parser.add_argument("--resolution", help = "The output resolution in px", type = int, default = 1024)
+    parser.add_argument("--render-samples", help = "The maximum number of render samples, more generally results in higher quality but longer render times", type = int, default = 256)
+    parser.add_argument("--rotations", help = "A list of rotations (in JSON) to rotate the molecule to a given alignment. The first item in each list item is the axis to rotate about (0=x, 1=y, 2=z), the second is the angle to rotate by (in radians)", nargs = "*", default = [])
+    parser.add_argument("--dipoles", help = "Draw dipoles from a list of the following data (in JSON): 0) start coord, 1) end coord, 2) RGBA color information", nargs = "*", default = [])
+    parser.add_argument("--alpha", help = "Override the opacity value for all molecule objects (but not dipoles) to this value, useful for showing dipole arrows more clearly", default = None, type = float)
+    parser.add_argument("--perspective", help = "The perspective mode, either orthographic or perspective", default = "perspective", choices = ["perspective", "orthographic"])
+    parser.add_argument("--padding", help = "Padding", type = float, default = 1.0)
+    
+    # Both blender and python share the same command line arguments.
+    # They are separated by double dash ('--'), everything before is for blender,
+    # everything afterwards is for python (except for the first argument, wich is
+    # the program name, which is for both).
+    if "--" in sys.argv:
+        python_argv = sys.argv[sys.argv.index("--") +1:]
+    else:
+        python_argv = []
+    
+    args = parser.parse_args(python_argv)
+    
+    # Batoms or blender will silently set the extension to png if it's not already.
+    # This is surprising, so stop now before that happens.
+    if Path(args.output).suffix.lower() != ".png":
+        raise ValueError("Output location must have a .png extension")
+    
+    if args.rotations is not None:
+        rotations = [yaml.safe_load(rotation) for rotation in args.rotations]
+    
+    # Load the input data.
+    mol = add_molecule(
+        args.cube_file,
+        name = "molecule",
+        visible = True,
+        rotations = rotations,
+        isovalues = args.isovalues,
+        isotype = args.isotype,
+        primary_color = args.primary_color,
+        secondary_color = args.secondary_color if args.isocolor == "sign" else args.primary_color,
+        style = args.style
+    )
+    
+    # Uncomment to show atom labels.
+    # Needs some tweaking to appear in render (viewport only by default).
+    #mol.show_label = 'species'
+    
+    # If we have a second cube, add that too.
+    if args.second_cube is not None:
+        mol2 = add_molecule(
+            args.second_cube,
+            name = "molecule2",
+            visible = True,
+            rotations = rotations,
+            isovalues = args.isovalues,
+            isotype = args.isotype,
+            primary_color = args.primary_color if args.isocolor == "sign" else args.secondary_color,
+            secondary_color = args.secondary_color,
+            style = args.style
+        )
+    
+    if args.alpha:
+        # Set all materials transparent.
+        for material in bpy.data.materials:
+            try:
+                material.node_tree.nodes['Principled BSDF'].inputs['Alpha'].default_value = args.alpha
+            except Exception as e:
+                pass
+        
+    
+    # Draw any dipoles.
+    if args.dipoles is not None:
+        
+        dipoles = [yaml.safe_load(dipole) for dipole in args.dipoles]
+        for start_coord, end_coord, rgba in dipoles:
+            draw_arrow(start_coord, end_coord, 0.08, rgba, collection = mol.coll)
+        
+    
+    # Setup rendering settings.
+#     mol.render.engine = 'workbench'
+#     mol.render.engine = 'eevee'
+    mol.render.engine = 'cycles'
+    mol.render.resolution = [args.resolution, args.resolution]
+    # Set up cycles for good quality rendering.
+    # Prevents early end to rendering (forces us to use the actual number of samples).
+    bpy.context.scene.cycles.use_adaptive_sampling = False
+    # Quality control, more = better and slower.
+    bpy.context.scene.cycles.samples = args.render_samples
+    # Post-processing to remove noise, works well for coloured backgrounds, useless for transparency.
+    bpy.context.scene.cycles.use_denoising = True
+    # Ray-tracing options
+    bpy.context.scene.cycles.max_bounces = 48
+    bpy.context.scene.cycles.transparent_max_bounces = 24
+    if args.use_gpu:
+        bpy.context.scene.cycles.device = "GPU"
+    
+    # Use maximum compression.
+    bpy.context.scene.render.image_settings.compression = 1000
+    
+    # Change light intensity.
+    bpy.data.lights["batoms_light_Default"].node_tree.nodes["Emission"].inputs[1].default_value = 0.5 #0.3
+    #mol.render.lights["Default"].energy=10
+    
+    # Change view mode.
+    if args.perspective == "perspective":
+        mol.render.camera.type = "PERSP"
+    
+    else:
+        mol.render.camera.type = "ORTHO"
+    
+    # Enable to add an outline.
+    #bpy.context.scene.render.use_freestyle = True
+    
+    # Performance options.
+    bpy.context.scene.render.threads_mode = 'FIXED'
+    bpy.context.scene.render.threads = args.cpus
+    
+    mol.get_image(viewport = args.orientation, output = args.output, padding = args.padding)
+#     # Move the camera.
+#     mol.render.camera.location = (100,0,0)
+#     mol.render.camera.look_at = mol.get_center_of_geometry()
+#     bpy.ops.object.select_all(action='DESELECT')
+#     for obj in mol.coll.objects[:]:
+#         obj.select_set(True)
+#     #bpy.ops.view3d.camera_to_view_selected()
+    
+    return 0
+    
+# If we've been invoked as a program, call main().    
+if __name__ == '__main__':
+    try:
+        sys.exit(main())
+    
+    except Exception as e:
+        logging.error("Erro", exc_info = True)
+        sys.exit(1)