def test_put_caching(mongo_client):
database_name = '_test_put_caching'
mongo_client.drop_database(database_name)
database = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
jsonizer=stk.MoleculeJsonizer(
key_makers=(
stk.InchiKey(),
),
),
)
molecule = stk.BuildingBlock('CCC')
database.put(molecule)
database.put(molecule)
cache_info = database._put.cache_info()
assert cache_info.hits == 1
assert cache_info.misses == 1
database.put(
molecule=molecule.with_position_matrix(
position_matrix=np.zeros((molecule.get_num_atoms(), 3)),
),
)
cache_info = database._put.cache_info()
assert cache_info.hits == 1
assert cache_info.misses == 2
def test_get_caching(mongo_client):
database_name = '_test_get_caching'
mongo_client.drop_database(database_name)
database = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
jsonizer=stk.MoleculeJsonizer(
key_makers=(
stk.InchiKey(),
),
),
)
molecule = stk.BuildingBlock('CCC')
database.put(molecule)
database.get({
stk.InchiKey().get_key_name():
stk.InchiKey().get_key(molecule),
})
database.get({
stk.InchiKey().get_key_name():
stk.InchiKey().get_key(molecule),
})
cache_info = database._get.cache_info()
assert cache_info.hits == 1
assert cache_info.misses == 1
def add_entries(client, database, key_makers, random_seed):
molecule_db = stk.MoleculeMongoDb(
mongo_client=client,
database=database,
molecule_collection='molecules',
position_matrix_collection='position_matrices',
jsonizer=stk.MoleculeJsonizer(key_makers=key_makers, ),
)
num_atoms_db = stk.ValueMongoDb(
mongo_client=client,
collection='numAtoms',
database=database,
key_makers=key_makers,
)
num_bonds_db = stk.ValueMongoDb(
mongo_client=client,
collection='numBonds',
database=database,
key_makers=key_makers,
)
add_value = True
for molecule in get_molecules(200, 5):
molecule_db.put(molecule)
num_bonds_db.put(molecule, molecule.get_num_bonds())
if add_value:
num_atoms_db.put(molecule, molecule.get_num_atoms())
add_value ^= 1
def molecule_db(mongo_client):
"""
A :class:`.MoleculeDatabase` instance.
"""
return stk.MoleculeMongoDb(
mongo_client=mongo_client,
database='_stk_pytest_database',
jsonizer=stk.MoleculeJsonizer(key_makers=(stk.Smiles(), ), ),
indices=(stk.Smiles().get_key_name(), ),
)
def get_database(
database_name: str,
mongo_client: pymongo.MongoClient,
key_makers: tuple[stk.MoleculeKeyMaker, ...],
indices: tuple[str, ...],
) -> stk.MoleculeMongoDb:
return stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
jsonizer=stk.MoleculeJsonizer(key_makers),
put_lru_cache_size=0,
get_lru_cache_size=0,
indices=indices,
)
def test_update_1(mongo_client):
"""
Test that existing entries are updated.
"""
database_name = '_test_update_1'
mongo_client.drop_database(database_name)
database = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
)
jsonizer = stk.MoleculeJsonizer()
molecule = stk.BuildingBlock('CCC').with_canonical_atom_ordering()
json = jsonizer.to_json(molecule)
database.put(molecule)
assert_database_state(
state1=get_database_state(database),
state2=DatabaseState({
DatabaseEntry(**json['molecule']): 1,
DatabaseEntry(**to_hashable(json['matrix'])): 1,
}),
)
molecule2 = molecule.with_position_matrix(position_matrix=np.zeros(
(molecule.get_num_atoms(), 3)), )
json2 = jsonizer.to_json(molecule2)
database.put(molecule2)
assert_database_state(
state1=get_database_state(database),
state2=DatabaseState({
# Molecule JSON should be unchanged.
DatabaseEntry(**json['molecule']):
1,
# Position matrix should be updated.
DatabaseEntry(**to_hashable(json2['matrix'])):
1,
}),
)
def test_get_all():
"""
Test iteration over all entries.
"""
database_name = '_test_get_entries_molecule'
client = pymongo.MongoClient()
client.drop_database(database_name)
key_maker = stk.Inchi()
jsonizer = stk.MoleculeJsonizer(key_makers=(key_maker, ))
database = stk.MoleculeMongoDb(
mongo_client=client,
database=database_name,
jsonizer=jsonizer,
put_lru_cache_size=0,
get_lru_cache_size=0,
)
molecules = (
stk.BuildingBlock('CCC').with_canonical_atom_ordering(),
stk.BuildingBlock('BrCCCBr').with_canonical_atom_ordering(),
stk.BuildingBlock('NCCN').with_canonical_atom_ordering(),
)
molecules_by_key = {
key_maker.get_key(molecule): molecule
for molecule in molecules
}
for molecule in molecules:
database.put(molecule)
for i, retrieved in enumerate(database.get_all()):
key = key_maker.get_key(retrieved)
molecule = molecules_by_key[key]
is_equivalent_molecule(
molecule1=molecule.with_canonical_atom_ordering(),
molecule2=retrieved.with_canonical_atom_ordering(),
)
# Check number of molecules.
assert i+1 == len(molecules)
'InChIKey':
rdkit.MolToInchiKey(rdkit.MolFromSmiles('BrCCBr'))
},
),
lambda: CaseDataData(
get_database=lambda mongo_client: stk.MoleculeMongoDb(
mongo_client=mongo_client,
database='_stk_test_database_for_testing',
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=stk.MoleculeJsonizer(
key_makers=(
stk.MoleculeKeyMaker(
key_name='SMILES',
get_key=lambda molecule:
rdkit.MolToSmiles(
molecule.to_rdkit_mol()
)
),
),
),
),
molecule=stk.BuildingBlock('BrBr'),
key={'SMILES': 'BrBr'},
),
lambda: CaseDataData(
get_database=lambda mongo_client: stk.MoleculeMongoDb(
mongo_client=mongo_client,
database='_stk_test_database_for_testing',
put_lru_cache_size=128,
get_lru_cache_size=128,
import numpy as np
import pytest
import stk
from .case_data import CaseData
@pytest.fixture(
scope='session',
params=(
lambda: CaseData(
jsonizer=stk.MoleculeJsonizer(key_makers=(
stk.Inchi(),
stk.InchiKey(),
)),
molecule=stk.BuildingBlock(smiles='Br[C+2][C+2]Br', ).
with_position_matrix(
np.array([
[0, 0, 0],
[1, 1, 1],
[2, 2, 2],
[3, 3, 3],
],
dtype=np.float64)),
json={
'molecule': {
'a': (
(35, 0),
(6, 2),
(6, 2),
def test_update_2(mongo_client):
"""
Test that existing entries are updated.
In this test, you first create two separate entries, using
different molecule keys. You then update both at the same time,
with a database which uses both molecule keys.
"""
database_name = '_test_update_2'
mongo_client.drop_database(database_name)
jsonizer1 = stk.MoleculeJsonizer()
database1 = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=jsonizer1,
)
jsonizer2 = stk.MoleculeJsonizer(
key_makers=(
stk.Smiles(),
),
)
database2 = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=jsonizer2,
)
jsonizer3 = stk.MoleculeJsonizer(
key_makers=(
stk.InchiKey(),
stk.Smiles(),
),
)
database3 = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=jsonizer3,
)
molecule1 = stk.BuildingBlock('CCC').with_canonical_atom_ordering()
json1 = jsonizer1.to_json(molecule1)
database1.put(molecule1)
assert_database_state(
state1=get_database_state(database1),
state2=DatabaseState({
DatabaseEntry(**json1['molecule']): 1,
DatabaseEntry(**to_hashable(json1['matrix'])): 1,
}),
)
# Should add another entry, as a different key maker is used.
molecule2 = molecule1.with_position_matrix(
position_matrix=np.zeros((molecule1.get_num_atoms(), 3)),
)
json2 = jsonizer2.to_json(molecule2)
database2.put(molecule2)
assert_database_state(
state1=get_database_state(database1),
state2=DatabaseState({
DatabaseEntry(**json1['molecule']): 1,
DatabaseEntry(**to_hashable(json1['matrix'])): 1,
DatabaseEntry(**json2['molecule']): 1,
DatabaseEntry(**to_hashable(json2['matrix'])): 1,
}),
)
# Should update both entries as both key makers are used.
molecule3 = molecule1.with_position_matrix(
position_matrix=np.ones((molecule1.get_num_atoms(), 3)),
)
json3 = jsonizer3.to_json(molecule3)
database3.put(molecule3)
assert_database_state(
state1=get_database_state(database1),
state2=DatabaseState({
DatabaseEntry(**json3['molecule']): 2,
DatabaseEntry(**to_hashable(json3['matrix'])): 2,
}),
)
def test_update_3(mongo_client):
"""
Test that existing entries are updated.
In this test, you first create one entry with two keys. Then
update the entry with databases, each using 1 different key.
No duplicate entries should be made in the database this way.
"""
database_name = '_test_update_3'
mongo_client.drop_database(database_name)
jsonizer1 = stk.MoleculeJsonizer(
key_makers=(
stk.InchiKey(),
stk.Smiles(),
),
)
database1 = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=jsonizer1,
)
jsonizer2 = stk.MoleculeJsonizer()
database2 = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=jsonizer2,
)
jsonizer3 = stk.MoleculeJsonizer(
key_makers=(
stk.Smiles(),
),
)
database3 = stk.MoleculeMongoDb(
mongo_client=mongo_client,
database=database_name,
put_lru_cache_size=0,
get_lru_cache_size=0,
jsonizer=jsonizer3,
)
molecule1 = stk.BuildingBlock('CCC').with_canonical_atom_ordering()
json1 = jsonizer1.to_json(molecule1)
database1.put(molecule1)
assert_database_state(
state1=get_database_state(database1),
state2=DatabaseState({
DatabaseEntry(**json1['molecule']): 1,
DatabaseEntry(**to_hashable(json1['matrix'])): 1,
}),
)
molecule2 = molecule1.with_position_matrix(
position_matrix=np.zeros((molecule1.get_num_atoms(), 3)),
)
json2 = jsonizer2.to_json(molecule2)
json2['matrix'] = dict(json1['matrix'])
json2['matrix']['m'] = jsonizer2.to_json(molecule2)['matrix']['m']
database2.put(molecule2)
assert_database_state(
state1=get_database_state(database1),
state2=DatabaseState({
DatabaseEntry(**json1['molecule']): 1,
DatabaseEntry(**to_hashable(json2['matrix'])): 1,
}),
)
molecule3 = molecule1.with_position_matrix(
position_matrix=np.zeros((molecule1.get_num_atoms(), 3)),
)
json3 = jsonizer3.to_json(molecule3)
json3['matrix'] = dict(json1['matrix'])
json3['matrix']['m'] = jsonizer3.to_json(molecule3)['matrix']['m']
database3.put(molecule3)
assert_database_state(
state1=get_database_state(database1),
state2=DatabaseState({
DatabaseEntry(**json1['molecule']): 1,
DatabaseEntry(**to_hashable(json3['matrix'])): 1,
}),
)
