# Data insertion

## Add structures
To add all the basic information of an ase.Atoms object (pbc, cell, numbers,
positions) you can use:
```python
from duckatoms import DB
# instantiate a database object

db = DB('<path_to_db_file>/<filename>.duckdb')
structure_ids = db.add_structures([atoms1, atoms2, ..])
```
The `add_structures`-method returns the automatically generated unique *id(s)* of
the inserted structures. If you want to add metadata to all of the added
structures, you can pass a dictionary, which has the structures-table
column names as keys and the corresponding values as values. For
example if all of the structures had the same Miller indices, you could use:
```python
structure_ids = db.add_structures(atoms, batch_metadata={'miller': '111'})
```
which sets the miller column value to '111' for all added structures.
Keep in mind that the column must be defined beforehand. If your atoms
objects contain additional information, which you'd like to store and/or
retrieve from the database, you can define custom converters, as shown in the
last section of this page.

## Add or update custom properties
If a property column exists, but a corresponding value has not been
specified, its value is initialized to NULL. To update this or an already
existing value you can use the following methods.

### Add custom structure properties
```python
# values_dict maps structure_id to property value
db.update_structure_properties(property_name='coordination_number', values_dict={1: 12, 2: 8})
```

### Add custom atoms properties
```python
# values_dict maps atom_id to property value
db.update_atom_properties(property_name='charge', values_dict={1: -0.5, 2: 0.3, 3: -0.2})
```
If you want to update the property values of all atoms of a single (or multiple
structures), you can also pass a list with the corresponding values for each
*structure_id*. The order of the list must correspond to the order of the
atom indices in the original ase.Atoms object.
```python
# structures_dict maps structure_id to list of values (one per atom in structure)
db.update_atom_properties_by_structure_id(property_name='coordination_number', structures_dict={1: [12, 8, 6], 2: [10, 9, 7]})
```

## Define custom converter
### CustomAtomsToDB converter
To define a custom converter for adding structures to the database, you have
to define a class that implements the *extract_structure_data* and
*extract_atom_data* functions of the abstract *AtomsToDB* class. Those define
how the data from a single atoms object is obtained and where it is inserted. The
example shows how to define the extraction of a label and the velocities of an
atoms object.
```python
from ase import Atoms
from duckatoms.converter import AtomsToDB

class CustomAtomsToDB(AtomsToDB):
    def extract_structure_data(self, atoms: Atoms) -> dict:
        """Extracts the label from an ase atoms object and passes it to be added to the `label` column of the `structures` table."""
        return {'label': atoms.info['label']}

    def extract_atom_data(self, atoms: Atoms) -> dict:
        """Extracts the velocities from ase atoms object and passes it to be added to the `velocity` column of the `atoms` table.

        Note: If the passed value is a 2D array, it has to be converted to a
        list."""
        velocities = atoms.get_velocities() # Shape: (n_atoms of one structure, 3)
        return {'velocity': velocities.tolist()}

# set the write-converter to the db instance
db.writer = CustomAtomsToDB()
```
### CustomDBToAtoms converter
Similar to the atoms insertion, you can define a custom converter to retrieve
atoms objects from the database. For that you need to define which property
values should be extracted and how one corresponding atoms object is
constructed. The example constructs atoms objects with the saved values for:
pbc, cell, label, positions, numbers, and velocities.
```python
import numpy as np
from ase import Atoms
from duckatoms.converter import DBToAtoms

class CustomDBToAtoms(DBToAtoms):

    def get_additional_structures_columns(self) -> list[str]:
        """Defines that the value of the `label` column of the `structures`
        table is made available in the values-dictionary."""
        return ['label']

    def get_additional_atoms_columns(self) -> list[str]:
        """Defines that the value of the `velocity` column of the `atoms`
        table is made available in the values-dictionary."""
        return ["velocity"]

    def convert_single_structure(self, **values) -> Atoms:
        """Create Atoms object from extracted data."""
        # Combine x, y, z into positions array
        positions = np.stack([values["x"], values["y"], values["z"]], axis=-1)
        velocities = np.array(values["velocity"])

        # Create and return Atoms object
        atoms = Atoms(
            numbers=values["number"],
            positions=positions,
            cell=values["cell"],
            pbc=values["pbc"],
            velocities=velocities
        )
        atoms.info = {"label": values["label"]}

        return atoms

# set the read-converter to the db instance
db.reader = CustomDBToAtoms()
```