Calculator instantiation inside Prefect task - #91
Merged
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughAdds a ChangesCalculator spec + dispersion propagation
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
mlip_arena/tasks/phonon.py (1)
97-114:⚠️ Potential issue | 🟠 Major | ⚡ Quick winCalculator is rebuilt for every displaced supercell.
Line 111 re-instantiates the calculator inside
_get_forces. In phonon runs with many displacements, this can dominate runtime. Build once inrunand reuse it.💡 Suggested fix
def _get_forces( phononpy_atoms: PhonopyAtoms, - calculator: str | MLIPEnum | BaseCalculator, - calculator_kwargs: dict | None = None, - dispersion: bool = False, - dispersion_kwargs: dict | None = None, + calculator: BaseCalculator, ) -> np.ndarray: @@ - atoms.calc = get_calculator(calculator, calculator_kwargs, dispersion, dispersion_kwargs) + atoms.calc = calculator @@ def run( @@ ): + calculator_obj = get_calculator(calculator, calculator_kwargs, dispersion, dispersion_kwargs) @@ phonon.forces = [ _get_forces( supercell, - calculator, - calculator_kwargs, - dispersion, - dispersion_kwargs, + calculator_obj, )Also applies to: 182-190
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@mlip_arena/tasks/phonon.py` around lines 97 - 114, The _get_forces function re-instantiates the calculator on every call via the get_calculator function, which is inefficient when processing many displaced supercells. Refactor by removing the get_calculator call from _get_forces, having it accept a pre-built calculator as a parameter instead, and then build the calculator once in the run function before calling _get_forces repeatedly. This same pattern appears in another function at lines 182-190 and should be refactored identically.mlip_arena/tasks/eos.py (1)
161-163:⚠️ Potential issue | 🔴 Critical
EOSwill crash when reading energies from atoms afterOPThas removed their calculator.At line 162,
r["atoms"].get_potential_energy()will raise aRuntimeErrorbecausemlip_arena/tasks/optimize.pyclears the calculator viaatoms.calc = None(line 137) before returning. ASE'sget_potential_energy()method requires an attached calculator and does not fall back to cached values.Compare with
mlip_arena/tasks/ev.py, which correctly reattaches the calculator before callingget_potential_energy().💡 Suggested fix
-from mlip_arena.tasks.utils import ARENA_TASK_CACHE_POLICY, resolve_calculator_name +from mlip_arena.tasks.utils import ARENA_TASK_CACHE_POLICY, get_calculator, resolve_calculator_name @@ - volumes = [r["atoms"].get_volume() for r in results] - energies = [r["atoms"].get_potential_energy() for r in results] + volumes = [r["atoms"].get_volume() for r in results] + energies = [] + for r in results: + atoms_i = r["atoms"] + atoms_i.calc = get_calculator(calculator, calculator_kwargs, dispersion, dispersion_kwargs) + energies.append(atoms_i.get_potential_energy()) + atoms_i.calc = None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@mlip_arena/tasks/eos.py` around lines 161 - 163, The EOS task will crash when attempting to read energies from atoms after OPT has removed their calculator. At the line where energies are extracted using get_potential_energy() on results, the calculator has been cleared by optimize.py and ASE's get_potential_energy() requires an attached calculator. Reattach the calculator to each atoms object before calling get_potential_energy(), following the same pattern used in ev.py which correctly handles this scenario. The fix should be applied to the energies list comprehension where it iterates over results and calls r["atoms"].get_potential_energy().
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@mlip_arena/flows/diatomics.py`:
- Around line 38-39: Optional output-directory parameters were introduced
without proper `None` handling, causing runtime crashes when tasks are invoked
without explicit directory arguments. Fix this across three locations: In
mlip_arena/flows/diatomics.py (lines 38-39), either make the `out_dir` parameter
required again or initialize a default `Path` object before any first use of
`out_dir`. In mlip_arena/flows/stability.py (lines 38-39) within the
`nvt_heat_one` function, initialize `run_dir` with a default `Path` value before
it is used to construct `traj_file`. In mlip_arena/flows/stability.py (lines
71-72) within the `npt_compress_one` function, initialize `run_dir` with a
default `Path` value before it is used to construct `traj_file`. Ensure all
three locations properly handle the `None` case by providing sensible defaults
before any path construction operations.
In `@mlip_arena/tasks/eos_alloy/flow.py`:
- Around line 57-67: Guard access to the REGISTRY dictionary to prevent silent
data loss when the calculator cannot be resolved. After calling
resolve_calculator_name() to get the calculator_name, add a check to verify the
calculator_name exists as a valid key in REGISTRY before attempting to access it
(at lines 89-90 where REGISTRY is accessed). If the calculator_name is not in
REGISTRY or is "Unknown", handle this case explicitly by logging a meaningful
warning or error message rather than allowing it to crash silently within the
broad except clause at line 102. Reference the guarded pattern used in
diatomics.py and stability.py as examples of the safe approach already
established in the codebase.
---
Outside diff comments:
In `@mlip_arena/tasks/eos.py`:
- Around line 161-163: The EOS task will crash when attempting to read energies
from atoms after OPT has removed their calculator. At the line where energies
are extracted using get_potential_energy() on results, the calculator has been
cleared by optimize.py and ASE's get_potential_energy() requires an attached
calculator. Reattach the calculator to each atoms object before calling
get_potential_energy(), following the same pattern used in ev.py which correctly
handles this scenario. The fix should be applied to the energies list
comprehension where it iterates over results and calls
r["atoms"].get_potential_energy().
In `@mlip_arena/tasks/phonon.py`:
- Around line 97-114: The _get_forces function re-instantiates the calculator on
every call via the get_calculator function, which is inefficient when processing
many displaced supercells. Refactor by removing the get_calculator call from
_get_forces, having it accept a pre-built calculator as a parameter instead, and
then build the calculator once in the run function before calling _get_forces
repeatedly. This same pattern appears in another function at lines 182-190 and
should be refactored identically.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 8bb3a6b7-9987-48db-b3a0-1d91b57d1b9e
📒 Files selected for processing (13)
mlip_arena/flows/diatomics.pymlip_arena/flows/eos_bulk.pymlip_arena/flows/stability.pymlip_arena/tasks/elasticity.pymlip_arena/tasks/eos.pymlip_arena/tasks/eos_alloy/flow.pymlip_arena/tasks/md.pymlip_arena/tasks/mof/flow.pymlip_arena/tasks/neb.pymlip_arena/tasks/optimize.pymlip_arena/tasks/phonon.pymlip_arena/tasks/utils.pytests/test_prefect_serialization.py
Comment on lines
+38
to
39
| out_dir: Path | None = None, | ||
| ): |
Contributor
There was a problem hiding this comment.
Optional output-directory parameters were introduced without matching None handling. This creates direct runtime crashes when these tasks are invoked without explicit directory arguments.
mlip_arena/flows/diatomics.py#L38-L39: either makeout_dirrequired again or initialize a defaultPathbefore first use.mlip_arena/flows/stability.py#L38-L39: initializerun_dirbefore buildingtraj_fileinnvt_heat_one.mlip_arena/flows/stability.py#L71-L72: initializerun_dirbefore buildingtraj_fileinnpt_compress_one.
📍 Affects 2 files
mlip_arena/flows/diatomics.py#L38-L39(this comment)mlip_arena/flows/stability.py#L38-L39mlip_arena/flows/stability.py#L71-L72
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/flows/diatomics.py` around lines 38 - 39, Optional
output-directory parameters were introduced without proper `None` handling,
causing runtime crashes when tasks are invoked without explicit directory
arguments. Fix this across three locations: In mlip_arena/flows/diatomics.py
(lines 38-39), either make the `out_dir` parameter required again or initialize
a default `Path` object before any first use of `out_dir`. In
mlip_arena/flows/stability.py (lines 38-39) within the `nvt_heat_one` function,
initialize `run_dir` with a default `Path` value before it is used to construct
`traj_file`. In mlip_arena/flows/stability.py (lines 71-72) within the
`npt_compress_one` function, initialize `run_dir` with a default `Path` value
before it is used to construct `traj_file`. Ensure all three locations properly
handle the `None` case by providing sensible defaults before any path
construction operations.
Comment on lines
+57
to
67
| # Retrieve calculator argument from parameters or task_inputs | ||
| calculator = None | ||
| if hasattr(run, "parameters") and run.parameters: | ||
| calculator = run.parameters.get("calculator") | ||
| if calculator is None and hasattr(run, "task_inputs"): | ||
| calculator = run.task_inputs.get("calculator") | ||
| if isinstance(calculator, list) and len(calculator) > 0: | ||
| calculator = calculator[0] | ||
|
|
||
| calculator_name = resolve_calculator_name(calculator) | ||
|
|
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
python - <<'PY'
from prefect.client.schemas.objects import TaskRun
field = TaskRun.model_fields.get("task_inputs")
print("TaskRun.task_inputs annotation:", field.annotation)
PYRepository: atomind-ai/mlip-arena
Length of output: 185
🏁 Script executed:
cat -n mlip_arena/tasks/eos_alloy/flow.py | sed -n '50,100p'Repository: atomind-ai/mlip-arena
Length of output: 2143
🏁 Script executed:
rg "resolve_calculator_name" -A 5 -B 5Repository: atomind-ai/mlip-arena
Length of output: 19442
🏁 Script executed:
git ls-files | grep -E "(prefect|task)" | head -20Repository: atomind-ai/mlip-arena
Length of output: 1108
🏁 Script executed:
cat -n mlip_arena/tasks/eos_alloy/flow.py | sed -n '40,110p'Repository: atomind-ai/mlip-arena
Length of output: 2786
🏁 Script executed:
cat -n mlip_arena/tasks/utils.py | grep -A 15 "def resolve_calculator_name"Repository: atomind-ai/mlip-arena
Length of output: 743
🏁 Script executed:
rg "def save_to_hdf" -A 80 mlip_arena/tasks/eos_alloy/flow.py | head -100Repository: atomind-ai/mlip-arena
Length of output: 2621
🏁 Script executed:
rg "task_inputs" mlip_arena/tasks/eos_alloy/ -B 2 -A 2Repository: atomind-ai/mlip-arena
Length of output: 894
🏁 Script executed:
rg "REGISTRY\[" mlip_arena/ -A 1 | head -50Repository: atomind-ai/mlip-arena
Length of output: 3136
🏁 Script executed:
grep -n "except" mlip_arena/tasks/eos_alloy/flow.pyRepository: atomind-ai/mlip-arena
Length of output: 96
🏁 Script executed:
rg "REGISTRY\s*=" mlip_arena/models/__init__.py -A 30 | head -50Repository: atomind-ai/mlip-arena
Length of output: 1005
🏁 Script executed:
grep -r '"Unknown"' mlip_arena/models/ | grep -i registryRepository: atomind-ai/mlip-arena
Length of output: 47
🏁 Script executed:
rg "registry\.yaml" mlip_arena/ | head -5Repository: atomind-ai/mlip-arena
Length of output: 604
🏁 Script executed:
fd "registry.yaml" mlip_arena/Repository: atomind-ai/mlip-arena
Length of output: 128
🏁 Script executed:
python3 - <<'PY'
import yaml
from pathlib import Path
# Try to find and load registry.yaml
registry_files = list(Path("mlip_arena").rglob("registry.yaml"))
print(f"Found {len(registry_files)} registry.yaml files:")
for f in registry_files:
print(f" {f}")
if registry_files:
with open(registry_files[0]) as f:
registry = yaml.safe_load(f)
print(f"\nRegistry keys: {list(registry.keys())[:10]}...")
print(f"Total registry entries: {len(registry)}")
print(f"'Unknown' in REGISTRY: {'Unknown' in registry}")
PYRepository: atomind-ai/mlip-arena
Length of output: 361
🏁 Script executed:
python3 - <<'PY'
import yaml
from pathlib import Path
with open("mlip_arena/models/registry.yaml") as f:
registry = yaml.safe_load(f)
# Simulate what happens when calculator_name is "Unknown"
calculator_name = "Unknown"
print(f"Is 'Unknown' in REGISTRY? {calculator_name in registry}")
try:
family = registry[calculator_name]["family"]
print(f"family_path would be: {family}")
except KeyError as e:
print(f"KeyError when accessing REGISTRY['{calculator_name}']: {e}")
print("This error would be caught by broad except and silently swallowed")
PYRepository: atomind-ai/mlip-arena
Length of output: 218
Guard REGISTRY access to prevent silent data loss when calculator cannot be resolved.
When run.parameters lacks a calculator and fallback to run.task_inputs yields None or an invalid value, resolve_calculator_name() returns "Unknown". Line 89 then crashes with KeyError since "Unknown" is not a valid REGISTRY key. The broad except at line 102 silently swallows this error, skipping HDF5 persistence without warning.
The codebase already handles this safely elsewhere (e.g., diatomics.py, stability.py) using a guarded pattern. Apply the same approach here:
Suggested fix
- family_path = Path(__file__).parent / REGISTRY[calculator_name]["family"]
+ family = REGISTRY.get(calculator_name, {}).get("family", "custom")
+ family_path = Path(__file__).parent / familyAlso applies to: 89-90
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/tasks/eos_alloy/flow.py` around lines 57 - 67, Guard access to the
REGISTRY dictionary to prevent silent data loss when the calculator cannot be
resolved. After calling resolve_calculator_name() to get the calculator_name,
add a check to verify the calculator_name exists as a valid key in REGISTRY
before attempting to access it (at lines 89-90 where REGISTRY is accessed). If
the calculator_name is not in REGISTRY or is "Unknown", handle this case
explicitly by logging a meaningful warning or error message rather than allowing
it to crash silently within the broad except clause at line 102. Reference the
guarded pattern used in diatomics.py and stability.py as examples of the safe
approach already established in the codebase.
…th unified signatures
…nified parameters
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
mlip_arena/flows/ev.py (1)
62-62: 💤 Low valueType hint missing
MLIPEnumfor consistency with the underlying task.The flow's
calculatorparameter acceptsstr | BaseCalculator, but the underlyingev_scantask acceptsstr | MLIPEnum | BaseCalculator | None. AddingMLIPEnumto this type hint would improve consistency and allow type checkers to validate callers passing enum values.- calculator: str | BaseCalculator, + calculator: str | MLIPEnum | BaseCalculator,Note: You'd also need to import
MLIPEnumfrommlip_arena.models.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@mlip_arena/flows/ev.py` at line 62, Update the type hint for the `calculator` parameter in the flow to include `MLIPEnum` alongside `str` and `BaseCalculator` to match the underlying `ev_scan` task's accepted types. Additionally, add an import statement for `MLIPEnum` from `mlip_arena.models` at the top of the file to enable the type checker to recognize and validate this type.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@mlip_arena/flows/conservation.py`:
- Around line 197-203: The parameters structures, out_dir, input_path, and
reference_path are defined with default values of None but are dereferenced
unconditionally throughout the code (e.g., out_dir.mkdir(), read(input_path,
...), parameters["reference_path"].stem), causing AttributeError or TypeError at
runtime. Either remove the None defaults to make these parameters required, or
add explicit None checks before each use of these variables. Apply this fix
across all affected locations in mlip_arena/flows/conservation.py (the function
definition at lines 197-203, and the dereferencing sites at lines 250-257,
266-274, and 308-308) to ensure the function contract matches its actual
requirements.
In `@mlip_arena/tasks/ev.py`:
- Line 35: The calculator parameter in the function signature allows None as a
default value and includes None in its type hint, but the get_calculator
function call does not accept None as a valid input and will raise a ValueError
at runtime. Either remove None from the type hint (str | MLIPEnum |
BaseCalculator) and remove the default value assignment, or add validation at
the start of the function to check if calculator is None and raise a clear error
message before get_calculator is invoked, making the constraint explicit to
callers.
---
Nitpick comments:
In `@mlip_arena/flows/ev.py`:
- Line 62: Update the type hint for the `calculator` parameter in the flow to
include `MLIPEnum` alongside `str` and `BaseCalculator` to match the underlying
`ev_scan` task's accepted types. Additionally, add an import statement for
`MLIPEnum` from `mlip_arena.models` at the top of the file to enable the type
checker to recognize and validate this type.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b785c1e6-5ff3-49d8-945b-1cdc6e61869a
📒 Files selected for processing (14)
benchmarks/c2db/run.pybenchmarks/eos_bulk/run.pybenchmarks/mof/classification/classification.pybenchmarks/submit.pybenchmarks/vacancy_migration/run.pymlip_arena/flows/conservation.pymlip_arena/flows/ev.pymlip_arena/tasks/ev.pytests/test_elasticity.pytests/test_eos.pytests/test_md.pytests/test_mof.pytests/test_neb.pytests/test_stability.py
Comment on lines
+197
to
203
| calculator: str | MLIPEnum | BaseCalculator | None = None, | ||
| calculator_kwargs: dict | None = None, | ||
| dispersion: bool = False, | ||
| dispersion_kwargs: dict | None = None, | ||
| structures: list[Atoms] = None, | ||
| out_dir: Path = None, | ||
| ): |
Contributor
There was a problem hiding this comment.
Optional parameters are treated as required and can crash at runtime.
structures, out_dir, input_path, and reference_path are defaulted to None, but the code dereferences them unconditionally (out_dir.mkdir(...), read(input_path, ...), parameters["reference_path"].stem). This creates AttributeError/TypeError paths and makes the flow contract misleading.
Suggested fix
def _generate_flow_run_name():
name = flow_run.flow_name
parameters = flow_run.parameters
calculator = parameters.get("calculator")
model_name = resolve_calculator_name(calculator)
- reference_path = parameters["reference_path"].stem
+ reference_path = parameters.get("reference_path")
+ reference_stem = Path(reference_path).stem if reference_path is not None else "unknown-reference"
- return f"{name}: {model_name} - {reference_path}"
+ return f"{name}: {model_name} - {reference_stem}" def run_simulations(
calculator: str | MLIPEnum | BaseCalculator | None = None,
calculator_kwargs: dict | None = None,
dispersion: bool = False,
dispersion_kwargs: dict | None = None,
structures: list[Atoms] = None,
out_dir: Path = None,
):
+ if structures is None or out_dir is None:
+ raise ValueError("`structures` and `out_dir` are required.") def differential_entropy_along_nve_trajectory(
@@
):
+ if input_path is None or reference_path is None:
+ raise ValueError("`input_path` and `reference_path` are required.")
+ if start_idx is None or end_idx is None or step is None:
+ raise ValueError("`start_idx`, `end_idx`, and `step` are required.")Also applies to: 250-257, 266-274, 308-308
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/flows/conservation.py` around lines 197 - 203, The parameters
structures, out_dir, input_path, and reference_path are defined with default
values of None but are dereferenced unconditionally throughout the code (e.g.,
out_dir.mkdir(), read(input_path, ...), parameters["reference_path"].stem),
causing AttributeError or TypeError at runtime. Either remove the None defaults
to make these parameters required, or add explicit None checks before each use
of these variables. Apply this fix across all affected locations in
mlip_arena/flows/conservation.py (the function definition at lines 197-203, and
the dereferencing sites at lines 250-257, 266-274, and 308-308) to ensure the
function contract matches its actual requirements.
| def run( | ||
| atoms: Atoms, | ||
| model: str | BaseCalculator, | ||
| calculator: str | MLIPEnum | BaseCalculator | None = None, |
Contributor
There was a problem hiding this comment.
calculator=None will fail at runtime in get_calculator.
The function signature allows calculator=None as default, but get_calculator at line 66 will raise ValueError("Invalid calculator: None") since None doesn't match any of its accepted types. Either remove None from the type hint and default, or add validation with a clear error message before calling get_calculator.
Suggested fix: Remove None default or add early validation
Option 1 - Remove None from signature:
def run(
atoms: Atoms,
- calculator: str | MLIPEnum | BaseCalculator | None = None,
+ calculator: str | MLIPEnum | BaseCalculator,
calculator_kwargs: dict | None = None,Option 2 - Add early validation:
def run(
atoms: Atoms,
calculator: str | MLIPEnum | BaseCalculator | None = None,
...
):
+ if calculator is None:
+ raise ValueError("calculator must be specified for E-V scan")
+
model_name = resolve_calculator_name(calculator)Also applies to: 66-66
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@mlip_arena/tasks/ev.py` at line 35, The calculator parameter in the function
signature allows None as a default value and includes None in its type hint, but
the get_calculator function call does not accept None as a valid input and will
raise a ValueError at runtime. Either remove None from the type hint (str |
MLIPEnum | BaseCalculator) and remove the default value assignment, or add
validation at the start of the function to check if calculator is None and raise
a clear error message before get_calculator is invoked, making the constraint
explicit to callers.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
mlip_arena/tasks/optimize.py (2)
136-140:⚠️ Potential issue | 🟠 Major | ⚡ Quick winClear calculator from returned atoms before returning the task result.
atomsis returned withatoms.calcstill attached. That increases serialization risk/payload size and breaks the intended “instantiate internally, then cleanup” behavior.Proposed fix
- return { + atoms.calc = None + return { "atoms": atoms, "steps": optimizer_instance.nsteps, "converged": converged, }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@mlip_arena/tasks/optimize.py` around lines 136 - 140, The atoms object being returned in the dictionary at the return statement still has its calculator attached via atoms.calc, which increases serialization overhead and breaks the cleanup pattern. Before the return statement, clear the calculator by setting atoms.calc to None, then return the dictionary with the cleaned atoms object.
65-99:⚠️ Potential issue | 🟠 Major | ⚡ Quick win
calculator=Noneis declared valid but fails at runtime.The task API allows
calculator=None, but the current execution path always callsget_calculator(...), which rejectsNone. This creates an avoidable runtime failure and an inconsistent contract.Proposed fix
-def run( +def run( atoms: Atoms, - calculator: str | MLIPEnum | BaseCalculator | None = None, + calculator: str | MLIPEnum | BaseCalculator, @@ - calculator_obj = get_calculator(calculator, calculator_kwargs, dispersion, dispersion_kwargs) + calculator_obj = get_calculator(calculator, calculator_kwargs, dispersion, dispersion_kwargs)If
Nonemust remain supported, add explicit fallback logic before callingget_calculator.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@mlip_arena/tasks/optimize.py` around lines 65 - 99, The optimize function signature declares calculator as an optional parameter that can be None, but the implementation unconditionally calls get_calculator with this parameter, which causes a runtime failure when calculator is None. Add explicit fallback logic before the get_calculator call to handle the case when calculator is None, either by setting a sensible default calculator or by skipping calculator assignment if None is intended to mean no calculator should be used.
🧹 Nitpick comments (1)
tests/test_coverage.py (1)
52-55: ⚡ Quick winPatch
run_simulationstoo to keep this test unit-scoped and fast.This test still runs real MD via
run_simulations; mocking it here will preserve coverage intent while avoiding expensive/flaky runtime in CI.Suggested change
with ( patch("mlip_arena.flows.conservation.read", return_value=[atoms]), + patch("mlip_arena.flows.conservation.run_simulations", return_value=[]), patch("mlip_arena.flows.conservation.get_trajectory_entropy", return_value=(0.0, [atoms])), ):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_coverage.py` around lines 52 - 55, The test for differential_entropy_along_nve_trajectory still executes real molecular dynamics through the run_simulations function, which is slow and potentially flaky in CI. Add another patch call for run_simulations (using the appropriate module path matching the pattern of the existing patches) within the same context manager to mock it with a reasonable return value. This will keep the test unit-scoped and fast while maintaining coverage of the differential_entropy_along_nve_trajectory function itself.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_coverage.py`:
- Around line 155-164: The test uses hasattr and del on MagicMock objects
(mock_run_params and mock_run_inputs) to simulate different branch conditions,
but MagicMock's dynamic attribute creation can make both attributes appear
present, invalidating the test coverage intent. Instead of relying on
hasattr/del, create separate mock objects with spec parameter or use spec_set to
explicitly restrict which attributes are present, ensuring mock_run_params only
has parameters and mock_run_inputs only has task_inputs without
cross-contamination.
---
Outside diff comments:
In `@mlip_arena/tasks/optimize.py`:
- Around line 136-140: The atoms object being returned in the dictionary at the
return statement still has its calculator attached via atoms.calc, which
increases serialization overhead and breaks the cleanup pattern. Before the
return statement, clear the calculator by setting atoms.calc to None, then
return the dictionary with the cleaned atoms object.
- Around line 65-99: The optimize function signature declares calculator as an
optional parameter that can be None, but the implementation unconditionally
calls get_calculator with this parameter, which causes a runtime failure when
calculator is None. Add explicit fallback logic before the get_calculator call
to handle the case when calculator is None, either by setting a sensible default
calculator or by skipping calculator assignment if None is intended to mean no
calculator should be used.
---
Nitpick comments:
In `@tests/test_coverage.py`:
- Around line 52-55: The test for differential_entropy_along_nve_trajectory
still executes real molecular dynamics through the run_simulations function,
which is slow and potentially flaky in CI. Add another patch call for
run_simulations (using the appropriate module path matching the pattern of the
existing patches) within the same context manager to mock it with a reasonable
return value. This will keep the test unit-scoped and fast while maintaining
coverage of the differential_entropy_along_nve_trajectory function itself.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 3c50d68f-225b-486c-9704-703130eb7d01
📒 Files selected for processing (6)
mlip_arena/tasks/eos.pymlip_arena/tasks/mof/flow.pymlip_arena/tasks/neb.pymlip_arena/tasks/optimize.pytests/test_coverage.pytests/test_md.py
💤 Files with no reviewable changes (3)
- mlip_arena/tasks/mof/flow.py
- mlip_arena/tasks/neb.py
- mlip_arena/tasks/eos.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/test_diatomics.py`:
- Around line 24-28: The assertion at line 28 only verifies that res has a
length of 1, but does not validate the actual content of the returned item.
Since the homonuclear_diatomics function uses raise_on_failure=False upstream, a
failed run could still return a result list with one item (even if that item is
malformed or empty). Add additional assertions after the length check to verify
the structure and content of the single returned item, such as checking for
expected dictionary keys or required fields that would be present in a
successful result. This ensures the test catches actual behavior regressions
beyond just the number of items returned.
In `@tests/test_eos_bulk.py`:
- Around line 20-21: The assertions at lines 20-21 and lines 30-36 only verify
that the result is a DataFrame type, which allows invalid or degraded results to
pass silently. Replace or enhance these assertions to also verify that the
returned DataFrame is non-empty (has rows) and contains all required columns/IDs
that indicate successful computation rather than fallback/error behavior. For
each test case, add assertions that check both len(res) > 0 and that expected
column names are present in res.columns.
In `@tests/test_ev.py`:
- Around line 21-22: The type-only assertions checking isinstance(res,
pd.DataFrame) at line 22 and the similar assertion pattern at lines 31-37 are
insufficient for regression testing because they pass even when EV calculations
fail and return empty or fallback DataFrames. Strengthen these assertions by
adding checks to verify the DataFrame is non-empty (using len() or shape check)
and contains the expected columns for EV calculations, ensuring the tests
properly catch calculation failures rather than just type checks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: fb5eeb12-7885-4d84-b800-f6b9ecf0a02f
📒 Files selected for processing (4)
.github/workflows/ci.yamltests/test_diatomics.pytests/test_eos_bulk.pytests/test_ev.py
Comment on lines
+24
to
+28
| res = homonuclear_diatomics( | ||
| calculator=calc, | ||
| run_dir=Path(tmpdir), | ||
| ) | ||
| assert len(res) == 1 |
Contributor
There was a problem hiding this comment.
Strengthen success assertions beyond list length.
At Line 28, assert len(res) == 1 can pass even when the single submitted run fails, because upstream uses raise_on_failure=False. Please assert the returned item shape/content (e.g., expected dict keys/fields) so this test actually catches behavior regressions.
Suggested assertion hardening
with patch("mlip_arena.flows.diatomics.chemical_symbols", ["", "Cu"]):
res = homonuclear_diatomics(
calculator=calc,
run_dir=Path(tmpdir),
)
assert len(res) == 1
+ assert res[0] is not None
+ assert isinstance(res[0], dict)
+ assert "symbol" in res[0] or "formula" in res[0]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| res = homonuclear_diatomics( | |
| calculator=calc, | |
| run_dir=Path(tmpdir), | |
| ) | |
| assert len(res) == 1 | |
| res = homonuclear_diatomics( | |
| calculator=calc, | |
| run_dir=Path(tmpdir), | |
| ) | |
| assert len(res) == 1 | |
| assert res[0] is not None | |
| assert isinstance(res[0], dict) | |
| assert "symbol" in res[0] or "formula" in res[0] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_diatomics.py` around lines 24 - 28, The assertion at line 28 only
verifies that res has a length of 1, but does not validate the actual content of
the returned item. Since the homonuclear_diatomics function uses
raise_on_failure=False upstream, a failed run could still return a result list
with one item (even if that item is malformed or empty). Add additional
assertions after the length check to verify the structure and content of the
single returned item, such as checking for expected dictionary keys or required
fields that would be present in a successful result. This ensures the test
catches actual behavior regressions beyond just the number of items returned.
Comment on lines
+20
to
+21
| res = run(atoms=atoms, calculator=calc) | ||
| assert isinstance(res, pd.DataFrame) |
Contributor
There was a problem hiding this comment.
Current assertions can pass on fallback/error paths.
At Line 21 and Line 36, checking only DataFrame type won’t fail when computation silently degrades (e.g., empty/fallback frames). Please assert minimum semantic expectations (non-empty + required columns/IDs).
Suggested assertion hardening
# Test run task
res = run(atoms=atoms, calculator=calc)
assert isinstance(res, pd.DataFrame)
+ assert not res.empty
+ assert {"method", "id", "eos"}.issubset(res.columns)
@@
df = run_db(
calculator=calc,
run_dir=Path(tmpdir),
dataset="dummy",
dataset_file="test.db",
)
assert isinstance(df, pd.DataFrame)
+ assert not df.empty
+ assert {"model", "structure", "missing"}.issubset(df.columns)Also applies to: 30-36
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_eos_bulk.py` around lines 20 - 21, The assertions at lines 20-21
and lines 30-36 only verify that the result is a DataFrame type, which allows
invalid or degraded results to pass silently. Replace or enhance these
assertions to also verify that the returned DataFrame is non-empty (has rows)
and contains all required columns/IDs that indicate successful computation
rather than fallback/error behavior. For each test case, add assertions that
check both len(res) > 0 and that expected column names are present in
res.columns.
Comment on lines
+21
to
+22
| res = ev_run(atoms=atoms, calculator=calc, npoints=3) | ||
| assert isinstance(res, pd.DataFrame) |
Contributor
There was a problem hiding this comment.
Type-only checks are too weak for EV regression coverage.
At Line 22 and Line 37, these assertions pass even when EV calculations fail and return fallback/empty frames. Please assert key output semantics (non-empty and expected columns).
Suggested assertion hardening
# Test tasks.ev.run
res = ev_run(atoms=atoms, calculator=calc, npoints=3)
assert isinstance(res, pd.DataFrame)
+ assert not res.empty
+ assert {"method", "id", "eos"}.issubset(res.columns)
@@
df = ev_run_db(
calculator=calc,
run_dir=Path(tmpdir),
dataset="dummy",
dataset_file="test.db",
)
assert isinstance(df, pd.DataFrame)
+ assert not df.empty
+ assert {"model", "structure", "missing"}.issubset(df.columns)Also applies to: 31-37
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_ev.py` around lines 21 - 22, The type-only assertions checking
isinstance(res, pd.DataFrame) at line 22 and the similar assertion pattern at
lines 31-37 are insufficient for regression testing because they pass even when
EV calculations fail and return empty or fallback DataFrames. Strengthen these
assertions by adding checks to verify the DataFrame is non-empty (using len() or
shape check) and contains the expected columns for EV calculations, ensuring the
tests properly catch calculation failures rather than just type checks.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
Refactor
Tests