from __future__ import print_function import collections import optparse import sys from random import randint import ase.io from ase.db import connect from ase.db.core import convert_str_to_int_float_or_str from ase.db.summary import Summary from ase.db.table import Table, all_columns from ase.calculators.calculator import get_calculator from ase.utils import plural, basestring try: input = raw_input # Python 2+3 compatibility except NameError: pass description = """Selecton is a comma-separated list of selections where each selection is of the type "ID", "key" or "key=value". Instead of "=", one can also use "<", "<=", ">=", ">" and "!=" (these must be protected from the shell by using quotes). Special keys: id, user, calculator, age, natoms, energy, magmom, and charge. Chemical symbols can also be used to select number of specific atomic species (H, He, Li, ...).""" examples = ['calculator=nwchem', 'age<1d', 'natoms=1', 'user=alice', '2.2=10'] def main(args=sys.argv[1:]): if isinstance(args, str): args = args.split(' ') parser = optparse.OptionParser( usage='Usage: %prog db-name [selection] [options]', description=description, epilog='Selection examples: ' + ', '.join(examples) + '.') add = parser.add_option add('-v', '--verbose', action='store_true', default=False) add('-q', '--quiet', action='store_true', default=False) add('-n', '--count', action='store_true', help='Count number of selected rows.') add('-l', '--long', action='store_true', help='Long description of selected row') add('-i', '--insert-into', metavar='db-name', help='Insert selected rows into another database.') add('-a', '--add-from-file', metavar='[type:]filename', help='Add results from file.') add('-k', '--add-key-value-pairs', metavar='key1=val1,key2=val2,...', help='Add key-value pairs to selected rows. Values must be numbers ' 'or strings and keys must follow the same rules as keywords.') add('-L', '--limit', type=int, default=20, metavar='N', help='Show only first N rows (default is 20 rows). Use --limit=0 ' 'to show all.') add('--offset', type=int, default=0, metavar='N', help='Skip first N rows. By default, no rows are skipped') add('--delete', action='store_true', help='Delete selected rows.') add('--delete-keys', metavar='key1,key2,...', help='Delete keys for selected rows.') add('-y', '--yes', action='store_true', help='Say yes.') add('--explain', action='store_true', help='Explain query plan.') add('-c', '--columns', metavar='col1,col2,...', help='Specify columns to show. Precede the column specification ' 'with a "+" in order to add columns to the default set of columns. ' 'Precede by a "-" to remove columns. Use "++" for all.') add('-s', '--sort', metavar='column', default='id', help='Sort rows using column. Use -column for a descendin sort. ' 'Default is to sort after id.') add('--cut', type=int, default=35, help='Cut keywords and key-value ' 'columns after CUT characters. Use --cut=0 to disable cutting. ' 'Default is 35 characters') add('-p', '--plot', metavar='[a,b:]x,y1,y2,...', help='Example: "-p x,y": plot y row against x row. Use ' '"-p a:x,y" to make a plot for each value of a.') add('-P', '--plot-data', metavar='name', help="Show plot from data['name'] from the selected row.") add('--csv', action='store_true', help='Write comma-separated-values file.') add('-w', '--open-web-browser', action='store_true', help='Open results in web-browser.') add('--no-lock-file', action='store_true', help="Don't use lock-files") add('--analyse', action='store_true', help='Gathers statistics about tables and indices to help make ' 'better query planning choices.') add('-j', '--json', action='store_true', help='Write json representation of selected row.') add('--unique', action='store_true', help='Give rows a new unique id when using --insert-into.') opts, args = parser.parse_args(args) if not args: parser.error('No database given') verbosity = 1 - opts.quiet + opts.verbose try: run(opts, args, verbosity) except Exception as x: if verbosity < 2: print('{0}: {1}'.format(x.__class__.__name__, x)) sys.exit(1) else: raise def run(opts, args, verbosity): filename = args.pop(0) query = ','.join(args) if query.isdigit(): query = int(query) add_key_value_pairs = {} if opts.add_key_value_pairs: for pair in opts.add_key_value_pairs.split(','): key, value = pair.split('=') add_key_value_pairs[key] = convert_str_to_int_float_or_str(value) if opts.delete_keys: delete_keys = opts.delete_keys.split(',') else: delete_keys = [] con = connect(filename, use_lock_file=not opts.no_lock_file) def out(*args): if verbosity > 0: print(*args) if opts.analyse: con.analyse() return if opts.add_from_file: filename = opts.add_from_file if ':' in filename: calculator_name, filename = filename.split(':') atoms = get_calculator(calculator_name)(filename).get_atoms() else: atoms = ase.io.read(filename) con.write(atoms, key_value_pairs=add_key_value_pairs) out('Added {0} from {1}'.format(atoms.get_chemical_formula(), filename)) return if opts.count: n = con.count(query) print('%s' % plural(n, 'row')) return if opts.explain: for dct in con.select(query, explain=True, verbosity=verbosity, limit=opts.limit, offset=opts.offset): print(dct['explain']) return if opts.insert_into: nkvp = 0 nrows = 0 with connect(opts.insert_into, use_lock_file=not opts.no_lock_file) as con2: for dct in con.select(query): kvp = dct.get('key_value_pairs', {}) nkvp -= len(kvp) kvp.update(add_key_value_pairs) nkvp += len(kvp) if opts.unique: dct['unique_id'] = '%x' % randint(16**31, 16**32 - 1) con2.write(dct, data=dct.get('data'), **kvp) nrows += 1 out('Added %s (%s updated)' % (plural(nkvp, 'key-value pair'), plural(len(add_key_value_pairs) * nrows - nkvp, 'pair'))) out('Inserted %s' % plural(nrows, 'row')) return if add_key_value_pairs or delete_keys: ids = [dct['id'] for dct in con.select(query)] m, n = con.update(ids, delete_keys, **add_key_value_pairs) out('Added %s (%s updated)' % (plural(m, 'key-value pair'), plural(len(add_key_value_pairs) * len(ids) - m, 'pair'))) out('Removed', plural(n, 'key-value pair')) return if opts.delete: ids = [dct['id'] for dct in con.select(query)] if ids and not opts.yes: msg = 'Delete %s? (yes/No): ' % plural(len(ids), 'row') if input(msg).lower() != 'yes': return con.delete(ids) out('Deleted %s' % plural(len(ids), 'row')) return if opts.plot_data: from ase.db.plot import dct2plot dct2plot(con.get(query).data, opts.plot_data) return if opts.plot: if ':' in opts.plot: tags, keys = opts.plot.split(':') tags = tags.split(',') else: tags = [] keys = opts.plot keys = keys.split(',') plots = collections.defaultdict(list) X = {} labels = [] for row in con.select(query, sort=opts.sort): name = ','.join(str(row[tag]) for tag in tags) x = row.get(keys[0]) if x is not None: if isinstance(x, basestring): if x not in X: X[x] = len(X) labels.append(x) x = X[x] plots[name].append([x] + [row.get(key) for key in keys[1:]]) import matplotlib.pyplot as plt for name, plot in plots.items(): xyy = zip(*plot) x = xyy[0] for y, key in zip(xyy[1:], keys[1:]): plt.plot(x, y, label=name + ':' + key) if X: plt.xticks(range(len(labels)), labels, rotation=90) plt.legend() plt.show() return if opts.long: dct = con.get(query) summary = Summary(dct) summary.write() elif opts.json: dct = con.get(query) con2 = connect(sys.stdout, 'json', use_lock_file=False) kvp = dct.get('key_value_pairs', {}) con2.write(dct, data=dct.get('data'), **kvp) else: if opts.open_web_browser: import ase.db.app as app app.db = con app.app.run(host='0.0.0.0', debug=True) else: columns = list(all_columns) c = opts.columns if c and c.startswith('++'): keys = set() for row in con.select(query, limit=opts.limit, offset=opts.offset): keys.update(row._keys) columns.extend(keys) if c[2:3] == ',': c = c[3:] else: c = '' if c: if c[0] == '+': c = c[1:] elif c[0] != '-': columns = [] for col in c.split(','): if col[0] == '-': columns.remove(col[1:]) else: columns.append(col.lstrip('+')) table = Table(con, verbosity, opts.cut) table.select(query, columns, opts.sort, opts.limit, opts.offset) if opts.csv: table.write_csv() else: table.write(query)