Constants + Utils updates

[vratins]

Small PR that adds graph creation constants to constants.py and makes changes to utils.py to use those constants (as well slightly improving the tqdm logging utility function.

Also makes formatting changes to util plotting functions.

[Copilot]

Pull request overview

This PR centralizes graph/feature hyperparameters into src/constants.py and updates utilities to consume those constants, while also improving tqdm-compatible logging and reformatting/clarifying plotting utilities.

Changes:

• Add feature-dimension and cutoff constants (NUM_RBF, RBF_CUTOFF, NODE_FEATURE_DIM) to src/constants.py.
• Update utils.rbf() defaults and expand docstrings; extend setup_logging_for_tqdm() with configurable level + optional file logging.
• Reformat/clarify visualization utilities and adjust minor test docstrings/whitespace.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
src/constants.py	Introduces centralized constants for RBF encoding and node features, plus normalizes edge-type literals.
src/utils.py	Uses new constants for RBF defaults, expands logging helper functionality, and reformats plotting helpers/docstrings.
tests/test_utils.py	Small docstring/whitespace cleanup in OT coupling tests.

Comments suppressed due to low confidence (2)

src/utils.py:279

• mate_pos is documented and used as optional (you pass None from create_trajectory_gif and guard with if mate_pos is not None), but the type annotation is np.ndarray. Update the annotation to accept None to match actual usage and avoid type-checker errors.

    ax,
    protein_pos: np.ndarray,
    mate_pos: np.ndarray,
    water_pred: np.ndarray,

src/utils.py:284

• xlim/ylim/zlim default to None but are annotated as tuple[float, float]. Adjust the annotations to be optional (e.g., tuple[float, float] | None) so the signature matches the defaults and intended calling pattern.

    title: str = "",
    xlim: tuple[float, float] = None,
    ylim: tuple[float, float] = None,
    zlim: tuple[float, float] = None,

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

setup_logging_for_tqdm now has new behavior (configurable log level + optional file sink/dir creation) but there are no unit tests covering these paths. Since tests/test_utils.py already exercises other utilities, consider adding a small test that calls setup_logging_for_tqdm(level=..., log_file=tmp_path/...) and asserts it doesn’t raise and creates/writes the log file.

[DorisMai]

Util updates are primarily stylistic changes, so merging shouldn’t harm the main branch. My personal preference is that the visualization-related functions might be cleaner as a separate script. They are not part of the core model logic and have quite a lot of hard-coded args like color and marker size that might be subject to quite some changes in the future. But this is really just a minor/non-blocking comment.

[marcuscollins]

keep all imports together at the top. Running ruff format and/or ruff check --fix can help do this automatically.

[marcuscollins]

This really applies to the code above: keep your code organized and put constants at the top. Generally it should go imports, then constants, then methods, then classes. It just makes for easier reading.

[DorisMai]

Style changes (organizing import order) should hopefully be addressed after merging PR #39 and #45 where pre-commit hooks run to check and fix these things (address issue #27 ). Approving for now to address the current batch of PRs.