Small enhancements for cathub organize (#48)

catkit/hub/cli.py

@@ -384,6 +384,14 @@ def connect(user):
     show_default=True,
     help="Specify the maximum density (#atoms/A^3) "
     " at which the structure are considered slabs and not bulk")
+@click.option(
+    '-t', '--traj-format',
+    type=bool,
+    is_flag=True,
+    default=False,
+    show_default=True,
+    help="Store intermediate filetype as traj"
+    "instead of json files")
 @click.option(
     '-v', '--verbose',
     is_flag=True,

catkit/hub/organize.py

@@ -74,7 +74,7 @@ def collect_structures(foldername, options):
             with open(posix_filename) as infile:
 
                 global PUBLICATION_TEMPLATE
-                PUBLICATION_TEMPLATE= infile.read()
+                PUBLICATION_TEMPLATE = infile.read()
         elif Path(posix_filename).is_file():
             try:
                 filetype = ase.io.formats.filetype(posix_filename)
@@ -86,7 +86,15 @@ def collect_structures(foldername, options):
                     structure.info['filename'] = posix_filename
                     structure.info['filetype'] = ase.io.formats.filetype(
                         posix_filename)
-                    structures.append(structure)
+                    try:
+                        structure.get_potential_energy()
+                        # ensure that the structure has an energy
+                        structures.append(structure)
+                    except RuntimeError:
+                        print("Did not add {posix_filename} since it has no energy"
+                              .format(
+                                  posix_filename=posix_filename,
+                              ))
                     print(structure)
                 except StopIteration:
                     print("Warning: StopIteration {posix_filename} hit."
@@ -107,12 +115,12 @@ def collect_structures(foldername, options):
                     print("Hit an assertion error with {posix_filename}: {e}".format(
                         posix_filename=posix_filename,
                         e=e,
-                        ))
+                    ))
                 except ValueError as e:
                     print("Trouble reading {posix_filename}: {e}".format(
                         posix_filename=posix_filename,
                         e=e,
-                        ))
+                    ))
 
     return structures
 
@@ -121,7 +129,7 @@ def fuzzy_match(structures, options):
     # filter out cell with ill-defined unit cells
     structures = [structure for structure in structures
                   if structure.number_of_lattice_vectors == 3
-            ]
+                  ]
     # sort by density
     structures = sorted(structures,
                         key=lambda x: len(x) / x.get_volume()
@@ -306,10 +314,23 @@ def fuzzy_match(structures, options):
                         for i, j in references:
                             atomic_references[i] = j
 
+                        adsorbates = []
+                        if additions:
+                            adsorbates.append(additions)
+                        if subtractions:
+                            adsorbates.append(subtractions)
+                        if options.verbose:
+                            print("    ADDITIONS " + str(additions))
+                            print("    SUBTRACTIONS " + str(subtractions))
+                            print("    ADSORBATES " + str(adsorbates))
+                            print("    REFERENCES " + str(references))
                         stoichiometry_factors =  \
                             gas_phase_references.get_stoichiometry_factors(
-                                additions + subtractions, references,
+                                adsorbates, references,
                             )
+                        if options.verbose:
+                            print("    STOICH FACTORS " +
+                                  str(stoichiometry_factors) + "\n\n")
 
                         formula = '* ->'
                         adsorbate = get_chemical_formula(
@@ -325,24 +346,15 @@ def fuzzy_match(structures, options):
 
                         gas_phase_corrections = {}
 
-                        for addition in additions:
-                            stoich_factors = stoichiometry_factors[addition]
+                        for adsorbate in adsorbates:
+                            stoich_factors = stoichiometry_factors[adsorbate]
                             for ref in stoich_factors:
                                 dE -= stoich_factors[ref] * \
                                     reference_energy[ref]
                                 gas_phase_corrections[ref] = \
                                     gas_phase_corrections.get(
                                         ref, 0) - stoich_factors[ref]
 
-                        for subtraction in subtractions:
-                            stoich_factors = stoichiometry_factors[subtraction]
-                            for ref in stoich_factors:
-                                dE += stoich_factors[ref] * \
-                                    reference_energy[ref]
-                                gas_phase_corrections[ref] = \
-                                    gas_phase_corrections.get(
-                                        ref, 0) + stoich_factors[ref]
-
                         for molecule, factor in gas_phase_corrections.items():
                             if factor != 0:
                                 sign = ' + ' if factor < 0 else ' +- '
@@ -432,6 +444,8 @@ def fuzzy_match(structures, options):
 
 
 def create_folders(options, structures, root=''):
+    out_format = 'json'
+
     for key in structures:
         if isinstance(structures[key], dict):
             d = Path(root).joinpath(key)
@@ -446,9 +460,9 @@ def create_folders(options, structures, root=''):
             create_folders(options, structures[key], root=d)
         else:
             ase.io.write(
-                str(Path(root).joinpath(key + '.traj')),
+                str(Path(root).joinpath(key + '.' + out_format)),
                 structures[key],
-                format='traj',
+                format=out_format,
             )
 
 

catkit/hub/tests/test_file_organization.py

@@ -43,8 +43,51 @@ def test_file_organization_module():
 
     })
 
+    subprocess.call(
+        ('python {path}/make_test_slabs.py'.format(path=path)).split())
+    catkit.hub.organize.main(options)
+
+
+def test_file_organization_module_non_keep():
+    options = Struct(**{
+        'adsorbates': 'O,H2',
+        'foldername': '{path}/unorganized'.format(path=path),
+        'verbose': True,
+        'include_pattern': '.',
+        'exclude_pattern': '%%$^#$',
+        'facet_name': '111',
+        'max_density_gas': 0.002,
+        'max_density_slab': 0.06,
+        'exclude_reference': '',
+        'max_energy': 10,
+        'keep_all_energies': False,
+
+    })
+
+    subprocess.call(
+        ('python {path}/make_test_slabs.py'.format(path=path)).split())
     catkit.hub.organize.main(options)
 
+def test_file_organization_module_collect_only():
+    options = Struct(**{
+        'adsorbates': 'O,H2',
+        'foldername': '{path}/unorganized'.format(path=path),
+        'verbose': True,
+        'include_pattern': '.',
+        'exclude_pattern': '%%$^#$',
+        'facet_name': '111',
+        'max_density_gas': 0.002,
+        'max_density_slab': 0.06,
+        'exclude_reference': '',
+        'max_energy': 10,
+        'keep_all_energies': False,
+
+    })
+
+    subprocess.call(
+        ('python {path}/make_test_slabs.py'.format(path=path)).split())
+    catkit.hub.organize.collect_structures(options.foldername, options)
+
 
 if __name__ == '__main__':
     test_file_organization()