Added SHEAP related tools

disp/cli/cmd_deploy.py

@@ -7,7 +7,7 @@
 
 import click
 from fireworks.core.firework import Workflow
-from fireworks.core.launchpad import LaunchPad, LAUNCHPAD_LOC
+from fireworks.core.launchpad import LaunchPad
 from ase.io import read
 
 from disp.fws.works import AirssSearchFW, RelaxFW
@@ -22,6 +22,7 @@
     'pp3': '.pp',
 }
 
+
 @click.group('deploy')
 @click.pass_context
 def deploy(ctx):
@@ -30,6 +31,7 @@ def deploy(ctx):
     ctx.obj = LaunchPad.from_file(lpad_file)
     click.echo(f'Using launchpad file at `{lpad_file}`')
 
+
 pass_lpad = click.make_pass_decorator(LaunchPad)
 
 
@@ -110,8 +112,9 @@ def info(lpad):
     'Alias for resolving the CASTEP executable, as define in the worker file.')
 @click.option('--cluster', is_flag=True, default=False)
 @pass_lpad
-def deploy_search(lpad, code, seed, project, num, exe, cycles, keep, wf_name, dryrun,
-                  priority, category, gzip, record_db, modcell, castep_code, cluster):
+def deploy_search(lpad, code, seed, project, num, exe, cycles, keep, wf_name,
+                  dryrun, priority, category, gzip, record_db, modcell,
+                  castep_code, cluster):
     """
     Deploy the search by uploading it to the Fireserver
     """
@@ -189,9 +192,9 @@ def deploy_search(lpad, code, seed, project, num, exe, cycles, keep, wf_name, dr
         click.echo('And parameters:')
         click.echo(param_content)
         if keep:
-            click.echo(f'The intermediate files will be kept.')
+            click.echo('The intermediate files will be kept.')
         else:
-            click.echo(f'The intermediate files will not be kept.')
+            click.echo('The intermediate files will not be kept.')
 
         click.echo(f'The default executable for tasks is: {exe}')
         click.echo(
@@ -215,12 +218,18 @@ def deploy_search(lpad, code, seed, project, num, exe, cycles, keep, wf_name, dr
     type=str,
     help=
     ('Base cell files to be used.'
-     'The "cell" file passed will be only used to define the crystal structure.'
+     'If supplied, the'
+     ' "cell" file passed will be only used to define the crystal structure as read using ASE.'
+     'Hence, any ASE-supported geometry file can be used for defining the structures to be relaxed.'
+     ))
+@click.option(
+    '--cell',
+    required=True,
+    type=str,
+    help=
+    ('Path the cell files - support globbing. If `--base-cell` is supplied, '
+     'any ase-support format may be used, otherwise only CASTEP .cell files are allowed.'
      ))
-@click.option('--cell',
-              required=True,
-              type=str,
-              help='Path the cell files - support globbing')
 @click.option('--priority',
               type=int,
               help='Priority to be used for the workflow')
@@ -270,9 +279,9 @@ def deploy_search(lpad, code, seed, project, num, exe, cycles, keep, wf_name, dr
               default='castep')
 @click.option('--cluster', is_flag=True, default=False)
 @pass_lpad
-def deploy_relax(lpad, code, seed, cell, base_cell, param, project, exe, cycles,
-                 keep, dryrun, priority, gzip, record_db, category, cluster,
-                 extra_cell_file, castep_code):
+def deploy_relax(lpad, code, seed, cell, base_cell, param, project, exe,
+                 cycles, keep, dryrun, priority, gzip, record_db, category,
+                 cluster, extra_cell_file, castep_code):
     """
     Deploy a workflow to do relaxation of a particular structure
     """
@@ -352,9 +361,9 @@ def deploy_relax(lpad, code, seed, cell, base_cell, param, project, exe, cycles,
         click.echo('And parameters:')
         click.echo(param_content)
         if keep:
-            click.echo(f'The intermediate files will be kept.')
+            click.echo('The intermediate files will be kept.')
         else:
-            click.echo(f'The intermediate files will not be kept.')
+            click.echo('The intermediate files will not be kept.')
 
         click.echo(f'The default executable for tasks is: {exe}')
         click.echo(

disp/cli/cmd_tools.py

@@ -1,9 +1,13 @@
 """
 Collection of useful tools
 """
+from json import dumps
 import click
 from ase.io import read
 from disp.tools.modcell import modify_cell
+from disp.tools.sheapio import sheap_to_dict, parse_sheap_output
+
+# pylint: disable=import-outside-toplevel
 
 
 @click.group('tools')
@@ -19,3 +23,107 @@ def tools(ctx):
 def modcell(base_cell, other_cell):
     """Modify the structure of a CELL file using another"""
     click.echo('\n'.join(modify_cell(base_cell, read(other_cell))))
+
+
+@tools.command('sheap2json')
+@click.argument('sheapout')
+@click.argument('path')
+def cmd_sheap2json(sheapout, path):
+    """Convert output from SHEAP to chemiscope format including the structures.
+
+    The results is printed to the STDOUT
+    """
+    with open(sheapout) as handle:
+        parsed_data = parse_sheap_output(handle)
+
+    outdict = sheap_to_dict(parsed_data, path)
+
+    # Include mapping settings
+    map_dict = {
+        'x': {
+            'property': 'sheap1'
+        },
+        'y': {
+            'property': 'sheap2'
+        },
+        'color': {
+            'max': min(parsed_data.enthalpy) + 0.5,
+            'min': min(parsed_data.enthalpy),
+            'property': 'energy',
+            'scale': 'linear'
+        },
+        'size': {
+            'factor': 20,
+            'mode': 'linear',
+            'property': 'size',
+            'reverse': False
+        }
+    }
+    settings = {
+        'map':
+        map_dict,
+        'structure': [{
+            'bonds': True,
+            'spaceFilling': False,
+            'atomLabels': False,
+            'unitCell': True,
+            'rotation': False,
+            'supercell': {
+                '0': 2,
+                '1': 2,
+                '2': 2
+            },
+            'axes': 'off',
+            'keepOrientation': False,
+            'environments': {
+                'activated': True,
+                'bgColor': 'grey',
+                'bgStyle': 'ball-stick',
+                'center': True,
+                'cutoff': 0
+            }
+        }]
+    }
+    outdict['settings'] = settings
+    click.echo(dumps(outdict))
+
+
+@tools.command('plot-sheap')
+@click.argument('sheapout')
+@click.option('--vmax',
+              help='Color scale maximum relative to the minimum value',
+              default=0.25)
+def plot_sheap(sheapout, vmax):
+    """
+    Plot the output of SHEAP as spheres, respecting the specification of radius output.
+    """
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import matplotlib.cm as cm
+
+    with open(sheapout) as handle:
+        parsed = parse_sheap_output(handle)
+    coords = np.array(parsed.coords)
+    radius = np.array(parsed.radius)
+
+    # Plot the output
+    _, axes = plt.subplots(1, 1)
+    axes.set_aspect('equal')
+
+    # Compute the colours
+    cmap = plt.get_cmap()
+    emin = min(parsed.enthalpy)
+    emax = emin + vmax
+    norm = plt.Normalize(emin, emax)
+
+    for i, coord in enumerate(coords):
+        axes.add_patch(
+            plt.Circle(coord,
+                       radius=radius[i],
+                       facecolor=cmap(norm(parsed.enthalpy[i]))))
+    xmin, xmax = coords[:, 0].min(), coords[:, 0].max()
+    ymin, ymax = coords[:, 1].min(), coords[:, 1].max()
+    axes.set_xlim(min(xmin, ymin) - 0.1, max(xmax, ymax) + 0.1)
+    axes.set_ylim(min(xmin, ymin) - 0.1, max(xmax, ymax) + 0.1)
+    plt.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap))
+    plt.show()