Esempio n. 1
0
    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
Esempio n. 2
0
    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
Esempio n. 3
0
    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
Esempio n. 4
0
    def _write(self, atoms, key_value_pairs, data):
        Database._write(self, atoms, key_value_pairs, data)

        con = self.connection or self._connect()
        self._initialize(con)
        cur = con.cursor()

        id = None

        if not isinstance(atoms, AtomsRow):
            row = AtomsRow(atoms)
            row.ctime = mtime = now()
            row.user = os.getenv('USER')
        else:
            row = atoms

            cur.execute('SELECT id FROM systems WHERE unique_id=?',
                        (row.unique_id,))
            results = cur.fetchall()
            if results:
                id = results[0][0]
                self._delete(cur, [id], ['keys', 'text_key_values',
                                         'number_key_values'])
            mtime = now()

        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:
            if not isinstance(row.calculator_parameters, basestring):
                row.calculator_parameters = encode(row.calculator_parameters)
            values += (row.calculator,
                       row.calculator_parameters)
        else:
            values += (None, None)

        if key_value_pairs is None:
            key_value_pairs = row.key_value_pairs

        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 ({0})'.format(q),
                        values)
        else:
            q = ', '.join(line.split()[0].lstrip() + '=?'
                          for line in init_statements[0].splitlines()[2:])
            cur.execute('UPDATE systems SET {0} WHERE id=?'.format(q),
                        values + (id,))

        if id is None:
            id = self.get_last_id(cur)

            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, (float, int)):
                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