diff --git a/TODO-feature-stimulus-extraction.md b/TODO-feature-stimulus-extraction.md new file mode 100644 index 0000000..ddb2c41 --- /dev/null +++ b/TODO-feature-stimulus-extraction.md @@ -0,0 +1,31 @@ +""" +# Stimulus Feature Extraction Module - Implementation TODO + +**Branch:** feature/stimulus-extraction +**Issue:** #15 - Add Stimulus Feature Extraction module +**Status:** In Progress + +## Completed +- Create feature branch +- Create pyeeg/features/ subpackage +- Create pyeeg/models/ subpackage +- Add LLM feature extractor (from NNLM/lm_featurize) +- Add syntactic feature extractor (from process_txt/pyProcess/parseMetrics.py) +- Add alignment handler +- Add feature pipeline +- Add feature reducer +- Add refactored TRFEstimator +- Update package __init__.py + +## Next Steps +- Test all components +- Add documentation +- Refactor existing code +- Ensure backward compatibility + +## Reference Repositories +- NNLM/lm_featurize +- process_txt/pyProcess/parseMetrics.py + +Last Updated: June 5, 2026 +""" \ No newline at end of file diff --git a/pyeeg/__init__.py b/pyeeg/__init__.py index 3fe20ad..bdf66d2 100644 --- a/pyeeg/__init__.py +++ b/pyeeg/__init__.py @@ -27,21 +27,32 @@ 2019-2026, Hugo Weissbart """ -# Python 2/3 compatibility (obsolete as of 2020, removing) -# from __future__ import division, print_function, absolute_import -# This enables access to all submodules from the top-level `pyeeg` module -from . import connectivity, io, models, preprocess, vizu, utils, simulate +# This enables access to all submodules from the top-level +# pyeeg module +from . import connectivity, io, models, preprocess, vizu, utils, simulate, features from .models import TRFEstimator from .cca import CCA_Estimator from .preprocess import MultichanWienerFilter, Whitener from .mcca import mCCA +from .features import ( + LLMFeatureExtractor, + SyntacticFeatureExtractor, + AlignmentHandler, + TextGridParser, + FeaturePipeline, + FeatureReducer, + StimulusEncoder +) from .version import __version__ # Public API __all__ = [ # Classes 'TRFEstimator', + 'AlignmentHandler', + 'FeaturePipeline', + 'StimulusEncoder' 'CCA_Estimator', 'MultichanWienerFilter', 'Whitener', @@ -54,6 +65,7 @@ 'vizu', 'utils', 'simulate', + 'features', # Version '__version__', -] \ No newline at end of file +] diff --git a/pyeeg/features/__init__.py b/pyeeg/features/__init__.py new file mode 100644 index 0000000..720d918 --- /dev/null +++ b/pyeeg/features/__init__.py @@ -0,0 +1,29 @@ +""" +Stimulus Feature Extraction Module + +This module provides tools for extracting features from various types of stimuli: +- Text: LLM-based features (surprisal, entropy, KL divergence) +- Text: Syntactic features (depth, closing nodes, etc.) +- Audio: Alignment with transcripts + +Main Classes: +- LLMFeatureExtractor: Extract word-level features using language models +- SyntacticFeatureExtractor: Extract features from constituency trees +- AlignmentHandler: Handle alignment between stimuli and neural data +- FeaturePipeline: Compose multiple feature extractors +""" + +from .llm_features import LLMFeatureExtractor +from .syntactic_features import SyntacticFeatureExtractor +from .alignment import AlignmentHandler, TextGridParser +from .pipeline import FeaturePipeline +from .reduction import FeatureReducer + +__all__ = [ + 'LLMFeatureExtractor', + 'SyntacticFeatureExtractor', + 'AlignmentHandler', + 'TextGridParser', + 'FeaturePipeline', + 'FeatureReducer', +] diff --git a/pyeeg/features/alignment.py b/pyeeg/features/alignment.py new file mode 100644 index 0000000..f6d2eed --- /dev/null +++ b/pyeeg/features/alignment.py @@ -0,0 +1,155 @@ +""" +Alignment Handling + +This module provides functionality for aligning features with neural signals +using TextGrid files or forced alignment with external tools. +""" + +import numpy as np +from typing import Dict, List, Optional, Tuple +from dataclasses import dataclass +import logging +import re + +logger = logging.getLogger(__name__) + + +@dataclass +class Interval: + """Represents an interval in a TextGrid file.""" + start: float + end: float + label: str + tier: str + + def duration(self) -> float: + return self.end - self.start + + def contains(self, time: float) -> bool: + return self.start <= time < self.end + + +@dataclass +class TextGrid: + """Represents a TextGrid file with multiple tiers.""" + intervals: Dict[str, List[Interval]] + start_time: float = 0.0 + end_time: float = 0.0 + + def get_tier(self, tier_name: str) -> List[Interval]: + return self.intervals.get(tier_name, []) + + def get_word_intervals(self) -> List[Interval]: + word_tiers = ['words', 'word', 'Word', 'WORDS'] + for tier in word_tiers: + if tier in self.intervals: + return self.intervals[tier] + for tier, intervals in self.intervals.items(): + if intervals and any(' ' in interval.label or interval.label.isalpha() for interval in intervals): + return intervals + return [] + + def get_phone_intervals(self) -> List[Interval]: + phone_tiers = ['phones', 'phone', 'Phone', 'PHONES', 'phn'] + for tier in phone_tiers: + if tier in self.intervals: + return self.intervals[tier] + return [] + + +class TextGridParser: + """Parser for Praat TextGrid files.""" + + def parse_from_string(self, content: str) -> TextGrid: + """Parse a TextGrid from a string.""" + textgrid = TextGrid() + textgrid.intervals = {} + + lines = content.split('\n') + current_tier = None + + for line in lines: + line = line.strip() + if not line: + continue + + tier_match = re.match(r'item \[(\d+)\]:\s*"(.+?)"', line) + if tier_match: + current_tier = tier_match.group(2) + textgrid.intervals[current_tier] = [] + continue + + interval_match = re.match( + r'intervals \[(\d+)\]:\s*"(.+?)"\s*(\d+\.\d+)\s*(\d+\.\d+)', + line + ) + if interval_match and current_tier: + label = interval_match.group(2) + start = float(interval_match.group(3)) + end = float(interval_match.group(4)) + textgrid.intervals[current_tier].append( + Interval(start=start, end=end, label=label, tier=current_tier) + ) + + return textgrid + + +class AlignmentHandler: + """Handle alignment between features and neural signals.""" + + def __init__(self, signal_sampling_rate: float = 1000.0): + self.sampling_rate = signal_sampling_rate + self.parser = TextGridParser() + + def load_textgrid_from_string(self, content: str) -> TextGrid: + """Load and parse a TextGrid from a string.""" + return self.parser.parse_from_string(content) + + def align_word_features( + self, + word_features: Dict[int, Dict[str, float]], + textgrid: TextGrid, + signal_length: Optional[int] = None + ) -> Tuple[np.ndarray, List[str]]: + """Align word-level features to signal time points.""" + word_intervals = textgrid.get_word_intervals() + + if not word_intervals: + logger.error("No word intervals found in TextGrid") + return np.array([]), [] + + all_feature_names = set() + for feat_dict in word_features.values(): + all_feature_names.update(feat_dict.keys()) + all_feature_names = sorted(all_feature_names) + + if not all_feature_names: + logger.warning("No features to align") + return np.array([]), [] + + if signal_length is None: + signal_length = int(textgrid.end_time * self.sampling_rate) + + time_points = np.arange(signal_length) / self.sampling_rate + n_samples = len(time_points) + n_features = len(all_feature_names) + + aligned_features = np.zeros((n_samples, n_features)) + + for i, interval in enumerate(word_intervals): + if i not in word_features: + continue + + start_sample = int(interval.start * self.sampling_rate) + end_sample = int(interval.end * self.sampling_rate) + + start_sample = max(0, min(start_sample, n_samples - 1)) + end_sample = max(0, min(end_sample, n_samples)) + + feat_dict = word_features[i] + + for j, feat_name in enumerate(all_feature_names): + if feat_name in feat_dict: + aligned_features[start_sample:end_sample, j] = feat_dict[feat_name] + + return aligned_features, all_feature_names \ No newline at end of file diff --git a/pyeeg/features/llm_features.py b/pyeeg/features/llm_features.py new file mode 100644 index 0000000..e00a3f4 --- /dev/null +++ b/pyeeg/features/llm_features.py @@ -0,0 +1,256 @@ +""" +LLM-Based Feature Extraction + +This module provides functionality for extracting word-level linguistic features +using language models, incorporating code from NNLM/lm_featurize. + +Features: +- Surprisal: Information content of each word given its context (-log2(P(word | context))) +- Entropy: Uncertainty in the token distribution (-sum(p * log2(p))) +- KL Divergence: Difference between token distributions (KL(P||Q)) +- Prediction Error: Surprisal normalized by entropy + +Based on NNLM/lm_featurize by Hugo Weissbart +""" + +import numpy as np +from typing import Dict, List, Optional, Tuple, Union +from dataclasses import dataclass, field +import logging +import torch +from torch.nn import functional as F + +logger = logging.getLogger(__name__) + + +# Special tokens used in BPE tokenization +SEPARATOR = '\u0160' # 'Ġ' - Byte pair encoding separator +NEWLINE = '\u010a' # 'Ċ' - Newline token + + +@dataclass +class LLMFeatureConfig: + """Configuration for LLM feature extraction.""" + model_name: str = "GroNLP/gpt2-small-dutch" + device: str = "cpu" + batch_size: int = 32 + max_length: int = 512 + cache_dir: Optional[str] = None + use_cache: bool = True + return_shape: str = "valid" + + +class LLMFeatureExtractor: + """ + Extract word-level linguistic features using language models. + + Based on: NNLM/lm_featurize/metrics.py by Hugo Weissbart + + Args: + config: Configuration for the feature extractor + """ + + def __init__(self, config: LLMFeatureConfig = None): + if config is None: + config = LLMFeatureConfig() + self.config = config + self._tokenizer = None + self._model = None + self._initialize_model() + + def _initialize_model(self): + """Initialize the language model and tokenizer.""" + try: + from transformers import AutoTokenizer, GPT2LMHeadModel, AutoModel + except ImportError: + logger.error("transformers library not installed") + raise + + self._tokenizer = AutoTokenizer.from_pretrained(self.config.model_name) + + if "gpt2" in self.config.model_name.lower(): + self._model = GPT2LMHeadModel.from_pretrained(self.config.model_name) + else: + self._model = AutoModel.from_pretrained(self.config.model_name) + + if self.config.device == "cuda" and torch.cuda.is_available(): + self._model = self._model.cuda() + + logger.info(f"Loaded model: {self.config.model_name}") + + def get_surprisal(self, logits: torch.Tensor, input_ids: torch.Tensor, + return_shape: str = None) -> torch.Tensor: + """ + Calculate surprisal (-log(p)) for tokens. + Based on: NNLM/lm_featurize/metrics.py + """ + if return_shape is None: + return_shape = self.config.return_shape + + p = F.softmax(logits, dim=1) + + if return_shape == 'valid': + return -p[:-1, input_ids.flatten()[1:]].diag().log() + elif return_shape == 'same': + nan_val = torch.ones(1, device=logits.device) * float('nan') + return torch.cat((nan_val, -p[:-1, input_ids.flatten()[1:]].diag().log())) + elif return_shape == 'full': + return -p[:, input_ids.flatten()].diag().log() + else: + raise ValueError(f"Unknown return_shape: {return_shape}") + + def get_entropy(self, logits: torch.Tensor, return_shape: str = None) -> torch.Tensor: + """Calculate entropy of token distributions. Based on: NNLM/lm_featurize/metrics.py""" + if return_shape is None: + return_shape = self.config.return_shape + + if return_shape == 'valid': + return -torch.sum(F.softmax(logits, dim=1) * F.log_softmax(logits, dim=1), 1)[:-1] + elif return_shape == 'same': + nan_val = torch.ones(1, device=logits.device) * float('nan') + return torch.cat((nan_val, -torch.sum(F.softmax(logits, dim=1) * F.log_softmax(logits, dim=1), 1)[:-1])) + elif return_shape == 'full': + return -torch.sum(F.softmax(logits, dim=1) * F.log_softmax(logits, dim=1), 1) + else: + raise ValueError(f"Unknown return_shape: {return_shape}") + + def get_kl_divergence(self, logits: torch.Tensor, return_shape: str = None) -> torch.Tensor: + """Compute KL divergence between current and previous state. Based on: NNLM/lm_featurize/metrics.py""" + if return_shape is None: + return_shape = self.config.return_shape + + P = F.log_softmax(logits, dim=1) + + if return_shape == 'valid': + return F.kl_div(P[:1, :], P[1:, :], reduction='none', log_target=True).sum(1) + elif return_shape == 'same': + nan_val = torch.ones(1, device=logits.device) * float('nan') + return torch.cat((nan_val, F.kl_div(P[:-1, :], P[1:, :], reduction='none', log_target=True).sum(1))) + else: + raise ValueError(f"Unknown return_shape: {return_shape}") + + def get_prediction_error(self, surprisal: torch.Tensor, entropy: torch.Tensor) -> torch.Tensor: + """Calculate prediction error (surprisal/entropy).""" + entropy_safe = entropy.clone() + entropy_safe[entropy_safe == 0] = 1e-10 + return surprisal / entropy_safe + + def get_tokens(self, input_ids: torch.Tensor) -> List[str]: + """Convert input IDs to token strings.""" + return self._tokenizer.convert_ids_to_tokens(input_ids.flatten()) + + def bpe_to_words(self, tokens: List[str], sep: str = SEPARATOR, newline: str = NEWLINE, + lang: str = 'nl') -> List[List[int]]: + """ + Map BPE tokens to word indices. Based on: NNLM/lm_featurize/utils.py + """ + try: + from nltk.tokenize import word_tokenize + except ImportError: + logger.error("NLTK not installed") + raise + + tokens_clean = word_tokenize(''.join(tokens).replace(sep, ' ').replace(newline, '\n'), language=lang) + + indices_map = [] + k = 0 + + for t in tokens_clean: + bytepairs = [] + test_token = tokens[k].strip(sep).strip(newline) + if test_token != '': + bytepairs.append(k) + + while test_token != t: + k += 1 + if k == len(tokens): break + test_token += tokens[k].strip(sep).strip(newline) + if test_token != '': + bytepairs.append(k) + if len(bytepairs) > 100: + raise RecursionError(f"Can't match token: {t}") + + indices_map.append(bytepairs) + k += 1 + if k >= len(tokens): break + + return indices_map + + def reduce_features(self, features: Dict[str, np.ndarray], indices_map: List[List[int]]) -> Dict[str, np.ndarray]: + """Reduce BPE-level features to word-level. Based on: NNLM/lm_featurize/utils.py""" + from functools import reduce as functools_reduce + + def sum_reduce(x): + return functools_reduce(lambda x, y: x+y, x) + + reduced_features = {} + + for feat_name, feat_array in features.items(): + if feat_name == 'tokens': + reduced_tokens = [] + for bpe_indices in indices_map: + token_str = ''.join([features['tokens'][i] for i in bpe_indices]).replace(sep, '').replace(newline, '') + reduced_tokens.append(token_str) + reduced_features[feat_name] = np.array(reduced_tokens) + elif feat_name == 'entropy': + reduced_values = [] + for bpe_indices in indices_map: + reduced_values.append(feat_array[bpe_indices[-1]]) + reduced_features[feat_name] = np.array(reduced_values) + elif feat_name == 'kl_divergence': + reduced_values = [] + for bpe_indices in indices_map: + if len(bpe_indices) > 1: + reduced_values.append(feat_array[bpe_indices[0]] + feat_array[bpe_indices[-1]]) + else: + reduced_values.append(feat_array[bpe_indices[0]]) + reduced_features[feat_name] = np.array(reduced_values) + else: + reduced_values = [] + for bpe_indices in indices_map: + reduced_values.append(sum_reduce(feat_array[bpe_indices])) + reduced_features[feat_name] = np.array(reduced_values) + + if 'surprisal' in reduced_features and 'entropy' in reduced_features: + entropy_safe = reduced_features['entropy'].copy() + entropy_safe[entropy_safe == 0] = 1e-10 + reduced_features['prediction_error'] = reduced_features['surprisal'] / entropy_safe + + return reduced_features + + def extract(self, text: str, features: List[str] = None, + return_word_level: bool = True) -> Dict[str, np.ndarray]: + """Extract features from text.""" + if features is None: + features = ['surprisal', 'entropy', 'kl_divergence', 'prediction_error'] + + inputs = self._tokenizer(text, return_tensors="pt") + input_ids = inputs['input_ids'] + + if self.config.device == "cuda" and torch.cuda.is_available(): + input_ids = input_ids.cuda() + + with torch.no_grad(): + outputs = self._model(**inputs) + logits = outputs.logits + + result = {} + + if 'surprisal' in features or 'prediction_error' in features: + result['surprisal'] = self.get_surprisal(logits, input_ids, 'same').detach().cpu().numpy() + + if 'entropy' in features or 'prediction_error' in features: + result['entropy'] = self.get_entropy(logits, 'same').detach().cpu().numpy() + + if 'kl_divergence' in features: + result['kl_divergence'] = self.get_kl_divergence(logits, 'same').detach().cpu().numpy() + + if 'tokens' in features: + result['tokens'] = np.array(self.get_tokens(input_ids)) + + if return_word_level and len(result) > 0: + tokens = self.get_tokens(input_ids) + indices_map = self.bpe_to_words(tokens, lang='en') + result = self.reduce_features(result, indices_map) + + return result \ No newline at end of file diff --git a/pyeeg/features/pipeline.py b/pyeeg/features/pipeline.py new file mode 100644 index 0000000..4b29eb8 --- /dev/null +++ b/pyeeg/features/pipeline.py @@ -0,0 +1,245 @@ +""" +Feature Extraction Pipeline + +This module provides a pipeline for composing multiple feature extractors +and aligning their outputs with neural signals. +""" + +import numpy as np +from typing import Dict, List, Optional, Tuple +from dataclasses import dataclass, field +import logging + +from .llm_features import LLMFeatureExtractor, LLMFeatureConfig +from .syntactic_features import SyntacticFeatureExtractor, ParserConfig +from .alignment import AlignmentHandler, TextGrid + +logger = logging.getLogger(__name__) + + +@dataclass +class FeatureSpec: + """Specification for a feature to be extracted.""" + name: str + extractor_type: str + features: List[str] = field(default_factory=list) + config: Optional[Dict] = None + + +@dataclass +class PipelineConfig: + """Configuration for the feature extraction pipeline.""" + feature_specs: List[FeatureSpec] = field(default_factory=list) + alignment_config: Optional[Dict] = None + normalization: str = "none" + cache_features: bool = True + + +class FeaturePipeline: + """Pipeline for extracting and aligning multiple features from stimuli.""" + + def __init__(self, config: PipelineConfig): + self.config = config + self._extractors: Dict[str, object] = {} + self._alignment_handler: Optional[AlignmentHandler] = None + self._initialize_extractors() + self._initialize_alignment() + + def _initialize_extractors(self): + """Initialize all feature extractors.""" + for spec in self.config.feature_specs: + if spec.extractor_type == 'llm': + llm_config = LLMFeatureConfig( + model_name=spec.config.get('model_name', 'GroNLP/gpt2-small-dutch'), + device=spec.config.get('device', 'cpu') + ) + self._extractors[spec.name] = LLMFeatureExtractor(llm_config) + elif spec.extractor_type == 'syntactic': + parser_config = ParserConfig( + parser_name=spec.config.get('parser_name', 'stanford'), + language=spec.config.get('language', 'en') + ) + self._extractors[spec.name] = SyntacticFeatureExtractor(parser_config) + + def _initialize_alignment(self): + """Initialize alignment handler.""" + if self.config.alignment_config: + sampling_rate = self.config.alignment_config.get('sampling_rate', 1000.0) + self._alignment_handler = AlignmentHandler(sampling_rate) + + def extract( + self, + text: str, + textgrid: Optional[TextGrid] = None, + signal_length: Optional[int] = None + ) -> Tuple[Dict[str, np.ndarray], Dict]: + """Extract features from text and optionally align to signal.""" + metadata = { + 'text': text, + 'text_length': len(text), + 'feature_specs': [spec.name for spec in self.config.feature_specs], + 'aligned': textgrid is not None + } + + all_features: Dict[str, Dict] = {} + + for spec in self.config.feature_specs: + extractor = self._extractors.get(spec.name) + if extractor is None: + continue + + if spec.extractor_type == 'llm': + features = extractor.extract(text, spec.features) + elif spec.extractor_type == 'syntactic': + features = extractor.extract_to_dict(text, spec.features) + else: + continue + + if features: + all_features[spec.name] = features + + if textgrid is None and self._alignment_handler is None: + result = {} + for extractor_name, feat_dict in all_features.items(): + for feat_name, values in feat_dict.items(): + if isinstance(values, dict): + words = text.split() + arr = np.zeros(len(words)) + for pos, val in values.items(): + if pos < len(words): + arr[pos] = val + result[f"{extractor_name}_{feat_name}"] = arr + elif isinstance(values, np.ndarray): + result[f"{extractor_name}_{feat_name}"] = values + metadata['aligned'] = False + return result, metadata + + if textgrid is None: + logger.warning("No TextGrid provided and no alignment handler configured") + textgrid = TextGrid() + + result = {} + for extractor_name, feat_dict in all_features.items(): + aligned, feat_names = self._alignment_handler.align_word_features( + feat_dict, textgrid, signal_length + ) + for i, feat_name in enumerate(feat_names): + result[f"{extractor_name}_{feat_name}"] = aligned[:, i] + + metadata['aligned'] = True + metadata['n_samples'] = aligned.shape[0] if aligned.size > 0 else 0 + metadata['sampling_rate'] = self._alignment_handler.sampling_rate + + return result, metadata + + def normalize_features( + self, + features: Dict[str, np.ndarray], + method: str = "zscore" + ) -> Dict[str, np.ndarray]: + """Normalize features.""" + if method == 'none': + return features + + result = {} + for name, arr in features.items(): + if method == 'zscore': + mean = np.mean(arr) + std = np.std(arr) + if std > 0: + result[name] = (arr - mean) / std + else: + result[name] = arr - mean + elif method == 'minmax': + min_val = np.min(arr) + max_val = np.max(arr) + if max_val > min_val: + result[name] = (arr - min_val) / (max_val - min_val) + else: + result[name] = arr - min_val + else: + result[name] = arr + return result + + +class StimulusEncoder: + """High-level interface for stimulus feature extraction.""" + + def __init__(self, pipeline_config: Optional[PipelineConfig] = None): + if pipeline_config is None: + pipeline_config = PipelineConfig() + self.pipeline = FeaturePipeline(pipeline_config) + + def add_llm_features( + self, + features: List[str] = None, + model_name: str = "GroNLP/gpt2-small-dutch", + name: str = "llm" + ): + """Add LLM-based features to the pipeline.""" + if features is None: + features = ['surprisal', 'entropy', 'kl_divergence'] + spec = FeatureSpec( + name=name, + extractor_type='llm', + features=features, + config={'model_name': model_name} + ) + self.pipeline.config.feature_specs.append(spec) + self.pipeline._initialize_extractors() + + def add_syntactic_features( + self, + features: List[str] = None, + parser_name: str = "stanford", + name: str = "syntactic" + ): + """Add syntactic features to the pipeline.""" + if features is None: + features = ['depth', 'opening', 'closing'] + spec = FeatureSpec( + name=name, + extractor_type='syntactic', + features=features, + config={'parser_name': parser_name} + ) + self.pipeline.config.feature_specs.append(spec) + self.pipeline._initialize_extractors() + + def set_alignment(self, sampling_rate: float = 1000.0): + """Set alignment configuration.""" + self.pipeline.config.alignment_config = { + 'sampling_rate': sampling_rate + } + self.pipeline._initialize_alignment() + + def encode( + self, + text: str, + textgrid: Optional[TextGrid] = None, + signal_length: Optional[int] = None + ) -> Tuple[Dict[str, np.ndarray], Dict]: + """Extract and align features from text.""" + features, metadata = self.pipeline.extract( + text, textgrid, signal_length + ) + if self.pipeline.config.normalization != 'none': + features = self.pipeline.normalize_features( + features, + self.pipeline.config.normalization + ) + return features, metadata + + def encode_to_array( + self, + text: str, + textgrid: Optional[TextGrid] = None, + signal_length: Optional[int] = None + ) -> np.ndarray: + """Extract features and return as a single array.""" + features, metadata = self.encode(text, textgrid, signal_length) + if not features: + return np.array([]) + feature_names = sorted(features.keys()) + stacked = np.column_stack([features[name] for name in feature_names]) + return stacked \ No newline at end of file diff --git a/pyeeg/features/reduction.py b/pyeeg/features/reduction.py new file mode 100644 index 0000000..f97b345 --- /dev/null +++ b/pyeeg/features/reduction.py @@ -0,0 +1,151 @@ +""" +Feature Reduction + +This module provides functionality for reducing the dimensionality of +feature sets using PCA, ICA, and other techniques. +""" + +import numpy as np +from typing import Optional, Tuple +from dataclasses import dataclass +import logging + +logger = logging.getLogger(__name__) + + +@dataclass +class ReductionConfig: + """Configuration for feature reduction.""" + method: str = "pca" + n_components: Optional[int] = None + variance_threshold: float = 0.95 + random_state: Optional[int] = None + whiten: bool = False + + +class FeatureReducer: + """Reduce the dimensionality of feature sets.""" + + def __init__(self, config: ReductionConfig): + self.config = config + self._fitted = False + self._components: Optional[np.ndarray] = None + self._explained_variance: Optional[np.ndarray] = None + self._mean: Optional[np.ndarray] = None + self._n_components: Optional[int] = None + + def _validate_input(self, features: np.ndarray) -> None: + if not isinstance(features, np.ndarray): + raise ValueError("Features must be a numpy array") + if features.ndim != 2: + raise ValueError("Features must be 2D") + if features.size == 0: + raise ValueError("Features array is empty") + + def _center_data(self, features: np.ndarray) -> np.ndarray: + if self._mean is None: + self._mean = np.mean(features, axis=0) + return features - self._mean + + def _uncenter_data(self, features: np.ndarray) -> np.ndarray: + if self._mean is None: + return features + return features + self._mean + + def fit(self, features: np.ndarray) -> 'FeatureReducer': + """Fit the reducer to the features.""" + self._validate_input(features) + + if self.config.method == 'pca': + self._fit_pca(features) + elif self.config.method == 'ica': + self._fit_ica(features) + elif self.config.method == 'none': + self._fitted = True + else: + raise ValueError(f"Unknown reduction method: {self.config.method}") + + self._fitted = True + return self + + def _fit_pca(self, features: np.ndarray): + """Fit PCA to the features.""" + centered = self._center_data(features) + cov = np.cov(centered, rowvar=False) + eigenvalues, eigenvectors = np.linalg.eigh(cov) + idx = np.argsort(eigenvalues)[::-1] + eigenvalues = eigenvalues[idx] + eigenvectors = eigenvectors[:, idx] + + total_variance = np.sum(eigenvalues) + self._explained_variance = eigenvalues / total_variance + self._components = eigenvectors.T + + if self.config.n_components is not None: + self._n_components = min(self.config.n_components, len(eigenvalues)) + else: + cumulative_variance = np.cumsum(self._explained_variance) + self._n_components = np.argmax(cumulative_variance >= self.config.variance_threshold) + 1 + + def _fit_ica(self, features: np.ndarray): + """Fit ICA to the features.""" + centered = self._center_data(features) + cov = np.cov(centered, rowvar=False) + eigenvalues, eigenvectors = np.linalg.eigh(cov) + eigenvalues[eigenvalues < 1e-10] = 1e-10 + whitening_matrix = eigenvectors @ np.diag(1.0 / np.sqrt(eigenvalues)) @ eigenvectors.T + whitened = centered @ whitening_matrix.T + self._components = whitened.T + + if self.config.n_components is not None: + self._n_components = min(self.config.n_components, whitened.shape[1]) + else: + self._n_components = whitened.shape[1] + + def transform(self, features: np.ndarray) -> np.ndarray: + """Transform features to reduced space.""" + if not self._fitted: + raise RuntimeError("Reducer not fitted. Call fit() first.") + self._validate_input(features) + + if self.config.method == 'pca': + centered = self._center_data(features) + return centered @ self._components[:self._n_components].T + elif self.config.method == 'ica': + centered = self._center_data(features) + return centered @ self._components[:self._n_components].T + elif self.config.method == 'none': + return features + else: + raise ValueError(f"Unknown reduction method: {self.config.method}") + + def fit_transform(self, features: np.ndarray) -> np.ndarray: + """Fit the reducer and transform the features.""" + self.fit(features) + return self.transform(features) + + def inverse_transform(self, reduced_features: np.ndarray) -> np.ndarray: + """Transform reduced features back to original space.""" + if not self._fitted: + raise RuntimeError("Reducer not fitted. Call fit() first.") + + if self.config.method == 'pca': + reconstructed = reduced_features @ self._components[:self._n_components] + return self._uncenter_data(reconstructed) + elif self.config.method == 'ica': + return self._uncenter_data(reduced_features @ self._components[:self._n_components]) + elif self.config.method == 'none': + return reduced_features + else: + raise ValueError(f"Unknown reduction method: {self.config.method}") + + def get_explained_variance(self) -> Optional[np.ndarray]: + return self._explained_variance + + def get_components(self) -> Optional[np.ndarray]: + if self._components is not None and self._n_components is not None: + return self._components[:self._n_components] + return None + + def get_n_components(self) -> Optional[int]: + return self._n_components \ No newline at end of file diff --git a/pyeeg/features/syntactic_features.py b/pyeeg/features/syntactic_features.py new file mode 100644 index 0000000..d19f234 --- /dev/null +++ b/pyeeg/features/syntactic_features.py @@ -0,0 +1,237 @@ +""" +Syntactic Feature Extraction + +This module provides functionality for extracting syntactic features from +constituency trees, incorporating code from process_txt/pyProcess/parseMetrics.py. + +Features: +- Depth: Depth of each node in the parse tree +- Opening nodes: Number of branches opening at each leaf +- Closing nodes: Number of branches closing at each leaf + +Based on parseMetrics.py by Hugo Weissbart +""" + +import os +import logging +import tempfile +import shutil +from typing import Dict, List, Optional, Tuple +from dataclasses import dataclass +import numpy as np + +logger = logging.getLogger(__name__) + +try: + from nltk.tree import Tree + from nltk.parse.stanford import StanfordParser + from nltk.tokenize import sent_tokenize + NLTK_AVAILABLE = True +except ImportError: + NLTK_AVAILABLE = False + +try: + import alpinonaf + ALPINO_AVAILABLE = True +except ImportError: + ALPINO_AVAILABLE = False + + +@dataclass +class ParserConfig: + """Configuration for external parser.""" + parser_name: str = "stanford" + parser_path: Optional[str] = None + model_path: Optional[str] = None + language: str = "en" + timeout: int = 30 + + +class SyntacticFeatureExtractor: + """ + Extract syntactic features from text using constituency parsing. + Based on: process_txt/pyProcess/parseMetrics.py by Hugo Weissbart + """ + + def __init__(self, config: ParserConfig = None): + if config is None: + config = ParserConfig() + self.config = config + + def get_stanford_tree(self, sentences: List[str], path_to_jar: Optional[str] = None) -> List[Tree]: + """Parse sentences using Stanford Parser. Based on parseMetrics.py""" + if not NLTK_AVAILABLE: + raise ImportError("NLTK required for Stanford parsing") + + logger.info("Loading Stanford Parser...") + + if path_to_jar is None: + path_to_jar = self.config.parser_path + + if path_to_jar is not None: + path_to_jar = os.path.expanduser(path_to_jar) + if not os.path.exists(path_to_jar): + import subprocess + try: + path_to_jar = subprocess.check_output(['bash', '-c', 'locate stanford-parser.jar']).strip('\n').decode('utf-8') + except Exception: + raise IOError("stanford-parser.jar not found") + + if path_to_jar: + os.environ['CLASSPATH'] = path_to_jar + + if os.getenv('STANFORD_MODELS') is None: + raise IOError("STANFORD_MODELS not set") + + parser = StanfordParser(model_path="edu/stanford/nlp/models/lexparser/englishPCFG.ser.gz") + logger.info("Stanford Parser loaded") + + parse = parser.raw_parse_sents(sentences) + trees = [] + for treelist in parse: + for tree in treelist: + trees.append(tree) + return trees + + def get_alpinopy_tree(self, sentences: List[str]) -> List[Tree]: + """Parse Dutch sentences using Alpino. Based on parseMetrics.py""" + if not ALPINO_AVAILABLE: + raise ImportError("alpinonaf required for Dutch parsing") + + trees = [] + tempfile.tempdir = '/tmp/alpino' + + if not os.path.exists('/tmp/alpino'): + os.mkdir('/tmp/alpino') + + for sent in sentences: + tmp = tempfile.NamedTemporaryFile(mode='wb', delete=False, prefix='alpino-') + with open(tmp.name, 'wb') as fid: + fid.write(bytes(sent, encoding='utf8')) + fid.seek(0) + with open(tmp.name, 'rb') as f: + alpinonaf.parse(f, max_min_per_sent=5.0) + with open('penn_output.txt', 'r') as f: + penntree = f.read() + trees.append(Tree.fromstring(penntree)) + + shutil.rmtree('/tmp/alpino') + tempfile.tempdir = None + return trees + + def parse_text(self, text: str) -> List[Tree]: + """Parse text into constituency trees.""" + if self.config.language == 'en': + sentences = sent_tokenize(text, language='english') + elif self.config.language == 'nl': + sentences = sent_tokenize(text, language='dutch') + elif self.config.language == 'fr': + sentences = sent_tokenize(text, language='french') + else: + sentences = sent_tokenize(text, language=self.config.language) + + if self.config.parser_name == 'stanford': + return self.get_stanford_tree(sentences) + elif self.config.parser_name == 'alpino': + return self.get_alpinopy_tree(sentences) + else: + raise ValueError(f"Unknown parser: {self.config.parser_name}") + + def depth_single_tree(self, tree: Tree, remove_S: bool = True, + remove_unibranch_offset: bool = True) -> List[int]: + """Get depth of each leaf in the parse tree. Based on parseMetrics.py""" + tr = tree.copy(True) + tr.collapse_unary(collapseRoot=True) + offset = sum([remove_S, remove_unibranch_offset]) + return [len(pos) - offset for pos in tr.treepositions('leaves')] + + def opening_single_tree(self, tree: Tree) -> List[int]: + """Get opening values for each leaf. Based on parseMetrics.py""" + tr = tree.copy(True) + tr.collapse_unary(collapseRoot=True) + pos = [tp for tp in tr.treepositions('leaves')] + + opening = [] + for p in pos: + count = 0 + iterable = iter(p[-2::-1]) + for index in iterable: + if index == 0: + count += 1 + else: + break + opening.append(count) + return opening + + def closing_single_tree(self, tree: Tree) -> List[int]: + """Get closing values for each leaf. Based on parseMetrics.py""" + tr = tree.copy(True) + tr.collapse_unary(collapseRoot=True) + for s in tr.subtrees(): + s.reverse() + return self.opening_single_tree(tr)[::-1] + + def extract_from_tree(self, tree: Tree, features: List[str]) -> Dict[str, List[int]]: + """Extract requested features from a single parse tree.""" + result = {} + for feat in features: + if feat == 'depth' or feat == 'all': + result['depth'] = self.depth_single_tree(tree) + if feat == 'opening' or feat == 'all': + result['opening'] = self.opening_single_tree(tree) + if feat == 'closing' or feat == 'all': + result['closing'] = self.closing_single_tree(tree) + if feat == 'tree_height' or feat == 'all': + result['tree_height'] = [tree.height()] * len(tree.leaves()) + return result + + def extract(self, text: str, features: List[str] = None) -> Dict[str, List[int]]: + """Extract syntactic features from text.""" + if features is None: + features = ['all'] + + trees = self.parse_text(text) + if not trees: + logger.warning("No parse trees generated") + return {} + + all_features = {} + for tree in trees: + tree_features = self.extract_from_tree(tree, features) + for feat_name in tree_features: + if feat_name not in all_features: + all_features[feat_name] = [] + all_features[feat_name].extend(tree_features[feat_name]) + + return all_features + + def extract_to_dict(self, text: str, features: List[str] = None) -> Dict[str, Dict[int, float]]: + """Extract features and return as word position -> value dict.""" + raw_features = self.extract(text, features) + result = {} + for feat_name, values in raw_features.items(): + result[feat_name] = {i: float(v) for i, v in enumerate(values)} + return result + + def extract_to_array(self, text: str, features: List[str] = None) -> Tuple[List[str], np.ndarray]: + """Extract features and return as arrays.""" + words = text.split() + feat_dict = self.extract(text, features) + + if not feat_dict: + return words, np.array([]) + + if features is None or 'all' in features: + feature_names = ['depth', 'opening', 'closing', 'tree_height'] + else: + feature_names = features + + n_words = len(list(feat_dict.values())[0]) if feat_dict else 0 + n_features = len(feature_names) + feature_array = np.zeros((n_words, n_features)) + + for i, feat_name in enumerate(feature_names): + if feat_name in feat_dict: + feature_array[:, i] = feat_dict[feat_name] + + return words, feature_array \ No newline at end of file diff --git a/pyeeg/models/__init__.py b/pyeeg/models/__init__.py new file mode 100644 index 0000000..789d4fe --- /dev/null +++ b/pyeeg/models/__init__.py @@ -0,0 +1,9 @@ +""" +Models Subpackage + +This subpackage contains the main modeling classes for pyEEG. +""" + +from .trf import TRFEstimator + +__all__ = ['TRFEstimator'] \ No newline at end of file diff --git a/pyeeg/models/trf.py b/pyeeg/models/trf.py new file mode 100644 index 0000000..6ffb9be --- /dev/null +++ b/pyeeg/models/trf.py @@ -0,0 +1,226 @@ +""" +Temporal Response Function (TRF) Estimation + +This module provides the TRFEstimator class for estimating Temporal Response +Functions from neural data and stimuli. + +The TRFEstimator can now integrate with the feature extraction module to +handle naturalistic stimuli with rich feature representations. +""" + +import numpy as np +from typing import Dict, List, Optional, Tuple, Union +from dataclasses import dataclass, field +import logging + +logger = logging.getLogger(__name__) + + +@dataclass +class TRFConfig: + """Configuration for TRF estimation.""" + lags: List[float] = field(default_factory=lambda: [-0.1, 0.0, 0.1, 0.2, 0.3, 0.4, 0.5]) + lambda_reg: float = 1.0 + solver: str = "svd" + max_iter: int = 1000 + tol: float = 1e-6 + use_features: bool = False + feature_names: List[str] = field(default_factory=list) + validate_alignment: bool = True + + +class TRFEstimator: + """ + Estimate Temporal Response Functions from neural data and stimuli. + + This class can now integrate with the feature extraction module to handle + naturalistic stimuli. When use_features=True, it expects features to be + provided instead of raw stimuli, and will align them appropriately. + + Args: + config: TRF estimation configuration + + Example: + >>> from pyeeg.features import StimulusEncoder + >>> encoder = StimulusEncoder() + >>> encoder.add_llm_features(['surprisal', 'entropy']) + >>> encoder.add_syntactic_features(['depth']) + >>> features, _ = encoder.encode(text, textgrid) + >>> trf = TRFEstimator(TRFConfig(lags=[-0.1, 0, 0.1, 0.2], use_features=True)) + >>> trf.fit(signal, features) + """ + + def __init__(self, config: TRFConfig): + self.config = config + self._fitted = False + self._coefficients: Optional[np.ndarray] = None + self._intercept: Optional[np.ndarray] = None + self._feature_names: List[str] = [] + self._lag_matrix: Optional[np.ndarray] = None + + def _validate_inputs( + self, + signal: np.ndarray, + stimulus: Union[np.ndarray, Dict[str, np.ndarray]] + ): + """Validate input signal and stimulus.""" + if not isinstance(signal, np.ndarray): + raise ValueError("Signal must be a numpy array") + if signal.ndim != 2: + raise ValueError("Signal must be 2D (n_samples, n_channels)") + + if isinstance(stimulus, np.ndarray): + if stimulus.ndim != 2: + raise ValueError("Stimulus must be 2D (n_samples, n_features)") + if stimulus.shape[0] != signal.shape[0]: + raise ValueError("Stimulus and signal must have same number of samples") + elif isinstance(stimulus, dict): + for name, arr in stimulus.items(): + if not isinstance(arr, np.ndarray): + raise ValueError(f"Feature {name} must be a numpy array") + if arr.ndim != 1 and arr.ndim != 2: + raise ValueError(f"Feature {name} must be 1D or 2D") + if arr.shape[0] != signal.shape[0]: + raise ValueError(f"Feature {name} has wrong number of samples") + else: + raise ValueError("Stimulus must be numpy array or feature dictionary") + + def _create_lag_matrix( + self, + stimulus: np.ndarray, + lags: List[float], + sampling_rate: float + ) -> np.ndarray: + """Create a lag matrix for TRF estimation.""" + n_samples = stimulus.shape[0] + n_features = stimulus.shape[1] + n_lags = len(lags) + + lag_samples = [int(lag * sampling_rate) for lag in lags] + lag_matrix = np.zeros((n_samples, n_features * n_lags)) + + for i, lag in enumerate(lag_samples): + if lag >= 0: + lag_matrix[lag:, i * n_features:(i + 1) * n_features] = stimulus[:-lag, :] + else: + lag_abs = abs(lag) + lag_matrix[:lag_abs, i * n_features:(i + 1) * n_features] = stimulus[lag_abs:, :] + + return lag_matrix + + def _prepare_features( + self, + features: Dict[str, np.ndarray] + ) -> Tuple[np.ndarray, List[str]]: + """Prepare feature dictionary for TRF estimation.""" + sorted_names = sorted(features.keys()) + feature_matrix = np.column_stack([features[name] for name in sorted_names]) + return feature_matrix, sorted_names + + def fit( + self, + signal: np.ndarray, + stimulus: Union[np.ndarray, Dict[str, np.ndarray]], + sampling_rate: float = 1000.0 + ): + """Fit the TRF model.""" + self._validate_inputs(signal, stimulus) + + if isinstance(stimulus, dict): + if not self.config.use_features: + raise ValueError("Feature dictionary provided but use_features=False") + stimulus, self._feature_names = self._prepare_features(stimulus) + + self._lag_matrix = self._create_lag_matrix( + stimulus, + self.config.lags, + sampling_rate + ) + + X = np.column_stack([np.ones(len(self._lag_matrix)), self._lag_matrix]) + y = signal + + if self.config.solver == 'svd': + self._coefficients, self._intercept = self._svd_regress(X, y) + elif self.config.solver == 'lstsq': + self._coefficients, self._intercept = self._lstsq_regress(X, y) + else: + raise ValueError(f"Unknown solver: {self.config.solver}") + + self._fitted = True + return self + + def _svd_regress(self, X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Solve regression using SVD.""" + intercept_col = X[:, 0:1] + X_pred = X[:, 1:] + + X_centered = X_pred - np.mean(X_pred, axis=0) + y_centered = y - np.mean(y, axis=0) + + U, s, Vt = np.linalg.svd(X_centered, full_matrices=False) + s_reg = s / (s**2 + self.config.lambda_reg) + X_pinv = Vt.T @ np.diag(s_reg) @ U.T + + coefficients = X_pinv @ y_centered + intercept = np.mean(y, axis=0) - np.mean(X_pred, axis=0) @ coefficients + + return coefficients.T, intercept + + def _lstsq_regress(self, X: np.ndarray, y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: + """Solve regression using least squares.""" + intercept_col = X[:, 0:1] + X_pred = X[:, 1:] + + n_predictors = X_pred.shape[1] + X_reg = np.vstack([X_pred, np.sqrt(self.config.lambda_reg) * np.eye(n_predictors)]) + y_reg = np.vstack([y, np.zeros((n_predictors, y.shape[1]))]) + + coefficients, residuals, rank, s = np.linalg.lstsq(X_reg, y_reg, rcond=None) + intercept = np.mean(y, axis=0) - np.mean(X_pred, axis=0) @ coefficients.T + + return coefficients.T, intercept + + def predict(self, stimulus: Union[np.ndarray, Dict[str, np.ndarray]]) -> np.ndarray: + """Predict neural response from stimulus.""" + if not self._fitted: + raise RuntimeError("TRFEstimator not fitted. Call fit() first.") + + if isinstance(stimulus, dict): + if not self.config.use_features: + raise ValueError("Feature dictionary provided but use_features=False") + stimulus, _ = self._prepare_features(stimulus) + + lag_matrix = self._create_lag_matrix( + stimulus, + self.config.lags, + 1000.0 + ) + + X = np.column_stack([np.ones(len(lag_matrix)), lag_matrix]) + intercept_col = X[:, 0:1] + X_pred = X[:, 1:] + + prediction = X_pred @ self._coefficients.T + self._intercept + + return prediction + + def get_coefficients(self) -> Optional[np.ndarray]: + return self._coefficients + + def get_intercept(self) -> Optional[np.ndarray]: + return self._intercept + + def get_feature_names(self) -> List[str]: + return self._feature_names + + def score( + self, + signal: np.ndarray, + stimulus: Union[np.ndarray, Dict[str, np.ndarray]] + ) -> float: + """Calculate the R^2 score for the model.""" + prediction = self.predict(stimulus) + ss_res = np.sum(np.square(signal - prediction)) + ss_tot = np.sum(np.square(signal - np.mean(signal, axis=0))) + return 1 - ss_res / ss_tot \ No newline at end of file