def _write(self, atoms, key_value_pairs, data, id):
Database._write(self, atoms, key_value_pairs, data)
bigdct = {}
ids = []
nextid = 1
if (isinstance(self.filename, basestring)
and os.path.isfile(self.filename)):
try:
bigdct, ids, nextid = self._read_json()
except (SyntaxError, ValueError):
pass
mtime = now()
if isinstance(atoms, AtomsRow):
row = atoms
else:
row = AtomsRow(atoms)
row.ctime = mtime
row.user = os.getenv('USER')
dct = {}
for key in row.__dict__:
if key[0] == '_' or key in row._keys or key == 'id':
continue
dct[key] = row[key]
dct['mtime'] = mtime
if key_value_pairs:
dct['key_value_pairs'] = key_value_pairs
if data:
dct['data'] = data
constraints = row.get('constraints')
if constraints:
dct['constraints'] = constraints
if id is None:
id = nextid
ids.append(id)
nextid += 1
else:
assert id in bigdct
bigdct[id] = dct
self._write_json(bigdct, ids, nextid)
return id
def _write(self, atoms, key_value_pairs, data, id):
ext_tables = key_value_pairs.pop("external_tables", {})
Database._write(self, atoms, key_value_pairs, data)
mtime = now()
encode = self.encode
blob = self.blob
if not isinstance(atoms, AtomsRow):
row = AtomsRow(atoms)
row.ctime = mtime
row.user = os.getenv('USER')
else:
row = atoms
# Extract the external tables from AtomsRow
names = self._get_external_table_names()
for name in names:
new_table = row.get(name, {})
if new_table:
ext_tables[name] = new_table
if not id and not key_value_pairs and not ext_tables:
key_value_pairs = row.key_value_pairs
for k, v in ext_tables.items():
dtype = self._guess_type(v)
self._create_table_if_not_exists(k, dtype)
constraints = row._constraints
if constraints:
if isinstance(constraints, list):
constraints = encode(constraints)
else:
constraints = None
values = (row.unique_id, row.ctime, mtime, row.user, blob(row.numbers),
blob(row.positions), blob(row.cell),
int(np.dot(row.pbc,
[1, 2, 4])), blob(row.get('initial_magmoms')),
blob(row.get('initial_charges')), blob(row.get('masses')),
blob(row.get('tags')), blob(row.get('momenta')), constraints)
if 'calculator' in row:
values += (row.calculator, encode(row.calculator_parameters))
else:
values += (None, None)
if not data:
data = row._data
with self.managed_connection() as con:
if not isinstance(data, (str, bytes)):
data = encode(data, binary=self.version >= 9)
values += (row.get('energy'), row.get('free_energy'),
blob(row.get('forces')), blob(row.get('stress')),
blob(row.get('dipole')), blob(row.get('magmoms')),
row.get('magmom'), blob(row.get('charges')),
encode(key_value_pairs), data, len(row.numbers),
float_if_not_none(row.get('fmax')),
float_if_not_none(row.get('smax')),
float_if_not_none(row.get('volume')), float(row.mass),
float(row.charge))
cur = con.cursor()
if id is None:
q = self.default + ', ' + ', '.join('?' * len(values))
cur.execute('INSERT INTO systems VALUES ({})'.format(q),
values)
id = self.get_last_id(cur)
else:
self._delete(cur, [id], [
'keys', 'text_key_values', 'number_key_values', 'species'
])
q = ', '.join(name + '=?' for name in self.columnnames[1:])
cur.execute('UPDATE systems SET {} WHERE id=?'.format(q),
values + (id, ))
count = row.count_atoms()
if count:
species = [(atomic_numbers[symbol], n, id)
for symbol, n in count.items()]
cur.executemany('INSERT INTO species VALUES (?, ?, ?)',
species)
text_key_values = []
number_key_values = []
for key, value in key_value_pairs.items():
if isinstance(value, (numbers.Real, np.bool_)):
number_key_values.append([key, float(value), id])
else:
assert isinstance(value, str)
text_key_values.append([key, value, id])
cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)',
text_key_values)
cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)',
number_key_values)
cur.executemany('INSERT INTO keys VALUES (?, ?)',
[(key, id) for key in key_value_pairs])
# Insert entries in the valid tables
for tabname in ext_tables.keys():
entries = ext_tables[tabname]
entries['id'] = id
self._insert_in_external_table(cur,
name=tabname,
entries=ext_tables[tabname])
return id
def _write(self, atoms, key_value_pairs, data, id):
Database._write(self, atoms, key_value_pairs, data)
encode = self.encode
con = self.connection or self._connect()
self._initialize(con)
cur = con.cursor()
mtime = now()
blob = self.blob
text_key_values = []
number_key_values = []
if not isinstance(atoms, AtomsRow):
row = AtomsRow(atoms)
row.ctime = mtime
row.user = os.getenv('USER')
else:
row = atoms
if id:
self._delete(cur, [id], ['keys', 'text_key_values',
'number_key_values', 'species'])
else:
if not key_value_pairs:
key_value_pairs = row.key_value_pairs
constraints = row._constraints
if constraints:
if isinstance(constraints, list):
constraints = encode(constraints)
else:
constraints = None
values = (row.unique_id,
row.ctime,
mtime,
row.user,
blob(row.numbers),
blob(row.positions),
blob(row.cell),
int(np.dot(row.pbc, [1, 2, 4])),
blob(row.get('initial_magmoms')),
blob(row.get('initial_charges')),
blob(row.get('masses')),
blob(row.get('tags')),
blob(row.get('momenta')),
constraints)
if 'calculator' in row:
values += (row.calculator, encode(row.calculator_parameters))
else:
values += (None, None)
if not data:
data = row._data
if not isinstance(data, basestring):
data = encode(data)
values += (row.get('energy'),
row.get('free_energy'),
blob(row.get('forces')),
blob(row.get('stress')),
blob(row.get('dipole')),
blob(row.get('magmoms')),
row.get('magmom'),
blob(row.get('charges')),
encode(key_value_pairs),
data,
len(row.numbers),
float_if_not_none(row.get('fmax')),
float_if_not_none(row.get('smax')),
float_if_not_none(row.get('volume')),
float(row.mass),
float(row.charge))
if id is None:
q = self.default + ', ' + ', '.join('?' * len(values))
cur.execute('INSERT INTO systems VALUES ({})'.format(q),
values)
id = self.get_last_id(cur)
else:
q = ', '.join(name + '=?' for name in self.columnnames[1:])
cur.execute('UPDATE systems SET {} WHERE id=?'.format(q),
values + (id,))
count = row.count_atoms()
if count:
species = [(atomic_numbers[symbol], n, id)
for symbol, n in count.items()]
cur.executemany('INSERT INTO species VALUES (?, ?, ?)',
species)
text_key_values = []
number_key_values = []
for key, value in key_value_pairs.items():
if isinstance(value, (numbers.Real, np.bool_)):
number_key_values.append([key, float(value), id])
else:
assert isinstance(value, basestring)
text_key_values.append([key, value, id])
cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)',
text_key_values)
cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)',
number_key_values)
cur.executemany('INSERT INTO keys VALUES (?, ?)',
[(key, id) for key in key_value_pairs])
if self.connection is None:
con.commit()
con.close()
return id
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
def _write(self, atoms, key_value_pairs, data, id):
ext_tables = key_value_pairs.pop("external_tables", {})
Database._write(self, atoms, key_value_pairs, data)
encode = self.encode
con = self.connection or self._connect()
self._initialize(con)
cur = con.cursor()
mtime = now()
blob = self.blob
text_key_values = []
number_key_values = []
if not isinstance(atoms, AtomsRow):
row = AtomsRow(atoms)
row.ctime = mtime
row.user = os.getenv('USER')
else:
row = atoms
# Extract the external tables from AtomsRow
names = self._get_external_table_names(db_con=con)
for name in names:
new_table = row.get(name, {})
if new_table:
ext_tables[name] = new_table
if id:
self._delete(
cur, [id],
['keys', 'text_key_values', 'number_key_values', 'species'])
else:
if not key_value_pairs:
key_value_pairs = row.key_value_pairs
constraints = row._constraints
if constraints:
if isinstance(constraints, list):
constraints = encode(constraints)
else:
constraints = None
values = (row.unique_id, row.ctime, mtime, row.user, blob(row.numbers),
blob(row.positions), blob(row.cell),
int(np.dot(row.pbc,
[1, 2, 4])), blob(row.get('initial_magmoms')),
blob(row.get('initial_charges')), blob(row.get('masses')),
blob(row.get('tags')), blob(row.get('momenta')), constraints)
if 'calculator' in row:
values += (row.calculator, encode(row.calculator_parameters))
else:
values += (None, None)
if not data:
data = row._data
if not isinstance(data, basestring):
data = encode(data)
values += (row.get('energy'), row.get('free_energy'),
blob(row.get('forces')), blob(row.get('stress')),
blob(row.get('dipole')), blob(row.get('magmoms')),
row.get('magmom'), blob(row.get('charges')),
encode(key_value_pairs), data, len(row.numbers),
float_if_not_none(row.get('fmax')),
float_if_not_none(row.get('smax')),
float_if_not_none(row.get('volume')), float(row.mass),
float(row.charge))
if id is None:
q = self.default + ', ' + ', '.join('?' * len(values))
cur.execute('INSERT INTO systems VALUES ({})'.format(q), values)
id = self.get_last_id(cur)
else:
q = ', '.join(name + '=?' for name in self.columnnames[1:])
cur.execute('UPDATE systems SET {} WHERE id=?'.format(q),
values + (id, ))
count = row.count_atoms()
if count:
species = [(atomic_numbers[symbol], n, id)
for symbol, n in count.items()]
cur.executemany('INSERT INTO species VALUES (?, ?, ?)', species)
text_key_values = []
number_key_values = []
for key, value in key_value_pairs.items():
if isinstance(value, (numbers.Real, np.bool_)):
number_key_values.append([key, float(value), id])
else:
assert isinstance(value, basestring)
text_key_values.append([key, value, id])
cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)',
text_key_values)
cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)',
number_key_values)
cur.executemany('INSERT INTO keys VALUES (?, ?)',
[(key, id) for key in key_value_pairs])
# Update external tables
valid_entries = []
for k, v in ext_tables.items():
try:
# Guess the type of the value
dtype = self._guess_type(v)
self._create_table_if_not_exists(k, dtype, db_con=con)
v["id"] = id
valid_entries.append(k)
except ValueError as exc:
# Close the connection without committing
if self.connection is None:
con.close()
# Raise error again
raise ValueError(exc)
# Insert entries in the valid tables
for tabname in valid_entries:
try:
self._insert_in_external_table(cur,
name=tabname,
entries=ext_tables[tabname])
except ValueError as exc:
# Close the connection without committing
if self.connection is None:
con.close()
# Raise the error again
raise ValueError(exc)
if self.connection is None:
con.commit()
con.close()
return id
