Charge density endpoint and client updates

src/mp_api/routes/charge_density/client.py

@@ -1,49 +1,72 @@
 from mp_api.core.client import BaseRester, MPRestError
 from mp_api.routes.tasks.client import TaskRester
+from mp_api.routes.materials.client import MaterialsRester
 
 
 class ChargeDensityRester(BaseRester):
 
     suffix = "charge_density"
 
-    def get_charge_density_from_material_id(self, material_id: str):
+    def get_charge_density_from_calculation_id(self, task_id: str):
         """
-        Get charge density data for a given Materials Project ID.
+        Get charge density data for a given calculation ID.
 
         Arguments:
-            material_id (str): Materials project ID
+            task_id (str): Calculation ID
 
         Returns:
             chgcar (dict): Pymatgen CHGCAR object.
         """
 
-        result = self._make_request(
-            "{}/?fields=data&all_fields=false".format(material_id)
-        )
+        result = self.get_document_by_id(document_id=task_id)
 
         if result.get("data", None) is not None:
             return result["data"]
         else:
-            raise MPRestError("No document found")
+            raise MPRestError("No charge density found")
 
-    def get_calculation_details(self, material_id: str):
+    def get_calculation_details(self, task_id: str):
         """
-        Get charge density calculations details for a given Materials Project ID.
+        Get charge density calculations details for a given calculation ID.
 
         Arguments:
-            material_id (str): Materials project ID
+            task_id (str): Calculation ID
 
         Returns:
             calc_details (dict): Dictionary containing INCAR, POSCAR, and KPOINTS data for the DFT calculation.
         """
 
-        base_endpoint = "/".join(self.endpoint.split("/")[0:3])
-        task_rester = TaskRester(api_key=self.api_key, endpoint=base_endpoint)  # type: ignore
+        task_rester = TaskRester(endpoint=self.base_endpoint, api_key=self.api_key)  # type: ignore
 
-        result = task_rester.get_task_from_material_id(  # type: ignore
-            material_id, fields=["task_id", "orig_inputs"]
-        ).get("data")[0]
-
-        task_rester.session.close()  # type: ignore
+        result = task_rester.get_document_by_id(
+            document_id=task_id,
+            fields=["orig_inputs.incar", "orig_inputs.poscar", "orig_inputs.kpoints"],
+        ).orig_inputs
 
         return result
+
+    def get_calculation_ids_from_material_id(self, material_id: str):
+        """
+        Get charge density calculation ids associated with a given Materials Project ID.
+
+        Arguments:
+            material_id (str): Materials Project ID
+
+        Returns:
+            calculation_ids (List[str]): List of calculation ids.
+        """
+
+        materials_rester = MaterialsRester(
+            endpoint=self.base_endpoint, api_key=self.api_key
+        )  # type: ignore
+
+        calculation_types = materials_rester.get_document_by_id(
+            document_id=material_id, fields=["calc_types"]
+        ).calc_types
+
+        calculation_ids = []
+        for calculation_id, calculation_type in calculation_types.items():
+            if "Static" in calculation_type:
+                calculation_ids.append(calculation_id)
+
+        return calculation_ids

src/mp_api/routes/charge_density/models.py

@@ -1,6 +1,7 @@
 from pydantic import BaseModel, Field, validator
 from datetime import datetime
 from monty.json import MontyDecoder
+from pymatgen.io.vasp import Chgcar
 
 
 class ChgcarDataDoc(BaseModel):
@@ -23,9 +24,4 @@ class ChgcarDataDoc(BaseModel):
             This comes in the form: mp-******",
     )
 
-    data: dict = Field(None, description="Pymatgen CHGCAR object.")
-
-    # Make sure that the datetime field is properly formatted
-    @validator("last_updated", pre=True)
-    def last_updated_dict_ok(cls, v):
-        return MontyDecoder().process_decoded(v)
+    data: Chgcar = Field(None, description="Pymatgen CHGCAR object.")

src/mp_api/routes/charge_density/resources.py

@@ -1,5 +1,6 @@
 from fastapi.param_functions import Path
 from mp_api.core.resource import GetResource
+from mp_api.core.models import Response
 from mp_api.routes.charge_density.models import ChgcarDataDoc
 from mp_api.core.query_operator import PaginationQuery, SparseFieldsQuery, SortQuery
 from mp_api.core.utils import STORE_PARAMS
@@ -47,15 +48,13 @@ async def get_chgcar_data(
                     status_code=404,
                     detail=f"Item with {key_name} = {material_id} not found",
                 )
-
-            response = item
-
-            return response
+            else:
+                return {"data": [item]}
 
         self.router.get(
             f"/{{{key_name}}}/",
             response_description=f"Get an {model_name} by {key_name}",
-            response_model=model,
+            response_model=Response[model],
             response_model_exclude_unset=True,
             tags=self.tags,
         )(get_chgcar_data)