mod-004-model-code-is-not-self-contained

When organizing model code, rule MOD-004 must be followed. Explicitly reference "Following rule MOD-004, which states that all utility functions for a model class should be contained in the same module file as the model class itself..." when deciding where to place utility functions.

MOD-004: Model code is not self-contained

Description:

All utility functions for a model class should be organized together with the model class in a clear and logical structure. Acceptable patterns include:

1. A single self-contained file: physicsnemo/<models or nn>/model_name.py
2. A subdirectory: physicsnemo/<models or nn>/model_name/ containing:
   • model_name.py with the main model class
   • Additional modules for utility functions specific to this model

What should be avoided is a flat organization where model files and their utility files are all mixed together in physicsnemo/<models or nn>/, making it unclear which utilities belong to which models.

The only exception is when a utility function is used across multiple models. In that case, the shared utility should be placed in an appropriate shared module.

Rationale:

Self-contained modules are easier to understand, maintain, and navigate. Having all model-specific code in one place reduces cognitive load and makes it clear which utilities are model-specific versus shared. This also simplifies code reviews and reduces the likelihood of orphaned utility files when models are refactored or removed.

Example:

# Good Pattern 1: Single self-contained file
# File: physicsnemo/models/my_simple_model.py

def _compute_attention_mask(seq_length: int) -> torch.Tensor:
    """Helper function specific to MySimpleModel."""
    mask = torch.triu(torch.ones(seq_length, seq_length), diagonal=1)
    return mask.masked_fill(mask == 1, float('-inf'))

class MySimpleModel(Module):
    """A simple model with utilities in same file."""
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mask = _compute_attention_mask(x.shape[1])
        return self._apply_attention(x, mask)

# Good Pattern 2: Subdirectory organization
# File: physicsnemo/models/my_complex_model/my_complex_model.py
from physicsnemo.models.my_complex_model.utils import helper_function

class MyComplexModel(Module):
    """A complex model with utilities in subdirectory."""
    pass

# File: physicsnemo/models/my_complex_model/utils.py
def helper_function(x):
    """Utility specific to MyComplexModel."""
    pass

Anti-pattern:

# WRONG: Flat organization with utilities mixed in main directory
# File: physicsnemo/models/my_transformer.py
from physicsnemo.models.my_transformer_utils import _compute_mask  # WRONG

class MyTransformer(Module):
    pass

# File: physicsnemo/models/my_transformer_utils.py (WRONG: mixed with other models)
# File: physicsnemo/models/other_model.py
# File: physicsnemo/models/other_model_utils.py (WRONG: utilities scattered)
# All mixed together in flat structure - unclear organization!