From 95703af0c630a249b51c68e97e10d782dedf4398 Mon Sep 17 00:00:00 2001 From: Yuchen Xiao Date: Sun, 27 Sep 2026 13:01:16 -0400 Subject: [PATCH 1/4] feat: add a tool to download cif file --- requirements/conda.txt | 2 + .../apps/refinebase/refinement_server.py | 21 +++++ src/diffpy/apps/refinebase/util.py | 80 +++++++++++++++++++ 3 files changed, 103 insertions(+) create mode 100644 src/diffpy/apps/refinebase/util.py diff --git a/requirements/conda.txt b/requirements/conda.txt index 38287db..9a93f82 100644 --- a/requirements/conda.txt +++ b/requirements/conda.txt @@ -8,3 +8,5 @@ diffpy.srreal diffpy.structure networkx mcp[cli] +mp-api +pymatgen diff --git a/src/diffpy/apps/refinebase/refinement_server.py b/src/diffpy/apps/refinebase/refinement_server.py index 8823104..4ddcec6 100644 --- a/src/diffpy/apps/refinebase/refinement_server.py +++ b/src/diffpy/apps/refinebase/refinement_server.py @@ -11,6 +11,7 @@ from mcp.server.mcpserver.exceptions import ToolError from diffpy.apps.refinebase.refinement_session import RefinementSession +from diffpy.apps.refinebase.util import download_mp_cifs session = RefinementSession() mcp = MCPServer("diffpy.apps") @@ -56,6 +57,26 @@ async def refinement_playbook(): return f.read() +@mcp.tool() +@tool_errors +async def download_mp_cifs_tool( + query: Annotated[str, "Query for the Materials Project structures"], + start_index: Annotated[int, "Start index of the structures to download"], + end_index: Annotated[int, "End index of the structures to download"], + output_dir: Annotated[str, "Directory to save the CIF files"] = "mp_cifs", + api_key: Annotated[str | None, "Materials Project API key"] = None, +) -> list[str]: + """ + Download a consecutive range of Materials Project structures as CIF files. + """ + return [ + str(path) + for path in download_mp_cifs( + query, start_index, end_index, output_dir, api_key + ) + ] + + @mcp.tool() @tool_errors async def add_profile_from_file( diff --git a/src/diffpy/apps/refinebase/util.py b/src/diffpy/apps/refinebase/util.py new file mode 100644 index 0000000..134a17b --- /dev/null +++ b/src/diffpy/apps/refinebase/util.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import os +from pathlib import Path + +from mp_api.client import MPRester +from pymatgen.io.cif import CifWriter + + +def download_mp_cifs( + query: str, + start_index: int, + end_index: int, + output_dir: str | Path = "mp_cifs", + api_key: str | None = None, +) -> list[Path]: + """ + Download a consecutive range of Materials Project structures as CIF files. + + Parameters + ---------- + query + One of: + - Materials Project ID: "mp-149" + - Formula: "SiO2" or "LiNbO3" + - Chemical system: "Li-Ti-Nb-O" + start_index, end_index + Zero-based inclusive indices in the returned MP result list. + output_dir + Directory in which CIF files will be written. + api_key + Materials Project API key. If omitted, reads MP_API_KEY from + the environment. + + Returns + ------- + list[Path] + Paths of the CIF files written. + """ + if start_index < 0 or end_index < start_index: + raise ValueError("Require 0 <= start_index <= end_index.") + + api_key = api_key or os.environ["MP_API_KEY"] + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + with MPRester(api_key) as mpr: + if query.startswith("mp-"): + docs = mpr.materials.summary.search( + material_ids=[query], + fields=["material_id", "structure"], + ) + elif "-" in query: + # Example: "Li-Ti-Nb-O" + docs = mpr.materials.summary.search( + chemsys=query, + fields=["material_id", "structure"], + ) + else: + # Example: "SiO2" or "LiNbO3" + docs = mpr.materials.summary.search( + formula=[query], + fields=["material_id", "structure"], + ) + + if end_index >= len(docs): + raise IndexError( + f"Requested index {end_index}, but MP returned only " + f"{len(docs)} structure(s)." + ) + + written_files = [] + + for index in range(start_index, end_index + 1): + doc = docs[index] + filename = output_dir / f"{index:04d}_{doc.material_id}.cif" + CifWriter(doc.structure).write_file(filename) + written_files.append(filename) + + return written_files From 98f0cfc6c964f3a42507b260fb6b1cc482cf0d5a Mon Sep 17 00:00:00 2001 From: Yuchen Xiao Date: Sun, 27 Sep 2026 13:06:19 -0400 Subject: [PATCH 2/4] feat: allow generator to run in parallel --- .../apps/refinebase/parametric_model.py | 51 +++++++++++++++++-- .../apps/refinebase/refinement_session.py | 4 ++ 2 files changed, 50 insertions(+), 5 deletions(-) diff --git a/src/diffpy/apps/refinebase/parametric_model.py b/src/diffpy/apps/refinebase/parametric_model.py index d61535c..1047425 100644 --- a/src/diffpy/apps/refinebase/parametric_model.py +++ b/src/diffpy/apps/refinebase/parametric_model.py @@ -4,6 +4,7 @@ from pathlib import Path import networkx as nx +import numpy from pyobjcryst import loadCrystal from diffpy.srfit.fitbase import FitContribution @@ -245,7 +246,14 @@ def evaluate(self): class ParametricModelPDF(ParametricModel): # NOTE: qmin, qmax, stype(scattering type) are meta handled # throughout the loaded profile in the refinement session - def __init__(self, name, structure, spacegroup_symbol="P1", finite=False): + def __init__( + self, + name, + structure, + spacegroup_symbol="P1", + finite=False, + run_parallel=True, + ): """ Create a ParametricModelPDF instance from a structure object. @@ -269,6 +277,25 @@ def __init__(self, name, structure, spacegroup_symbol="P1", finite=False): self.space_group_symbol = spacegroup_symbol self.sgpar_names = [] self._rebuild_graph() + if run_parallel: + try: + import multiprocessing + from multiprocessing import Pool + + import psutil + except ImportError: + print( + "\nYou don't appear to have the " + "necessary packages for parallelization" + ) + syst_cores = multiprocessing.cpu_count() + cpu_percent = psutil.cpu_percent() + avail_cores = numpy.floor( + (100 - cpu_percent) / (100.0 / syst_cores) + ) + ncpu = int(numpy.max([1, avail_cores])) + pool = Pool(processes=ncpu) + self.calc_obj.parallel(ncpu=ncpu, mapfunc=pool.map) def _hide_dependent_parameters(self, use_uiso=True): if use_uiso: @@ -545,7 +572,11 @@ def residual(self): def create_pdf_model_from_file( - name, structure_file_path, library="Diffpy", finite=False + name, + structure_file_path, + library="Diffpy", + finite=False, + run_parallel=True, ): """Create a ParametricModelPDF by parsing a structure file.""" stru_parser = get_parser("auto") @@ -557,16 +588,21 @@ def create_pdf_model_from_file( ) or library == "ObjCryst": structure = loadCrystal(structure_file_path) return ParametricModelPDF( - name, structure, spacegroup_symbol=spacegroup_symbol, finite=finite + name, + structure, + spacegroup_symbol=spacegroup_symbol, + finite=finite, + run_parallel=run_parallel, ) -def create_pdf_model_from_model(name, from_model): +def create_pdf_model_from_model(name, from_model, run_parallel=True): """Create a ParametricModelPDF sharing the phase of from_model.""" return ParametricModelPDF( name, from_model.calc_obj.phase, spacegroup_symbol=from_model.space_group_symbol, + run_parallel=run_parallel, ) @@ -577,6 +613,7 @@ def create_pdf_model_from_code( global_namespace={}, local_structure_name="structure", finite=False, + run_parallel=True, ): """Create a ParametricModelPDF by executing code that builds a structure. @@ -618,5 +655,9 @@ def create_pdf_model_from_code( ) return ParametricModelPDF( - name, structure, spacegroup_symbol=spacegroup_symbol, finite=finite + name, + structure, + spacegroup_symbol=spacegroup_symbol, + finite=finite, + run_parallel=run_parallel, ) diff --git a/src/diffpy/apps/refinebase/refinement_session.py b/src/diffpy/apps/refinebase/refinement_session.py index 9676be8..7aba9b6 100644 --- a/src/diffpy/apps/refinebase/refinement_session.py +++ b/src/diffpy/apps/refinebase/refinement_session.py @@ -128,6 +128,7 @@ def add_pdf_model( global_namespace={}, local_structure_name="structure", finite=False, + run_parallel=True, ): from diffpy.apps.refinebase.parametric_model import ( create_pdf_model_from_code, @@ -143,6 +144,7 @@ def add_pdf_model( structure_file_path, library=library, finite=finite, + run_parallel=run_parallel, ) elif from_model_name is not None: if from_model_name not in self.models_dict: @@ -153,6 +155,7 @@ def add_pdf_model( pdf_model = create_pdf_model_from_model( model_name, from_model, + run_parallel=run_parallel, ) elif code is not None: pdf_model = create_pdf_model_from_code( @@ -161,6 +164,7 @@ def add_pdf_model( global_namespace=global_namespace, local_structure_name=local_structure_name, finite=finite, + run_parallel=run_parallel, ) else: raise ValueError( From 8dae7fa488a61c48a957f21f8cdb8ca09d786e6f Mon Sep 17 00:00:00 2001 From: Yuchen Xiao Date: Sun, 27 Sep 2026 13:12:53 -0400 Subject: [PATCH 3/4] feat: allow soft bound restraints --- .../apps/refinebase/refinement_server.py | 26 ++++++++--- .../apps/refinebase/refinement_session.py | 27 +++++++---- tests/test_refinement_session.py | 45 +++++++++++++++++++ 3 files changed, 85 insertions(+), 13 deletions(-) diff --git a/src/diffpy/apps/refinebase/refinement_server.py b/src/diffpy/apps/refinebase/refinement_server.py index 4ddcec6..3a1019e 100644 --- a/src/diffpy/apps/refinebase/refinement_server.py +++ b/src/diffpy/apps/refinebase/refinement_server.py @@ -489,8 +489,8 @@ async def solve( "and the second dict is variable-constraint_equation pair." ), ] = None, - restraints: Annotated[ - list[str], "List of restraints to apply during the refinement" + bounds: Annotated[ + dict, "Dictionary of bounds for variables or equations" ] = None, name: Annotated[str, "Name of the refinement session"] = None, weights: Annotated[ @@ -519,8 +519,24 @@ async def solve( constraints : list[dict], optional First dict is new_variable-initial value pair, and the second dict is variable-constraint_equation pair. - restraints : list[str], optional - List of restraints to apply during the refinement. + bounds : dict, optional + Dictionary of bounds for variables or equations. + e.g. {"variable_name": + { + "lower_bound": 0, + "upper_bound": 10, + "uncertainty": 1, + "scaled": False + }} + # start copied from diffpy.srfit docstring + scaled : bool, optional + If True, the restraint penalty is scaled by the unrestrained + point-average chi^2 (chi^2/numpoints) (default is False). + params : dict, optional + The dictionary of Parameters, indexed by name, that are used in + `param_or_eq` (if an equation string is used) but are not part + of the RecipeOrganizer (default is {}). + # end copied from diffpy.srfit docstring name : str, optional Name of the refinement session. weights : list[float], optional @@ -547,7 +563,7 @@ async def solve( variable_names, residual_equations=residual_equations, constraints=constraints, - restraints=restraints, + bounds=bounds, name=name, weights=weights, metas=metas, diff --git a/src/diffpy/apps/refinebase/refinement_session.py b/src/diffpy/apps/refinebase/refinement_session.py index 7aba9b6..5960a21 100644 --- a/src/diffpy/apps/refinebase/refinement_session.py +++ b/src/diffpy/apps/refinebase/refinement_session.py @@ -447,19 +447,19 @@ def _solve( models, variable_names, constraints=None, - restraints=None, + bounds=None, weights=None, residual_equations=None, metas=None, - verbose_iterations=0, ): - # NOTE: restraints to be implemented recipe = FitRecipe() self.recipes_dict[name] = recipe if weights is None: weights = numpy.ones(len(profiles)) if residual_equations is None: residual_equations = ["chiv"] * len(profiles) + if bounds is None: + bounds = {} if metas is not None: for i in range(len(metas)): profiles[i].meta.update(metas[i]) @@ -494,7 +494,20 @@ def _solve( if var in recipe._parameters.values(): continue recipe.add_variable(var, name=variable_names[i]) - + for eq_or_var_name, arg_dict in bounds.items(): + lb = arg_dict.get("lower_bound", -numpy.inf) + ub = arg_dict.get("upper_bound", numpy.inf) + use_soft_bounds = arg_dict.get("use_soft_bounds", True) + if use_soft_bounds: + uncertainty = arg_dict.get("uncertainty", 1) + scaled = arg_dict.get("scaled", False) + eq_or_var_name = eq_or_var_name.replace(".", "_") + recipe.add_soft_bounds( + eq_or_var_name, lb, ub, sig=uncertainty, scaled=scaled + ) + else: + par = self.get_variable(eq_or_var_name)["obj"] + par.bound_range(lb, ub) recipe.free("all") leastsq(recipe.residual, recipe.getValues()) # NOTE: non-scalar value will raise error in `get_results_string` @@ -511,12 +524,11 @@ def solve( variable_names=[], residual_equations=None, constraints=None, - restraints=None, + bounds=None, name=uuid.uuid4(), weights=None, metas=None, include_sgpars=False, - verbose_iterations=0, ): profiles = [] for profile_name in profile_names: @@ -555,11 +567,10 @@ def solve( variable_names=variable_names, residual_equations=residual_equations, constraints=constraints, - restraints=restraints, + bounds=bounds, name=name, weights=weights, metas=metas, - verbose_iterations=verbose_iterations, ) def plot(self): diff --git a/tests/test_refinement_session.py b/tests/test_refinement_session.py index af4436b..e72f71e 100644 --- a/tests/test_refinement_session.py +++ b/tests/test_refinement_session.py @@ -13,6 +13,51 @@ _DATA_DIR = Path(__file__).parent / "data" +def test_bonds(): + # C1: Create a soft bound with an upper limit and high penalty + # Expect the variable to be constrained by the upper bound + xarray = numpy.linspace(0, 10, 100) + yarray = 2 * xarray + session = RefinementSession() + session.add_profile_from_arrays( + profile_name="linear", xarray=xarray, yarray=yarray + ) + session.add_equation_model(model_name="linear_model", equation_str="m*x") + session.set_variables_value( + { + "linear_model.m": 1, + } + ) + session._solve( + name="linear_upper_bound", + profiles=[session.profiles_dict["linear"]], + models=[session.models_dict["linear_model"]], + variable_names=["linear_model.m"], + bounds={"linear_model.m": {"upper_bound": 1.8, "uncertainty": 1e-4}}, + ) + assert numpy.isclose( + session.get_variable("linear_model.m")["value"], + 1.8, + rtol=1e-2, + ) + # C2: Create a soft bound with a lower limit and high penalty + # Expect the variable to be constrained by the lower bound + session._solve( + name="linear_lower_bound", + profiles=[session.profiles_dict["linear"]], + models=[session.models_dict["linear_model"]], + variable_names=["linear_model.m"], + bounds={"linear_model.m": {"lower_bound": 2.2, "uncertainty": 1e-4}}, + ) + assert numpy.isclose( + session.get_variable("linear_model.m")["value"], + 2.2, + rtol=1e-2, + ) + # NOTE: hard bounds are defined but not used by diffpy.srfit + # Skip testing hard bounds. + + def test_refine_sine(sine_profile): # C1: Refinement session without additional calculator or functions session = RefinementSession() From 92b8fe465fb399fac8d691d0517ace2527c95d96 Mon Sep 17 00:00:00 2001 From: Yuchen Xiao Date: Sun, 27 Sep 2026 13:13:34 -0400 Subject: [PATCH 4/4] chore: update argument names --- news/misc_functionalities.rst | 25 +++++++++++ .../apps/refinebase/refinement_server.py | 41 ++++++++++++++++++- .../apps/refinebase/refinement_session.py | 3 +- 3 files changed, 65 insertions(+), 4 deletions(-) create mode 100644 news/misc_functionalities.rst diff --git a/news/misc_functionalities.rst b/news/misc_functionalities.rst new file mode 100644 index 0000000..06f6a02 --- /dev/null +++ b/news/misc_functionalities.rst @@ -0,0 +1,25 @@ +**Added:** + +* Add generator parallelism. +* Add soft bounds. +* add a tool to download cif file from database. + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/src/diffpy/apps/refinebase/refinement_server.py b/src/diffpy/apps/refinebase/refinement_server.py index 3a1019e..9fff0b1 100644 --- a/src/diffpy/apps/refinebase/refinement_server.py +++ b/src/diffpy/apps/refinebase/refinement_server.py @@ -102,7 +102,6 @@ async def add_profile_from_file( async def add_profile_from_arrays( xarray: Annotated[list, "X-values of the profile"], yarray: Annotated[list, "Y-values of the profile"], - dx: Annotated[list, "Uncertainties in the x-values"] = None, dy: Annotated[list, "Uncertainties in the y-values"] = None, profile_name: Annotated[str, "Unique name for the profile"] = None, xname: Annotated[str, "Name of the x-axis for the profile"] = "x", @@ -113,7 +112,6 @@ async def add_profile_from_arrays( session.add_profile_from_arrays( xarray, yarray, - dx=dx, dy=dy, profile_name=profile_name, xname=xname, @@ -203,6 +201,24 @@ async def add_pdf_model( from_model_name: Annotated[ str, "Name of the existing model to base the new model on" ] = None, + code: Annotated[ + str, "Code that builds a structure to base the new model on" + ] = None, + library: Annotated[ + str, "Structure library to use ('Diffpy' or 'ObjCryst')" + ] = "Diffpy", + global_namespace: Annotated[ + dict, "Global namespace to execute 'code' in" + ] = {}, + local_structure_name: Annotated[ + str, "Name of the structure variable assigned by 'code'" + ] = "structure", + finite: Annotated[ + bool, "Whether to use DebyePDFGenerator instead of PDFGenerator" + ] = False, + run_parallel: Annotated[ + bool, "Whether to run the PDF model in parallel" + ] = True, ) -> str: """ Add a structure-file-based parametric model to the refinement session. @@ -214,6 +230,18 @@ async def add_pdf_model( structure_file_path : str, optional Path to the structure file. from_model_name : str, optional + code : str, optional + Code that builds a structure to base the new model on. + library : str, optional + Structure library to use ('Diffpy' or 'ObjCryst'). + global_namespace : dict, optional + Global namespace to execute 'code' in. + local_structure_name : str, optional + Name of the structure variable assigned by 'code'. + finite : bool, optional + Whether to use DebyePDFGenerator instead of PDFGenerator. + run_parallel : bool, optional + Whether to run the PDF model in parallel. Notes ----- @@ -221,11 +249,20 @@ async def add_pdf_model( existing computation object. For example, when the same phase's signal is observed in multiple profiles, 'from_model_name' allows the refinement backend to modify the same structure across multiple profiles. + + Exactly one of 'structure_file_path', 'from_model_name', or 'code' + must be provided. """ session.add_pdf_model( model_name=model_name, structure_file_path=structure_file_path, from_model_name=from_model_name, + code=code, + library=library, + global_namespace=global_namespace, + local_structure_name=local_structure_name, + finite=finite, + run_parallel=run_parallel, ) return f"Model {model_name} added successfully." diff --git a/src/diffpy/apps/refinebase/refinement_session.py b/src/diffpy/apps/refinebase/refinement_session.py index 5960a21..6e820a4 100644 --- a/src/diffpy/apps/refinebase/refinement_session.py +++ b/src/diffpy/apps/refinebase/refinement_session.py @@ -56,7 +56,6 @@ def add_profile_from_arrays( self, xarray, yarray, - dx=None, dy=None, profile_name: str = None, xname: str = "x", @@ -68,7 +67,7 @@ def add_profile_from_arrays( if profile_name is None: profile_name = str(uuid.uuid4()) profile = Profile() - profile.setObservedProfile(xarray, yarray, dx=dx, dy=dy) + profile.setObservedProfile(xobs=xarray, yobs=yarray, dyobs=dy) profile.xpar.name = xname profile._xname = xname profile.ypar.name = yname