FileBatchModelDisk¶
FileBatchModelDisk 是 AutoParser 返回的批次容器，保存成功解析的文件，并提供批量过滤、摘要、
frame 选择、格式转换和图像预览操作。
Python
from molop import AutoParser
batch = AutoParser("water_mp2.out", n_jobs=1)
summary = batch.to_summary_df(brief=False, flatten_columns=True)
print(summary.shape)
FileBatchModelDisk
¶
Bases: BatchFormatTransformMixin, MutableMapping, Generic[TFileDisk]
A class to store a batch of files and their corresponding models.
Source code in src/molop/io/FileBatchModelDisk.py
Python
class FileBatchModelDisk(BatchFormatTransformMixin, MutableMapping, Generic[TFileDisk]):
"""
A class to store a batch of files and their corresponding models.
"""
def __init__(self, diskfiles: Iterable[TFileDisk] | None = None):
self.__diskfiles: OrderedDict[str, TFileDisk] = OrderedDict()
if diskfiles is not None:
self.add_diskfiles(diskfiles)
@overload
def parallel_execute(
self,
func: Callable[..., R],
desc: str = "",
n_jobs: int = -1,
return_as: Literal["list", "generator", "generator_unordered"] = "list",
*args: Any,
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
return_results: Literal[False],
**kwargs: Any,
) -> None: ...
@overload
def parallel_execute(
self,
func: Callable[..., R],
desc: str = "",
n_jobs: int = -1,
return_as: Literal["list", "generator", "generator_unordered"] = "list",
*args: Any,
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
return_results: Literal[True],
**kwargs: Any,
) -> Iterable[R]: ...
@overload
def parallel_execute(
self,
func: Callable[..., R],
desc: str = "",
n_jobs: int = -1,
return_as: Literal["list", "generator", "generator_unordered"] = "list",
*args: Any, # do not use this parameter
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
return_results: None = None,
**kwargs: Any,
) -> Iterable[R] | None: ...
def parallel_execute(
self,
func: Callable[..., R],
desc: str = "",
n_jobs: int = -1,
return_as: Literal["list", "generator", "generator_unordered"] = "list",
*args: Any, # do not use this parameter
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
return_results: bool | None = None,
**kwargs: Any,
) -> Iterable[R] | None:
"""
Internal helper to execute a function over the batch in parallel or serial.
Refactored to support Adaptive UI (Rich/Ipywidgets) and Type Hints.
Accepts extra keyword arguments to pass to the function.
Parameters:
func (Callable[..., R]): The function to execute.
desc (str, optional): The description to show in the progress bar. Defaults to "".
n_jobs (int, optional): The number of jobs to use. Defaults to -1.
return_as (Literal["list", "generator", "generator_unordered"], optional): The return mode to use. Defaults to "list".
_diskfiles_snapshot (Sequence[TFileDisk] | None, optional): Internal precomputed
batch snapshot used by callers that need the same item order for zipping results.
return_results (bool | None, optional): When False, exhaust execution and return None.
When None, list results containing only implicit None values are returned as None.
**kwargs: Additional keyword arguments to pass to the function.
Returns:
Iterable[R] | None: The function results, or None for side-effect-only calls.
"""
diskfiles = (
_diskfiles_snapshot if _diskfiles_snapshot is not None else self._snapshot_diskfiles()
)
return parallel_map(
lambda diskfile: func(diskfile, *args, **kwargs),
diskfiles,
n_jobs=molopconfig.set_n_jobs(n_jobs),
desc=desc,
total=len(diskfiles),
disable=not molopconfig.show_progress_bar,
return_as=return_as,
return_results=return_results,
maxtasks_per_child=50,
max_nbytes=molopconfig.parallel_max_size,
)
def _prewarm_topologies(
self,
*,
frame: FrameSelector = "all",
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
backend: Literal["cpp", "python"] | None = None,
make_dative_bonds: bool | None = None,
make_stereochemistry: bool | None = None,
max_workers: int | None = None,
queue_size: int = 16,
ordered: bool = False,
raise_on_error: bool = False,
) -> list[ReconstructionBatchResult]:
"""Prewarm coordinate-only frames with MolGR's native batch API.
Frames are gathered in this process and submitted as one or more
configuration-homogeneous batches. This keeps MolGR's bounded native
worker pool at the outermost reconstruction boundary instead of
nesting it inside :meth:`parallel_execute` workers.
"""
from molop.io.base_models.Molecule import reconstruct_topologies_batch
diskfiles = (
_diskfiles_snapshot if _diskfiles_snapshot is not None else self._snapshot_diskfiles()
)
def iter_frames() -> Iterable[Any]:
for diskfile in diskfiles:
file_frames = getattr(diskfile, "frames", None)
if file_frames is None:
file_frames = diskfile
frame_ids = normalize_frame_selector(
frame,
len(file_frames),
parameter_name="frame",
)
yield from (file_frames[frame_id] for frame_id in frame_ids)
return reconstruct_topologies_batch(
iter_frames(),
backend=backend,
make_dative_bonds=make_dative_bonds,
make_stereochemistry=make_stereochemistry,
max_workers=max_workers,
queue_size=queue_size,
ordered=ordered,
raise_on_error=raise_on_error,
retain_results=False,
)
def _snapshot_diskfiles(self) -> list[TFileDisk]:
return list(self.__diskfiles.values())
def _filter_diskfiles(
self, diskfiles: Sequence[TFileDisk], keep_flags: Iterable[bool]
) -> FileBatchModelDisk[TFileDisk]:
return self.new_batch(
diskfile for diskfile, keep in zip(diskfiles, keep_flags, strict=True) if keep
)
def _add_sorted_diskfiles(self, diskfiles: Iterable[TFileDisk]) -> None:
for diskfile in diskfiles:
if not _looks_like_disk_file(diskfile):
raise TypeError(f"diskfiles must be MolOP file-like objects, got {type(diskfile)}")
if diskfile.file_path not in self.__diskfiles:
self.__diskfiles[diskfile.file_path] = diskfile
else:
moloplogger.warning(
f"File {diskfile.file_path} already exists in the batch, skipped"
)
def add_diskfiles(self, diskfiles: Iterable[TFileDisk]) -> None:
"""
Add disk files to the batch.
Parameters:
diskfiles (Iterable[FileDiscObj]): The disk files to add.
"""
self._add_sorted_diskfiles(sorted(diskfiles, key=lambda d: d.file_path))
@classmethod
def new_batch(cls, diskfiles: Iterable[TFileDisk]):
new_batch: FileBatchModelDisk[TFileDisk] = cls()
new_batch.add_diskfiles(diskfiles)
return new_batch
@classmethod
def _new_batch_from_sorted_diskfiles(
cls, diskfiles: Iterable[TFileDisk]
) -> FileBatchModelDisk[TFileDisk]:
new_batch: FileBatchModelDisk[TFileDisk] = cls()
new_batch._add_sorted_diskfiles(diskfiles)
return new_batch
def __repr__(self) -> str:
return f"{self.__class__.__name__}({len(self)})"
def __contains__(self, key: object) -> bool:
if isinstance(key, str):
return key in self.__diskfiles
else:
return key in self.__diskfiles.values()
@overload
def __getitem__(self, key: int | str) -> TFileDisk: ...
@overload
def __getitem__(self, key: slice | Sequence[int | str]) -> FileBatchModelDisk[TFileDisk]: ...
def __getitem__(
self,
key: int | str | slice | Sequence[int | str],
) -> TFileDisk | FileBatchModelDisk[TFileDisk]:
if isinstance(key, int):
return list(self.__diskfiles.values())[key]
if isinstance(key, str):
return self.__diskfiles[key]
if isinstance(key, slice):
return self.new_batch(list(self.__diskfiles.values())[key])
if isinstance(key, Sequence):
return self.new_batch([self[k] for k in key])
raise TypeError(f"Invalid key type: {type(key)}")
@overload
def require(self, key: int, expected_cls: type[TFile]) -> TFile: ...
@overload
def require(self, key: str, expected_cls: type[TFile]) -> TFile: ...
def require(self, key: int | str, expected_cls: type[TFile]) -> TFile:
"""Return a file with precise type hints, validating at runtime.
This avoids hardcoding a global list of concrete file classes while still
allowing callers to opt into format-specific APIs.
"""
value = self[key]
if not isinstance(value, expected_cls):
raise TypeError(
f"Expected {expected_cls.__name__} for key {key!r}, got {type(value).__name__}"
)
return value
def __setitem__(self, key: str, value: FileDiskObj) -> None:
if not isinstance(key, str):
raise TypeError(f"Key should be a string, got {type(key)}")
if not _looks_like_disk_file(value):
raise TypeError(f"Value must be a MolOP file-like object, got {type(value)}")
if key in self.__diskfiles:
moloplogger.warning(
f"File {key} already exists in the batch, replaced with the new one"
)
self.__diskfiles[key] = cast(TFileDisk, value)
def __delitem__(self, key: str | int) -> None:
if isinstance(key, int):
key = list(self.__diskfiles.keys())[key]
del self.__diskfiles[key]
def __len__(self) -> int:
return len(self.__diskfiles)
def __iter__(self) -> Iterator[TFileDisk]: # type: ignore[override]
return iter(self.__diskfiles.values())
def keys(self):
return self.__diskfiles.keys()
def values(self):
return self.__diskfiles.values()
def items(self):
return self.__diskfiles.items()
def __add__(self, other: FileBatchModelDisk[TFileDisk]) -> FileBatchModelDisk[TFileDisk]:
"""
Operator +: Merge two batches (Union).
Supports passing another FileBatchModelDisk instance.
Parameters:
other (FileBatchModelDisk): The other batch to merge with.
Returns:
FileBatchModelDisk: A new batch containing all unique files from both batches.
"""
if not isinstance(other, FileBatchModelDisk):
return NotImplemented
new_batch = self.new_batch(self.__diskfiles.values())
new_batch.add_diskfiles(other.__diskfiles.values())
return new_batch
def __sub__(
self, other: FileBatchModelDisk[TFileDisk] | Iterable[str]
) -> FileBatchModelDisk[TFileDisk]:
"""
Operator -: Remove files from the batch (Difference).
Supports passing another batch OR a simple list of file paths.
This avoids the IO overhead of parsing files just to exclude them.
Parameters:
other (Union[FileBatchModelDisk, Iterable[str]]): The other batch or list of file paths to exclude.
Returns:
FileBatchModelDisk: A new batch containing files not present in the other batch.
"""
other_paths: set[str]
if isinstance(other, FileBatchModelDisk):
other_paths = set(other.file_paths)
elif isinstance(other, str):
other_paths = {other}
elif isinstance(other, Iterable):
other_paths = set(other)
else:
return NotImplemented
filtered_files = [f for f in self.__diskfiles.values() if f.file_path not in other_paths]
return self.new_batch(filtered_files)
def __and__(
self, other: FileBatchModelDisk[TFileDisk] | Iterable[str]
) -> FileBatchModelDisk[TFileDisk]:
"""
Operator &: Keep only files present in both (Intersection).
Also supports passing a list of file paths.
Parameters:
other (Union[FileBatchModelDisk, Iterable[str]]): The other batch or list of file paths to intersect with.
Returns:
FileBatchModelDisk: A new batch containing files present in both batches.
"""
other_paths: set[str]
if isinstance(other, FileBatchModelDisk):
other_paths = set(other.file_paths)
elif isinstance(other, str):
other_paths = {other}
elif isinstance(other, Iterable):
other_paths = set(other)
else:
return NotImplemented
common_files = [f for f in self.__diskfiles.values() if f.file_path in other_paths]
return self.new_batch(common_files)
@property
def file_paths(self) -> list[str]:
"""
Get the file paths of the batch.
Returns:
List[str]: The file paths.
"""
return [diskfile.file_path for diskfile in self]
@property
def file_names(self) -> list[str]:
"""
Get the file names of the batch.
Returns:
List[str]: The file names.
"""
return [diskfile.filename for diskfile in self]
def filter_state(
self,
state: Literal["ts", "error", "opt", "normal", "thermal", "no-img"],
negate: bool = False,
n_jobs: int = -1,
) -> FileBatchModelDisk[TFileDisk]:
"""
Filter the files based on their state using the new execution engine.
Parameters:
state (Literal["ts", "error", "opt", "normal", "thermal", "no-img"]): The state to filter.
negate (bool): Whether to negate the filter.
n_jobs (int): Number of parallel jobs to filter.
Returns:
FileBatchModelDisk: The filtered batch.
"""
def has_frequency_data(vibrations: Any) -> bool:
try:
return len(vibrations) > 0
except TypeError:
frequencies = getattr(vibrations, "frequencies", None)
if frequencies is None:
return True
try:
return len(frequencies) > 0
except TypeError:
return True
def judge_func(diskfile: TFileDisk) -> bool:
res: bool = False
if state == "ts":
res = any(getattr(frame, "is_TS", None) for frame in diskfile)
elif state == "error":
res = any(getattr(frame, "is_error", None) for frame in diskfile)
elif state == "opt":
if (
any(
getattr(frame, "geometry_optimization_status", None) is not None
for frame in diskfile
)
and (opt_frame := getattr(diskfile, "closest_optimized_frame", None))
is not None
):
res = (
all(getattr(frame, "is_normal", None) for frame in diskfile)
and opt_frame.is_optimized
)
else:
res = False
elif state == "normal":
res = all(getattr(frame, "is_normal", None) for frame in diskfile)
elif state == "thermal":
res = any(
getattr(frame, "thermal_informations", None) is not None for frame in diskfile
)
elif state == "no-img":
vibration_frames: list[HasVibrations] = [
cast(HasVibrations, frame)
for frame in diskfile
if (vibrations := getattr(frame, "vibrations", None)) is not None
and has_frequency_data(vibrations)
]
res = bool(vibration_frames) and all(
frame.vibrations.num_imaginary == 0 for frame in vibration_frames
)
else:
raise ValueError(f"Invalid state: {state}")
return res != negate
desc = f"Filtering {state} files{' negated' if negate else ''} with {molopconfig.set_n_jobs(n_jobs)} jobs"
diskfiles = self._snapshot_diskfiles()
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
judge_func,
desc,
n_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
def filter_value(
self,
target: Literal["charge", "multiplicity", "format"],
value: str | float | int,
compare: Literal["==", "!=", ">", "<", ">=", "<="] = "==",
n_jobs: int = -1,
) -> FileBatchModelDisk[TFileDisk]:
"""
Filter with updated operator module and parallel execution.
Parameters:
target (Literal["charge", "multiplicity", "format"]): The target property to filter.
value (Union[str, float, int]): The value to compare against.
compare (Literal["==", "!=", ">", "<", ">=", "<="]): The comparison operator.
n_jobs (int): Number of parallel jobs to filter.
Returns:
FileBatchModelDisk: The filtered batch.
"""
compare_ops = {
"==": operator.eq,
"!=": operator.ne,
">": operator.gt,
"<": operator.lt,
">=": operator.ge,
"<=": operator.le,
}
op_func = compare_ops[compare]
def judge_func(diskfile: TFileDisk) -> bool:
if target == "charge":
return op_func(diskfile[-1].charge, value)
elif target == "multiplicity":
return op_func(diskfile[-1].multiplicity, value)
elif target == "format":
return op_func(diskfile.file_format, value)
return False
desc = f"Filtering files with {target} {compare} {value} with {molopconfig.set_n_jobs(n_jobs)} jobs"
diskfiles = self._snapshot_diskfiles()
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
judge_func,
desc,
n_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
def filter_custom(
self,
condition: Callable[[TFileDisk], bool],
n_jobs: int = -1,
) -> FileBatchModelDisk[TFileDisk]:
"""
Filter the files based on a custom callable condition.
This allows for complex filtering logic that goes beyond simple state or value checks.
Parameters:
condition (Callable[[FileDiscObj], bool]): A function that accepts a diskfile object and returns True (keep) or False (discard).
n_jobs (int): Number of parallel jobs.
Returns:
FileBatchModelDisk: A new batch containing only the files that satisfy the condition.
"""
effective_jobs = molopconfig.set_molgr_n_jobs(n_jobs)
desc = f"Filtering with custom function with {effective_jobs} jobs"
diskfiles = self._snapshot_diskfiles()
# A user callback may inspect any frame-level graph (``rdmol``,
# canonical SMILES, descriptors, ...). Materialize those graphs in
# the parent process before the callback can run in a loky worker.
if molopconfig.prewarm_topologies:
self._prewarm_topologies(
_diskfiles_snapshot=diskfiles,
max_workers=effective_jobs,
)
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
condition,
desc,
effective_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
def groupby(
self, key_func: Callable[[TFileDisk], str], n_jobs: int = -1
) -> dict[str, FileBatchModelDisk[TFileDisk]]:
"""
Group the files into multiple batches based on a key function.
Parameters:
key_func (Callable[[FileDiscObj], str]): A function that generates a group key for each file.
n_jobs (int): Number of parallel jobs to calculate keys.
Returns:
(dict[str, FileBatchModelDisk]): A dictionary where keys are group names and values are new batch objects.
"""
effective_jobs = molopconfig.set_molgr_n_jobs(n_jobs)
desc = f"Grouping files with {effective_jobs} jobs"
diskfiles = self._snapshot_diskfiles()
# ``key_func`` is intentionally arbitrary, so treat it as graph
# dependent even when the current callback only uses metadata. This
# keeps future callbacks from accidentally entering MolGR in workers.
if molopconfig.prewarm_topologies:
self._prewarm_topologies(
_diskfiles_snapshot=diskfiles,
max_workers=effective_jobs,
)
keys = self.parallel_execute(
key_func,
desc,
effective_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
)
temp_groups: dict[str, list[TFileDisk]] = {}
for diskfile, key in zip(diskfiles, keys, strict=True):
if key not in temp_groups:
temp_groups[key] = []
temp_groups[key].append(diskfile)
return {group_key: self.new_batch(files) for group_key, files in temp_groups.items()}
def filter_by_codec_id(
self,
codec_id: str,
*,
negate: bool = False,
n_jobs: int = -1,
on_missing: Literal["keep", "drop", "error"] = "drop",
) -> FileBatchModelDisk[TFileDisk]:
"""Filter batch by detected reader codec id.
This differs from `filter_value(target="format", ...)` which filters by file suffix.
Parameters:
codec_id: Reader format id (e.g. "xyz", "g16log"). Compared against each diskfile's
`detected_format_id` populated during parsing.
negate: If True, invert the predicate.
n_jobs: Parallel jobs.
on_missing: Behavior when `detected_format_id` is missing/None on an item.
- "drop": drop the item
- "keep": keep the item
- "error": raise ValueError
"""
wanted = codec_id.strip().lower()
if not wanted:
raise ValueError("codec_id must be non-empty")
def judge_func(diskfile: TFileDisk) -> bool:
detected = getattr(diskfile, "detected_format_id", None)
if detected is None:
if on_missing == "keep":
return True
if on_missing == "drop":
return False
raise ValueError(
f"Missing detected_format_id for {getattr(diskfile, 'file_path', '?')}"
)
hit = str(detected).strip().lower() == wanted
return (not hit) if negate else hit
desc = (
f"Filtering files with codec_id == {wanted} with {molopconfig.set_n_jobs(n_jobs)} jobs"
)
diskfiles = self._snapshot_diskfiles()
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
judge_func,
desc,
n_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
def to_summary_df(
self,
mode: Literal["file", "frame"] = "frame",
frame: FrameSelector = -1,
n_jobs: int = -1,
*,
brief: bool = True,
flatten_columns: bool = False,
on_missing_frame: Literal["skip", "error"] = "skip",
**kwargs: Any,
) -> pd.DataFrame:
"""
Generate summary DataFrame using the parallel execution engine.
Parameters:
mode (Literal["file", "frame"]): Whether to generate file-level or frame-level summary.
frame (FrameSelector): Frame IDs to process.
Use -1 for the last frame or "all" for every frame.
n_jobs (int): Number of parallel jobs.
brief (bool): Whether to use the compact summary payload.
flatten_columns (bool): Whether to flatten MultiIndex columns with dot-separated names.
on_missing_frame (Literal["skip", "error"]): How to handle frame indices that are
out of range for a file.
**kwargs: Additional arguments for to_summary_series.
Returns:
pd.DataFrame: A DataFrame containing the summary information.
"""
if mode not in ["file", "frame"]:
raise ValueError(f"Invalid mode: {mode}")
if on_missing_frame not in ["skip", "error"]:
raise ValueError(f"Invalid on_missing_frame: {on_missing_frame}")
def selected_frame_ids(diskfile: FileDiskObj) -> Sequence[int]:
try:
frame_ids = normalize_frame_selector(
frame,
len(diskfile),
parameter_name="frame",
validate_range=on_missing_frame == "error",
)
except IndexError as exc:
raise IndexError(
f"{exc} for {getattr(diskfile, 'file_path', '<unknown>')} "
f"with {len(diskfile)} frames"
) from exc
if on_missing_frame == "skip":
return [fid for fid in frame_ids if 0 <= fid < len(diskfile)]
return frame_ids
diskfiles = self._snapshot_diskfiles()
summary_tasks: list[tuple[FileDiskObj, Sequence[int]]] = []
if mode == "file":
summary_tasks = [(diskfile, ()) for diskfile in diskfiles]
else:
selected_frames: list[Any] = []
for diskfile in diskfiles:
frame_ids = list(selected_frame_ids(diskfile))
selected_frames.extend(diskfile[fid] for fid in frame_ids)
summary_tasks.append((diskfile, frame_ids))
if molopconfig.prewarm_topologies:
from molop.io.base_models.Molecule import reconstruct_topologies_batch
# The selected source frames have a stable, enumerable input
# set, so optionally warm their ordinary graph cache in one
# parent-process native batch before process workers start.
reconstruct_topologies_batch(
selected_frames,
max_workers=molopconfig.set_molgr_n_jobs(n_jobs),
retain_results=False,
)
def process_file_summary(task: tuple[FileDiskObj, Sequence[int]]) -> list[pd.Series]:
diskfile, frame_ids = task
if mode == "file":
return [diskfile.to_summary_series(brief=brief, **kwargs)]
summaries: list[pd.Series] = []
for fid in frame_ids:
summaries.append(diskfile[fid].to_summary_series(brief=brief, **kwargs))
return summaries
effective_jobs = (
molopconfig.set_molgr_n_jobs(n_jobs)
if mode == "frame"
else molopconfig.set_n_jobs(n_jobs)
)
desc = f"MolOP processing {mode} summary with {effective_jobs} jobs"
nested_results = parallel_map(
process_file_summary,
summary_tasks,
n_jobs=effective_jobs,
desc=desc,
total=len(summary_tasks),
disable=not molopconfig.show_progress_bar,
return_results=True,
maxtasks_per_child=50,
max_nbytes=molopconfig.parallel_max_size,
)
return build_summary_df(
(series for sublist in nested_results for series in sublist),
flatten_columns=flatten_columns,
)
def release_file_content(self) -> None:
"""
Release the content of the files in the batch.
"""
for diskfile in self:
diskfile.release_file_content()
def save_pre_post_ts(
self,
output_dir: os.PathLike[str] | str,
*,
format: Literal["xyz", "sdf"] = "xyz",
min_ratio: float = 0.6,
max_ratio: float = 1.4,
steps: int = 8,
sampling_method: VibrationSamplingMethod = "harmonic_potential",
n_jobs: int = -1,
) -> dict[str, dict[int, tuple[Path, Path]]]:
"""Export endpoint candidates for every TS frame in every calculation file.
Each supported source file receives a separate directory. Unique filename
stems are retained; duplicate stems gain a stable source-path digest to
prevent overwrites. Files without endpoint-export support are skipped with
a warning. The result maps supported source paths to the frame-ID-to-endpoint
mappings returned by each file. Endpoint inference samples ``steps``
amplitudes on each side from ``min_ratio`` through ``max_ratio``
according to ``sampling_method``. Defaults are ``min_ratio=0.6``,
``max_ratio=1.4``, ``steps=8``, and
``sampling_method="harmonic_potential"``.
"""
destination = Path(output_dir)
destination.mkdir(parents=True, exist_ok=True)
diskfiles = self._snapshot_diskfiles()
exportable_diskfiles: list[TFileDisk] = []
for diskfile in diskfiles:
if callable(getattr(diskfile, "save_pre_post_ts", None)):
exportable_diskfiles.append(diskfile)
else:
moloplogger.warning(
f"Skipping {diskfile.file_path}: "
f"{type(diskfile).__name__} does not support TS endpoint export"
)
output_directories = _endpoint_output_directories(
cast(Sequence[DiskFileLike], exportable_diskfiles), destination
)
endpoint_tasks: list[tuple[TFileDisk, Path]] = [
(diskfile, output_directories[str(diskfile.file_path)])
for diskfile in exportable_diskfiles
]
def export_file(
task: tuple[TFileDisk, Path],
) -> tuple[str, dict[int, tuple[Path, Path]]]:
diskfile, output_directory = task
save_endpoints = getattr(diskfile, "save_pre_post_ts", None)
assert callable(save_endpoints)
source_path = str(diskfile.file_path)
exports = cast(
dict[int, tuple[Path, Path]],
save_endpoints(
output_directory,
format=format,
min_ratio=min_ratio,
max_ratio=max_ratio,
steps=steps,
sampling_method=sampling_method,
),
)
return source_path, exports
effective_jobs = molopconfig.set_molgr_n_jobs(n_jobs)
desc = f"Exporting TS endpoint candidates with {effective_jobs} jobs"
results = parallel_map(
export_file,
endpoint_tasks,
n_jobs=effective_jobs,
desc=desc,
total=len(endpoint_tasks),
disable=not molopconfig.show_progress_bar,
return_results=True,
maxtasks_per_child=50,
max_nbytes=molopconfig.parallel_max_size,
)
return dict(results)
def draw_grid_image(
self,
molsPerRow: int = 4,
subImgSize: tuple[int, int] = (200, 200),
maxMols: int = 16,
useSVG: bool = True,
n_jobs: int = -1,
**kwargs: Any,
):
"""
Quick preview of the structures of the first N molecules in a Jupyter Notebook.
Parameters:
molsPerRow (int): Number of molecules per row in the grid.
subImgSize (tuple[int, int]): Size of each sub-image in the grid.
maxMols (int): Maximum number of molecules to visualize.
useSVG (bool): Whether to use SVG format for visualization.
n_jobs (int): Number of parallel jobs.
**kwargs: Additional arguments for the configured RDKit drawer.
Returns:
(PIL.Image.Image | None): A grid image of the molecules, or None if RDKit is not installed.
"""
try:
from rdkit import Chem
from rdkit.Chem import Draw
except ImportError:
moloplogger.warning("RDKit is not installed, visualization skipped.")
return None
mols = []
legends = []
for file_path, smi in (
self[:maxMols].format_transform("smi", output_dir=None, n_jobs=n_jobs).items()
):
smi = smi[-1] if isinstance(smi, list) and smi else smi
if isinstance(smi, str) and smi:
mol = Chem.MolFromSmiles(smi.strip())
if mol:
mols.append(mol)
legends.append(os.path.basename(file_path))
else:
mols.append(Chem.MolFromSmiles(""))
legends.append(os.path.basename(file_path))
if not mols:
moloplogger.warning("No valid molecules found for visualization.")
return None
if molopconfig.use_dof_effect_drawer:
try:
from rdkit_dof import MolsToGridDofImage
except ImportError:
moloplogger.warning("rdkit-dof is not installed; falling back to the RDKit drawer.")
else:
dof_kwargs = dict(kwargs)
dof_kwargs.pop("use_svg", None)
return MolsToGridDofImage(
mols,
molsPerRow=molsPerRow,
subImgSize=subImgSize,
legends=legends,
use_svg=useSVG,
return_image=not useSVG,
**dof_kwargs,
)
return Draw.MolsToGridImage(
mols,
molsPerRow=molsPerRow,
subImgSize=subImgSize,
legends=legends,
useSVG=useSVG,
**kwargs,
)
def sample(self, n: int = 10, seed: int | None = None) -> FileBatchModelDisk[TFileDisk]:
"""
Randomly sample n files from the batch.
Parameters:
n (int): Number of files to sample.
seed (int): Random seed for reproducibility.
Returns:
FileBatchModelDisk: A new batch containing the sampled files.
"""
if seed is not None:
random.seed(seed)
n = min(n, len(self))
sampled_files = random.sample(list(self.__diskfiles.values()), n)
return self.new_batch(sampled_files)
def copy_to(self, output_dir: str, n_jobs: int = -1):
if not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
def copy(diskfile: TFileDisk):
"""
Copy a file to the output directory.
"""
try:
shutil.copy(diskfile.file_path, output_dir)
except Exception as e:
return e
return f"successfully copied {diskfile.file_path} to {output_dir}"
desc = f"Copying files to {output_dir} with {molopconfig.set_n_jobs(n_jobs)} jobs"
return self.parallel_execute(copy, desc, n_jobs, return_as="list", return_results=True)
def move_to(self, output_dir: str, n_jobs: int = -1):
if not os.path.exists(output_dir):
os.makedirs(output_dir, exist_ok=True)
def move(diskfile: TFileDisk):
"""
Move a file to the output directory.
"""
try:
shutil.move(diskfile.file_path, output_dir)
except Exception as e:
return e
return f"successfully moved {diskfile.file_path} to {output_dir}"
desc = f"Moving files to {output_dir} with {molopconfig.set_n_jobs(n_jobs)} jobs"
return self.parallel_execute(move, desc, n_jobs, return_as="list", return_results=True)
file_names
property
¶
Get the file names of the batch.
Returns:
| Type | Description |
|---|---|
list[str]
|
List[str]: The file names. |
file_paths
property
¶
Get the file paths of the batch.
Returns:
| Type | Description |
|---|---|
list[str]
|
List[str]: The file paths. |
__add__(other)
¶
Operator +: Merge two batches (Union).
Supports passing another FileBatchModelDisk instance.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
other
|
FileBatchModelDisk
|
The other batch to merge with. |
required |
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
A new batch containing all unique files from both batches. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def __add__(self, other: FileBatchModelDisk[TFileDisk]) -> FileBatchModelDisk[TFileDisk]:
"""
Operator +: Merge two batches (Union).
Supports passing another FileBatchModelDisk instance.
Parameters:
other (FileBatchModelDisk): The other batch to merge with.
Returns:
FileBatchModelDisk: A new batch containing all unique files from both batches.
"""
if not isinstance(other, FileBatchModelDisk):
return NotImplemented
new_batch = self.new_batch(self.__diskfiles.values())
new_batch.add_diskfiles(other.__diskfiles.values())
return new_batch
__and__(other)
¶
Operator &: Keep only files present in both (Intersection).
Also supports passing a list of file paths.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
other
|
Union[FileBatchModelDisk, Iterable[str]]
|
The other batch or list of file paths to intersect with. |
required |
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
A new batch containing files present in both batches. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def __and__(
self, other: FileBatchModelDisk[TFileDisk] | Iterable[str]
) -> FileBatchModelDisk[TFileDisk]:
"""
Operator &: Keep only files present in both (Intersection).
Also supports passing a list of file paths.
Parameters:
other (Union[FileBatchModelDisk, Iterable[str]]): The other batch or list of file paths to intersect with.
Returns:
FileBatchModelDisk: A new batch containing files present in both batches.
"""
other_paths: set[str]
if isinstance(other, FileBatchModelDisk):
other_paths = set(other.file_paths)
elif isinstance(other, str):
other_paths = {other}
elif isinstance(other, Iterable):
other_paths = set(other)
else:
return NotImplemented
common_files = [f for f in self.__diskfiles.values() if f.file_path in other_paths]
return self.new_batch(common_files)
__sub__(other)
¶
Operator -: Remove files from the batch (Difference).
Supports passing another batch OR a simple list of file paths. This avoids the IO overhead of parsing files just to exclude them.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
other
|
Union[FileBatchModelDisk, Iterable[str]]
|
The other batch or list of file paths to exclude. |
required |
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
A new batch containing files not present in the other batch. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def __sub__(
self, other: FileBatchModelDisk[TFileDisk] | Iterable[str]
) -> FileBatchModelDisk[TFileDisk]:
"""
Operator -: Remove files from the batch (Difference).
Supports passing another batch OR a simple list of file paths.
This avoids the IO overhead of parsing files just to exclude them.
Parameters:
other (Union[FileBatchModelDisk, Iterable[str]]): The other batch or list of file paths to exclude.
Returns:
FileBatchModelDisk: A new batch containing files not present in the other batch.
"""
other_paths: set[str]
if isinstance(other, FileBatchModelDisk):
other_paths = set(other.file_paths)
elif isinstance(other, str):
other_paths = {other}
elif isinstance(other, Iterable):
other_paths = set(other)
else:
return NotImplemented
filtered_files = [f for f in self.__diskfiles.values() if f.file_path not in other_paths]
return self.new_batch(filtered_files)
add_diskfiles(diskfiles)
¶
Add disk files to the batch.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
diskfiles
|
Iterable[FileDiscObj]
|
The disk files to add. |
required |
Source code in src/molop/io/FileBatchModelDisk.py
draw_grid_image(molsPerRow=4, subImgSize=(200, 200), maxMols=16, useSVG=True, n_jobs=-1, **kwargs)
¶
Quick preview of the structures of the first N molecules in a Jupyter Notebook.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
molsPerRow
|
int
|
Number of molecules per row in the grid. |
4
|
subImgSize
|
tuple[int, int]
|
Size of each sub-image in the grid. |
(200, 200)
|
maxMols
|
int
|
Maximum number of molecules to visualize. |
16
|
useSVG
|
bool
|
Whether to use SVG format for visualization. |
True
|
n_jobs
|
int
|
Number of parallel jobs. |
-1
|
**kwargs
|
Any
|
Additional arguments for the configured RDKit drawer. |
{}
|
Returns:
| Type | Description |
|---|---|
Image | None
|
A grid image of the molecules, or None if RDKit is not installed. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def draw_grid_image(
self,
molsPerRow: int = 4,
subImgSize: tuple[int, int] = (200, 200),
maxMols: int = 16,
useSVG: bool = True,
n_jobs: int = -1,
**kwargs: Any,
):
"""
Quick preview of the structures of the first N molecules in a Jupyter Notebook.
Parameters:
molsPerRow (int): Number of molecules per row in the grid.
subImgSize (tuple[int, int]): Size of each sub-image in the grid.
maxMols (int): Maximum number of molecules to visualize.
useSVG (bool): Whether to use SVG format for visualization.
n_jobs (int): Number of parallel jobs.
**kwargs: Additional arguments for the configured RDKit drawer.
Returns:
(PIL.Image.Image | None): A grid image of the molecules, or None if RDKit is not installed.
"""
try:
from rdkit import Chem
from rdkit.Chem import Draw
except ImportError:
moloplogger.warning("RDKit is not installed, visualization skipped.")
return None
mols = []
legends = []
for file_path, smi in (
self[:maxMols].format_transform("smi", output_dir=None, n_jobs=n_jobs).items()
):
smi = smi[-1] if isinstance(smi, list) and smi else smi
if isinstance(smi, str) and smi:
mol = Chem.MolFromSmiles(smi.strip())
if mol:
mols.append(mol)
legends.append(os.path.basename(file_path))
else:
mols.append(Chem.MolFromSmiles(""))
legends.append(os.path.basename(file_path))
if not mols:
moloplogger.warning("No valid molecules found for visualization.")
return None
if molopconfig.use_dof_effect_drawer:
try:
from rdkit_dof import MolsToGridDofImage
except ImportError:
moloplogger.warning("rdkit-dof is not installed; falling back to the RDKit drawer.")
else:
dof_kwargs = dict(kwargs)
dof_kwargs.pop("use_svg", None)
return MolsToGridDofImage(
mols,
molsPerRow=molsPerRow,
subImgSize=subImgSize,
legends=legends,
use_svg=useSVG,
return_image=not useSVG,
**dof_kwargs,
)
return Draw.MolsToGridImage(
mols,
molsPerRow=molsPerRow,
subImgSize=subImgSize,
legends=legends,
useSVG=useSVG,
**kwargs,
)
filter_by_codec_id(codec_id, *, negate=False, n_jobs=-1, on_missing='drop')
¶
Filter batch by detected reader codec id.
This differs from filter_value(target="format", ...) which filters by file suffix.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
codec_id
|
str
|
Reader format id (e.g. “xyz”, “g16log”). Compared against each diskfile’s
|
required |
negate
|
bool
|
If True, invert the predicate. |
False
|
n_jobs
|
int
|
Parallel jobs. |
-1
|
on_missing
|
Literal['keep', 'drop', 'error']
|
Behavior when |
'drop'
|
Source code in src/molop/io/FileBatchModelDisk.py
Python
def filter_by_codec_id(
self,
codec_id: str,
*,
negate: bool = False,
n_jobs: int = -1,
on_missing: Literal["keep", "drop", "error"] = "drop",
) -> FileBatchModelDisk[TFileDisk]:
"""Filter batch by detected reader codec id.
This differs from `filter_value(target="format", ...)` which filters by file suffix.
Parameters:
codec_id: Reader format id (e.g. "xyz", "g16log"). Compared against each diskfile's
`detected_format_id` populated during parsing.
negate: If True, invert the predicate.
n_jobs: Parallel jobs.
on_missing: Behavior when `detected_format_id` is missing/None on an item.
- "drop": drop the item
- "keep": keep the item
- "error": raise ValueError
"""
wanted = codec_id.strip().lower()
if not wanted:
raise ValueError("codec_id must be non-empty")
def judge_func(diskfile: TFileDisk) -> bool:
detected = getattr(diskfile, "detected_format_id", None)
if detected is None:
if on_missing == "keep":
return True
if on_missing == "drop":
return False
raise ValueError(
f"Missing detected_format_id for {getattr(diskfile, 'file_path', '?')}"
)
hit = str(detected).strip().lower() == wanted
return (not hit) if negate else hit
desc = (
f"Filtering files with codec_id == {wanted} with {molopconfig.set_n_jobs(n_jobs)} jobs"
)
diskfiles = self._snapshot_diskfiles()
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
judge_func,
desc,
n_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
filter_custom(condition, n_jobs=-1)
¶
Filter the files based on a custom callable condition.
This allows for complex filtering logic that goes beyond simple state or value checks.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
condition
|
Callable[[FileDiscObj], bool]
|
A function that accepts a diskfile object and returns True (keep) or False (discard). |
required |
n_jobs
|
int
|
Number of parallel jobs. |
-1
|
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
A new batch containing only the files that satisfy the condition. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def filter_custom(
self,
condition: Callable[[TFileDisk], bool],
n_jobs: int = -1,
) -> FileBatchModelDisk[TFileDisk]:
"""
Filter the files based on a custom callable condition.
This allows for complex filtering logic that goes beyond simple state or value checks.
Parameters:
condition (Callable[[FileDiscObj], bool]): A function that accepts a diskfile object and returns True (keep) or False (discard).
n_jobs (int): Number of parallel jobs.
Returns:
FileBatchModelDisk: A new batch containing only the files that satisfy the condition.
"""
effective_jobs = molopconfig.set_molgr_n_jobs(n_jobs)
desc = f"Filtering with custom function with {effective_jobs} jobs"
diskfiles = self._snapshot_diskfiles()
# A user callback may inspect any frame-level graph (``rdmol``,
# canonical SMILES, descriptors, ...). Materialize those graphs in
# the parent process before the callback can run in a loky worker.
if molopconfig.prewarm_topologies:
self._prewarm_topologies(
_diskfiles_snapshot=diskfiles,
max_workers=effective_jobs,
)
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
condition,
desc,
effective_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
filter_state(state, negate=False, n_jobs=-1)
¶
Filter the files based on their state using the new execution engine.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
state
|
Literal['ts', 'error', 'opt', 'normal', 'thermal', 'no-img']
|
The state to filter. |
required |
negate
|
bool
|
Whether to negate the filter. |
False
|
n_jobs
|
int
|
Number of parallel jobs to filter. |
-1
|
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
The filtered batch. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def filter_state(
self,
state: Literal["ts", "error", "opt", "normal", "thermal", "no-img"],
negate: bool = False,
n_jobs: int = -1,
) -> FileBatchModelDisk[TFileDisk]:
"""
Filter the files based on their state using the new execution engine.
Parameters:
state (Literal["ts", "error", "opt", "normal", "thermal", "no-img"]): The state to filter.
negate (bool): Whether to negate the filter.
n_jobs (int): Number of parallel jobs to filter.
Returns:
FileBatchModelDisk: The filtered batch.
"""
def has_frequency_data(vibrations: Any) -> bool:
try:
return len(vibrations) > 0
except TypeError:
frequencies = getattr(vibrations, "frequencies", None)
if frequencies is None:
return True
try:
return len(frequencies) > 0
except TypeError:
return True
def judge_func(diskfile: TFileDisk) -> bool:
res: bool = False
if state == "ts":
res = any(getattr(frame, "is_TS", None) for frame in diskfile)
elif state == "error":
res = any(getattr(frame, "is_error", None) for frame in diskfile)
elif state == "opt":
if (
any(
getattr(frame, "geometry_optimization_status", None) is not None
for frame in diskfile
)
and (opt_frame := getattr(diskfile, "closest_optimized_frame", None))
is not None
):
res = (
all(getattr(frame, "is_normal", None) for frame in diskfile)
and opt_frame.is_optimized
)
else:
res = False
elif state == "normal":
res = all(getattr(frame, "is_normal", None) for frame in diskfile)
elif state == "thermal":
res = any(
getattr(frame, "thermal_informations", None) is not None for frame in diskfile
)
elif state == "no-img":
vibration_frames: list[HasVibrations] = [
cast(HasVibrations, frame)
for frame in diskfile
if (vibrations := getattr(frame, "vibrations", None)) is not None
and has_frequency_data(vibrations)
]
res = bool(vibration_frames) and all(
frame.vibrations.num_imaginary == 0 for frame in vibration_frames
)
else:
raise ValueError(f"Invalid state: {state}")
return res != negate
desc = f"Filtering {state} files{' negated' if negate else ''} with {molopconfig.set_n_jobs(n_jobs)} jobs"
diskfiles = self._snapshot_diskfiles()
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
judge_func,
desc,
n_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
filter_value(target, value, compare='==', n_jobs=-1)
¶
Filter with updated operator module and parallel execution.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
target
|
Literal['charge', 'multiplicity', 'format']
|
The target property to filter. |
required |
value
|
Union[str, float, int]
|
The value to compare against. |
required |
compare
|
Literal['==', '!=', '>', '<', '>=', '<=']
|
The comparison operator. |
'=='
|
n_jobs
|
int
|
Number of parallel jobs to filter. |
-1
|
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
The filtered batch. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def filter_value(
self,
target: Literal["charge", "multiplicity", "format"],
value: str | float | int,
compare: Literal["==", "!=", ">", "<", ">=", "<="] = "==",
n_jobs: int = -1,
) -> FileBatchModelDisk[TFileDisk]:
"""
Filter with updated operator module and parallel execution.
Parameters:
target (Literal["charge", "multiplicity", "format"]): The target property to filter.
value (Union[str, float, int]): The value to compare against.
compare (Literal["==", "!=", ">", "<", ">=", "<="]): The comparison operator.
n_jobs (int): Number of parallel jobs to filter.
Returns:
FileBatchModelDisk: The filtered batch.
"""
compare_ops = {
"==": operator.eq,
"!=": operator.ne,
">": operator.gt,
"<": operator.lt,
">=": operator.ge,
"<=": operator.le,
}
op_func = compare_ops[compare]
def judge_func(diskfile: TFileDisk) -> bool:
if target == "charge":
return op_func(diskfile[-1].charge, value)
elif target == "multiplicity":
return op_func(diskfile[-1].multiplicity, value)
elif target == "format":
return op_func(diskfile.file_format, value)
return False
desc = f"Filtering files with {target} {compare} {value} with {molopconfig.set_n_jobs(n_jobs)} jobs"
diskfiles = self._snapshot_diskfiles()
return self._filter_diskfiles(
diskfiles,
self.parallel_execute(
judge_func,
desc,
n_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
),
)
groupby(key_func, n_jobs=-1)
¶
Group the files into multiple batches based on a key function.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
key_func
|
Callable[[FileDiscObj], str]
|
A function that generates a group key for each file. |
required |
n_jobs
|
int
|
Number of parallel jobs to calculate keys. |
-1
|
Returns:
| Type | Description |
|---|---|
dict[str, FileBatchModelDisk]
|
A dictionary where keys are group names and values are new batch objects. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def groupby(
self, key_func: Callable[[TFileDisk], str], n_jobs: int = -1
) -> dict[str, FileBatchModelDisk[TFileDisk]]:
"""
Group the files into multiple batches based on a key function.
Parameters:
key_func (Callable[[FileDiscObj], str]): A function that generates a group key for each file.
n_jobs (int): Number of parallel jobs to calculate keys.
Returns:
(dict[str, FileBatchModelDisk]): A dictionary where keys are group names and values are new batch objects.
"""
effective_jobs = molopconfig.set_molgr_n_jobs(n_jobs)
desc = f"Grouping files with {effective_jobs} jobs"
diskfiles = self._snapshot_diskfiles()
# ``key_func`` is intentionally arbitrary, so treat it as graph
# dependent even when the current callback only uses metadata. This
# keeps future callbacks from accidentally entering MolGR in workers.
if molopconfig.prewarm_topologies:
self._prewarm_topologies(
_diskfiles_snapshot=diskfiles,
max_workers=effective_jobs,
)
keys = self.parallel_execute(
key_func,
desc,
effective_jobs,
return_as="generator",
_diskfiles_snapshot=diskfiles,
return_results=True,
)
temp_groups: dict[str, list[TFileDisk]] = {}
for diskfile, key in zip(diskfiles, keys, strict=True):
if key not in temp_groups:
temp_groups[key] = []
temp_groups[key].append(diskfile)
return {group_key: self.new_batch(files) for group_key, files in temp_groups.items()}
parallel_execute(func, desc='', n_jobs=-1, return_as='list', *args, _diskfiles_snapshot=None, return_results=None, **kwargs)
¶
Python
parallel_execute(
func: Callable[..., R],
desc: str = "",
n_jobs: int = -1,
return_as: Literal[
"list", "generator", "generator_unordered"
] = "list",
*args: Any,
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
return_results: Literal[False],
**kwargs: Any,
) -> None
Internal helper to execute a function over the batch in parallel or serial. Refactored to support Adaptive UI (Rich/Ipywidgets) and Type Hints.
Accepts extra keyword arguments to pass to the function.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
func
|
Callable[..., R]
|
The function to execute. |
required |
desc
|
str
|
The description to show in the progress bar. Defaults to “”. |
''
|
n_jobs
|
int
|
The number of jobs to use. Defaults to -1. |
-1
|
return_as
|
Literal['list', 'generator', 'generator_unordered']
|
The return mode to use. Defaults to “list”. |
'list'
|
_diskfiles_snapshot
|
Sequence[TFileDisk] | None
|
Internal precomputed batch snapshot used by callers that need the same item order for zipping results. |
None
|
return_results
|
bool | None
|
When False, exhaust execution and return None. When None, list results containing only implicit None values are returned as None. |
None
|
**kwargs
|
Any
|
Additional keyword arguments to pass to the function. |
{}
|
Returns: Iterable[R] | None: The function results, or None for side-effect-only calls.
Source code in src/molop/io/FileBatchModelDisk.py
Python
def parallel_execute(
self,
func: Callable[..., R],
desc: str = "",
n_jobs: int = -1,
return_as: Literal["list", "generator", "generator_unordered"] = "list",
*args: Any, # do not use this parameter
_diskfiles_snapshot: Sequence[TFileDisk] | None = None,
return_results: bool | None = None,
**kwargs: Any,
) -> Iterable[R] | None:
"""
Internal helper to execute a function over the batch in parallel or serial.
Refactored to support Adaptive UI (Rich/Ipywidgets) and Type Hints.
Accepts extra keyword arguments to pass to the function.
Parameters:
func (Callable[..., R]): The function to execute.
desc (str, optional): The description to show in the progress bar. Defaults to "".
n_jobs (int, optional): The number of jobs to use. Defaults to -1.
return_as (Literal["list", "generator", "generator_unordered"], optional): The return mode to use. Defaults to "list".
_diskfiles_snapshot (Sequence[TFileDisk] | None, optional): Internal precomputed
batch snapshot used by callers that need the same item order for zipping results.
return_results (bool | None, optional): When False, exhaust execution and return None.
When None, list results containing only implicit None values are returned as None.
**kwargs: Additional keyword arguments to pass to the function.
Returns:
Iterable[R] | None: The function results, or None for side-effect-only calls.
"""
diskfiles = (
_diskfiles_snapshot if _diskfiles_snapshot is not None else self._snapshot_diskfiles()
)
return parallel_map(
lambda diskfile: func(diskfile, *args, **kwargs),
diskfiles,
n_jobs=molopconfig.set_n_jobs(n_jobs),
desc=desc,
total=len(diskfiles),
disable=not molopconfig.show_progress_bar,
return_as=return_as,
return_results=return_results,
maxtasks_per_child=50,
max_nbytes=molopconfig.parallel_max_size,
)
release_file_content()
¶
require(key, expected_cls)
¶
Return a file with precise type hints, validating at runtime.
This avoids hardcoding a global list of concrete file classes while still allowing callers to opt into format-specific APIs.
Source code in src/molop/io/FileBatchModelDisk.py
Python
def require(self, key: int | str, expected_cls: type[TFile]) -> TFile:
"""Return a file with precise type hints, validating at runtime.
This avoids hardcoding a global list of concrete file classes while still
allowing callers to opt into format-specific APIs.
"""
value = self[key]
if not isinstance(value, expected_cls):
raise TypeError(
f"Expected {expected_cls.__name__} for key {key!r}, got {type(value).__name__}"
)
return value
sample(n=10, seed=None)
¶
Randomly sample n files from the batch.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
n
|
int
|
Number of files to sample. |
10
|
seed
|
int
|
Random seed for reproducibility. |
None
|
Returns:
| Name | Type | Description |
|---|---|---|
FileBatchModelDisk |
FileBatchModelDisk[TFileDisk]
|
A new batch containing the sampled files. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def sample(self, n: int = 10, seed: int | None = None) -> FileBatchModelDisk[TFileDisk]:
"""
Randomly sample n files from the batch.
Parameters:
n (int): Number of files to sample.
seed (int): Random seed for reproducibility.
Returns:
FileBatchModelDisk: A new batch containing the sampled files.
"""
if seed is not None:
random.seed(seed)
n = min(n, len(self))
sampled_files = random.sample(list(self.__diskfiles.values()), n)
return self.new_batch(sampled_files)
save_pre_post_ts(output_dir, *, format='xyz', min_ratio=0.6, max_ratio=1.4, steps=8, sampling_method='harmonic_potential', n_jobs=-1)
¶
Export endpoint candidates for every TS frame in every calculation file.
Each supported source file receives a separate directory. Unique filename
stems are retained; duplicate stems gain a stable source-path digest to
prevent overwrites. Files without endpoint-export support are skipped with
a warning. The result maps supported source paths to the frame-ID-to-endpoint
mappings returned by each file. Endpoint inference samples steps
amplitudes on each side from min_ratio through max_ratio
according to sampling_method. Defaults are min_ratio=0.6,
max_ratio=1.4, steps=8, and
sampling_method="harmonic_potential".
Source code in src/molop/io/FileBatchModelDisk.py
Python
def save_pre_post_ts(
self,
output_dir: os.PathLike[str] | str,
*,
format: Literal["xyz", "sdf"] = "xyz",
min_ratio: float = 0.6,
max_ratio: float = 1.4,
steps: int = 8,
sampling_method: VibrationSamplingMethod = "harmonic_potential",
n_jobs: int = -1,
) -> dict[str, dict[int, tuple[Path, Path]]]:
"""Export endpoint candidates for every TS frame in every calculation file.
Each supported source file receives a separate directory. Unique filename
stems are retained; duplicate stems gain a stable source-path digest to
prevent overwrites. Files without endpoint-export support are skipped with
a warning. The result maps supported source paths to the frame-ID-to-endpoint
mappings returned by each file. Endpoint inference samples ``steps``
amplitudes on each side from ``min_ratio`` through ``max_ratio``
according to ``sampling_method``. Defaults are ``min_ratio=0.6``,
``max_ratio=1.4``, ``steps=8``, and
``sampling_method="harmonic_potential"``.
"""
destination = Path(output_dir)
destination.mkdir(parents=True, exist_ok=True)
diskfiles = self._snapshot_diskfiles()
exportable_diskfiles: list[TFileDisk] = []
for diskfile in diskfiles:
if callable(getattr(diskfile, "save_pre_post_ts", None)):
exportable_diskfiles.append(diskfile)
else:
moloplogger.warning(
f"Skipping {diskfile.file_path}: "
f"{type(diskfile).__name__} does not support TS endpoint export"
)
output_directories = _endpoint_output_directories(
cast(Sequence[DiskFileLike], exportable_diskfiles), destination
)
endpoint_tasks: list[tuple[TFileDisk, Path]] = [
(diskfile, output_directories[str(diskfile.file_path)])
for diskfile in exportable_diskfiles
]
def export_file(
task: tuple[TFileDisk, Path],
) -> tuple[str, dict[int, tuple[Path, Path]]]:
diskfile, output_directory = task
save_endpoints = getattr(diskfile, "save_pre_post_ts", None)
assert callable(save_endpoints)
source_path = str(diskfile.file_path)
exports = cast(
dict[int, tuple[Path, Path]],
save_endpoints(
output_directory,
format=format,
min_ratio=min_ratio,
max_ratio=max_ratio,
steps=steps,
sampling_method=sampling_method,
),
)
return source_path, exports
effective_jobs = molopconfig.set_molgr_n_jobs(n_jobs)
desc = f"Exporting TS endpoint candidates with {effective_jobs} jobs"
results = parallel_map(
export_file,
endpoint_tasks,
n_jobs=effective_jobs,
desc=desc,
total=len(endpoint_tasks),
disable=not molopconfig.show_progress_bar,
return_results=True,
maxtasks_per_child=50,
max_nbytes=molopconfig.parallel_max_size,
)
return dict(results)
to_summary_df(mode='frame', frame=-1, n_jobs=-1, *, brief=True, flatten_columns=False, on_missing_frame='skip', **kwargs)
¶
Generate summary DataFrame using the parallel execution engine.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
mode
|
Literal['file', 'frame']
|
Whether to generate file-level or frame-level summary. |
'frame'
|
frame
|
FrameSelector
|
Frame IDs to process. Use -1 for the last frame or “all” for every frame. |
-1
|
n_jobs
|
int
|
Number of parallel jobs. |
-1
|
brief
|
bool
|
Whether to use the compact summary payload. |
True
|
flatten_columns
|
bool
|
Whether to flatten MultiIndex columns with dot-separated names. |
False
|
on_missing_frame
|
Literal['skip', 'error']
|
How to handle frame indices that are out of range for a file. |
'skip'
|
**kwargs
|
Any
|
Additional arguments for to_summary_series. |
{}
|
Returns:
| Type | Description |
|---|---|
DataFrame
|
pd.DataFrame: A DataFrame containing the summary information. |
Source code in src/molop/io/FileBatchModelDisk.py
Python
def to_summary_df(
self,
mode: Literal["file", "frame"] = "frame",
frame: FrameSelector = -1,
n_jobs: int = -1,
*,
brief: bool = True,
flatten_columns: bool = False,
on_missing_frame: Literal["skip", "error"] = "skip",
**kwargs: Any,
) -> pd.DataFrame:
"""
Generate summary DataFrame using the parallel execution engine.
Parameters:
mode (Literal["file", "frame"]): Whether to generate file-level or frame-level summary.
frame (FrameSelector): Frame IDs to process.
Use -1 for the last frame or "all" for every frame.
n_jobs (int): Number of parallel jobs.
brief (bool): Whether to use the compact summary payload.
flatten_columns (bool): Whether to flatten MultiIndex columns with dot-separated names.
on_missing_frame (Literal["skip", "error"]): How to handle frame indices that are
out of range for a file.
**kwargs: Additional arguments for to_summary_series.
Returns:
pd.DataFrame: A DataFrame containing the summary information.
"""
if mode not in ["file", "frame"]:
raise ValueError(f"Invalid mode: {mode}")
if on_missing_frame not in ["skip", "error"]:
raise ValueError(f"Invalid on_missing_frame: {on_missing_frame}")
def selected_frame_ids(diskfile: FileDiskObj) -> Sequence[int]:
try:
frame_ids = normalize_frame_selector(
frame,
len(diskfile),
parameter_name="frame",
validate_range=on_missing_frame == "error",
)
except IndexError as exc:
raise IndexError(
f"{exc} for {getattr(diskfile, 'file_path', '<unknown>')} "
f"with {len(diskfile)} frames"
) from exc
if on_missing_frame == "skip":
return [fid for fid in frame_ids if 0 <= fid < len(diskfile)]
return frame_ids
diskfiles = self._snapshot_diskfiles()
summary_tasks: list[tuple[FileDiskObj, Sequence[int]]] = []
if mode == "file":
summary_tasks = [(diskfile, ()) for diskfile in diskfiles]
else:
selected_frames: list[Any] = []
for diskfile in diskfiles:
frame_ids = list(selected_frame_ids(diskfile))
selected_frames.extend(diskfile[fid] for fid in frame_ids)
summary_tasks.append((diskfile, frame_ids))
if molopconfig.prewarm_topologies:
from molop.io.base_models.Molecule import reconstruct_topologies_batch
# The selected source frames have a stable, enumerable input
# set, so optionally warm their ordinary graph cache in one
# parent-process native batch before process workers start.
reconstruct_topologies_batch(
selected_frames,
max_workers=molopconfig.set_molgr_n_jobs(n_jobs),
retain_results=False,
)
def process_file_summary(task: tuple[FileDiskObj, Sequence[int]]) -> list[pd.Series]:
diskfile, frame_ids = task
if mode == "file":
return [diskfile.to_summary_series(brief=brief, **kwargs)]
summaries: list[pd.Series] = []
for fid in frame_ids:
summaries.append(diskfile[fid].to_summary_series(brief=brief, **kwargs))
return summaries
effective_jobs = (
molopconfig.set_molgr_n_jobs(n_jobs)
if mode == "frame"
else molopconfig.set_n_jobs(n_jobs)
)
desc = f"MolOP processing {mode} summary with {effective_jobs} jobs"
nested_results = parallel_map(
process_file_summary,
summary_tasks,
n_jobs=effective_jobs,
desc=desc,
total=len(summary_tasks),
disable=not molopconfig.show_progress_bar,
return_results=True,
maxtasks_per_child=50,
max_nbytes=molopconfig.parallel_max_size,
)
return build_summary_df(
(series for sublist in nested_results for series in sublist),
flatten_columns=flatten_columns,
)