digichem-core 6.1.0__py3-none-any.whl → 6.10.1__py3-none-any.whl

digichem/data/batoms/batoms_renderer.py

@@ -0,0 +1,500 @@
+# 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 debugpy
+# debugpy.listen(5678)
+# debugpy.wait_for_client()
+
+import addon_utils
+def handle_error(exception):
+    raise exception
+addon_utils.enable("batoms", handle_error = handle_error, default_set=True)
+
+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
+from batoms.render import Render
+
+class Digichem_render(Render):
+    def set_viewport_distance_center(self, center=None, padding=0, canvas=None):
+        """
+        Calculate canvas and direction
+        """
+        batoms = self.batoms
+        if padding is None:
+            padding = max(batoms.size) + 0.5
+        if center is None:
+            center = batoms.get_center_of_geometry()
+        self.center = center
+        if canvas is None:
+            width, height, depth = batoms.get_canvas_box(
+                direction=self.viewport, padding=padding
+            )
+        else:
+            width = canvas[0]
+            height = canvas[1]
+            depth = canvas[2]
+        if self.distance < 0:
+            self.distance = max(10, depth)
+
+        self.update_camera(width, height, depth / 2)
+
+        # To auto centre the camera, we need to select the molecule as well as all isosurfaces that might be present.
+        # Select the molecule.
+        self.batoms.obj.select_set(True)
+
+        # Isosurfaces.
+        for obj in self.batoms.coll.all_objects:
+            if obj.batoms.type == "ISOSURFACE":
+                obj.select_set(True)
+
+        # Set camera as active.
+        bpy.context.scene.camera = self.camera.obj
+
+        # Manually set a focal point.
+        self.camera.lens = 50
+        # Auto centre.
+        bpy.ops.view3d.camera_to_view_selected()
+
+        self.update_light()
+
+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.    
+    
+    # Hide cell boundaries.
+    mol.cell.hide = True
+    
+    # Colour tuning.
+    # Carbon to 'black'.
+    new_colors = {
+        "C": (0.095, 0.095, 0.095, 1),
+        "B": (1.0, 0.396, 0.468, 1)
+    }
+
+    for atom, color in new_colors.items():
+        try:
+            mol[atom].color = color
+        except AttributeError:
+            pass
+
+    # And bonds.
+    for key, value in mol.bond.settings.items():
+        for atom, color in new_colors.items():
+            if key[0] == atom:
+                value['color1'] = color
+
+            if key[1] == "C":
+                value['color2'] = color
+    
+    # Change molecule style.
+    mol.model_style = 1
+
+    # Slightly increase volume of all atoms.
+    for atom in mol.species.keys():
+        mol[atom].scale *= 1.25
+    
+    # Increase volume of H atoms
+    mol['H'].scale = 0.75
+    
+    # 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.obj.select_set(True)
+    bpy.ops.object.origin_set(type="ORIGIN_GEOMETRY", center='MEDIAN')
+    bpy.context.object.location = [0,0,0]
+    
+    # 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))
+    
+    if not visible:
+        mol.hide = True
+    
+    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
+    try:
+        # Blener 3.x
+        bsdf.inputs["Specular"].default_value = 0.2
+    except KeyError:
+        # Blender 4.x
+        bsdf.inputs["Specular IOR Level"].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)
+
+    return 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])
+    cylinder = draw_primitive(start, join, radius, "cylinder", color, collection = collection)
+    cone = draw_primitive(join, end, radius*2, "cone", color, collection = collection)
+
+    # Select the two objects and join them together.
+    bpy.ops.object.select_all(action='DESELECT')
+    cylinder.select_set(True)
+    cone.select_set(True)
+    bpy.ops.object.join()
+
+    arrow = cone
+
+    # Set the origin of the new combined object to the origin of the arrow.
+    bpy.context.scene.cursor.location = start
+    bpy.ops.object.origin_set(type='ORIGIN_CURSOR')
+
+    return arrow
+
+
+
+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", nargs="?", default = None)
+    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.5])
+    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.5])
+    parser.add_argument("--style", help = "Material style for isosurfaces", choices = ('default', 'metallic', 'plastic', 'ceramic', 'mirror'), default = "ceramic")
+    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, 0])
+    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 = 64)# 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  1- 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)
+    parser.add_argument("--multi", help = "Render multiple images, each of a different angle of the scene. Each argument should consist of 6 parts, the x y z position, the resolution, the samples, and the filename (which is appended to 'output')", nargs = 6, default =[], action="append")
+    # 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 args.output is not None and 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]
+
+    if args.multi != []:
+        if args.orientation != [0, 0, 0]:
+            raise ValueError("You cannot set both --orientation and --multi!")
+
+        if args.resolution != 1024:
+            raise ValueError("You cannot set both --resolution and --multi!")
+        
+        if args.output is not None:
+            raise ValueError("You cannot set both 'output' and --multi!")
+    
+    # Remove the starting cube object.
+    bpy.ops.object.select_all(action='SELECT')
+    bpy.ops.object.delete()
+
+    # Load the input data.
+    mol = add_molecule(
+        args.cube_file,
+        name = "first_mol",
+        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'
+    mol2 = None
+    
+    # If we have a second cube, add that too.
+    if args.second_cube is not None:
+        mol2 = add_molecule(
+            args.second_cube,
+            name = "second_mol",
+            visible = False,
+            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 = (1 - args.alpha)
+            except Exception as e:
+                pass
+        
+    
+    # Draw any dipoles.
+    arrows = []
+    if args.dipoles is not None:
+        
+        dipoles = [yaml.safe_load(dipole) for dipole in args.dipoles]
+        for start_coord, end_coord, rgba in dipoles:
+            arrows.append(draw_arrow(start_coord, end_coord, 0.1, rgba, collection = mol.coll))
+        
+    
+    # Setup rendering settings.
+#     mol.render.engine = 'workbench'
+#     mol.render.engine = 'eevee'
+    mol.render.engine = 'cycles'
+    # 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
+    # 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.
+    mol.render.lights["Default"].direction = [0.1, 0.1, 1]
+    mol.render.lights["Default"].obj.data.node_tree.nodes["Emission"].inputs[1].default_value = 0.2
+    mol.render.lights["Default"].obj.data.angle = 0.174533
+
+    # Add a second light for depth.
+    mol.render.lights.add("Accent1", direction = [1,0.5,0.75])
+    mol.render.lights.add("Accent2", direction = [0.5,1,0.75])
+
+    mol.render.lights["Accent1"].obj.data.angle = 0.0872665
+    mol.render.lights["Accent1"].obj.data.node_tree.nodes["Emission"].inputs[1].default_value = 0.25
+    mol.render.lights["Accent2"].obj.data.angle = 0.0872665
+    mol.render.lights["Accent2"].obj.data.node_tree.nodes["Emission"].inputs[1].default_value = 0.25
+
+    # bpy.data.lights["batoms_light_Default"].node_tree.nodes["Emission"].inputs[1].default_value = 0.45
+    # bpy.data.lights["batoms_light_Default"].angle
+    #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
+    
+    # We have plenty of memory to play with, use one tile.
+    bpy.context.scene.cycles.tile_x = args.resolution
+    bpy.context.scene.cycles.tile_y = args.resolution
+    bpy.context.scene.cycles.tile_size = args.resolution
+
+    # Performance options.
+    bpy.context.scene.render.threads_mode = 'FIXED'
+    bpy.context.scene.render.threads = args.cpus
+    
+    # Set our custom renderer so we can modify zoom etc.
+    mol.render = Digichem_render()
+
+    # We have two ways we can change which angle we render from.
+    # 1) the viewport keyword arg (which places the camera in a certain location).
+    # 2) rotate the molecule.
+    #
+    # We use option 2, because this gives us more control.
+
+    # Work out how many angles we're rendering from.
+    if args.multi == []:
+        # Just one.
+        targets = [[args.orientation[0], args.orientation[1], args.orientation[2], args.resolution, args.render_samples, args.output]]
+    
+    else:
+        # More than one.
+        targets = args.multi
+
+    for x, y, z, resolution, samples, full_file_name in targets:
+        # Add args.output and mini_file_name together (useful for --multi).
+        orientation = (float(x), float(y), float(z))
+
+        mol.render.resolution = [resolution, resolution]
+        # Quality control, more = better and slower.
+        bpy.context.scene.cycles.samples = int(samples)
+        mol.obj.delta_rotation_euler = orientation
+        
+        if mol2 is not None:
+            mol2.obj.delta_rotation_euler = orientation
+
+        for arrow in arrows:
+            arrow.delta_rotation_euler = orientation
+
+        mol.get_image(viewport = [0,0,1], output = full_file_name, padding = args.padding)
+    
+    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)

digichem/file/base.py

@@ -45,6 +45,20 @@ class File_maker_ABC():
         """
         raise NotImplementedError("delete() is not implemented in this ABC")
     
+    def get_parallel(self, name = 'file', cpus = 1):
+        """
+        Generate the file represented by this object in a parallel context.
+
+        get_parallel() will be called by a higher level function as an argument to ThreadPoolExecutor.map() or similar,
+        to generate many files simultaneously. This is useful for slow operations (such as cube generation) that are
+        difficult to parallelise individually, but easy to parallelise across multiple files.
+
+        :param name: The file to generate.
+        :param cpus: The number of CPUs this operation should use, nearly always 1.
+        """
+        # This default implementation does nothing.
+        pass
+    
     #################################
     # Implemented in this sub-class #
     #################################