def test_mongostore_newer_in(mongostore):
target = MongoStore("maggma_test", "test_target")
target.connect()
# make sure docs are newer in mongostore then target and check updated_keys
target.update(
[
{mongostore.key: i, mongostore.last_updated_field: datetime.utcnow()}
for i in range(10)
]
)
# Update docs in source
mongostore.update(
[
{mongostore.key: i, mongostore.last_updated_field: datetime.utcnow()}
for i in range(10)
]
)
assert len(target.newer_in(mongostore)) == 10
assert len(target.newer_in(mongostore, exhaustive=True)) == 10
assert len(mongostore.newer_in(target)) == 0
target._collection.drop()
class ElasticAnalysisBuilderTest(unittest.TestCase):
@classmethod
def setUp(self):
# Set up test db, set up mpsft, etc.
self.test_tasks = MongoStore("test_emmet", "tasks")
self.test_tasks.connect()
docs = loadfn(test_tasks, cls=None)
self.test_tasks.update(docs)
self.test_elasticity = MongoStore("test_emmet", "elasticity")
self.test_elasticity.connect()
if PROFILE_MODE:
self.pr = cProfile.Profile()
self.pr.enable()
print("\n<<<---")
@classmethod
def tearDown(self):
if not DEBUG_MODE:
self.test_elasticity.collection.drop()
self.test_tasks.collection.drop()
if PROFILE_MODE:
p = Stats(self.pr)
p.strip_dirs()
p.sort_stats('cumtime')
p.print_stats()
print("\n--->>>")
def test_builder(self):
ec_builder = ElasticAnalysisBuilder(self.test_tasks,
self.test_elasticity,
incremental=False)
ec_builder.connect()
for t in ec_builder.get_items():
processed = ec_builder.process_item(t)
self.assertTrue(bool(processed))
runner = Runner([ec_builder])
runner.run()
# Test warnings
doc = ec_builder.elasticity.query_one(
criteria={"pretty_formula": "NaN3"})
self.assertEqual(doc['warnings'], None)
self.assertAlmostEqual(doc['compliance_tensor'][0][0], 41.576072, 6)
def test_grouping_functions(self):
docs1 = list(
self.test_tasks.query(criteria={"formula_pretty": "NaN3"}))
docs_grouped1 = group_by_parent_lattice(docs1)
self.assertEqual(len(docs_grouped1), 1)
grouped_by_opt = group_deformations_by_optimization_task(docs1)
self.assertEqual(len(grouped_by_opt), 1)
docs2 = self.test_tasks.query(
criteria={"task_label": "elastic deformation"})
sgroup2 = group_by_parent_lattice(docs2)
def test_get_distinct_rotations(self):
struct = PymatgenTest.get_structure("Si")
conv = SpacegroupAnalyzer(struct).get_conventional_standard_structure()
rots = get_distinct_rotations(conv)
ops = SpacegroupAnalyzer(conv).get_symmetry_operations()
for op in ops:
self.assertTrue(
any([np.allclose(op.rotation_matrix, r) for r in rots]))
self.assertEqual(len(rots), 48)
def test_process_elastic_calcs(self):
test_struct = PymatgenTest.get_structure('Sn') # use cubic test struct
dss = DeformedStructureSet(test_struct)
# Construct test task set
opt_task = {
"output": {
"structure": test_struct.as_dict()
},
"input": {
"structure": test_struct.as_dict()
}
}
defo_tasks = []
for n, (struct, defo) in enumerate(zip(dss, dss.deformations)):
strain = defo.green_lagrange_strain
defo_task = {
"output": {
"structure": struct.as_dict(),
"stress": (strain * 5).tolist()
},
"input": None,
"task_id": n,
"completed_at": datetime.utcnow()
}
defo_task.update({
"transmuter": {
"transformation_params": [{
"deformation": defo
}]
}
})
defo_tasks.append(defo_task)
defo_tasks.pop(0)
explicit, derived = process_elastic_calcs(opt_task, defo_tasks)
self.assertEqual(len(explicit), 23)
self.assertEqual(len(derived), 1)
def test_process_elastic_calcs_toec(self):
# Test TOEC tasks
test_struct = PymatgenTest.get_structure('Sn') # use cubic test struct
strain_states = get_default_strain_states(3)
# Default stencil in atomate, this maybe shouldn't be hard-coded
stencil = np.linspace(-0.075, 0.075, 7)
strains = [
Strain.from_voigt(s * np.array(strain_state))
for s, strain_state in product(stencil, strain_states)
]
strains = [s for s in strains if not np.allclose(s, 0)]
sym_reduced = symmetry_reduce(strains, test_struct)
opt_task = {
"output": {
"structure": test_struct.as_dict()
},
"input": {
"structure": test_struct.as_dict()
}
}
defo_tasks = []
for n, strain in enumerate(sym_reduced):
defo = strain.get_deformation_matrix()
new_struct = defo.apply_to_structure(test_struct)
defo_task = {
"output": {
"structure": new_struct.as_dict(),
"stress": (strain * 5).tolist()
},
"input": None,
"task_id": n,
"completed_at": datetime.utcnow()
}
defo_task.update({
"transmuter": {
"transformation_params": [{
"deformation": defo
}]
}
})
defo_tasks.append(defo_task)
explicit, derived = process_elastic_calcs(opt_task, defo_tasks)
self.assertEqual(len(explicit), len(sym_reduced))
self.assertEqual(len(derived), len(strains) - len(sym_reduced))
for calc in derived:
self.assertTrue(
np.allclose(calc['strain'], calc['cauchy_stress'] / -0.5))
class ElasticAggregateBuilderTest(unittest.TestCase):
def setUp(self):
# Empty aggregated collection
self.test_elasticity_agg = MongoStore("test_emmet", "elasticity_agg")
self.test_elasticity_agg.connect()
# Generate test materials collection
self.test_materials = MongoStore("test_emmet", "materials")
self.test_materials.connect()
mat_docs = []
for n, formula in enumerate(['Si', 'BaNiO3', 'Li2O2', 'TiO2']):
structure = PymatgenTest.get_structure(formula)
structure.add_site_property("magmoms", [0.0] * len(structure))
mat_docs.append({
"task_id": "mp-{}".format(n),
"structure": structure.as_dict(),
"pretty_formula": formula
})
self.test_materials.update(mat_docs, update_lu=False)
# Create elasticity collection and add docs
self.test_elasticity = MongoStore("test_emmet",
"elasticity",
key="optimization_task_id")
self.test_elasticity.connect()
si = PymatgenTest.get_structure("Si")
si.add_site_property("magmoms", [0.0] * len(si))
et = ElasticTensor.from_voigt([[50, 25, 25, 0, 0, 0],
[25, 50, 25, 0, 0, 0],
[25, 25, 50, 0, 0, 0],
[0, 0, 0, 75, 0,
0], [0, 0, 0, 0, 75, 0],
[0, 0, 0, 0, 0, 75]])
doc = {
"input_structure": si.copy().as_dict(),
"order": 2,
"magnetic_type": "non-magnetic",
"optimization_task_id": "mp-1",
"last_updated": datetime.utcnow(),
"completed_at": datetime.utcnow(),
"optimized_structure": si.copy().as_dict(),
"pretty_formula": "Si",
"state": "successful"
}
doc['elastic_tensor'] = et.voigt
doc.update(et.property_dict)
self.test_elasticity.update([doc])
# Insert second doc with diff params
si.perturb(0.005)
doc.update({
"optimized_structure": si.copy().as_dict(),
"updated_at": datetime.utcnow(),
"optimization_task_id": "mp-5"
})
self.test_elasticity.update([doc])
self.builder = self.get_a_new_builder()
def tearDown(self):
if not DEBUG_MODE:
self.test_elasticity.collection.drop()
self.test_elasticity_agg.collection.drop()
self.test_materials.collection.drop()
def test_materials_aggregator(self):
materials_dict = generate_formula_dict(self.test_materials)
docs = []
grouped_by_mpid = group_by_material_id(
materials_dict['Si'],
[{
'structure': PymatgenTest.get_structure('Si').as_dict(),
'magnetic_type': "non-magnetic"
}])
self.assertEqual(len(grouped_by_mpid), 1)
materials_dict = generate_formula_dict(self.test_materials)
def test_get_items(self):
iterator = self.builder.get_items()
for item in iterator:
self.assertIsNotNone(item)
def test_process_items(self):
docs = list(
self.test_elasticity.query(criteria={"pretty_formula": "Si"}))
formula_dict = generate_formula_dict(self.test_materials)
processed = self.builder.process_item((docs, formula_dict['Si']))
self.assertEqual(len(processed), 1)
self.assertEqual(len(processed[0]['all_elastic_fits']), 2)
def test_update_targets(self):
processed = [
self.builder.process_item(item)
for item in self.builder.get_items()
]
self.builder.update_targets(processed)
def test_aggregation(self):
runner = Runner([self.builder])
runner.run()
all_agg_docs = list(self.test_elasticity_agg.query())
self.assertTrue(bool(all_agg_docs))
def get_a_new_builder(self):
return ElasticAggregateBuilder(self.test_elasticity,
self.test_materials,
self.test_elasticity_agg)
class TestCopyBuilder(TestCase):
@classmethod
def setUpClass(cls):
cls.dbname = "test_" + uuid4().hex
s = MongoStore(cls.dbname, "test")
s.connect()
cls.client = s.collection.database.client
@classmethod
def tearDownClass(cls):
cls.client.drop_database(cls.dbname)
def setUp(self):
tic = datetime.now()
toc = tic + timedelta(seconds=1)
keys = list(range(20))
self.old_docs = [{"lu": tic, "k": k, "v": "old"} for k in keys]
self.new_docs = [{"lu": toc, "k": k, "v": "new"} for k in keys[:10]]
kwargs = dict(key="k", lu_field="lu")
self.source = MongoStore(self.dbname, "source", **kwargs)
self.target = MongoStore(self.dbname, "target", **kwargs)
self.builder = CopyBuilder(self.source, self.target)
self.source.connect()
self.source.collection.create_index("lu")
self.target.connect()
self.target.collection.create_index("lu")
self.target.collection.create_index("k")
def tearDown(self):
self.source.collection.drop()
self.target.collection.drop()
def test_get_items(self):
self.source.collection.insert_many(self.old_docs)
self.assertEqual(len(list(self.builder.get_items())),
len(self.old_docs))
self.target.collection.insert_many(self.old_docs)
self.assertEqual(len(list(self.builder.get_items())), 0)
self.source.update(self.new_docs, update_lu=False)
self.assertEqual(len(list(self.builder.get_items())),
len(self.new_docs))
def test_process_item(self):
self.source.collection.insert_many(self.old_docs)
items = list(self.builder.get_items())
self.assertCountEqual(items, map(self.builder.process_item, items))
def test_update_targets(self):
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
self.target.collection.insert_many(self.old_docs)
items = list(map(self.builder.process_item, self.builder.get_items()))
self.builder.update_targets(items)
self.assertEqual(self.target.query_one(criteria={"k": 0})["v"], "new")
self.assertEqual(self.target.query_one(criteria={"k": 10})["v"], "old")
def test_confirm_lu_field_index(self):
self.source.collection.drop_index("lu_1")
with self.assertRaises(Exception) as cm:
self.builder.get_items()
self.assertTrue(cm.exception.args[0].startswith("Need index"))
self.source.collection.create_index("lu")
def test_runner(self):
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
self.target.collection.insert_many(self.old_docs)
runner = Runner([self.builder])
runner.run()
self.assertEqual(self.target.query_one(criteria={"k": 0})["v"], "new")
self.assertEqual(self.target.query_one(criteria={"k": 10})["v"], "old")
def test_query(self):
self.builder.query = {"k": {"$gt": 5}}
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
runner = Runner([self.builder])
runner.run()
all_docs = list(self.target.query(criteria={}))
self.assertEqual(len(all_docs), 14)
self.assertTrue(min([d['k'] for d in all_docs]), 6)
class TestCopyBuilder(TestCase):
@classmethod
def setUpClass(cls):
cls.dbname = "test_" + uuid4().hex
s = MongoStore(cls.dbname, "test")
s.connect()
cls.client = s.collection.database.client
@classmethod
def tearDownClass(cls):
cls.client.drop_database(cls.dbname)
def setUp(self):
tic = datetime.now()
toc = tic + timedelta(seconds=1)
keys = list(range(20))
self.old_docs = [{"lu": tic, "k": k, "v": "old"} for k in keys]
self.new_docs = [{"lu": toc, "k": k, "v": "new"} for k in keys[:10]]
kwargs = dict(key="k", lu_field="lu")
self.source = MongoStore(self.dbname, "source", **kwargs)
self.target = MongoStore(self.dbname, "target", **kwargs)
self.builder = CopyBuilder(self.source, self.target)
self.source.connect()
self.source.ensure_index(self.source.key)
self.source.ensure_index(self.source.lu_field)
self.target.connect()
self.target.ensure_index(self.target.key)
self.target.ensure_index(self.target.lu_field)
def tearDown(self):
self.source.collection.drop()
self.target.collection.drop()
def test_get_items(self):
self.source.collection.insert_many(self.old_docs)
self.assertEqual(len(list(self.builder.get_items())),
len(self.old_docs))
self.target.collection.insert_many(self.old_docs)
self.assertEqual(len(list(self.builder.get_items())), 0)
self.source.update(self.new_docs, update_lu=False)
self.assertEqual(len(list(self.builder.get_items())),
len(self.new_docs))
def test_process_item(self):
self.source.collection.insert_many(self.old_docs)
items = list(self.builder.get_items())
self.assertCountEqual(items, map(self.builder.process_item, items))
def test_update_targets(self):
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
self.target.collection.insert_many(self.old_docs)
items = list(map(self.builder.process_item, self.builder.get_items()))
self.builder.update_targets(items)
self.assertEqual(self.target.query_one(criteria={"k": 0})["v"], "new")
self.assertEqual(self.target.query_one(criteria={"k": 10})["v"], "old")
@unittest.skip(
"Have to refactor how we force read-only so a warning will get thrown")
def test_index_warning(self):
"""Should log warning when recommended store indexes are not present."""
self.source.collection.drop_index([(self.source.key, 1)])
with self.assertLogs(level=logging.WARNING) as cm:
list(self.builder.get_items())
self.assertIn("Ensure indices", "\n".join(cm.output))
def test_run(self):
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
self.target.collection.insert_many(self.old_docs)
self.builder.run()
self.assertEqual(self.target.query_one(criteria={"k": 0})["v"], "new")
self.assertEqual(self.target.query_one(criteria={"k": 10})["v"], "old")
def test_query(self):
self.builder.query = {"k": {"$gt": 5}}
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
self.builder.run()
all_docs = list(self.target.query(criteria={}))
self.assertEqual(len(all_docs), 14)
self.assertTrue(min([d['k'] for d in all_docs]), 6)
def test_delete_orphans(self):
self.builder = CopyBuilder(self.source,
self.target,
delete_orphans=True)
self.source.collection.insert_many(self.old_docs)
self.source.update(self.new_docs, update_lu=False)
self.target.collection.insert_many(self.old_docs)
deletion_criteria = {"k": {"$in": list(range(5))}}
self.source.collection.delete_many(deletion_criteria)
self.builder.run()
self.assertEqual(
self.target.collection.count_documents(deletion_criteria), 0)
self.assertEqual(self.target.query_one(criteria={"k": 5})["v"], "new")
self.assertEqual(self.target.query_one(criteria={"k": 10})["v"], "old")
def test_incremental_false(self):
tic = datetime.now()
toc = tic + timedelta(seconds=1)
keys = list(range(20))
earlier = [{"lu": tic, "k": k, "v": "val"} for k in keys]
later = [{"lu": toc, "k": k, "v": "val"} for k in keys]
self.source.collection.insert_many(earlier)
self.target.collection.insert_many(later)
query = {"k": {"$gt": 5}}
self.builder = CopyBuilder(self.source,
self.target,
incremental=False,
query=query)
self.builder.run()
docs = sorted(self.target.query(), key=lambda d: d["k"])
self.assertTrue(all(d["lu"] == tic) for d in docs[5:])
self.assertTrue(all(d["lu"] == toc) for d in docs[:5])
class JointStoreTest(unittest.TestCase):
def setUp(self):
self.jointstore = JointStore("maggma_test", ["test1", "test2"])
self.jointstore.connect()
self.jointstore.collection.drop()
self.jointstore.collection.insert_many([{
"task_id":
k,
"my_prop":
k + 1,
"last_updated":
datetime.utcnow(),
"category":
k // 5
} for k in range(10)])
self.jointstore.collection.database["test2"].drop()
self.jointstore.collection.database["test2"].insert_many([{
"task_id":
2 * k,
"your_prop":
k + 3,
"last_updated":
datetime.utcnow(),
"category2":
k // 3
} for k in range(5)])
self.test1 = MongoStore("maggma_test", "test1")
self.test1.connect()
self.test2 = MongoStore("maggma_test", "test2")
self.test2.connect()
def test_query(self):
# Test query all
docs = list(self.jointstore.query())
self.assertEqual(len(docs), 10)
docs_w_field = [d for d in docs if "test2" in d]
self.assertEqual(len(docs_w_field), 5)
docs_w_field = sorted(docs_w_field, key=lambda x: x['task_id'])
self.assertEqual(docs_w_field[0]['test2']['your_prop'], 3)
self.assertEqual(docs_w_field[0]['task_id'], 0)
self.assertEqual(docs_w_field[0]['my_prop'], 1)
def test_query_one(self):
doc = self.jointstore.query_one()
self.assertEqual(doc['my_prop'], doc['task_id'] + 1)
# Test limit properties
doc = self.jointstore.query_one(properties=['test2', 'task_id'])
self.assertEqual(doc['test2']['your_prop'], doc['task_id'] + 3)
self.assertIsNone(doc.get("my_prop"))
# Test criteria
doc = self.jointstore.query_one(criteria={"task_id": {"$gte": 10}})
self.assertIsNone(doc)
doc = self.jointstore.query_one(
criteria={"test2.your_prop": {
"$gt": 6
}})
self.assertEqual(doc['task_id'], 8)
# Test merge_at_root
self.jointstore.merge_at_root = True
# Test merging is working properly
doc = self.jointstore.query_one(criteria={"task_id": 2})
self.assertEqual(doc['my_prop'], 3)
self.assertEqual(doc['your_prop'], 4)
# Test merging is allowing for subsequent match
doc = self.jointstore.query_one(criteria={"your_prop": {"$gt": 6}})
self.assertEqual(doc['task_id'], 8)
def test_distinct(self):
dyour_prop = self.jointstore.distinct("test2.your_prop")
self.assertEqual(set(dyour_prop), {k + 3 for k in range(5)})
dmy_prop = self.jointstore.distinct("my_prop")
self.assertEqual(set(dmy_prop), {k + 1 for k in range(10)})
dmy_prop_cond = self.jointstore.distinct(
"my_prop", {"test2.your_prop": {
"$gte": 5
}})
self.assertEqual(set(dmy_prop_cond), {5, 7, 9})
def test_last_updated(self):
doc = self.jointstore.query_one({"task_id": 0})
test1doc = self.test1.query_one({"task_id": 0})
test2doc = self.test2.query_one({"task_id": 0})
self.assertEqual(test2doc['last_updated'], doc['last_updated'])
self.assertNotEqual(test1doc['last_updated'], doc['last_updated'])
# Swap the two
test2date = test2doc['last_updated']
test2doc['last_updated'] = test1doc['last_updated']
test1doc['last_updated'] = test2date
self.test1.update([test1doc], update_lu=False)
self.test2.update([test2doc], update_lu=False)
doc = self.jointstore.query_one({"task_id": 0})
test1doc = self.test1.query_one({"task_id": 0})
test2doc = self.test2.query_one({"task_id": 0})
self.assertEqual(test1doc['last_updated'], doc['last_updated'])
self.assertNotEqual(test2doc['last_updated'], doc['last_updated'])
# Check also that still has a field if no task2 doc
doc = self.jointstore.query_one({"task_id": 1})
self.assertIsNotNone(doc['last_updated'])
def test_groupby(self):
docs = list(self.jointstore.groupby("category"))
self.assertEqual(len(docs[0]['docs']), 5)
self.assertEqual(len(docs[1]['docs']), 5)
docs = list(self.jointstore.groupby("test2.category2"))
docs_by_id = {get(d, '_id.test2.category2'): d['docs'] for d in docs}
self.assertEqual(len(docs_by_id[None]), 5)
self.assertEqual(len(docs_by_id[0]), 3)
self.assertEqual(len(docs_by_id[1]), 2)
