A brief introduction to the Materials Project API

The Materials Project API provides programmatic access to a massive store of computed materials data: from crystal structures (CIFs), to band structures, to thermodynamic quantities like formation energy and energy above hull. The official Python client is mp_api. Typical usage (non-MCP) might look like:

from mp_api.client import MPRester

with MPRester("your_api_key_here") as mpr:
    docs = mpr.materials.summary.search(elements=["Si", "O"], band_gap=(0.5, 1.0))

Although it’s straightforward to call mp_api from your local Python environment, hooking it up to an LLM for agent-like usage typically requires additional bridging code.

Why not just call the MP API directly?

If you’re only using one AI platform or a single environment (like a Jupyter notebook), you can just call mp_api functions directly. However, this approach doesn’t scale well if you want:

1. Multiple AI apps: Tools like Cursor, Windsurf, Claude Desktop, Emacs clients, etc., all want to discover and call your “tools” in a standard, stable way.
2. One integration for many LLMs: With MCP, you write your code once. Any LLM app that supports the protocol can discover your server, see its docs, and call your exposed tools.
3. Granular permission control: The protocol ensures the user must approve a tool call if the LLM attempts to use it. This keeps you in the loop on usage.

Using MCP is especially convenient if you want to combine the Materials Project data with other advanced data sources or local resources, e.g. local quantum chemistry codes, Python scripts, or HPC job management tools. AI applications that speak MCP can chain all of them together seamlessly.

Step-by-step: The Materials Project MCP server

Below is the main code of our server (a Python script). You can see how we define two tools:

# File: materials_project_plugin.py
#
# An MCP server exposing tools for querying the Materials Project database
# via mp_api. Once running, tools can be invoked through any MCP-compatible
# client (e.g. Claude Desktop).
#
# Dependencies:
#   pip install "mcp[cli]" aiohttp pydantic mp_api 
#
# Environment:
#   MP_API_KEY=<your_materials_project_api_key>
#

import os
import logging
from typing import Optional, List
from emmet.core.electronic_structure import BSPathType
from mcp.server.fastmcp import FastMCP
from pydantic import Field

# Materials Project client
from mp_api.client import MPRester

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("materials_project_mcp")

# Obtain your Materials Project API key from env var
API_KEY = os.environ.get("MP_API_KEY")

# Create the MCP server instance
mcp = FastMCP(
    name="Materials Project Plugin",
    version="0.0.1",
    description=(
        "A Model Context Protocol (MCP) server that exposes query tools "
        "for the Materials Project database using the mp_api client."
    ),
)


def _get_mp_rester() -> MPRester:
    """
    Helper function to initialize a MPRester session with the user's API key.
    """
    if not API_KEY:
        logger.warning(
            "No MP_API_KEY found in environment. Attempting MPRester() without key."
        )
        return MPRester()
    return MPRester(API_KEY)


@mcp.tool()
async def search_materials(
    elements: Optional[List[str]] = Field(
        default=None,
        description="List of element symbols to filter by (e.g. ['Si', 'O']).",
    ),
    band_gap_min: float = Field(
        default=0.0, description="Lower bound for band gap filtering in eV."
    ),
    band_gap_max: float = Field(
        default=10.0, description="Upper bound for band gap filtering in eV."
    ),
    is_stable: bool = Field(
        default=False,
        description="Whether to only retrieve stable materials (True) or all (False).",
    ),
    max_results: int = Field(
        default=20, ge=1, le=200, description="Maximum number of results to return."
    ),
) -> str:
    """
    Search for materials in the Materials Project database using basic filters:
    - elements (list of elements to include)
    - band_gap (range in eV)
    - is_stable
    Returns a formatted list of matches with their material_id, formula, and band gap.
    """
    logger.info("Starting search_materials query...")
    with _get_mp_rester() as mpr:
        docs = mpr.materials.summary.search(
            elements=elements,
            band_gap=(band_gap_min, band_gap_max),
            is_stable=is_stable,
            fields=["material_id", "formula_pretty", "band_gap", "energy_above_hull"],
        )

    # Truncate results to max_results
    docs = list(docs)[:max_results]

    if not docs:
        return "No materials found matching your criteria."

    results_md = (
        f"## Materials Search Results\n\n"
        f"- **Elements**: {elements or 'Any'}\n"
        f"- **Band gap range**: {band_gap_min} eV to {band_gap_max} eV\n"
        f"- **Stable only**: {is_stable}\n\n"
        f"**Showing up to {max_results} matches**\n\n"
    )
    for i, mat in enumerate(docs, 1):
        results_md += (
            f"**{i}.** ID: `{mat.material_id}` | Formula: **{mat.formula_pretty}** | "
            f"Band gap: {mat.band_gap:.3f} eV | E above hull: {mat.energy_above_hull:.3f} eV\n"
        )
    return results_md


@mcp.tool()
async def get_structure_by_id(
    material_id: str = Field(..., description="Materials Project ID (e.g. 'mp-149')")
) -> str:
    """
    Retrieve the final computed structure for a given material_id from the Materials Project.
    Returns a plain text summary of the structure (lattice, sites, formula).
    """
    logger.info(f"Fetching structure for {material_id}...")
    with _get_mp_rester() as mpr:
        # Shortcut method to get just the final structure
        structure = mpr.get_structure_by_material_id(material_id)

    if not structure:
        return f"No structure found for {material_id}."

    # Summarize the structure
    formula = structure.composition.reduced_formula
    lattice = structure.lattice
    sites_count = len(structure)
    text_summary = (
        f"## Structure for {material_id}\n\n"
        f"- **Formula**: {formula}\n"
        f"- **Lattice**:\n"
        f"   a = {lattice.a:.3f} Å, b = {lattice.b:.3f} Å, c = {lattice.c:.3f} Å\n"
        f"   α = {lattice.alpha:.2f}°, β = {lattice.beta:.2f}°, γ = {lattice.gamma:.2f}°\n"
        f"- **Number of sites**: {sites_count}\n"
        f"- **Reduced formula**: {structure.composition.reduced_formula}\n"
    )
    return text_summary


if __name__ == "__main__":
    logger.info("Starting Materials Project MCP server...")
    mcp.run()

1. search_materials: Queries the Materials Project for stable/instable materials with certain elements and band gap constraints.
2. get_structure_by_id: Retrieves the final computed crystal structure for a known material_id (like mp-149 for silicon).

When run, it starts an MCP server over STDIO. Any MCP-compatible client can spawn this script, discover the tools, and use them interactively with user approval.

Wrap-up and additional links

That’s it! You now have an MCP server that can call Materials Project queries. You can start from this basic example and adapt it to your needs by:

• Experiment with advanced queries (like stable phases in multi-element chemical systems).
• Add or refine tools in your server for more specialized searches (e.g., stable doping compositions).
• Combine with other servers in your environment, for a multi-functional AI agent.

Understanding Nodeology’s Core Concepts

Defining Our Simulation State

We’ll begin by specifying which data we expect to carry throughout the workflow. This includes physical parameters (e.g., mass, charge, initial_velocity, etc.), plus placeholders for the user’s analysis or plot results.

from typing import List, Dict
import numpy as np
from nodeology.state import State

class TrajectoryState(State):
    """
    State for our particle trajectory workflow.
    It holds everything from initial parameters,
    user confirmations, and final analysis.
    """
    mass: float                   # Particle mass (kg)
    charge: float                 # Particle charge (C)
    initial_velocity: np.ndarray  # [vx, vy, vz]
    E_field: np.ndarray           # [Ex, Ey, Ez]
    B_field: np.ndarray           # [Bx, By, Bz]
    
    # A boolean for user confirmation
    confirm_parameters: bool

    # For storing the raw JSON from an LLM-based node
    parameters_updater_output: str

    # The computed 3D positions of the particle
    positions: List[np.ndarray]

    # We'll store a path or figure for the plotted trajectory
    trajectory_plot: str
    trajectory_plot_path: str

    # Where the LLM’s analysis result will land
    analysis_result: Dict

    # A boolean for continuing or finishing
    continue_simulation: bool

Key Takeaway: If you want a new value to persist or be shared across nodes, define it here.

Writing Our Nodes

import chainlit as cl
from chainlit import Message, run_sync
from nodeology.node import as_node

@as_node(sink=[])
def display_parameters(
    mass: float,
    charge: float,
    initial_velocity: np.ndarray,
    E_field: np.ndarray,
    B_field: np.ndarray,
):
    """
    Display the current simulation parameters using
    Chainlit's custom UI element.
    """
    parameters = {
        "Mass (kg)": mass,
        "Charge (C)": charge,
        "Initial Velocity (m/s)": initial_velocity.tolist(),
        "Electric Field (N/C)": E_field.tolist(),
        "Magnetic Field (T)": B_field.tolist(),
    }

    run_sync(
        Message(
            content="Below are the current simulation parameters:",
            elements=[
                cl.CustomElement(
                    name="DataDisplay",
                    props={
                        "data": parameters,
                        "title": "Particle Parameters",
                        "badge": "Configured",
                    },
                )
            ],
        ).send()
    )
    # No return value needed, so sink=[] above

from chainlit import AskActionMessage

@as_node(sink="confirm_parameters")
def ask_confirm_parameters():
    """
    Prompt user with 2 actions: 'Yes' or 'No'.
    Return True if 'Yes', False otherwise.
    """
    res = run_sync(
        AskActionMessage(
            content="Are you happy with the parameters?",
            actions=[
                cl.Action(name="yes", payload={"value": "yes"}, label="Yes"),
                cl.Action(name="no",  payload={"value": "no"},  label="No"),
            ],
        ).send()
    )
    # We store the boolean in 'confirm_parameters' in the State
    if res and res["payload"]["value"] == "yes":
        return True
    else:
        return False

from chainlit import AskUserMessage

@as_node(sink=["human_input"])
def ask_parameters_input():
    """
    Wait for a text message from the user specifying
    how they want to update the parameters.
    """
    user_msg = run_sync(
        AskUserMessage(
            content="Please let me know how you want to change any of the parameters :)",
        ).send()
    )
    return user_msg["output"]  # We'll store this in 'human_input'

from nodeology.node import Node
import json
import numpy as np

parameters_updater = Node(
    node_type="parameters_updater",
    prompt_template="""Update the parameters based on the user's input.

Current parameters:
mass: {mass}
charge: {charge}
initial_velocity: {initial_velocity}
E_field: {E_field}
B_field: {B_field}

User input:
{human_input}

Please return the updated parameters in JSON format.
{
    "mass": float,
    "charge": float,
    "initial_velocity": list[float],
    "E_field": list[float],
    "B_field": list[float]
}
""",
    sink="parameters_updater_output",  # We'll store the LLM response here (raw text/JSON)
    sink_format="json",
)

# We'll need a post-processing function to interpret that JSON
def parameters_updater_transform(state, client, **kwargs):
    # Convert the LLM's output from text to Python objects
    params_dict = json.loads(state["parameters_updater_output"])
    state["mass"] = params_dict["mass"]
    state["charge"] = params_dict["charge"]
    state["initial_velocity"] = np.array(params_dict["initial_velocity"])
    state["E_field"] = np.array(params_dict["E_field"])
    state["B_field"] = np.array(params_dict["B_field"])
    return state

parameters_updater.post_process = parameters_updater_transform

import tempfile
from scipy.integrate import solve_ivp
from typing import List

@as_node(sink=["positions"])
def calculate_trajectory(
    mass: float,
    charge: float,
    initial_velocity: np.ndarray,
    E_field: np.ndarray,
    B_field: np.ndarray,
) -> List[np.ndarray]:
    """
    Numerically integrate the trajectory of a particle under
    electric & magnetic fields, returning the positions over time.
    """
    # Estimate cyclotron period if B is non-zero:
    B_magnitude = np.linalg.norm(B_field)
    if B_magnitude == 0 or charge == 0:
        cyclotron_period = 1e-6
    else:
        cyclotron_frequency = abs(charge) * B_magnitude / mass
        cyclotron_period    = 2 * np.pi / cyclotron_frequency

    # We'll simulate 5 cycles, each with 100 steps
    num_periods = 5
    num_points_per_period = 100
    total_time = num_periods * cyclotron_period
    total_points = num_periods * num_points_per_period

    time_points = np.linspace(0, total_time, total_points)

    def lorentz_force(t, state):
        vel = state[3:]
        force = charge * (E_field + np.cross(vel, B_field))
        acc = force / mass
        return np.concatenate([vel, acc])

    initial_position = np.array([0.0, 0.0, 0.0])
    initial_state = np.concatenate([initial_position, initial_velocity])

    sol = solve_ivp(
        lorentz_force,
        (time_points[0], time_points[-1]),
        initial_state,
        t_eval=time_points,
        method="RK45",
        rtol=1e-8,
    )

    if not sol.success:
        # Return zeros if something fails
        return [np.zeros(3) for _ in range(len(time_points))]

    return [sol.y[:3, i] for i in range(len(time_points))]

import plotly.graph_objects as go

@as_node(sink=["trajectory_plot", "trajectory_plot_path"])
def plot_trajectory(positions: List[np.ndarray]) -> str:
    """
    Generate a 3D plot of the positions and
    save it to a temporary file.
    """
    arr = np.array(positions)
    fig = go.Figure(
        data=[
            go.Scatter3d(
                x=arr[:, 0],
                y=arr[:, 1],
                z=arr[:, 2],
                mode="lines",
                line=dict(width=4, color="green"),
            )
        ]
    )
    fig.update_layout(
        scene=dict(xaxis_title="X (m)", yaxis_title="Y (m)", zaxis_title="Z (m)")
    )

    image_path = tempfile.mktemp(suffix=".png")
    fig.write_image(image_path)

    # Display in UI
    run_sync(
        Message(
            content="Below is the trajectory plot:",
            elements=[cl.Plotly(figure=fig)],
        ).send()
    )

    # Return figure + path
    return fig, image_path

trajectory_analyzer = Node(
    node_type="trajectory_analyzer",
    prompt_template="""Analyze this particle trajectory plot.

Please determine:
1. The type of motion (linear, circular, helical, or chaotic)
2. Key physical features (radius, period, pitch angle if applicable)
3. Explanation of the motion
4. Anomalies in the motion

Output in JSON format:
{
  "trajectory_type": "type_name",
  "key_features": {
     "feature1": value,
     "feature2": value
  },
  "explanation": "detailed explanation",
  "anomalies": "anomaly description"
}""",
    sink="analysis_result",
    sink_format="json",
    image_keys=["trajectory_plot_path"],  # for image-based context, if the LLM can handle it
)

def display_trajectory_analyzer_result(state, client, **kwargs):
    """
    A post-process function that picks up the LLM's
    output, converts it to a dict, and displays it nicely.
    """
    import json
    state["analysis_result"] = json.loads(state["analysis_result"])
    run_sync(
        Message(
            content="Here is the trajectory analysis:",
            elements=[
                cl.CustomElement(
                    name="DataDisplay",
                    props={
                        "data": state["analysis_result"],
                        "title": "Trajectory Analysis",
                        "badge": state["analysis_result"].get("trajectory_type", "Unknown"),
                    },
                )
            ],
        ).send()
    )
    return state

trajectory_analyzer.post_process = display_trajectory_analyzer_result

@as_node(sink="continue_simulation")
def ask_continue_simulation():
    res = run_sync(
        AskActionMessage(
            content="Would you like to continue the simulation?",
            actions=[
                cl.Action(name="continue", payload={"value": "continue"}, label="Continue Simulation"),
                cl.Action(name="finish",   payload={"value": "finish"},   label="Finish"),
            ],
        ).send()
    )
    return (res and res["payload"]["value"] == "continue")

Orchestrating Everything in a Workflow

Now we piece all Nodes together in a Workflow. Let’s define it:

from langgraph.graph import END
from nodeology.workflow import Workflow

class TrajectoryWorkflow(Workflow):
    def create_workflow(self):
        # 1) Register our Nodes
        self.add_node("display_parameters", display_parameters)
        self.add_node("ask_confirm_parameters", ask_confirm_parameters)
        self.add_node("ask_parameters_input", ask_parameters_input)
        self.add_node("update_parameters", parameters_updater)
        self.add_node("calculate_trajectory", calculate_trajectory)
        self.add_node("plot_trajectory", plot_trajectory)
        self.add_node("analyze_trajectory", trajectory_analyzer)
        self.add_node("ask_continue_simulation", ask_continue_simulation)

        # 2) Define the flow (edges)
        self.add_flow("display_parameters", "ask_confirm_parameters")
        self.add_conditional_flow(
            "ask_confirm_parameters",
            "confirm_parameters",       # If user is happy
            then="calculate_trajectory",
            otherwise="ask_parameters_input",  # If not
        )
        self.add_flow("ask_parameters_input", "update_parameters")
        self.add_flow("update_parameters", "display_parameters")
        self.add_flow("calculate_trajectory", "plot_trajectory")
        self.add_flow("plot_trajectory", "analyze_trajectory")
        self.add_flow("analyze_trajectory", "ask_continue_simulation")
        self.add_conditional_flow(
            "ask_continue_simulation",
            "continue_simulation",  # If user wants to keep going
            then="display_parameters",
            otherwise=END,         # If user says “finish”
        )

        # 3) Set the first Node to run
        self.set_entry("display_parameters")

        # 4) Compile
        self.compile()

Note: Once compiled, Nodeology internally organizes everything as a state machine with branching logic, enabling flexible re-runs and so on.

Launching the Interactive Chat

At last, we bring it all together. We instantiate our workflow with an initial set of parameters and run it in UI mode:

# 1) Create the workflow
workflow = TrajectoryWorkflow(
    state_defs=TrajectoryState,       # Our custom state
    llm_name="gemini/gemini-2.0-flash",
    vlm_name="gemini/gemini-2.0-flash",
    debug_mode=False
)

# 2) Provide initial data
initial_state = {
    "mass": 9.1093837015e-31,  # electron mass
    "charge": -1.602176634e-19,
    "initial_velocity": np.array([1e6, 1e6, 1e6]),
    "E_field": np.array([5e6, 1e6, 5e6]),
    "B_field": np.array([0.0, 0.0, 50000.0]),
}

# 3) Run with a user interface
result = workflow.run(init_values=initial_state, ui=True)

When you run this file, a Chainlit web server will pop up in your console logs, typically at http://localhost:8000. Open that address in your browser and you now have a complete, AI-driven pipeline that merges human interaction, classical simulation, and LLM-based analysis. The user sees a straightforward chat interface, while under the hood, Nodeology coordinates each Node in a robust workflow. Below is a short video showing the workflow in action:

What is “Deep Research”?

There is not a rigorous definition of “Deep Research”, but it generally refers to an AI-powered agent that autonomously conducts multi-step investigations by gathering, analyzing, and synthesizing diverse online data—from text and images to PDFs—to generate comprehensive, cited reports. This concept has origins in frameworks such as STORM and has been further developed by large players:

• OpenAI integrated a variant into ChatGPT for Pro subscribers.
• Google’s Gemini includes a “Deep Research” component for Gemini Advanced subscribers.
• Some open-source alternatives are also emerging.

Excitement has grown around these capabilities, fueled by user testimonials and demos. Researchers like me are wondering if these tools can expedite their literature reviews while maintaining rigor and reliability.

Deep Research for Scientific Studies

In academia, literature reviews are fundamental to scientific research. Scientists regularly summarize current research findings to identify trends, gaps, and future directions. If Deep Research could automate large portions of the review process, that would be a significant boon.

I chose generative models for inorganic crystal structures as my test case for several reasons. During my PhD research in computational materials design, I focused extensively on inorganic crystal materials and authored a review paper on crystal structure prediction (CSP). At that time, generative models were a relatively small subfield. However, there has since been an explosion of research in this area, with the emergence of new techniques like diffusion models, flow matching models, and foundation models, along with more application-oriented studies focusing on conditional generation. This topic is particularly relevant to me as I recently led a seed LDRD project on using crystal generative models for X-ray diffraction (XRD) structure determination. Given both my background and current research interests, I’m eager to understand how this field has evolved.

Simple Prompt Engineering for “Deep Research”

Rather than expecting a single, monolithic prompt to produce a full-blown review, I started by asking each system a guiding question:

I am a computational material scientist trying to do a literature review on generative models for inorganic crystal structure prediction/generation. What do you think I should pay attention to or focus on when reading through the papers?

This initial prompt helped the AI define an outline or review structure. Based on the responses, I then formulated a final prompt:

Can you help me research on generative models for inorganic crystal structure prediction/generation/design? Return me a report with three sections:

1. A table with seven columns: authors/method, generative modeling approach, representation of crystal structures, data and training protocols, validation & evaluation metrics, interpretability, computational efficiency.
2. A comparison of these methods and how the field has evolved.
3. Discussion of gaps, challenges, and future directions.

OpenAI vs. Gemini vs. Semi-Manual

1. OpenAI’s Deep Research
   • After choosing “Deep Research,” ChatGPT typically asks follow-up questions to refine search criteria.
   • An internal version of O3 model is used.
   • It will output a rendered markdown file that the user can copy.
     
   • I ran the same prompt (and minor clarifications) three times to see if results varied.
   • I also merged the three outputs to see if it would improve the quality of the report.
2. Google Gemini
   • Gemini generates a research plan, requests user review, and then “executes” the plan.
   • Currently only Gemini 1.5 Pro can be used for this feature.
   • The final step yields a Google Doc-like report (which can be exported to other formats).
     
   • I repeated the entire process three times as well.
   • I also merged the three outputs to see if it would improve the quality of the report.
3. Semi-Manual Approach
   • I collected 27 relevant papers from my own personal research library (PDF format).
   • Using Marker-PDF, I converted PDFs into markdown files. This will likely overcome paywalls and limited web-scraping issues.
     
   • I fed these markdown documents as a context to LLM, prompting for the same final structure.
   • I used Gemini 2 Pro for this approach due to context window constraints.
   • I repeated the process three times as well.
   • I also merged the three outputs to see if it would improve the quality of the report.

Finally, I also experimented with “merging” all individual reports from each approach to see if combining them could reduce errors or improve coverage. I ended up with 13 total reports (3 from each method, plus 4 merges). All the reports can be found in the Google Drive Folder

Conclusion

From this experiment, it’s clear that Deep Research tools can provide valuable and well-structured summaries, but they aren’t yet a drop-in replacement for a thorough, human-curated literature review—particularly in cutting-edge scientific fields.

Key Observations:

• Information retrieval is not yet up to par due to paywalls, limited web-scraping, and domain specificity.
• Merging multiple AI outputs can substantially improve final report quality
• The advanced models can generate interesting and novel research insights.

Overall, while the promise of “Deep Research” is highly appealing, scientists should remain actively involved for a thorough literature review. Nonetheless, the strides made in foundation models are both impressive and indicative of a rapidly evolving landscape. Deep research as for now is more like a starting point for a literature review. But it can be envisioned that future advancements in data sources, model capabilities, and workflow automation will make it a truly robust autonomous tool for literature review.

Highest Scored Final Report

Below is the final merged output from my semi-manual approach, which scored highest overall.

Code Description

Given the significance of local doping patterns, this blog presents a Python code designed to:

1. Identify Neighbors: Determine atomic species within a specified radius of a target site, forming the essential first step in analyzing the local atomic environment.

2. Utilize Symmetry: Leverage crystallographic site-symmetry operations to identify equivalent neighbor positions. This method eliminates redundancy by consolidating symmetrically equivalent doping configurations.

3. Generate Configurations with Controlled Composition: Perform exhaustive enumeration of all possible local doping arrangements in manageable systems. For more complex systems with extensive possibilities, employ quasi-random Sobol sampling to efficiently explore the configuration space while maintaining targeted dopant concentrations.

4. Represent and Visualize Patterns: Produce structured JSON outputs detailing unique doping patterns and neighbor orbits obtained through site-symmetry operations. Interactive visualization capabilities are provided via the crystaltoolkit package, allowing intuitive inspection of atomic arrangements.

Code Walkthrough

Below, I highlight the essential Python constructs that make the workflow possible.

structure = Structure.from_file("cubic_batio3.cif")
site_index = 0
cutoff = 7.0

target_site = structure[site_index]
neighbors = structure.get_sites_in_sphere(
    pt=target_site.coords,
    r=cutoff,
    include_index=True
)

rel_positions = []
neighbor_indices = []
for (site, dist, idx) in neighbors:
    if idx != site_index:
        rel_positions.append(
            structure.frac_coords[idx] - structure.frac_coords[site_index]
        )
        neighbor_indices.append(idx)
rel_positions = np.mod(rel_positions, 1.0)  # ensure within [0,1) for fractional coords

sym_ops = get_site_symmetries(structure, site_index=site_index)

def get_unique_rotations(sym_ops, decimals: int = 6):
    unique_ops = []
    seen = set()
    for op in sym_ops:
        rot_flat = tuple(np.round(op.rotation_matrix.ravel(), decimals))
        if rot_flat not in seen:
            seen.add(rot_flat)
            unique_ops.append(op)
    return unique_ops

sym_ops = get_unique_rotations(sym_ops)

def build_kdtree(rel_positions: np.ndarray) -> cKDTree:
    return cKDTree(rel_positions)

tree = build_kdtree(rel_positions)

def find_permutation(sym_op, rel_positions: np.ndarray, tree, rtol: float = 1e-2) -> List[int]:
    N = len(rel_positions)
    perm = [None] * N
    transformed = np.array([sym_op.operate(pos) for pos in rel_positions])
    transformed = np.mod(transformed, 1.0)  # keep fractional coords in [0,1)

    for i, tpos in enumerate(transformed):
        dist, idx = tree.query(tpos)
        if dist < rtol:
            perm[i] = idx
        else:
            # handle errors or no match cases
            pass
    return perm

def get_orbits_under_group(N: int, permutations: List[List[int]]) -> List[List[int]]:
    visited = [False] * N
    orbits = []

    for start_idx in range(N):
        if not visited[start_idx]:
            orbit = set()
            queue = [start_idx]
            while queue:
                current = queue.pop()
                if current not in orbit:
                    orbit.add(current)
                    visited[current] = True
                    for perm in permutations:
                        neighbor = perm[current]
                        if not visited[neighbor]:
                            queue.append(neighbor)
            orbits.append(sorted(orbit))
    return orbits

doping_dict = {
  "Ba": ["Ba", "Sr"],
  "Ti": ["Ti", "Zr"]
}

"doping_fraction_constraints": {
  "Sr": [0.0, 0.3],
  "Zr": [0.0, 0.5]
}

How This Code Could Be Used

1. DFT Cluster Setup
   By enumerating unique local doping patterns, you can generate input structures for small-cluster or supercell-based ab initio calculations, ensuring that you capture all relevant local doping motifs (without duplicates) near a site of interest.

2. Machine Learning Training Data
   When training ML potentials or neural network force fields, coverage of doped environments is essential. This code systematically creates local doping configurations to feed into high-fidelity calculations (e.g., DFT) for label generation.

3. High-Throughput Screening
   Researchers exploring doping across multiple materials or sites can automate the pipeline: for each material, each site, and each doping fraction range, the script enumerates or samples doping arrangements.

4. Local Defect Complex Analysis
   In ion conductors or oxide-based electrolytes (like ceria or zirconia), dopants often cluster with oxygen vacancies. With minor modifications (e.g., including "Vac" as a valid species), you can systematically explore dopant–vacancy complexes.

5. Interfacing with Cluster Expansion
   Even if you eventually plan a global cluster expansion, having a library of local dopant environments is valuable for building an initial pool of training structures that sufficiently represent local doping motifs.

Equilibrium Theory for PSA

Code Overview

The Python implementation follows the equilibrium theory equations described by Knaebel & Hill for simulating wave propagation in a PSA (Pressure Swing Adsorption) column. The core data structures (Region, WaveFront, Characteristic) capture the state of the bed and any traveling fronts or characteristics. The main computation logic resides in the WaveTracker class, which calculates wave speeds at both constant and changing pressures. A higher-level controller, WaveDiagram, orchestrates time-stepping (or pressure-stepping) through multiple PSA steps, updates the column profile, and manages plotting.

Examples and Usage

Below are simplified examples illustrating how to set up and run PSA steps. The bed length, inlet compositions, pressures, and velocities are all adjustable, making it easy to experiment with different cycle designs.

# Initialize the WaveTracker and WaveDiagram
wave_tracker = WaveTracker(betaA=0.3, betaB=0.7, beta=0.5)
wave_diagram = WaveDiagram(wave_tracker)

# Define the bed and step parameters
L = 1.0
initial_profile = [Region(y_value=0.0, z_left=0.0, z_right=L)]
T_pressurization = 2.0
T_feed = 2.0
T_blowdown = 4.0
T_purge = 2.0
P_low = 1.0
P_high = 5.0
y_feed = 0.8
y_product = 0.0  # pure light

# Pressurization step: increases pressure from P_low to P_high,
# with flow entering from the 'reverse' end (z=L).
wave_diagram.run_changing_pressure_step(
    step_name="Pressurization",
    y_in=y_product,
    p_start=P_low,
    p_end=P_high,
    t_step=T_pressurization,
    z_profile_in=initial_profile,
    flow_dir="reverse",
)

# Feed step: keeps pressure constant, injects heavier component from z=0
wave_diagram.run_constant_pressure_step(
    step_name="Feed",
    u_in=1.8,
    y_in=y_feed,
    t_step=T_feed,
    flow_dir="default",
)

# Blowdown step: decreases pressure from P_high to P_low
wave_diagram.run_changing_pressure_step(
    step_name="Blowdown",
    y_in=None,
    p_start=P_high,
    p_end=P_low,
    t_step=T_blowdown,
    flow_dir="default",
)

# Purge step: at low pressure, flow enters from z=L to desorb heavy component
wave_diagram.run_constant_pressure_step(
    step_name="Purge",
    u_in=2.0,
    y_in=0.0,
    t_step=T_purge,
    flow_dir="reverse",
)

# Generate and display the wave diagram
wave_diagram.plot_steps(["Pressurization","Feed","Blowdown","Purge"])

How to Run and Adapt

1. Dependencies and Setup. You need Python 3, NumPy, Scipy, and Matplotlib. Save the code as wave_tracking.py (or any name) and run it directly (python wave_tracking.py).
   
2. Modifying the Bed. Adjust or extend initial_profile to reflect different initial loading conditions (multiple Region objects with various y_values).
   
3. Adjusting Equilibrium Parameters. The isotherm parameters kA, kB, and bed porosity epsilon determine the values of betaA, betaB, and beta. Changing these inputs reflects different adsorbent materials.
   
4. Defining Custom Steps. In your own PSA script, call run_constant_pressure_step or run_changing_pressure_step with the times, pressures, flows, and compositions you want. The code automatically spawns shocks or simple waves where discontinuities appear.
   
5. Plotting and Analysis. The plot_steps method creates a stitched diagram of composition fronts across multiple steps. It can be customized or replaced with your own post-processing routines if desired.

By experimenting with different inlet conditions, velocities, and pressure profiles, you can simulate and visualize how composition waves and shocks evolve in a PSA system under various operating strategies.

Cluster Expansions

A cluster expansion approximates a property () of a crystal structure by decomposing that structure into clusters of sites (singlets, pairs, triplets, etc.) and summing up their contributions. Mathematically, if () is the configuration specifying which species occupy each site, we write

where:

• indexes distinct symmetry orbits of clusters (e.g., all pairs of sites separated by the same distance/orientation in the crystal),
• is the effective cluster interaction (ECI) for orbit ,
• is the average of certain basis functions (often orthonormal) over the individual clusters in orbit .

To train a cluster expansion, we first sample a set of training structures/configurations and compute their property values via DFT or other methods. Then we do a linear regression (often with regularization) to obtain the ECIs ({}). The typical workflow looks like this:

• Picking a prototype structure (primitive cell).
• Defining cutoffs for 2-body, 3-body (and higher) distances so we only keep clusters within certain radii.
• Computing the DFT energies of various sampled structures/configurations of the same underlying lattice.
• Solving a linear regression problem to learn those “effective cluster interactions” (ECIs).

These ECIs can then be used to predict the property of new configurations cheaply. Tools like icet automate and streamline the above workflow.

MILP Reformulation

To effectively utilize a trained cluster expansion in an optimization problem, we need to reformulate/transform the parameters. Below we show how to do this for a Mixed-Integer Linear Programming (MILP) formulation, where we have explicit linear expressions and binary variables.

An End-to-end Example

In the sections below, we show how to:

1. Load DFT data for Cu–Ni–Pd–Ag structures
2. Build a ClusterSpace and train the expansion with LassoCV
3. Fully expand the learned cluster expansion onto a “design supercell”
4. Build a matopt MILP and solve for an optimal arrangement.

import itertools
import pandas as pd
import ase
from collections import Counter
from sklearn.linear_model import LassoCV
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_absolute_error

# icet imports
from icet import ClusterSpace, ClusterExpansion, StructureContainer
from icet.core.local_orbit_list_generator import LocalOrbitListGenerator
from icet.core.structure import Structure

# MatOpt imports
from matopt.materials.atom import Atom
from matopt.materials.canvas import Canvas
from matopt.aml.expr import SumSites, SumClusters
from matopt.aml.rule import EqualTo, FixedTo
from matopt.aml.model import MatOptModel

def parameterize_cluster_expansion(cs, ce, design_atoms, weight=1.0, tol=1e-15):
    """
    Expand a trained ClusterExpansion over a 'design_atoms' structure into
    fully enumerated cluster specs and coefficients for use in MatOpt.
    """
    cluster_specs = []
    cluster_coeffs = []

    # (0) Zero-let (constant term)
    zero_coeff = ce.parameters[0] * weight
    if abs(zero_coeff) > tol:
        cluster_specs.append([])
        cluster_coeffs.append(zero_coeff)

    # (1) Generate orbit list for the design structure
    lolg = LocalOrbitListGenerator(
        cs.orbit_list,
        Structure.from_atoms(design_atoms),
        ce.fractional_position_tolerance,
    )
    orbit_list = lolg.generate_full_orbit_list().get_orbit_list()

    # (2) We assume chemical_symbols used in cs match the "elements" list
    elements = list(cs.species_maps[0].keys())

    # (3) Expand each orbit's cluster-vector elements
    param_idx = 1  # param[0] is zero-let
    for orbit in orbit_list:
        cves = orbit.cluster_vector_elements
        if not cves:
            continue
        multiplicity = cves[0]["multiplicity"]

        for cve in cves:
            param_eci = ce.parameters[param_idx]
            param_idx += 1
            outer_coeff = (param_eci / multiplicity) * weight

            if abs(outer_coeff) < tol:
                continue

            # (4) For each actual cluster in that orbit
            for cluster in orbit.clusters:
                site_indices = [site.index for site in cluster.lattice_sites]

                # (5) Loop over all species combos for those sites
                for species_combo in itertools.product(range(len(elements)), repeat=len(site_indices)):
                    final_coeff = outer_coeff
                    if abs(final_coeff) > tol:
                        cluster_spec = list(zip(site_indices, species_combo))
                        cluster_specs.append(cluster_spec)
                        cluster_coeffs.append(final_coeff)

    return cluster_specs, cluster_coeffs

# Reference energies for pure elements (per atom)
reference_energies = {
    "Cu": -240.11121086,
    "Ni": -352.62417749,
    "Pd": -333.69496589,
    "Ag": -173.55506507,
}
reference_structure_indice = 155
cutoffs = [5, 5]

# Read from JSON database
db = ase.io.read("structures_CuNiPdAg.json", ":")
dft_data = pd.read_csv("properties_CuNiPdAg.csv", index_col=0)
indices = dft_data.index.values

df = pd.DataFrame(columns=["atoms", "MixingEnergy", "TotalEnergy"])
for idx, i in enumerate(indices):
    st = db[i]
    energy = dft_data.loc[i, "final_energy"]
    dict_representation = dict(Counter(st.get_chemical_symbols()))
    NumAtoms = len(st)

    # Mixing energy calculation
    mixing_energy = (
        energy
        - sum(
            [
                v / NumAtoms * reference_energies[k]
                for k, v in dict_representation.items()
            ]
        )
    ) / NumAtoms

    df.loc[idx] = [st, mixing_energy, energy]

# 4-species expansion: Cu, Ni, Pd, Ag
cs = ClusterSpace(
    structure=df.loc[reference_structure_indice, "atoms"],  # pick a reference
    cutoffs=cutoffs,
    chemical_symbols=list(reference_energies.keys()),
)

# Collect training data
sc = StructureContainer(cluster_space=cs)
for i in df.index:
    sc.add_structure(
        structure=df.loc[i, "atoms"],
        properties={"mixing_energy": df.loc[i, "MixingEnergy"]},
    )

# Extract design matrix (X) and target (y)
X, y = sc.get_fit_data(key="mixing_energy")

# Fit with LassoCV
x_train, x_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
lasso_model = LassoCV(cv=5, random_state=42)
lasso_model.fit(x_train, y_train)

print("Train R2:", lasso_model.score(x_train, y_train))
print("Test  R2:", lasso_model.score(x_test, y_test))
print("Test MAE:", mean_absolute_error(lasso_model.predict(x_test), y_test))

# Build the cluster expansion object
ce = ClusterExpansion(cluster_space=cs, parameters=lasso_model.coef_)

# Intercept goes into the zero-let
ce.parameters[0] += lasso_model.intercept_

design_atoms = df.loc[0, "atoms"]  # for demonstration
cluster_specs, cluster_coeffs = parameterize_cluster_expansion(
    cs, ce, design_atoms, weight=1.0
)

# Example of what we get:
# cluster_specs might look like [ [(40,3),(47,0),(44,3)], [(44,3),(57,1)], ... ]
# cluster_coeffs like [5.8e-05, -1.85e-04, ...]

points_list = design_atoms.get_positions().tolist()
num_sites = len(design_atoms)
Canv = Canvas(points_list)

# Define the candidate atoms in the same order as "reference_energies.keys()"
AllAtoms = [Atom(s) for s in reference_energies.keys()]

# Convert cluster_specs to "site-Atom" pairs
cluster_specs_expanded = []
for c in cluster_specs:
    c_expanded = []
    for st, spidx in c:
        c_expanded.append((st, AllAtoms[spidx]))
    cluster_specs_expanded.append(c_expanded)

# Create the MatOpt model
m = MatOptModel(Canv, AllAtoms, clusters=cluster_specs_expanded[1:])

# Optional: add composition constraints 
CompBounds = {
    AllAtoms[0]: (0, num_sites),
    AllAtoms[1]: (0, num_sites),
    AllAtoms[2]: (5, 10),  # e.g., require 5–10 Pd
    AllAtoms[3]: (5, 10),  # e.g., require 5–10 Ag
}
m.addGlobalTypesDescriptor("Composition", bounds=CompBounds, rules=EqualTo(SumSites(desc=m.Yik)))

# Force the structure to be fully occupied (or any other site-based constraint)
m.Yi.rules.append(FixedTo(1.0))

# Build the objective from the cluster coefficients
obj_expr = SumClusters(desc=m.Zn, coefs=cluster_coeffs[1:])

# Solve the MILP
D = m.minimize(obj_expr, solver="neos-cplex", tilim=60)

D.toPDB("result.pdb")
D.toXYZ("result.xyz")
D.toPOSCAR("result.vasp")

from pymatgen.core import Structure
structure = D.to_pymatgen()
structure.to("cif", "result.cif")

Remarks

Using icet to train a cluster expansion and matopt to solve the resulting discrete optimization problem provides a powerful path toward rational materials design. You can rapidly propose candidate structures that minimize or maximize a desired property, subject to constraints (composition, geometry, etc.). While we showed a small demonstration on a single Cu–Ni–Pd–Ag cell, you can scale up to larger supercells, more complex constraints, or different properties.