def search_resource(eos_store):
resource = Resource(
eos_store,
SearchDoc,
query_operators=[
FormulaQuery(),
MinMaxQuery(),
SymmetryQuery(),
ThermoEnergySearchQuery(),
IsStableQuery(),
SearchBandGapQuery(),
BulkModulusQuery(),
ShearModulusQuery(),
PoissonQuery(),
DielectricQuery(),
PiezoelectricQuery(),
SurfaceMinMaxQuery(),
SearchTaskIDsQuery(),
HasPropsQuery(),
DeprecationQuery(),
PaginationQuery(),
SparseFieldsQuery(SearchDoc, default_fields=["task_id"]),
],
tags=["Search"],
)
return resource
def xas_resource(xas_store):
resource = Resource(
xas_store,
XASDoc,
query_operators=[
FormulaQuery(),
XASQuery(),
PaginationQuery(),
SparseFieldsQuery(
XASDoc,
default_fields=[
"xas_id",
"task_id",
"edge",
"absorbing_element",
"formula_pretty",
"spectrum_type",
"last_updated",
],
),
],
tags=["XAS"],
)
return resource
def task_deprecation_resource(materials_store):
def custom_deprecation_prep(self):
async def check_deprecation(task_id: str = Path(
...,
alias="task_id",
title="Task id to check for deprecation.",
), ):
"""
Checks whether a task_id is deprecated.
Returns:
Dictionary containing deprecation information
"""
crit = {"deprecated_tasks": task_id}
self.store.connect()
deprecation = {"deprecated": False, "deprecation_reason": None}
for r in self.store.query(criteria=crit,
properties=["deprecated_tasks"]):
if r != {}:
deprecation = {
"deprecated": True,
"deprecation_reason": None
}
break
response = {"data": deprecation}
return response
self.router.get(
"/{task_id}/",
response_model_exclude_unset=True,
response_description="Check deprecation of a specific task_id",
tags=self.tags,
)(check_deprecation)
resource = Resource(
materials_store,
MaterialsCoreDoc,
query_operators=[
PaginationQuery(),
],
tags=["Tasks"],
custom_endpoint_funcs=[custom_deprecation_prep],
enable_get_by_key=False,
enable_default_search=False,
)
return resource
def fermi_resource(fermi_store):
resource = Resource(
fermi_store,
FermiDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(FermiDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Electronic Structure"],
)
return resource
def dois_resource(dois_store):
resource = Resource(
dois_store,
DOIDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(DOIDoc, default_fields=["task_id", "doi"]),
],
tags=["DOIs"],
enable_default_search=False,
)
return resource
def wulff_resource(wulff_store):
resource = Resource(
wulff_store,
WulffDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(WulffDoc, default_fields=["task_id"]),
],
tags=["Surface Properties"],
enable_default_search=False,
)
return resource
def similarity_resource(similarity_store):
resource = Resource(
similarity_store,
SimilarityDoc,
query_operators=[
PaginationQuery(),
SparseFieldsQuery(SimilarityDoc, default_fields=["task_id"]),
],
tags=["Similarity"],
enable_default_search=False,
)
return resource
def eos_resource(eos_store):
resource = Resource(
eos_store,
EOSDoc,
query_operators=[
EnergyVolumeQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(EOSDoc, default_fields=["task_id"]),
],
tags=["EOS"],
)
return resource
def magnetism_resource(magnetism_store):
resource = Resource(
magnetism_store,
MagnetismDoc,
query_operators=[
MagneticQuery(),
PaginationQuery(),
SparseFieldsQuery(MagnetismDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Magnetism"],
)
return resource
def piezo_resource(piezo_store):
resource = Resource(
piezo_store,
PiezoDoc,
query_operators=[
PiezoelectricQuery(),
PaginationQuery(),
SparseFieldsQuery(PiezoDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Piezoelectric"],
)
return resource
def dielectric_resource(dielectric_store):
resource = Resource(
dielectric_store,
DielectricDoc,
query_operators=[
DielectricQuery(),
PaginationQuery(),
SparseFieldsQuery(DielectricDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Dielectric"],
)
return resource
def surface_props_resource(surface_prop_store):
resource = Resource(
surface_prop_store,
SurfacePropDoc,
query_operators=[
SurfaceMinMaxQuery(),
ReconstructedQuery(),
PaginationQuery(),
SparseFieldsQuery(SurfacePropDoc, default_fields=["task_id"]),
],
tags=["Surface Properties"],
)
return resource
def substrates_resource(substrates_store):
resource = Resource(
substrates_store,
SubstratesDoc,
query_operators=[
SubstrateStructureQuery(),
EnergyAreaQuery(),
PaginationQuery(),
SparseFieldsQuery(SubstratesDoc,
default_fields=["film_id", "sub_id"]),
],
tags=["Substrates"],
enable_get_by_key=False,
)
return resource
def task_resource(task_store):
resource = Resource(
task_store,
TaskDoc,
query_operators=[
FormulaQuery(),
PaginationQuery(),
SparseFieldsQuery(
TaskDoc,
default_fields=["task_id", "formula_pretty", "last_updated"],
),
],
tags=["Tasks"],
)
return resource
def molecules_resource(molecules_store):
resource = Resource(
molecules_store,
MoleculesDoc,
query_operators=[
MoleculeBaseQuery(),
MoleculeElementsQuery(),
MoleculeFormulaQuery(),
SearchTaskIDsQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(MoleculesDoc, default_fields=["task_id"]),
],
tags=["Molecules"],
)
return resource
def thermo_resource(thermo_store):
resource = Resource(
thermo_store,
ThermoDoc,
query_operators=[
VersionQuery(),
ThermoChemicalQuery(),
IsStableQuery(),
ThermoEnergyQuery(),
PaginationQuery(),
SparseFieldsQuery(ThermoDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Thermo"],
)
return resource
def gb_resource(gb_store):
resource = Resource(
gb_store,
GBDoc,
query_operators=[
GBTaskIDQuery(),
GBEnergyQuery(),
GBStructureQuery(),
PaginationQuery(),
SparseFieldsQuery(GBDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Grain Boundaries"],
enable_get_by_key=False,
)
return resource
def elasticity_resource(elasticity_store):
resource = Resource(
elasticity_store,
ElasticityDoc,
query_operators=[
ChemsysQuery(),
BulkModulusQuery(),
ShearModulusQuery(),
PoissonQuery(),
PaginationQuery(),
SparseFieldsQuery(
ElasticityDoc,
default_fields=["task_id", "pretty_formula"],
),
],
tags=["Elasticity"],
)
return resource
def insertion_electrodes_resource(insertion_electrodes_store):
resource = Resource(
insertion_electrodes_store,
InsertionElectrodeDoc,
query_operators=[
VoltageStepQuery(),
InsertionVoltageStepQuery(),
InsertionElectrodeQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
InsertionElectrodeDoc,
default_fields=["task_id", "last_updated"],
),
],
tags=["Electrodes"],
)
return resource
def trajectory_resource(task_store):
class TrajectoryProcess(APIRoute):
def get_route_handler(self) -> Callable:
original_route_handler = super().get_route_handler()
async def custom_route_handler(request: Request) -> Response:
response: Response = await original_route_handler(request)
d = json.loads(response.body, encoding=response.charset)
trajectories = []
for entry in d["data"]:
trajectories.append(
calcs_reversed_to_trajectory(entry["calcs_reversed"]))
trajectories = jsanitize(trajectories)
response.body = json.dumps(
trajectories,
ensure_ascii=False,
allow_nan=False,
indent=None,
separators=(",", ":"),
).encode(response.charset)
traj_len = str(len(response.body))
response.headers["content-length"] = traj_len
return response
return custom_route_handler
resource = Resource(
task_store,
TaskDoc,
query_operators=[FormulaQuery(), PaginationQuery()],
route_class=TrajectoryProcess,
key_fields=["calcs_reversed"],
tags=["Tasks"],
)
return resource
def bs_resource(bs_store, s3_store):
def custom_bs_endpoint_prep(self):
self.s3 = s3_store
model = BSObjectReturn
model_name = model.__name__
key_name = self.s3.key
field_input = SparseFieldsQuery(
model, [self.s3.key, self.s3.last_updated_field]).query
async def get_object(
key: str = Query(
...,
alias=key_name,
title=f"The {key_name} of the {model_name} to get",
),
path_type: BSPathType = Query(
...,
title=
"The k-path convention type for the band structure object",
),
fields: STORE_PARAMS = Depends(field_input),
):
f"""
Get's a document by the primary key in the store
Args:
{key_name}: the id of a single {model_name}
Returns:
a single {model_name} document
"""
self.store.connect()
self.s3.connect()
bs_entry = self.store.query_one(
criteria={self.store.key: key},
properties=[f"{str(path_type.name)}.task_id"],
)
bs_task = bs_entry.get(str(path_type.name)).get("task_id", None)
if bs_task is None:
raise HTTPException(
status_code=404,
detail=
f"Band structure with {self.store.key} = {key} not found",
)
item = self.s3.query_one({"task_id": bs_task},
properties=fields["properties"])
response = item
return response
self.router.get(
"/object/",
response_description=f"Get an {model_name} by {key_name}",
response_model=model,
response_model_exclude_unset=True,
tags=self.tags,
)(get_object)
resource = Resource(
bs_store,
BSDoc,
query_operators=[
BSDataQuery(),
FormulaQuery(),
MinMaxQuery(),
PaginationQuery(),
SparseFieldsQuery(BSDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Electronic Structure"],
custom_endpoint_funcs=[custom_bs_endpoint_prep],
)
return resource
def materials_resource(materials_store):
def custom_version_prep(self):
model_name = self.model.__name__
async def get_versions():
f"""
Obtains the database versions for the data in {model_name}
Returns:
A list of database versions one can use to query on
"""
try:
conn = MongoClient(self.store.host, self.store.port)
db = conn[self.store.database]
if self.core.username != "":
db.authenticate(self.username, self.password)
except AttributeError:
conn = MongoClient(self.store.uri)
db = conn[self.store.database]
col_names = db.list_collection_names()
d = [
name.replace("_", ".")[15:]
for name in col_names
if "materials" in name
if name != "materials.core"
]
response = {"data": d}
return response
self.router.get(
"/versions/",
response_model_exclude_unset=True,
response_description=f"Get versions of {model_name}",
tags=self.tags,
)(get_versions)
def custom_findstructure_prep(self):
model_name = self.model.__name__
async def find_structure(
structure: Structure = Body(
...,
title="Pymatgen structure object to query with",
),
ltol: float = Query(
0.2,
title="Fractional length tolerance. Default is 0.2.",
),
stol: float = Query(
0.3,
title="Site tolerance. Defined as the fraction of the average free \
length per atom := ( V / Nsites ) ** (1/3). Default is 0.3.",
),
angle_tol: float = Query(
5,
title="Angle tolerance in degrees. Default is 5 degrees.",
),
limit: int = Query(
1,
title="Maximum number of matches to show. Defaults to 1, only showing the best match.",
),
):
"""
Obtains material structures that match a given input structure within some tolerance.
Returns:
A list of Material IDs for materials with matched structures alongside the associated RMS values
"""
try:
s = PS.from_dict(structure.dict())
except Exception:
raise HTTPException(
status_code=404,
detail="Body cannot be converted to a pymatgen structure object.",
)
m = StructureMatcher(
ltol=ltol,
stol=stol,
angle_tol=angle_tol,
primitive_cell=True,
scale=True,
attempt_supercell=False,
comparator=ElementComparator(),
)
crit = {"composition_reduced": dict(s.composition.to_reduced_dict)}
self.store.connect()
matches = []
for r in self.store.query(
criteria=crit, properties=["structure", "task_id"]
):
s2 = PS.from_dict(r["structure"])
matched = m.fit(s, s2)
if matched:
rms = m.get_rms_dist(s, s2)
matches.append(
{
"task_id": r["task_id"],
"normalized_rms_displacement": rms[0],
"max_distance_paired_sites": rms[1],
}
)
response = {
"data": sorted(
matches[:limit],
key=lambda x: (
x["normalized_rms_displacement"],
x["max_distance_paired_sites"],
),
)
}
return response
self.router.post(
"/find_structure/",
response_model_exclude_unset=True,
response_description=f"Get matching structures using data from {model_name}",
tags=self.tags,
)(find_structure)
resource = Resource(
materials_store,
MaterialsCoreDoc,
query_operators=[
VersionQuery(),
FormulaQuery(),
MultiTaskIDQuery(),
SymmetryQuery(),
DeprecationQuery(),
MinMaxQuery(),
PaginationQuery(),
SparseFieldsQuery(
MaterialsCoreDoc,
default_fields=["task_id", "formula_pretty", "last_updated"],
),
],
tags=["Materials"],
custom_endpoint_funcs=[custom_version_prep, custom_findstructure_prep],
)
return resource
def materials_resource(materials_store):
def custom_version_prep(self):
model_name = self.model.__name__
async def get_versions():
f"""
Obtains the database versions for the data in {model_name}
Returns:
A list of database versions one can use to query on
"""
try:
conn = MongoClient(self.store.host, self.store.port)
db = conn[self.store.database]
if self.core.username != "":
db.authenticate(self.username, self.password)
except AttributeError:
conn = MongoClient(self.store.uri)
db = conn[self.store.database]
col_names = db.list_collection_names()
d = [
name.replace("_", ".")[15:] for name in col_names
if "materials" in name if name != "materials.core"
]
response = {"data": d}
return response
self.router.get(
"/versions/",
response_model_exclude_unset=True,
response_description=f"Get versions of {model_name}",
tags=self.tags,
)(get_versions)
def custom_findstructure_prep(self):
model_name = self.model.__name__
async def find_structure(
structure: Structure = Body(
...,
title="Pymatgen structure object to query with",
),
ltol: float = Query(
0.2,
title="Fractional length tolerance. Default is 0.2.",
),
stol: float = Query(
0.3,
title=
"Site tolerance. Defined as the fraction of the average free \
length per atom := ( V / Nsites ) ** (1/3). Default is 0.3.",
),
angle_tol: float = Query(
5,
title="Angle tolerance in degrees. Default is 5 degrees.",
),
limit: int = Query(
1,
title=
"Maximum number of matches to show. Defaults to 1, only showing the best match.",
),
):
"""
Obtains material structures that match a given input structure within some tolerance.
Returns:
A list of Material IDs for materials with matched structures alongside the associated RMS values
"""
try:
s = PS.from_dict(structure.dict())
except Exception:
raise HTTPException(
status_code=404,
detail=
"Body cannot be converted to a pymatgen structure object.",
)
m = StructureMatcher(
ltol=ltol,
stol=stol,
angle_tol=angle_tol,
primitive_cell=True,
scale=True,
attempt_supercell=False,
comparator=ElementComparator(),
)
crit = {"composition_reduced": dict(s.composition.to_reduced_dict)}
self.store.connect()
matches = []
for r in self.store.query(criteria=crit,
properties=["structure", "task_id"]):
s2 = PS.from_dict(r["structure"])
matched = m.fit(s, s2)
if matched:
rms = m.get_rms_dist(s, s2)
matches.append({
"task_id": r["task_id"],
"normalized_rms_displacement": rms[0],
"max_distance_paired_sites": rms[1],
})
response = {
"data":
sorted(
matches[:limit],
key=lambda x: (
x["normalized_rms_displacement"],
x["max_distance_paired_sites"],
),
)
}
return response
self.router.post(
"/find_structure/",
response_model_exclude_unset=True,
response_description=
f"Get matching structures using data from {model_name}",
tags=self.tags,
)(find_structure)
def custom_autocomplete_prep(self):
async def formula_autocomplete(
text: str = Query(
...,
description="Text to run against formula autocomplete",
),
limit: int = Query(
10,
description="Maximum number of matches to show. Defaults to 10",
),
):
comp = Composition(text)
ind_str = []
if len(comp) == 1:
d = comp.get_integer_formula_and_factor()
s = d[0] + str(int(d[1])) if d[1] != 1 else d[0]
ind_str.append(s)
else:
comp_red = comp.reduced_composition.items()
for (i, j) in comp_red:
if j != 1:
ind_str.append(i.name + str(int(j)))
else:
ind_str.append(i.name)
final_terms = ["".join(entry) for entry in permutations(ind_str)]
pipeline = [
{
"$search": {
"index": "formula_autocomplete",
"autocomplete": {
"path": "formula_pretty",
"query": final_terms,
"tokenOrder": "any",
},
}
},
{
"$group": {
"_id": "$formula_pretty",
}
},
{
"$project": {
"score": {
"$strLenCP": "$_id"
}
}
},
{
"$sort": {
"score": 1
}
},
{
"$limit": limit
},
]
self.store.connect()
data = list(
self.store._collection.aggregate(pipeline, allowDiskUse=True))
response = {"data": data}
return response
self.router.get(
"/formula_autocomplete/",
response_model_exclude_unset=True,
response_description="Get autocomplete results for a formula",
tags=self.tags,
)(formula_autocomplete)
resource = Resource(
materials_store,
MaterialsCoreDoc,
query_operators=[
VersionQuery(),
FormulaQuery(),
MultiTaskIDQuery(),
SymmetryQuery(),
DeprecationQuery(),
MinMaxQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(
MaterialsCoreDoc,
default_fields=["task_id", "formula_pretty", "last_updated"],
),
],
tags=["Materials"],
custom_endpoint_funcs=[
custom_version_prep,
custom_findstructure_prep,
custom_autocomplete_prep,
],
)
return resource
def charge_density_resource(s3_store):
def custom_charge_density_endpoint_prep(self):
self.s3 = s3_store
model = ChgcarDataDoc
model_name = model.__name__
key_name = "task_id"
field_input = SparseFieldsQuery(
model, [key_name, self.s3.last_updated_field]).query
async def get_chgcar_data(
material_id:
str = Path(
...,
alias=key_name,
title=
f"The Material ID ({key_name}) associated with the {model_name}",
),
fields: STORE_PARAMS = Depends(field_input),
):
f"""
Get's a document by the primary key in the store
Args:
material_id: The Materials Project ID ({key_name}) of a single {model_name}
Returns:
a single {model_name} document
"""
self.s3.connect()
chgcar_key = self.s3.query_one(
criteria={
key_name: material_id
},
properties=[key_name],
).get(key_name, None)
if chgcar_key is None:
raise HTTPException(
status_code=404,
detail=
f"Charge density data with {key_name} = {chgcar_key} not found",
)
item = self.s3.query_one({key_name: chgcar_key},
properties=fields["properties"])
response = item
return response
self.router.get(
f"/{{{key_name}}}/",
response_description=f"Get an {model_name} by {key_name}",
response_model=model,
response_model_exclude_unset=True,
tags=self.tags,
)(get_chgcar_data)
resource = Resource(
s3_store,
ChgcarDataDoc,
query_operators=[
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(ChgcarDataDoc,
default_fields=["task_id", "last_updated"]),
],
tags=["Charge Density"],
custom_endpoint_funcs=[custom_charge_density_endpoint_prep],
enable_default_search=False,
enable_get_by_key=False,
)
return resource
def synth_resource(synth_store):
def custom_synth_prep(self):
async def query_synth_text(
keywords: str = Query(
...,
description=
"Comma delimited string keywords to search synthesis description text with",
),
skip: int = Query(
0, description="Number of entries to skip in the search"),
limit: int = Query(
100,
description=
"Max number of entries to return in a single query. Limited to 100",
),
):
pipeline = [
{
"$search": {
"index": "synth_descriptions",
"regex": {
"query":
[word + ".*" for word in keywords.split(",")],
"path": "text",
"allowAnalyzedField": True,
},
}
},
{
"$project": {
"_id": 0,
"doi": 1,
"formula": 1,
"text": 1,
"search_score": {
"$meta": "searchScore"
},
}
},
{
"$sort": {
"search_score": -1
}
},
{
"$skip": skip
},
{
"$limit": limit
},
]
self.store.connect()
data = list(
self.store._collection.aggregate(pipeline, allowDiskUse=True))
response = {"data": data}
return response
self.router.get(
"/text_search/",
response_model=self.response_model,
response_model_exclude_unset=True,
response_description=
"Find synthesis description documents through text search.",
tags=self.tags,
)(query_synth_text)
resource = Resource(
synth_store,
SynthesisDoc,
query_operators=[
SynthFormulaQuery(),
SortQuery(),
PaginationQuery(),
SparseFieldsQuery(SynthesisDoc, default_fields=["formula", "doi"]),
],
tags=["Synthesis"],
custom_endpoint_funcs=[custom_synth_prep],
enable_default_search=True,
enable_get_by_key=False,
)
return resource