"""Workspace unit-cell histogram aggregation and Gaussian fitting helpers."""
from __future__ import annotations
from dataclasses import dataclass
import glob
import os
import re
from typing import Iterable
import numpy as np
from coseda.io import config_to_paths
from coseda.crystfel_cell import (
UNIQUE_AXIS_LATTICE_TYPES,
normalize_crystfel_lattice_type,
)
PARAMS = ("a", "b", "c", "alpha", "beta", "gamma")
LENGTH_PARAMS = {"a", "b", "c"}
ANGLE_PARAMS = {"alpha", "beta", "gamma"}
CELL_KEY_BY_PARAM = {
"a": "a",
"b": "b",
"c": "c",
"alpha": "al",
"beta": "be",
"gamma": "ga",
}
CELL_UNIT_BY_PARAM = {
"a": "A",
"b": "A",
"c": "A",
"alpha": "deg",
"beta": "deg",
"gamma": "deg",
}
WORKSPACE_REFINED_CELL_DIR = "coseda_refined_cells"
WORKSPACE_REFINED_MARKER = "Cell.refined"
WORKSPACE_REFINED_SOURCE = "workspace_unit_cell_fit"
[docs]
@dataclass(frozen=True)
class WorkspaceStreamSource:
ini_path: str
run_dir: str
stream_path: str
cell_path: str
def _existing_dirs(paths: Iterable[str | None]) -> list[str]:
seen = set()
out = []
for path in paths:
if not path:
continue
norm = os.path.abspath(path)
if norm in seen or not os.path.isdir(norm):
continue
seen.add(norm)
out.append(norm)
return out
[docs]
def run_base_candidates_for_ini(ini_path: str) -> list[str]:
"""Return plausible folders containing indexing runs for a workspace INI."""
ini_dir = os.path.dirname(os.path.abspath(ini_path))
candidates: list[str | None] = [
os.path.join(ini_dir, "output"),
os.path.join(ini_dir, "runs"),
ini_dir,
]
try:
_, outputfolder_path, _, _, _, _ = config_to_paths(ini_path)
candidates.insert(0, outputfolder_path)
except Exception:
pass
return _existing_dirs(candidates)
def _is_ici_run_dir(run_dir: str) -> bool:
name = os.path.basename(os.path.abspath(run_dir).rstrip(os.sep))
return (
name.startswith("ici_")
or os.path.exists(os.path.join(run_dir, "ici_settings.json"))
or os.path.isdir(os.path.join(run_dir, "run_000"))
)
[docs]
def resolve_stream_path_for_run(run_dir: str) -> str | None:
"""Resolve the best stream path for a simple indexamajig or ICI run."""
if not run_dir or not os.path.isdir(run_dir):
return None
candidates: list[str | None] = []
if _is_ici_run_dir(run_dir):
try:
from coseda.ici.run_contract import canonical_stream_path, resolve_final_stream_path
candidates.extend(
[
resolve_final_stream_path(run_dir),
canonical_stream_path(run_dir),
]
)
except Exception:
pass
candidates.extend(
[
os.path.join(run_dir, "done.stream"),
os.path.join(run_dir, "early_break.stream"),
]
)
run_name = os.path.basename(os.path.abspath(run_dir).rstrip(os.sep))
candidates.append(os.path.join(run_dir, f"{run_name}.stream"))
candidates.extend(sorted(glob.glob(os.path.join(run_dir, "*.stream"))))
seen = set()
for candidate in candidates:
if not candidate:
continue
norm = os.path.abspath(candidate)
if norm in seen:
continue
seen.add(norm)
if os.path.isfile(norm):
return norm
return None
[docs]
def discover_workspace_unit_cell_streams(ini_paths: Iterable[str]) -> list[WorkspaceStreamSource]:
"""
Discover one usable unit-cell stream per workspace INI.
The newest run folder with an existing stream is selected for each INI, matching the
indexing UI's "latest run per file" behavior.
"""
sources: list[WorkspaceStreamSource] = []
for ini_path in ini_paths:
if not ini_path or not os.path.isfile(ini_path):
continue
run_dirs = set()
ini_dir = os.path.dirname(os.path.abspath(ini_path))
for base in run_base_candidates_for_ini(ini_path):
run_dirs.update(glob.glob(os.path.join(base, "indexingintegration_*")))
run_dirs.update(glob.glob(os.path.join(base, "ici_*")))
sorted_runs = sorted(
(rd for rd in run_dirs if os.path.isdir(rd)),
key=lambda rd: (os.path.getmtime(rd), os.path.basename(rd)),
reverse=True,
)
for run_dir in sorted_runs:
stream_path = resolve_stream_path_for_run(run_dir)
if stream_path:
sources.append(
WorkspaceStreamSource(
ini_path=os.path.abspath(ini_path),
run_dir=os.path.abspath(run_dir),
stream_path=stream_path,
cell_path=os.path.join(os.path.abspath(run_dir), "cellfile.cell"),
)
)
break
return sources
[docs]
def gaussian_curve(x, amplitude, mean, sigma, baseline):
sigma = max(float(sigma), np.finfo(float).eps)
return baseline + amplitude * np.exp(-0.5 * ((x - mean) / sigma) ** 2)
def _weighted_moments(bin_centers: np.ndarray, counts: np.ndarray) -> tuple[float, float]:
weight_sum = float(np.sum(counts))
if weight_sum <= 0:
raise ValueError("Selected histogram range has no counts.")
mean = float(np.sum(bin_centers * counts) / weight_sum)
var = float(np.sum(counts * (bin_centers - mean) ** 2) / weight_sum)
sigma = max(var**0.5, np.finfo(float).eps)
return mean, sigma
[docs]
def fit_gaussian_to_values(values, value_range: tuple[float, float], bins: int = 60) -> dict:
"""Fit a Gaussian to unit-cell values inside ``value_range``."""
arr = np.asarray(values, dtype=float)
arr = arr[np.isfinite(arr)]
if arr.size == 0:
raise ValueError("No finite values available for fitting.")
xmin, xmax = sorted((float(value_range[0]), float(value_range[1])))
selected = arr[(arr >= xmin) & (arr <= xmax)]
if selected.size < 3:
raise ValueError("Select at least three values for a Gaussian fit.")
bins = max(5, int(bins))
counts, edges = np.histogram(selected, bins=bins, range=(xmin, xmax))
centers = 0.5 * (edges[:-1] + edges[1:])
nonzero = counts > 0
if np.count_nonzero(nonzero) < 3:
mean = float(np.mean(selected))
sigma = max(float(np.std(selected)), np.finfo(float).eps)
amplitude = float(np.max(counts)) if counts.size else float(selected.size)
baseline = 0.0
method = "moments"
else:
fit_x = centers[nonzero]
fit_y = counts[nonzero].astype(float)
moment_mean, moment_sigma = _weighted_moments(centers, counts)
p0 = [
max(float(np.max(fit_y) - np.min(fit_y)), 1.0),
moment_mean,
max(moment_sigma, (xmax - xmin) / max(bins, 1)),
max(float(np.min(fit_y)), 0.0),
]
try:
from scipy.optimize import curve_fit
lower = [0.0, xmin, np.finfo(float).eps, 0.0]
upper = [np.inf, xmax, max(xmax - xmin, np.finfo(float).eps), np.inf]
popt, _ = curve_fit(
gaussian_curve,
fit_x,
fit_y,
p0=p0,
bounds=(lower, upper),
maxfev=10000,
)
amplitude, mean, sigma, baseline = [float(v) for v in popt]
method = "curve_fit"
except Exception:
mean, sigma = moment_mean, moment_sigma
amplitude = p0[0]
baseline = p0[3]
method = "moments"
curve_x = np.linspace(xmin, xmax, 256)
curve_y = gaussian_curve(curve_x, amplitude, mean, sigma, baseline)
return {
"range": (xmin, xmax),
"selected_count": int(selected.size),
"mean": float(mean),
"sigma": float(abs(sigma)),
"fwhm": float(2.354820045 * abs(sigma)),
"amplitude": float(amplitude),
"baseline": float(baseline),
"curve_x": curve_x,
"curve_y": curve_y,
"method": method,
}
[docs]
def update_cell_file_parameter(cell_path: str, param: str, value: float) -> bool:
"""Update one CrystFEL cell parameter in-place, preserving the line suffix."""
if param not in CELL_KEY_BY_PARAM:
raise ValueError(f"Unsupported cell parameter: {param}")
if not cell_path or not os.path.isfile(cell_path):
raise FileNotFoundError(cell_path)
key = CELL_KEY_BY_PARAM[param]
unit = CELL_UNIT_BY_PARAM[param]
value_text = f"{float(value):.6g}"
pattern = re.compile(rf"^(\s*{re.escape(key)}\s*=\s*)([-+0-9.eE]+)(.*)$")
with open(cell_path, "r", encoding="utf-8", errors="ignore") as handle:
lines = handle.readlines()
changed = False
for idx, line in enumerate(lines):
match = pattern.match(line)
if not match:
continue
suffix = match.group(3).rstrip("\n") or f" {unit}"
lines[idx] = f"{match.group(1)}{value_text}{suffix}\n"
changed = True
break
if not changed:
if lines and not lines[-1].endswith("\n"):
lines[-1] = lines[-1] + "\n"
lines.append(f"{key} = {value_text} {unit}\n")
changed = True
with open(cell_path, "w", encoding="utf-8") as handle:
handle.writelines(lines)
return changed
def _parse_comment_kv_file(path: str) -> dict[str, str]:
kv = {}
if not path or not os.path.isfile(path):
return kv
with open(path, "r", encoding="utf-8", errors="replace") as handle:
for line in handle:
stripped = line.strip()
if not stripped.startswith("#") or "=" not in stripped:
continue
key, value = stripped[1:].split("=", 1)
kv[key.strip()] = value.strip()
return kv
[docs]
def parse_workspace_cell(path: str) -> dict[str, str]:
"""Return workspace cell header values, including ``Cell.*`` keys."""
return _parse_comment_kv_file(path)
[docs]
def workspace_has_refined_cell(path: str) -> bool:
kv = parse_workspace_cell(path)
if kv.get(WORKSPACE_REFINED_MARKER, "").strip().lower() in {"1", "true", "yes"}:
return all(kv.get(f"Cell.{key}") for key in ("a", "b", "c", "al", "be", "ga"))
return False
[docs]
def update_workspace_refined_cell_parameter(workspace_path: str, param: str, value: float) -> bool:
"""Update one fitted cell parameter in the workspace header and mark it refined."""
if param not in CELL_KEY_BY_PARAM:
raise ValueError(f"Unsupported cell parameter: {param}")
if not workspace_path:
raise FileNotFoundError(workspace_path)
key = CELL_KEY_BY_PARAM[param]
header_updates = {
f"Cell.{key}": f"{float(value):.6g}",
WORKSPACE_REFINED_MARKER: "true",
"Cell.refined_source": WORKSPACE_REFINED_SOURCE,
}
if os.path.exists(workspace_path):
with open(workspace_path, "r", encoding="utf-8", errors="replace") as handle:
lines = handle.readlines()
else:
lines = []
kv_indices = {}
for idx, line in enumerate(lines):
stripped = line.strip()
if stripped.startswith("#") and "=" in stripped:
raw_key, _ = stripped[1:].split("=", 1)
clean_key = raw_key.strip()
if clean_key in header_updates and clean_key not in kv_indices:
kv_indices[clean_key] = idx
insert_at = 0
for idx, line in enumerate(lines):
stripped = line.strip()
if stripped and not stripped.startswith("#"):
insert_at = idx
break
insert_at = idx + 1
for update_key, update_value in header_updates.items():
new_line = f"# {update_key}={update_value}\n"
if update_key in kv_indices:
lines[kv_indices[update_key]] = new_line
else:
lines.insert(insert_at, new_line)
insert_at += 1
with open(workspace_path, "w", encoding="utf-8") as handle:
handle.writelines(lines)
return True
[docs]
def workspace_refined_cell_file_path(workspace_path: str) -> str:
root = os.path.dirname(os.path.abspath(workspace_path)) if workspace_path else os.getcwd()
stem = os.path.splitext(os.path.basename(workspace_path))[0] if workspace_path else "workspace"
return os.path.join(root, WORKSPACE_REFINED_CELL_DIR, f"{stem}_refined.cell")
[docs]
def write_workspace_refined_cell_file(workspace_path: str) -> str | None:
"""Materialize the refined workspace cell header into a CrystFEL .cell file."""
kv = parse_workspace_cell(workspace_path)
required = ["a", "b", "c", "al", "be", "ga"]
if not workspace_has_refined_cell(workspace_path):
return None
if not all(kv.get(f"Cell.{key}") for key in required):
return None
path = workspace_refined_cell_file_path(workspace_path)
os.makedirs(os.path.dirname(path), exist_ok=True)
centering = kv.get("Cell.centering", kv.get("centering", "P"))
lattice_type = normalize_crystfel_lattice_type(
kv.get("Cell.lattice_type", kv.get("lattice_type", "triclinic")),
centering,
)
unique_axis = kv.get("Cell.unique_axis", kv.get("unique_axis", "")).strip()
if lattice_type not in UNIQUE_AXIS_LATTICE_TYPES:
unique_axis = ""
with open(path, "w", encoding="utf-8") as handle:
handle.write("CrystFEL unit cell file version 1.0\n\n")
handle.write(f"lattice_type = {lattice_type}\n\n")
handle.write(f"centering = {centering}\n")
if unique_axis:
handle.write(f"unique_axis = {unique_axis}\n\n")
else:
handle.write("\n")
handle.write(f"a = {kv['Cell.a'].split()[0]} A\n")
handle.write(f"b = {kv['Cell.b'].split()[0]} A\n")
handle.write(f"c = {kv['Cell.c'].split()[0]} A\n")
handle.write(f"al = {kv['Cell.al'].split()[0]} deg\n")
handle.write(f"be = {kv['Cell.be'].split()[0]} deg\n")
handle.write(f"ga = {kv['Cell.ga'].split()[0]} deg\n")
return path