Source code for ase.db.core

import functools
import json
import numbers
import operator
import os
import re
import warnings
from time import time
from typing import List, Dict, Any

import numpy as np

from ase.atoms import Atoms
from ase.calculators.calculator import all_properties, all_changes
from ase.data import atomic_numbers
from ase.db.row import AtomsRow
from ase.formula import Formula
from ase.io.jsonio import create_ase_object
from ase.parallel import world, DummyMPI, parallel_function, parallel_generator
from ase.utils import Lock, PurePath


T2000 = 946681200.0  # January 1. 2000
YEAR = 31557600.0  # 365.25 days


# Format of key description: ('short', 'long', 'unit')
default_key_descriptions = {
    'id': ('ID', 'Uniqe row ID', ''),
    'age': ('Age', 'Time since creation', ''),
    'formula': ('Formula', 'Chemical formula', ''),
    'pbc': ('PBC', 'Periodic boundary conditions', ''),
    'user': ('Username', '', ''),
    'calculator': ('Calculator', 'ASE-calculator name', ''),
    'energy': ('Energy', 'Total energy', 'eV'),
    'natoms': ('Number of atoms', '', ''),
    'fmax': ('Maximum force', '', 'eV/Ang'),
    'smax': ('Maximum stress', 'Maximum stress on unit cell',
             '`\\text{eV/Ang}^3`'),
    'charge': ('Charge', 'Net charge in unit cell', '|e|'),
    'mass': ('Mass', 'Sum of atomic masses in unit cell', 'au'),
    'magmom': ('Magnetic moment', '', 'au'),
    'unique_id': ('Unique ID', 'Random (unique) ID', ''),
    'volume': ('Volume', 'Volume of unit cell', '`\\text{Ang}^3`')}


def now():
    """Return time since January 1. 2000 in years."""
    return (time() - T2000) / YEAR


seconds = {'s': 1,
           'm': 60,
           'h': 3600,
           'd': 86400,
           'w': 604800,
           'M': 2629800,
           'y': YEAR}

longwords = {'s': 'second',
             'm': 'minute',
             'h': 'hour',
             'd': 'day',
             'w': 'week',
             'M': 'month',
             'y': 'year'}

ops = {'<': operator.lt,
       '<=': operator.le,
       '=': operator.eq,
       '>=': operator.ge,
       '>': operator.gt,
       '!=': operator.ne}

invop = {'<': '>=', '<=': '>', '>=': '<', '>': '<=', '=': '!=', '!=': '='}

word = re.compile('[_a-zA-Z][_0-9a-zA-Z]*$')

reserved_keys = set(all_properties +
                    all_changes +
                    list(atomic_numbers) +
                    ['id', 'unique_id', 'ctime', 'mtime', 'user',
                     'fmax', 'smax',
                     'momenta', 'constraints', 'natoms', 'formula', 'age',
                     'calculator', 'calculator_parameters',
                     'key_value_pairs', 'data'])

numeric_keys = set(['id', 'energy', 'magmom', 'charge', 'natoms'])


def check(key_value_pairs):
    for key, value in key_value_pairs.items():
        if key == "external_tables":
            # Checks for external_tables are not
            # performed
            continue

        if not word.match(key) or key in reserved_keys:
            raise ValueError('Bad key: {}'.format(key))
        try:
            Formula(key, strict=True)
        except ValueError:
            pass
        else:
            warnings.warn(
                'It is best not to use keys ({0}) that are also a '
                'chemical formula.  If you do a "db.select({0!r})",'
                'you will not find rows with your key.  Instead, you wil get '
                'rows containing the atoms in the formula!'.format(key))
        if not isinstance(value, (numbers.Real, str, np.bool_)):
            raise ValueError('Bad value for {!r}: {}'.format(key, value))
        if isinstance(value, str):
            for t in [int, float]:
                if str_represents(value, t):
                    raise ValueError(
                        'Value ' + value + ' is put in as string ' +
                        'but can be interpreted as ' +
                        '{}! Please convert '.format(t.__name__) +
                        'to {} using '.format(t.__name__) +
                        '{}(value) before '.format(t.__name__) +
                        'writing to the database OR change ' +
                        'to a different string.')


def str_represents(value, t=int):
    try:
        t(value)
    except ValueError:
        return False
    return True


[docs]def connect(name, type='extract_from_name', create_indices=True, use_lock_file=True, append=True, serial=False): """Create connection to database. name: str Filename or address of database. type: str One of 'json', 'db', 'postgresql', (JSON, SQLite, PostgreSQL). Default is 'extract_from_name', which will guess the type from the name. use_lock_file: bool You can turn this off if you know what you are doing ... append: bool Use append=False to start a new database. """ if isinstance(name, PurePath): name = str(name) if type == 'extract_from_name': if name is None: type = None elif not isinstance(name, str): type = 'json' elif (name.startswith('postgresql://') or name.startswith('postgres://')): type = 'postgresql' elif name.startswith('mysql://') or name.startswith('mariadb://'): type = 'mysql' else: type = os.path.splitext(name)[1][1:] if type == '': raise ValueError('No file extension or database type given') if type is None: return Database() if not append and world.rank == 0: if isinstance(name, str) and os.path.isfile(name): os.remove(name) if type not in ['postgresql', 'mysql'] and isinstance(name, str): name = os.path.abspath(name) if type == 'json': from ase.db.jsondb import JSONDatabase return JSONDatabase(name, use_lock_file=use_lock_file, serial=serial) if type == 'db': from ase.db.sqlite import SQLite3Database return SQLite3Database(name, create_indices, use_lock_file, serial=serial) if type == 'postgresql': from ase.db.postgresql import PostgreSQLDatabase return PostgreSQLDatabase(name) if type == 'mysql': from ase.db.mysql import MySQLDatabase return MySQLDatabase(name) raise ValueError('Unknown database type: ' + type)
def lock(method): """Decorator for using a lock-file.""" @functools.wraps(method) def new_method(self, *args, **kwargs): if self.lock is None: return method(self, *args, **kwargs) else: with self.lock: return method(self, *args, **kwargs) return new_method def convert_str_to_int_float_or_str(value): """Safe eval()""" try: return int(value) except ValueError: try: value = float(value) except ValueError: value = {'True': True, 'False': False}.get(value, value) return value def parse_selection(selection, **kwargs): if selection is None or selection == '': expressions = [] elif isinstance(selection, int): expressions = [('id', '=', selection)] elif isinstance(selection, list): expressions = selection else: expressions = [w.strip() for w in selection.split(',')] keys = [] comparisons = [] for expression in expressions: if isinstance(expression, (list, tuple)): comparisons.append(expression) continue if expression.count('<') == 2: value, expression = expression.split('<', 1) if expression[0] == '=': op = '>=' expression = expression[1:] else: op = '>' key = expression.split('<', 1)[0] comparisons.append((key, op, value)) for op in ['!=', '<=', '>=', '<', '>', '=']: if op in expression: break else: # no break if expression in atomic_numbers: comparisons.append((expression, '>', 0)) else: try: count = Formula(expression).count() except ValueError: keys.append(expression) else: comparisons.extend((symbol, '>', n - 1) for symbol, n in count.items()) continue key, value = expression.split(op) comparisons.append((key, op, value)) cmps = [] for key, value in kwargs.items(): comparisons.append((key, '=', value)) for key, op, value in comparisons: if key == 'age': key = 'ctime' op = invop[op] value = now() - time_string_to_float(value) elif key == 'formula': if op != '=': raise ValueError('Use fomula=...') f = Formula(value) count = f.count() cmps.extend((atomic_numbers[symbol], '=', n) for symbol, n in count.items()) key = 'natoms' value = len(f) elif key in atomic_numbers: key = atomic_numbers[key] value = int(value) elif isinstance(value, str): value = convert_str_to_int_float_or_str(value) if key in numeric_keys and not isinstance(value, (int, float)): msg = 'Wrong type for "{}{}{}" - must be a number' raise ValueError(msg.format(key, op, value)) cmps.append((key, op, value)) return keys, cmps
[docs]class Database: """Base class for all databases.""" def __init__(self, filename=None, create_indices=True, use_lock_file=False, serial=False): """Database object. serial: bool Let someone else handle parallelization. Default behavior is to interact with the database on the master only and then distribute results to all slaves. """ if isinstance(filename, str): filename = os.path.expanduser(filename) self.filename = filename self.create_indices = create_indices if use_lock_file and isinstance(filename, str): self.lock = Lock(filename + '.lock', world=DummyMPI()) else: self.lock = None self.serial = serial # Decription of columns and other stuff: self._metadata: Dict[str, Any] = None @property def metadata(self) -> Dict[str, Any]: raise NotImplementedError
[docs] @parallel_function @lock def write(self, atoms, key_value_pairs={}, data={}, id=None, **kwargs): """Write atoms to database with key-value pairs. atoms: Atoms object Write atomic numbers, positions, unit cell and boundary conditions. If a calculator is attached, write also already calculated properties such as the energy and forces. key_value_pairs: dict Dictionary of key-value pairs. Values must be strings or numbers. data: dict Extra stuff (not for searching). id: int Overwrite existing row. Key-value pairs can also be set using keyword arguments:: connection.write(atoms, name='ABC', frequency=42.0) Returns integer id of the new row. """ if atoms is None: atoms = Atoms() kvp = dict(key_value_pairs) # modify a copy kvp.update(kwargs) id = self._write(atoms, kvp, data, id) return id
def _write(self, atoms, key_value_pairs, data, id=None): check(key_value_pairs) return 1
[docs] @parallel_function @lock def reserve(self, **key_value_pairs): """Write empty row if not already present. Usage:: id = conn.reserve(key1=value1, key2=value2, ...) Write an empty row with the given key-value pairs and return the integer id. If such a row already exists, don't write anything and return None. """ for dct in self._select([], [(key, '=', value) for key, value in key_value_pairs.items()]): return None atoms = Atoms() calc_name = key_value_pairs.pop('calculator', None) if calc_name: # Allow use of calculator key assert calc_name.lower() == calc_name # Fake calculator class: class Fake: name = calc_name def todict(self): return {} def check_state(self, atoms): return ['positions'] atoms.calc = Fake() id = self._write(atoms, key_value_pairs, {}, None) return id
def __delitem__(self, id): self.delete([id])
[docs] def get_atoms(self, selection=None, attach_calculator=False, add_additional_information=False, **kwargs): """Get Atoms object. selection: int, str or list See the select() method. attach_calculator: bool Attach calculator object to Atoms object (default value is False). add_additional_information: bool Put key-value pairs and data into Atoms.info dictionary. In addition, one can use keyword arguments to select specific key-value pairs. """ row = self.get(selection, **kwargs) return row.toatoms(attach_calculator, add_additional_information)
def __getitem__(self, selection): return self.get(selection)
[docs] def get(self, selection=None, **kwargs): """Select a single row and return it as a dictionary. selection: int, str or list See the select() method. """ rows = list(self.select(selection, limit=2, **kwargs)) if not rows: raise KeyError('no match') assert len(rows) == 1, 'more than one row matched' return rows[0]
[docs] @parallel_generator def select(self, selection=None, filter=None, explain=False, verbosity=1, limit=None, offset=0, sort=None, include_data=True, columns='all', **kwargs): """Select rows. Return AtomsRow iterator with results. Selection is done using key-value pairs and the special keys: formula, age, user, calculator, natoms, energy, magmom and/or charge. selection: int, str or list Can be: * an integer id * a string like 'key=value', where '=' can also be one of '<=', '<', '>', '>=' or '!='. * a string like 'key' * comma separated strings like 'key1<value1,key2=value2,key' * list of strings or tuples: [('charge', '=', 1)]. filter: function A function that takes as input a row and returns True or False. explain: bool Explain query plan. verbosity: int Possible values: 0, 1 or 2. limit: int or None Limit selection. offset: int Offset into selected rows. sort: str Sort rows after key. Prepend with minus sign for a decending sort. include_data: bool Use include_data=False to skip reading data from rows. columns: 'all' or list of str Specify which columns from the SQL table to include. For example, if only the row id and the energy is needed, queries can be speeded up by setting columns=['id', 'energy']. """ if sort: if sort == 'age': sort = '-ctime' elif sort == '-age': sort = 'ctime' elif sort.lstrip('-') == 'user': sort += 'name' keys, cmps = parse_selection(selection, **kwargs) for row in self._select(keys, cmps, explain=explain, verbosity=verbosity, limit=limit, offset=offset, sort=sort, include_data=include_data, columns=columns): if filter is None or filter(row): yield row
[docs] def count(self, selection=None, **kwargs): """Count rows. See the select() method for the selection syntax. Use db.count() or len(db) to count all rows. """ n = 0 for row in self.select(selection, **kwargs): n += 1 return n
def __len__(self): return self.count()
[docs] @parallel_function @lock def update(self, id, atoms=None, delete_keys=[], data=None, **add_key_value_pairs): """Update and/or delete key-value pairs of row(s). id: int ID of row to update. atoms: Atoms object Optionally update the Atoms data (positions, cell, ...). data: dict Data dict to be added to the existing data. delete_keys: list of str Keys to remove. Use keyword arguments to add new key-value pairs. Returns number of key-value pairs added and removed. """ if not isinstance(id, numbers.Integral): if isinstance(id, list): err = ('First argument must be an int and not a list.\n' 'Do something like this instead:\n\n' 'with db:\n' ' for id in ids:\n' ' db.update(id, ...)') raise ValueError(err) raise TypeError('id must be an int') check(add_key_value_pairs) row = self._get_row(id) kvp = row.key_value_pairs n = len(kvp) for key in delete_keys: kvp.pop(key, None) n -= len(kvp) m = -len(kvp) kvp.update(add_key_value_pairs) m += len(kvp) moredata = data data = row.get('data', {}) if moredata: data.update(moredata) if not data: data = None if atoms: oldrow = row row = AtomsRow(atoms) # Copy over data, kvp, ctime, user and id row._data = oldrow._data row.__dict__.update(kvp) row._keys = list(kvp) row.ctime = oldrow.ctime row.user = oldrow.user row.id = id if atoms or os.path.splitext(self.filename)[1] == '.json': self._write(row, kvp, data, row.id) else: self._update(row.id, kvp, data) return m, n
[docs] def delete(self, ids): """Delete rows.""" raise NotImplementedError
def time_string_to_float(s): if isinstance(s, (float, int)): return s s = s.replace(' ', '') if '+' in s: return sum(time_string_to_float(x) for x in s.split('+')) if s[-2].isalpha() and s[-1] == 's': s = s[:-1] i = 1 while s[i].isdigit(): i += 1 return seconds[s[i:]] * int(s[:i]) / YEAR def float_to_time_string(t, long=False): t *= YEAR for s in 'yMwdhms': x = t / seconds[s] if x > 5: break if long: return '{:.3f} {}s'.format(x, longwords[s]) else: return '{:.0f}{}'.format(round(x), s) def object_to_bytes(obj: Any) -> bytes: """Serialize Python object to bytes.""" parts = [b'12345678'] obj = o2b(obj, parts) offset = sum(len(part) for part in parts) x = np.array(offset, np.int64) if not np.little_endian: x.byteswap(True) parts[0] = x.tobytes() parts.append(json.dumps(obj, separators=(',', ':')).encode()) return b''.join(parts) def bytes_to_object(b: bytes) -> Any: """Deserialize bytes to Python object.""" x = np.frombuffer(b[:8], np.int64) if not np.little_endian: x = x.byteswap() offset = x.item() obj = json.loads(b[offset:].decode()) return b2o(obj, b) def o2b(obj: Any, parts: List[bytes]): if isinstance(obj, (int, float, bool, str, type(None))): return obj if isinstance(obj, dict): return {key: o2b(value, parts) for key, value in obj.items()} if isinstance(obj, (list, tuple)): return [o2b(value, parts) for value in obj] if isinstance(obj, np.ndarray): assert obj.dtype != object, \ 'Cannot convert ndarray of type "object" to bytes.' offset = sum(len(part) for part in parts) if not np.little_endian: obj = obj.byteswap() parts.append(obj.tobytes()) return {'__ndarray__': [obj.shape, obj.dtype.name, offset]} if isinstance(obj, complex): return {'__complex__': [obj.real, obj.imag]} objtype = getattr(obj, 'ase_objtype') if objtype: dct = o2b(obj.todict(), parts) dct['__ase_objtype__'] = objtype return dct raise ValueError('Objects of type {type} not allowed' .format(type=type(obj))) def b2o(obj: Any, b: bytes) -> Any: if isinstance(obj, (int, float, bool, str, type(None))): return obj if isinstance(obj, list): return [b2o(value, b) for value in obj] assert isinstance(obj, dict) x = obj.get('__complex__') if x is not None: return complex(*x) x = obj.get('__ndarray__') if x is not None: shape, name, offset = x dtype = np.dtype(name) size = dtype.itemsize * np.prod(shape).astype(int) a = np.frombuffer(b[offset:offset + size], dtype) a.shape = shape if not np.little_endian: a = a.byteswap() return a dct = {key: b2o(value, b) for key, value in obj.items()} objtype = dct.pop('__ase_objtype__', None) if objtype is None: return dct return create_ase_object(objtype, dct)