diff --git a/material-db-tools/material_db_tools.py b/material-db-tools/material_db_tools.py index 7a5144d..375dcd3 100644 --- a/material-db-tools/material_db_tools.py +++ b/material-db-tools/material_db_tools.py @@ -1,6 +1,9 @@ +from typing import Dict, Union, Optional from pyne import material from pyne.material import Material, MultiMaterial from pyne.material_library import MaterialLibrary +from pyne import nucname +from decimal import Decimal, getcontext def make_mat(nucvec, density, citation, molecular_mass = None): mat = Material(nucvec, density = density, metadata = {'citation' : citation}) @@ -53,3 +56,39 @@ def mix_by_volume(material_library, vol_fracs, citation, density_factor=1): ) return mat + +def enrich( + material_composition: Dict[Union[int, str], float], + element_to_enrich: Union[int, str], + isotope_enrichments: Dict[Union[int, str], float], +) -> Dict[int, float]: + """ + Enrich a specific element in a material composition by replacing it with its isotopes. + + Args: + material_composition (Dict[Union[int, str], float]): The original material composition. + Keys can be nuclide IDs (int) or element symbols (str), and values are their fractions (float). + element_to_enrich (Union[int, str]): The ID or symbol of the element to be enriched. + isotope_enrichments (Dict[Union[int, str], float]): The isotopic composition for element_to_enrich. + Keys can be isotope IDs (int) or isotope symbols (str), and values are their enrichment fractions (float). + + Returns: + Dict[int, float]: The updated material composition with the enriched element. + Keys are nuclide IDs (int) and values are their fractions (float). + """ + # Convert all keys to PyNE nuclide IDs + material_composition = { + nucname.id(k): Decimal(str(v)) for k, v in material_composition.items() + } + isotope_enrichments = { + nucname.id(k): Decimal(str(v)) for k, v in isotope_enrichments.items() + } + element_to_enrich_id = nucname.id(element_to_enrich) + + if element_to_enrich_id in material_composition: + fract = material_composition.pop(element_to_enrich_id) + for nuclide, enrich_frac in isotope_enrichments.items(): + material_composition[nuclide] = enrich_frac * fract + + # Convert back to float for consistency with the original function signature + return {k: float(v) for k, v in material_composition.items()} \ No newline at end of file