Source code for pleiades.sammy.backends.local

#!/usr/bin/env python
"""Local backend implementation for SAMMY execution."""

import os
import subprocess
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Union
from uuid import uuid4

from pleiades.sammy.config import LocalSammyConfig
from pleiades.sammy.interface import (
    EnvironmentPreparationError,
    SammyExecutionError,
    SammyExecutionResult,
    SammyFiles,
    SammyFilesMultiMode,
    SammyRunner,
)
from pleiades.utils.logger import loguru_logger

logger = loguru_logger.bind(name=__name__)


def _move_broadening_inp_to_par(inp_file: Path, par_file: Path) -> None:
    """Move fitted broadening parameters from SAMNDF.INP to SAMNDF.PAR.

    After a JSON-mode pass, SAMMY writes broadening parameters into
    ``SAMNDF.INP`` (one or two ``BROADENING PARAMETERS FOLLOW`` data
    sections, plus the ``BROADENING IS WANTED`` command).  ``SAMNDF.PAR``
    does **not** contain a broadening section.

    Running a follow-up traditional-mode pass with broadening in the INP
    triggers ``STOP [Broadening parameters in both INPut & PAR files?]``
    because SAMMY interprets broadening in the INP data section as
    conflicting with the PAR file format.

    This helper:

    1. Extracts the **last** ``BROADENING PARAMETERS FOLLOW`` section from
       the INP (which contains the fitted values from pass 1).
    2. Appends that section to the PAR file.
    3. Strips the ``BROADENING IS WANTED`` command and all broadening data
       sections from the INP.

    This ensures pass 2 reads the fitted broadening from the PAR file only.

    Args:
        inp_file: Path to the SAMMY INP file (modified in-place).
        par_file: Path to the SAMMY PAR file (broadening section appended).
    """
    # --- Step 1: Extract broadening sections from INP ---
    with open(inp_file, encoding="utf-8") as f:
        lines = f.readlines()

    in_broadening = False
    filtered = []
    broadening_sections: list[list[str]] = []  # each section = [header, data..., blank]
    current_section: list[str] = []

    for line in lines:
        upper = line.upper().strip()

        # Remove the command
        if upper.startswith("BROADENING IS WANTED"):
            continue

        # Capture broadening data sections
        if upper.startswith("BROAD"):
            in_broadening = True
            current_section = [line]
            continue

        if in_broadening:
            current_section.append(line)
            if not line.strip():
                in_broadening = False
                broadening_sections.append(current_section)
                current_section = []
            continue

        filtered.append(line)

    # Handle broadening block that runs to EOF without a trailing blank line
    if in_broadening and current_section:
        broadening_sections.append(current_section)

    # Write cleaned INP
    with open(inp_file, "w", encoding="utf-8") as f:
        f.writelines(filtered)

    logger.debug(
        f"Stripped {len(broadening_sections)} broadening section(s) "
        f"and BROADENING IS WANTED command from {inp_file.name}"
    )

    # --- Step 2: Append last broadening section to PAR ---
    if broadening_sections:
        # Use the LAST section (contains fitted values from pass 1)
        fitted_section = broadening_sections[-1]

        with open(par_file, "a", encoding="utf-8") as f:
            # Ensure we start on a new line
            f.write("\n")
            f.writelines(fitted_section)

        logger.debug(f"Appended fitted broadening ({len(fitted_section)} lines) to {par_file.name}")
    else:
        logger.warning(f"No broadening sections found in {inp_file.name}")


def _enable_abundance_fitting_in_par(par_file: Path) -> None:
    """Modify a SAMMY parameter file in-place to enable per-isotope abundance fitting.

    Finds the Card Set 10 section (``ISOTOpic abundances and masses`` or
    ``NUCLIde abundances and masses``) and sets the IFLISO flag from 0 to 1
    for every isotope line.  IFLISO occupies columns 31-32 in the fixed-width
    Card-10 format (0 = fixed, 1 = vary).

    Args:
        par_file: Path to the SAMMY parameter file to modify in-place.
    """
    with open(par_file, encoding="utf-8") as f:
        lines = f.readlines()

    in_card10 = False
    in_continuation = False
    modified = []
    isotopes_modified = 0

    for line in lines:
        upper = line.upper().rstrip()

        # Detect Card 10 header
        if upper.startswith("ISOTO") or upper.startswith("NUCLI"):
            in_card10 = True
            in_continuation = False
            modified.append(line)
            continue

        if in_card10:
            # Blank line terminates Card 10
            if not line.strip():
                in_card10 = False
                in_continuation = False
                modified.append(line)
                continue

            stripped = line.strip()

            # "-1" marks the start of a spin-group continuation block
            if stripped == "-1":
                in_continuation = True
                modified.append(line)
                continue

            # Positively identify isotope data lines by atomic mass in cols 1-10.
            # Atomic masses are floats >= 1.0; continuation lines contain only
            # small integers for spin-group indices.
            is_isotope_line = False
            try:
                mass = float(line[:10])
                if mass >= 1.0:
                    is_isotope_line = True
                    in_continuation = False
            except (ValueError, IndexError):
                pass

            # Only modify IFLISO (columns 31-32) on isotope data lines
            if is_isotope_line and len(line) >= 32:
                line = line[:30] + " 1" + line[32:]
                isotopes_modified += 1

            modified.append(line)
        else:
            modified.append(line)

    with open(par_file, "w", encoding="utf-8") as f:
        f.writelines(modified)

    logger.debug(f"Enabled abundance fitting for {isotopes_modified} isotope(s) in {par_file.name}")


[docs] class LocalSammyRunner(SammyRunner): """Implementation of SAMMY runner for local installation.""" def __init__(self, config: LocalSammyConfig): super().__init__(config) self.config: LocalSammyConfig = config self._moved_files: List[Path] = []
[docs] def prepare_environment(self, files: Union[SammyFiles, SammyFilesMultiMode]) -> None: """Prepare environment for local SAMMY execution.""" try: logger.debug("Validating input files") files.validate() # Additional validation for JSON mode if isinstance(files, SammyFilesMultiMode): logger.debug("Performing JSON-ENDF mapping validation") self._validate_json_endf_mapping(files) # Move files to working directory logger.debug("Moving files to working directory") files.move_to_working_dir(self.config.working_dir) # No need to validate directories as this is done in config validation logger.debug("Environment preparation complete") except Exception as e: raise EnvironmentPreparationError(f"Environment preparation failed: {str(e)}")
def _validate_json_endf_mapping(self, files: SammyFilesMultiMode) -> None: """ Validate that JSON configuration references existing ENDF files. Args: files: SammyFilesMultiMode containing JSON config and ENDF directory Raises: ValueError: If JSON references missing ENDF files """ import json try: # Parse JSON to find referenced ENDF files with open(files.json_config_file, "r", encoding="utf-8") as f: json_data = json.load(f) # Find isotope entries (lists in JSON) - keys are ENDF filenames endf_files_referenced = [] for key, value in json_data.items(): if isinstance(value, list) and len(value) > 0 and isinstance(value[0], dict): # Key is the ENDF filename (e.g., "079-Au-197.B-VIII.0.par") endf_files_referenced.append(key) # Check each referenced ENDF file exists in ENDF directory missing_files = [] for endf_filename in endf_files_referenced: endf_path = files.endf_directory / endf_filename if not endf_path.exists(): missing_files.append(endf_filename) if missing_files: raise ValueError( f"JSON references missing ENDF files: {missing_files}. " f"Expected in directory: {files.endf_directory}" ) logger.debug(f"JSON-ENDF validation passed: {len(endf_files_referenced)} ENDF files verified") except json.JSONDecodeError as e: raise ValueError(f"Invalid JSON configuration file: {e}") def _build_env(self) -> dict: """Build environment variables for SAMMY subprocess.""" env = dict(os.environ) env.update(self.config.env_vars) if "LD_LIBRARY_PATH" in env: env["LD_LIBRARY_PATH"] = f"/usr/lib64:{env['LD_LIBRARY_PATH']}" else: env["LD_LIBRARY_PATH"] = "/usr/lib64" return env def _run_sammy_once(self, sammy_input: str, env: dict) -> subprocess.CompletedProcess: """Run a single SAMMY subprocess invocation.""" return subprocess.run( [str(self.config.sammy_executable)], input=sammy_input, shell=False, text=True, env=env, cwd=str(self.config.working_dir), capture_output=True, )
[docs] def execute_sammy(self, files: Union[SammyFiles, SammyFilesMultiMode]) -> SammyExecutionResult: """Execute SAMMY using local installation. For ``SammyFilesMultiMode`` with ``fit_abundances=True``, a transparent two-pass strategy is used: 1. **Pass 1 (JSON mode)** – SAMMY reads the JSON config and ENDF files, generates an internal parameter file (``SAMNDF.PAR``), and fits global parameters (thickness, normalization) with fixed per-isotope abundances. 2. **Pass 2 (traditional mode)** – The generated ``SAMNDF.PAR`` is modified to set Card-10 IFLISO flags to 1 (vary abundance), then SAMMY re-runs in traditional mode to fit per-isotope abundances. The caller sees a single ``SammyExecutionResult``; the two passes are an implementation detail. """ execution_id = str(uuid4()) start_time = datetime.now() logger.info(f"Starting SAMMY execution {execution_id}") logger.debug(f"Working directory: {self.config.working_dir}") # Prepare input text for SAMMY based on mode if isinstance(files, SammyFilesMultiMode): sammy_input = f"{files.input_file.name}\n#file {files.json_config_file.name}\n{files.data_file.name}\n\n" logger.debug("Using JSON mode input format") else: sammy_input = f"{files.input_file.name}\n{files.parameter_file.name}\n{files.data_file.name}\n\n" logger.debug("Using traditional mode input format") try: env = self._build_env() # --- Pass 1 --- process = self._run_sammy_once(sammy_input, env) console_output = process.stdout + process.stderr success = " Normal finish to SAMMY" in console_output if not success: logger.error(f"SAMMY execution failed for {execution_id}") return SammyExecutionResult( success=False, execution_id=execution_id, start_time=start_time, end_time=datetime.now(), console_output=console_output, error_message=f"SAMMY execution failed with return code {process.returncode}. Check console output for details.", ) # --- Pass 2 (abundance fitting) --- error_message: Optional[str] = None if isinstance(files, SammyFilesMultiMode) and files.fit_abundances: samndf_par = self.config.working_dir / "SAMNDF.PAR" samndf_inp = self.config.working_dir / "SAMNDF.INP" if samndf_par.exists() and samndf_inp.exists(): logger.info("Pass 2: enabling per-isotope abundance fitting (IFLISO=1)") _enable_abundance_fitting_in_par(samndf_par) # Move fitted broadening from INP → PAR to avoid # "Broadening parameters in both INPut & PAR files?" error. _move_broadening_inp_to_par(samndf_inp, samndf_par) # Resolve the data file name as it exists in working_dir data_name = files.data_file.name sammy_input_2 = f"{samndf_inp.name}\n{samndf_par.name}\n{data_name}\n\n" process2 = self._run_sammy_once(sammy_input_2, env) console_output_2 = process2.stdout + process2.stderr success = " Normal finish to SAMMY" in console_output_2 console_output += "\n--- Pass 2 (abundance fitting) ---\n" + console_output_2 if not success: error_message = ( f"SAMMY pass 2 (abundance fitting) failed " f"(return code={process2.returncode}). Check console output." ) logger.error( f"SAMMY pass 2 (abundance) failed for {execution_id} (return code={process2.returncode})" ) else: error_message = "SAMMY pass 2 (abundance fitting) failed: SAMNDF.PAR/INP not found after pass 1." logger.error( "fit_abundances=True but SAMNDF.PAR/INP not found after pass 1; " "cannot perform abundance fitting" ) success = False end_time = datetime.now() logger.info(f"SAMMY execution completed for {execution_id} (success={success})") return SammyExecutionResult( success=success, execution_id=execution_id, start_time=start_time, end_time=end_time, console_output=console_output, error_message=error_message, ) except Exception as e: logger.exception(f"SAMMY execution failed for {execution_id}") raise SammyExecutionError(f"SAMMY execution failed: {str(e)}")
[docs] def cleanup(self) -> None: """Clean up after execution.""" logger.debug("Performing cleanup for local backend") self._moved_files = []
[docs] def validate_config(self) -> bool: """Validate the configuration.""" return self.config.validate()
if __name__ == "__main__": from pathlib import Path # Setup paths sammy_executable = Path.home() / "code.ornl.gov/SAMMY/build/bin/sammy" test_data_dir = Path(__file__).parents[4] / "tests/data/ex012" working_dir = Path.home() / "tmp/pleiades_test" output_dir = working_dir / "output" # Create config config = LocalSammyConfig(sammy_executable=sammy_executable, working_dir=working_dir, output_dir=output_dir) config.validate() # Create files container files = SammyFiles( input_file=test_data_dir / "ex012a.inp", parameter_file=test_data_dir / "ex012a.par", data_file=test_data_dir / "ex012a.dat", ) try: # Create and use runner runner = LocalSammyRunner(config) # Prepare environment runner.prepare_environment(files) # Execute SAMMY result = runner.execute_sammy(files) # Process results if result.success: print(f"SAMMY execution successful (runtime: {result.runtime_seconds:.2f}s)") runner.collect_outputs(result) else: print("SAMMY execution failed:") print(result.error_message) print("\nConsole output:") print(result.console_output) except Exception as e: print(f"Error running SAMMY: {str(e)}") finally: # Cleanup runner.cleanup()