Add Helper Classes for Parameter and Label Management

[BalzaniEdoardo]

Problem

The parameter mapping and label management code in _composition_utils.py and _basis_mixin.py is complex due to:

1. String manipulation scattered throughout - Regex patterns for parsing labels like "BSplineEval_2" appear in multiple functions
2. Bidirectional dict management - Two separate dicts (label_to_attr and attr_to_label) must be kept in sync manually
3. Hard to maintain - Logic for parsing, formatting, and mapping is spread across many functions

Example of current complexity:

# Regex patterns repeated in multiple places
match = re.match(rf"^{cls_name}(_\d+)?$", bas._label)
if match:
    id_str = match.group(1)
    new_id = int(id_str[1:]) if id_str else 0
    bas._label = f"{cls_name}_{new_id}" if new_id else cls_name

Proposed Solution

Add two helper classes to encapsulate the complexity:

1. DefaultLabel - Label Parsing/Formatting

@dataclass
class DefaultLabel:
    """
    Parse and format default basis labels like 'BSplineEval' or 'BSplineEval_2'.

    Examples
    --------
    >>> DefaultLabel.parse("BSplineEval")
    DefaultLabel(class_name='BSplineEval', id=0)

    >>> DefaultLabel.parse("BSplineEval_3")
    DefaultLabel(class_name='BSplineEval', id=3)

    >>> DefaultLabel("RaisedCosineLinear", 2).format()
    'RaisedCosineLinear_2'
    """
    class_name: str
    id: int = 0

    @classmethod
    def parse(cls, label: str) -> Optional['DefaultLabel']:
        """Parse a label string. Returns None if not a default label."""
        match = re.match(r'^([A-Za-z]\w+?)(?:_(\d+))?$', label)
        if not match:
            return None

        class_name = match.group(1)
        if class_name not in __PUBLIC_BASES__:
            return None

        id_str = match.group(2)
        return cls(class_name, int(id_str) if id_str else 0)

    def format(self) -> str:
        """Format back to string representation."""
        if self.id == 0:
            return self.class_name
        return f"{self.class_name}_{self.id}"

Usage:

# Before (scattered regex):
match = re.match(rf"^{cls_name}(_\d+)?$", label)
if match:
    id_str = match.group(1)
    id_val = int(id_str[1:]) if id_str else 0

# After (clean):
parsed = DefaultLabel.parse(label)
if parsed:
    id_val = parsed.id
    new_label = DefaultLabel(parsed.class_name, id_val + 1).format()

2. ParameterMapping - Bidirectional Mapping

class ParameterMapping:
    """
    Bidirectional mapping between label-based and attribute-based parameter names.

    Maps between:
    - User-facing: "BSplineEval__order" (label-based)
    - Internal:    "basis1__order" (attribute-based)

    Examples
    --------
    >>> mapping = ParameterMapping()
    >>> mapping.add("BSplineEval__order", "basis1__order")
    >>> mapping.get_attr("BSplineEval__order")
    'basis1__order'
    >>> mapping.get_label("basis1__order")
    'BSplineEval__order'
    """

    def __init__(self):
        self._label_to_attr: dict[str, str] = {}
        self._attr_to_label: dict[str, str] = {}

    def add(self, label_key: str, attr_key: str):
        """Add a bidirectional mapping."""
        self._label_to_attr[label_key] = attr_key
        self._attr_to_label[attr_key] = label_key

    def get_attr(self, label_key: str) -> str:
        """Get attribute name from label-based key."""
        return self._label_to_attr[label_key]

    def get_label(self, attr_key: str) -> str:
        """Get label-based key from attribute name."""
        return self._attr_to_label[attr_key]

    def __repr__(self):
        return f"ParameterMapping({len(self._label_to_attr)} entries)"

Usage:

# Before (two dicts, manual sync):
parameter_dict = {}
key_map = {}
key_map[label_key] = attr_key  # Which direction?
parameter_dict[label_key] = value

# After (clear direction):
params = {}
mapping = ParameterMapping()
mapping.add(label_key="BSpline__order", attr_key="basis1__order")
params[label_key] = value
# ... later ...
attr_key = mapping.get_attr(label_key)  # Direction is explicit!

Benefits

1. Centralized logic - All regex patterns in one place (DefaultLabel.parse())
2. Self-documenting - mapping.add(label_key=..., attr_key=...) is clear
3. Easier to test - Each helper can be unit tested independently
4. Maintainability - Change label format? Update DefaultLabel.format() only
5. Debuggability - print(mapping) shows all parameter mappings

Implementation Plan

Phase 1: Add Helper Classes

• Add DefaultLabel class to _composition_utils.py
• Add ParameterMapping class to _composition_utils.py
• Add unit tests for both classes

Phase 2: Incremental Refactoring

Replace usage one function at a time:

1. Update _has_default_label() to use DefaultLabel.parse()
2. Update update_default_label_id() to use DefaultLabel
3. Update _recompute_class_default_labels() to use DefaultLabel
4. Update _get_params_and_key_map() to return ParameterMapping
5. Update set_params() decorator to use ParameterMapping

After each step: Run tests to ensure no behavior changes.

Files to Modify

• src/nemos/basis/_composition_utils.py - Add helper classes
• src/nemos/basis/_basis_mixin.py - Update _get_params_and_key_map(), set_params()
• tests/test_basis.py - Add tests for helper classes

Acceptance Criteria

• DefaultLabel class added with parse() and format() methods
• ParameterMapping class added with bidirectional lookup
• At least _has_default_label() and one other function refactored to use helpers
• Unit tests for both helper classes
• All existing tests pass without modification
• No behavior changes (purely internal refactoring)

References

• Current regex usage: _composition_utils.py lines 90-148 (label management)
• Current parameter mapping: _basis_mixin.py lines 1226-1328 (_get_params_and_key_map)