add crystallographic data to fetch metadata

Author: esoteric-ephemera

mp_api/mcp/_schemas.py

@@ -3,7 +3,9 @@
 
 from typing import Any
 
+from emmet.core.summary import SummaryDoc
 from pydantic import BaseModel, Field, model_validator
+from typing_extensions import Self
 
 from mp_api.client.core.utils import validate_ids
 
@@ -264,6 +266,120 @@ class MaterialMetadata(BaseModel):
         ),
     )
 
+    cell_vectors: list[list[float]] | None = Field(
+        None,
+        description=(
+            "The 3x3 array of (lattice) cell vectors, all values in Å."
+            "The first row is the a vector, the second the b vector, "
+            "and the third the c vector."
+        ),
+    )
+
+    atoms: list[str] | None = Field(
+        None,
+        description=(
+            "A list of atom symbols on each site. Should have length `nsites`."
+        ),
+    )
+
+    cartesian_coordinates: list[list[float]] | None = Field(
+        None,
+        description=(
+            "A `nsites` x 3 array of floats, all values in Å, "
+            "representing the Cartesian coordinates of the atoms."
+            "The order is the same as in `atoms`."
+        ),
+    )
+    magnetic_moments: list[float] | None = Field(
+        None,
+        description=(
+            "A `nsites` array of floats, all values in μB, representing "
+            "the on-site magnetic moments found by integrating the "
+            "electronic spin density in a sphere surrounding each site "
+            "in the structure."
+        ),
+    )
+
+    @staticmethod
+    def _summary_fields() -> list[str]:
+        """Get a list of the fields needed in a SummaryDoc to populate this document."""
+        return [
+            *(set(MaterialMetadata.model_fields) & set(SummaryDoc.model_fields)),
+            # The following fields get renamed and flattened in `MaterialMetadata`
+            "structure",
+            "bulk_modulus",
+            "shear_modulus",
+            "database_IDs",
+            "symmetry",
+        ]
+
+    @classmethod
+    def from_summary_data(cls, summary_data: dict[str, Any], **kwargs) -> Self:
+        """Create a MaterialMetadata document from materials summary data.
+
+        Args:
+            summary_data : dict of str to Any
+                The dict representation of an `emmet.core.summary.SummaryDoc`
+                document (i.e., its `model_dump_json`)
+            **kwargs : to pass to `MaterialMetadata`
+        """
+        metadata = {
+            **kwargs,
+            **{
+                k: summary_data[k]
+                for k in MaterialMetadata.model_fields
+                if summary_data.get(k) is not None
+            },
+        }
+        for k in {"bulk", "shear"}:
+            if summary_data.get(f"{k}_modulus"):
+                metadata.update(
+                    {
+                        f"{k}_modulus_{v}": summary_data[f"{k}_modulus"].get(v)
+                        for v in ("voigt", "reuss", "hill")
+                    }
+                )
+
+        # Augment with experimental database id information
+        if summary_data.get("database_IDs"):
+            metadata.update(
+                {
+                    f"linked_{database}_ids": ", ".join(matched_ids)
+                    for database, matched_ids in summary_data["database_IDs"].items()
+                }
+            )
+
+        if (symm_meta := summary_data.get("symmetry")) is not None:
+            metadata.update(
+                {
+                    k: symm_meta.get(v)
+                    for k, v in {
+                        "space_group_number": "number",
+                        "space_group_symbol": "symbol",
+                        "crystal_system": "crystal_system",
+                        "point_group": "point_group",
+                    }.items()
+                }
+            )
+
+        # flatten structure data
+        if struct_dict := summary_data.get("structure"):
+            magnetic_moments = [
+                site["properties"].get("magmom") for site in struct_dict["sites"]
+            ]
+            metadata.update(
+                cell_vectors=struct_dict["lattice"]["matrix"],
+                atoms=[site["species"][0]["element"] for site in struct_dict["sites"]],
+                cartesian_coordinates=[site["xyz"] for site in struct_dict["sites"]],
+                magnetic_moments=(
+                    None
+                    if any(magmom is None for magmom in magnetic_moments)
+                    else magnetic_moments
+                ),
+            )
+
+        return cls(**metadata)
+
 
 class FetchResult(BaseModel):
     """Schematize result of the `fetch` MCP tool.

mp_api/mcp/tools.py

@@ -167,57 +167,7 @@ def fetch(self, idx: str) -> FetchResult:
         if not robo_desc:
             return FetchResult(id=idx)
 
-        metadata: dict[str, str] | None = None
-        if (
-            len(
-                summary_docs := self.client.materials.summary.search(
-                    material_ids=[idx],
-                    fields=list(
-                        set(MaterialMetadata.model_fields)
-                        & set(self.client.materials.summary.document_model.model_fields)
-                    ),
-                )
-            )
-            > 0
-        ):
-            # Try to avoid more nested fields, just provide things with
-            # simple str or numeric type
-            summary_doc = summary_docs[0]
-            metadata = {
-                k: summary_doc[k]
-                for k in MaterialMetadata.model_fields
-                if summary_doc.get(k) is not None
-            }
-            for k in {"bulk", "shear"}:
-                if summary_doc.get(f"{k}_modulus"):
-                    metadata.update(
-                        {
-                            f"{k}_modulus_{v}": summary_doc[f"{k}_modulus"].get(v)
-                            for v in ("voigt", "reuss", "hill")
-                        }
-                    )
-
-            # Augment with experimental database id information
-            if summary_doc.get("database_IDs"):
-                metadata.update(
-                    {
-                        f"linked_{database}_ids": ", ".join(matched_ids)
-                        for database, matched_ids in summary_doc["database_IDs"].items()
-                    }
-                )
-
-            if (symm_meta := summary_doc.get("symmetry")) is not None:
-                metadata.update(
-                    {
-                        k: symm_meta.get(v)
-                        for k, v in {
-                            "space_group_number": "number",
-                            "space_group_symbol": "symbol",
-                            "crystal_system": "crystal_system",
-                            "point_group": "point_group",
-                        }.items()
-                    }
-                )
+        metadata: dict[str, str] = {}
 
         if len(sim_docs := self.client.materials.similarity.find_similar(idx, top=10)):
             if not isinstance(sim_docs[0], dict):
@@ -231,7 +181,25 @@ def fetch(self, idx: str) -> FetchResult:
                 )
             )
 
-        return FetchResult(id=idx, text=robo_desc, metadata=metadata)
+        summary_doc = {}
+        if (
+            len(
+                summary_docs := self.client.materials.summary.search(
+                    material_ids=[idx],
+                    fields=MaterialMetadata._summary_fields(),
+                )
+            )
+            > 0
+        ):
+            # Try to avoid more nested fields, just provide things with
+            # simple str or numeric type
+            summary_doc = summary_docs[0]
+
+        return FetchResult(
+            id=idx,
+            text=robo_desc,
+            metadata=MaterialMetadata.from_summary_data(summary_doc, **metadata),
+        )
 
     def get_phase_diagram_from_elements(
         self,

tests/mcp/test_tools.py

@@ -1,5 +1,6 @@
-from re import search
+import numpy as np
 from pymatgen.core import Composition
+import pytest
 
 from mp_api.mcp._schemas import SearchOutput, FetchResult, MaterialMetadata
 from mp_api.mcp.tools import MPCoreMCP
@@ -23,6 +24,7 @@ def test_core_tools():
         robo_desc_docs = mcp_tools.client.materials.robocrys.search_docs(
             material_ids=[*[doc.id for doc in search_results.results], fetch_results.id]
         )
+        ref_struct = mcp_tools.client.get_structure_by_material_id(fetch_results.id)
 
     robo_descs = {doc["material_id"]: doc["description"] for doc in robo_desc_docs}
 
@@ -41,3 +43,19 @@ def test_core_tools():
     assert isinstance(fetch_results.metadata, MaterialMetadata)
     assert isinstance(fetch_results.metadata.structurally_similar_materials, str)
     assert fetch_results.text == robo_descs[fetch_results.id]
+
+    assert np.allclose(
+        ref_struct.lattice.matrix,
+        fetch_results.metadata.cell_vectors,
+    )
+    assert np.allclose(
+        ref_struct.cart_coords,
+        fetch_results.metadata.cartesian_coordinates,
+    )
+    assert fetch_results.metadata.atoms == [
+        str(site.species.elements[0]) for site in ref_struct
+    ]
+    if magmoms := ref_struct.site_properties.get("magmom"):
+        assert fetch_results.metadata.magnetic_moments == pytest.approx(magmoms)
+    else:
+        assert fetch_results.metadata.magnetic_moments is None