class TestDiffractionBuilder(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):
kwargs = dict(key="k", lu_field="lu")
self.source = MongoStore(self.dbname, "source", **kwargs)
self.target = MongoStore(self.dbname, "target", **kwargs)
self.source.connect()
self.source.collection.create_index("lu")
self.source.collection.create_index("k", unique=True)
self.target.connect()
self.target.collection.create_index("lu")
self.target.collection.create_index("k", unique=True)
def tearDown(self):
self.source.collection.drop()
self.target.collection.drop()
def test_get_xrd_from_struct(self):
builder = DiffractionBuilder(self.source, self.target)
structure = PymatgenTest.get_structure("Si")
self.assertIn("Cu", builder.get_xrd_from_struct(structure))
def test_serialization(self):
builder = DiffractionBuilder(self.source, self.target)
self.assertIsNone(builder.as_dict()["xrd_settings"])
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_mongostore_from_collection(mongostore, db_json):
ms = MongoStore.from_db_file(db_json)
ms.connect()
other_ms = MongoStore.from_collection(ms._collection)
assert ms._collection.full_name == other_ms._collection.full_name
assert ms.database == other_ms.database
def setUp(self):
self.ml_strucs = MongoStore("emmet_test", "ml_strucs", key="entry_id")
self.ml_strucs.connect()
self.ml_strucs.collection.drop()
self.mlbuilder = MLStructuresBuilder(
self.tasks,
self.ml_strucs,
task_types=("Structure Optimization", "Static"))
def setUp(self):
self.materials = MongoStore("emmet_test", "materials")
self.materials.connect()
self.materials.collection.drop()
self.mbuilder = MaterialsBuilder(self.tasks,
self.materials,
mat_prefix="",
chunk_size=1)
class TestMaterials(BuilderTest):
def setUp(self):
self.materials = MongoStore("emmet_test", "materials")
self.materials.connect()
self.materials.collection.drop()
self.mbuilder = MaterialsBuilder(self.tasks,
self.materials,
mat_prefix="",
chunk_size=1)
def test_get_items(self):
to_process = list(self.mbuilder.get_items())
to_process_forms = {tasks[0]["formula_pretty"] for tasks in to_process}
self.assertEqual(len(to_process), 12)
self.assertEqual(len(to_process_forms), 12)
self.assertEqual(len(list(chain.from_iterable(to_process))), 197)
self.assertTrue("Sr" in to_process_forms)
self.assertTrue("Hf" in to_process_forms)
self.assertTrue("O2" in to_process_forms)
self.assertFalse("H" in to_process_forms)
def test_process_item(self):
tasks = list(self.tasks.query(criteria={"chemsys": "Sr"}))
mats = self.mbuilder.process_item(tasks)
self.assertEqual(len(mats), 7)
tasks = list(self.tasks.query(criteria={"chemsys": "Hf"}))
mats = self.mbuilder.process_item(tasks)
self.assertEqual(len(mats), 4)
tasks = list(self.tasks.query(criteria={"chemsys": "O"}))
mats = self.mbuilder.process_item(tasks)
self.assertEqual(len(mats), 6)
tasks = list(self.tasks.query(criteria={"chemsys": "O-Sr"}))
mats = self.mbuilder.process_item(tasks)
self.assertEqual(len(mats), 5)
tasks = list(self.tasks.query(criteria={"chemsys": "Hf-O-Sr"}))
mats = self.mbuilder.process_item(tasks)
self.assertEqual(len(mats), 13)
def test_update_targets(self):
tasks = list(self.tasks.query(criteria={"chemsys": "Sr"}))
mats = self.mbuilder.process_item(tasks)
self.assertEqual(len(mats), 7)
self.mbuilder.update_targets([mats])
self.assertEqual(len(self.materials.distinct("task_id")), 7)
self.assertEqual(len(list(self.materials.query())), 7)
def tearDown(self):
self.materials.collection.drop()
def test_mongostore_connect_via_ssh():
mongostore = MongoStore("maggma_test", "test")
class fake_pipe:
remote_bind_address = ("localhost", 27017)
local_bind_address = ("localhost", 37017)
server = fake_pipe()
mongostore.connect(ssh_tunnel=server)
assert isinstance(mongostore._collection, pymongo.collection.Collection)
def setUp(self):
kwargs = dict(key="k", lu_field="lu")
self.source = MongoStore(self.dbname, "source", **kwargs)
self.target = MongoStore(self.dbname, "target", **kwargs)
self.source.connect()
self.source.collection.create_index("lu")
self.source.collection.create_index("k", unique=True)
self.target.connect()
self.target.collection.create_index("lu")
self.target.collection.create_index("k", unique=True)
def setUpClass(cls):
cls.dbname = "test_" + uuid4().hex
cls.source = MongoStore(cls.dbname, "source")
cls.target = MongoStore(cls.dbname, "target")
cls.stores = [cls.source, cls.target]
for store in cls.stores:
store.connect()
store.ensure_index(store.key)
store.ensure_index(store.lu_field)
cls.client = cls.stores[0].collection.database.client
def setUp(self):
coords = list()
coords.append([0, 0, 0])
coords.append([0.75, 0.5, 0.75])
lattice = Lattice([[3.8401979337, 0.00, 0.00], [1.9200989668, 3.3257101909, 0.00],
[0.00, -2.2171384943, 3.1355090603]])
self.structure = Structure(lattice, ["Si", "Si"], coords)
tasks = MongoStore("emmet_test", "tasks")
materials = MongoStore("emmet_test", "materials")
self.mbuilder = MaterialsBuilder(tasks, materials, mat_prefix="", chunk_size=1)
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<<<---")
def mongostore():
store = MongoStore("maggma_test", "test")
store.connect()
store.remove_docs({})
yield store
store.remove_docs({})
store._collection.drop()
def reporting_store():
store = MongoStore("maggma_test", "reporting")
store.connect()
store.remove_docs({})
yield store
store.remove_docs({})
store._collection.drop()
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()
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 setUp(self):
coords = list()
coords.append([0, 0, 0])
coords.append([0.75, 0.5, 0.75])
lattice = Lattice([[3.8401979337, 0.00, 0.00],
[1.9200989668, 3.3257101909, 0.00],
[0.00, -2.2171384943, 3.1355090603]])
self.structure = Structure(lattice, ["Si", "Si"], coords)
materials = MongoStore("emmet_test", "materials")
electronic_structure = MongoStore("emmet_test", "electronic_structure")
self.builder = ElectronicStructureImageBuilder(materials,
electronic_structure)
def setUpClass(cls):
root = logging.getLogger()
root.setLevel(logging.DEBUG)
ch = logging.StreamHandler(sys.stdout)
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter(
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
ch.setFormatter(formatter)
root.addHandler(ch)
cls.tasks = MongoStore("emmet_test", "tasks", lu_field="last_updated")
cls.tasks.connect()
cleardb(cls.tasks.collection.database)
vaspdb = VaspCalcDb(database="emmet_test")
tasks_dir = os.path.join(test_dir, "tasks")
raw_tasks = glob.glob(os.path.join(test_dir, "tasks", "*.json.gz"))
for task_path in raw_tasks:
with zopen(task_path) as f:
data = f.read().decode()
task = json.loads(data)
vaspdb.insert_task(task, parse_dos=True, parse_bs=True)
def mongo_jobstore(database):
from maggma.stores import MongoStore
from jobflow import JobStore
store = JobStore(MongoStore(database, "outputs"))
store.connect()
return store
def setUp(self):
self.materials = MongoStore("emmet_test", "materials")
self.thermo = MongoStore("emmet_test", "thermo")
self.materials.connect()
self.thermo.connect()
self.mbuilder = MaterialsBuilder(self.tasks,
self.materials,
mat_prefix="",
chunk_size=1)
self.tbuilder = ThermoBuilder(self.materials,
self.thermo,
chunk_size=1)
runner = Runner([self.mbuilder])
runner.run()
class TestMaterials(BuilderTest):
def setUp(self):
self.ml_strucs = MongoStore("emmet_test", "ml_strucs", key="entry_id")
self.ml_strucs.connect()
self.ml_strucs.collection.drop()
self.mlbuilder = MLStructuresBuilder(
self.tasks,
self.ml_strucs,
task_types=("Structure Optimization", "Static"))
def test_get_items(self):
to_process = list(self.mlbuilder.get_items())
to_process_forms = {task["formula_pretty"] for task in to_process}
self.assertEqual(len(to_process), 197)
self.assertEqual(len(to_process_forms), 12)
self.assertTrue("Sr" in to_process_forms)
self.assertTrue("Hf" in to_process_forms)
self.assertTrue("O2" in to_process_forms)
self.assertFalse("H" in to_process_forms)
def test_process_item(self):
for task in self.tasks.query():
ml_strucs = self.mlbuilder.process_item(task)
t_type = task_type(get(task, 'input.incar'))
if not any([t in t_type for t in self.mlbuilder.task_types]):
self.assertEqual(len(ml_strucs), 0)
else:
self.assertEqual(
len(ml_strucs),
sum([
len(t["output"]["ionic_steps"])
for t in task["calcs_reversed"]
]))
def test_update_targets(self):
for task in self.tasks.query():
ml_strucs = self.mlbuilder.process_item(task)
self.mlbuilder.update_targets([ml_strucs])
self.assertEqual(len(self.ml_strucs.distinct("task_id")), 102)
self.assertEqual(len(list(self.ml_strucs.query())), 1012)
def tearDown(self):
self.ml_strucs.collection.drop()
def setUpClass(cls):
# Set up test db, set up mpsft, etc.
cls.test_tasks = MongoStore("test_emmet", "tasks")
cls.test_tasks.connect()
docs = loadfn(test_tasks, cls=None)
cls.test_tasks.update(docs)
cls.test_elasticity = MongoStore("test_emmet", "elasticity")
# Generate test materials collection
cls.test_materials = MongoStore("test_emmet", "materials")
cls.test_materials.connect()
opt_docs = cls.test_tasks.query(
["output.structure", "formula_pretty"],
{"task_label": "structure optimization"})
mat_docs = [{
"material_id": "mp-{}".format(n),
"structure": opt_doc['output']['structure'],
"pretty_formula": opt_doc['formula_pretty']
} for n, opt_doc in enumerate(opt_docs)]
cls.test_materials.update(mat_docs, key='material_id', update_lu=False)
def init_drone(test_dir):
"""
Initialize the drone, do not initialize the connection with the database
:return:
initialized drone
"""
mongo_store = MongoStore(database="drone_test",
collection_name="drone_test",
key="record_key")
simple_path = test_dir / "simple_bib_example_data"
assert simple_path.exists(), f"{simple_path} not found"
simple_bib_drone = SimpleBibDrone(store=mongo_store, path=simple_path)
return simple_bib_drone
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 __init__(self,
pj_store,
lpad,
incremental=True,
query=None,
base_priority=2500,
site_penalty=10,
vote_weight=10,
**kwargs):
"""
Takes a propjockey collection and sets the priority
of a fireworks in a fireworks collection from a LaunchPad
Args:
pj_store (Store): store corresponding to propjockey collection
lpad (LaunchPad): fireworks launchpad
query (dict): query to filter the propjockey store
base_priority (int): base priority to assign to fireworks
site_penalty (int): per-site penalty to priority
vote_weight (int): priority boost per vote
**kwargs (kwargs): kwargs for builder
"""
self.pj_store = pj_store
if isinstance(lpad, dict):
self.lpad = LaunchPad.from_dict(lpad)
else:
self.lpad = lpad
self.fws_store = MongoStore.from_collection(self.lpad.fireworks)
self.fws_store.lu_field = "_pj_lu"
self.incremental = incremental
self.start_date = datetime.utcnow()
self.query = query or {}
self.base_priority = base_priority
self.site_penalty = site_penalty
self.vote_weight = vote_weight
super().__init__(sources=[self.pj_store],
targets=[self.fws_store],
**kwargs)
def _get_s3_store(self, store_name):
"""
Add a maggma store to this object for storage of large chunk data
The maggma store will be stored to self.maggma_store[store_name]
For aws store, all documents will be stored to the same bucket
and the store_name will double as the sub_dir name.
Args:
store_name: correspond to the the key within calcs_reversed.0 that will be stored
"""
if self.host_uri is not None:
index_store_ = MongoURIStore(
uri=self.host_uri,
database=self.db_name,
collection_name=
f"{self.maggma_store_prefix}_{store_name}_index",
key="fs_id",
)
else:
index_store_ = MongoStore(
database=self.db_name,
collection_name=
f"{self.maggma_store_prefix}_{store_name}_index",
host=self.host,
port=self.port,
username=self.user,
password=self.password,
key="fs_id",
)
store = S3Store(
index=index_store_,
sub_dir=f"{self.maggma_store_prefix}_{store_name}",
key="fs_id",
**self._maggma_store_kwargs,
)
return store
def test_mongostore_connect():
mongostore = MongoStore("maggma_test", "test")
assert mongostore._collection is None
mongostore.connect()
assert isinstance(mongostore._collection, pymongo.collection.Collection)
def test_mongostore_from_db_file(mongostore, db_json):
ms = MongoStore.from_db_file(db_json)
ms.connect()
assert ms._collection.full_name == "maggma_tests.tmp"
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)
def setUpClass(cls):
cls.dbname = "test_" + uuid4().hex
s = MongoStore(cls.dbname, "test")
s.connect()
cls.client = s.collection.database.client
