LSTM System for Frequency Extraction from Mixed Signals

A PyTorch implementation of an LSTM-based frequency filter that extracts pure sinusoidal components from noisy mixed signals through conditional regression.

Project Overview

This project demonstrates LSTM's ability to act as a frequency-selective filter by learning to isolate specific frequency components from a mixed, noisy signal. The system learns to extract pure sinusoids despite random amplitude and phase variations at each time sample.

Key Features

Frequency Extraction: Isolates 4 different frequency components (1Hz, 3Hz, 5Hz, 7Hz)
Noise Robustness: Handles amplitude variations (±20%) and random phase shifts
Conditional Regression: Uses one-hot encoding to select target frequency
State Management: Demonstrates LSTM with sequence length L=1 and proper state reset
Generalization: Tests on completely different noise patterns (separate random seeds)

Problem Statement

Given a mixed signal S(t) composed of 4 sinusoidal frequencies with varying noise:

S(t) = (1/4) Σ A_i(t) · sin(2π·f_i·t + φ_i(t))

where:
    A_i(t) ~ Uniform(0.8, 1.2)  (varies at each sample)
    φ_i(t) ~ Uniform(0, 2π)      (varies at each sample)
    f_i ∈ {1, 3, 5, 7} Hz

Goal: Train an LSTM to extract pure frequency components:

Target_i(t) = sin(2π·f_i·t)

Quick Start

Prerequisites

Python 3.10+
UV package manager

Installation

Clone the repository:

git clone https://github.com/TalHibner/llmcourse-hw2-lstm.git
cd llmcourse-hw2-lstm

Install UV (if not already installed):

curl -LsSf https://astral.sh/uv/install.sh | sh

Install dependencies:

uv sync

Running the Project

Option 1: Jupyter Notebook (Recommended)

uv run jupyter notebook notebooks/lstm_frequency_extraction.ipynb

Then run all cells to:

Generate datasets
Train the LSTM model
Evaluate performance
Generate visualizations

Option 2: Python Scripts

# Generate datasets
uv run python -c "from src.data_generator import create_datasets; create_datasets()"

# Train model
uv run python -c "from src.trainer import train_model; train_model()"

# Evaluate
uv run python -c "from src.evaluator import evaluate_and_visualize; evaluate_and_visualize()"

Project Structure

llmcourse-hw2-lstm/
│
├── src/                        # Source code modules
│   ├── __init__.py
│   ├── data_generator.py       # Signal generation and dataset creation
│   ├── model.py                # LSTM model definition
│   ├── trainer.py              # Training loop and state management
│   ├── evaluator.py            # Metrics and evaluation
│   └── visualizer.py           # Plotting functions
│
├── notebooks/                  # Jupyter notebooks
│   └── lstm_frequency_extraction.ipynb
│
├── data/                       # Generated datasets
│   ├── train_dataset.npz       # Training set (seed #1)
│   └── test_dataset.npz        # Test set (seed #2)
│
├── models/                     # Trained models
│   └── best_model.pth
│
├── results/                    # Results and visualizations
│   ├── metrics.json
│   └── plots/
│       ├── training_curve.png
│       ├── freq_comparison.png
│       └── all_frequencies.png
│
├── prompts/                    # Claude CLI documentation
│   └── *.png                   # Development screenshots
│
├── tests/                      # Comprehensive test suite (121 tests)
│   ├── test_data_generator.py  # Dataset generation tests (28 tests)
│   ├── test_model.py           # LSTM model tests (27 tests)
│   ├── test_trainer.py         # Training loop tests (19 tests)
│   ├── test_evaluator.py       # Evaluation tests (26 tests)
│   └── test_visualizer.py      # Visualization tests (21 tests)
│
├── PRD.md                      # Product Requirements Document
├── DESIGN.md                   # Technical Design Document
├── TASKS.md                    # Implementation Task Breakdown
├── README.md                   # This file
├── pyproject.toml              # UV package configuration
├── .gitignore
└── L2-homework.pdf             # Assignment specification

Technical Details

Dataset Specifications

Parameter	Value
Frequencies	f1=1Hz, f2=3Hz, f3=5Hz, f4=7Hz
Time Domain	0-10 seconds
Sampling Rate	1000 Hz
Total Samples	10,000 per frequency
Dataset Size	40,000 rows (10,000 × 4 frequencies)

Model Architecture

Input: [S[t], C1, C2, C3, C4]  (5 features)
   ↓
LSTM Layer (64 hidden units)
   ↓
Fully Connected Layer
   ↓
Output: Target_i[t]  (scalar)

Critical Implementation Detail: Internal state (h_t, c_t) is reset between consecutive samples (sequence length L=1), forcing the LSTM to learn frequency patterns through internal memory management alone.

Training Configuration

Loss Function: Mean Squared Error (MSE)
Optimizer: Adam (learning rate = 0.001)
Batch Size: 64
Epochs: 50 (with early stopping)
State Management: Reset hidden state for each batch

Computational Requirements

Hardware Requirements

Requirement	Minimum	Recommended
CPU	Intel i5 / Ryzen 3	Intel i7 / Ryzen 5+
RAM	8 GB	16 GB
GPU	None (CPU-only)	NVIDIA GTX 1060+ / RTX 2060+
Storage	200 MB	500 MB
Python	3.10+	3.11+

Performance Benchmarks

Hardware Configuration	Training Time	Memory Usage	Cost Estimate
CPU (Intel i7-10700)	4-5 minutes	2-3 GB RAM	$0 (local)
GPU (NVIDIA RTX 3060)	45-60 seconds	1.5 GB VRAM	$0 (local)
Cloud CPU (AWS t3.large)	~5 minutes	4 GB RAM	~$0.02/run
Cloud GPU (AWS g4dn.xlarge)	~1 minute	8 GB VRAM	~$0.05/run

Resource Usage Breakdown

Dataset Generation: < 10 seconds, ~50 MB disk space
Model Size: ~127 KB (64 hidden units LSTM)
Training: 50 epochs × ~5 seconds/epoch = ~4-5 minutes (CPU)
Evaluation: < 5 seconds
Visualization: < 10 seconds
Total Storage: ~100-150 MB (data + models + results + plots)

Cost Considerations

Local Development:

Zero API costs (no external services)
Electricity cost: negligible (~$0.001 per training run)
One-time setup: ~15-20 minutes

Cloud Deployment (Optional):

AWS EC2 t3.large: ~$0.0832/hour → ~$0.007/run (5 min)
AWS EC2 g4dn.xlarge (GPU): ~$0.526/hour → ~$0.009/run (1 min)
Data storage (S3): ~$0.023/GB/month → negligible for this project

Development Time:

Initial implementation: ~10-15 hours
Testing and documentation: ~5-8 hours
Total effort: ~15-23 hours

Optimization Strategies

Automatic GPU Detection: Code automatically uses GPU if available
```
device = 'cuda' if torch.cuda.is_available() else 'cpu'
```
Efficient Batch Processing: Batch size of 64 balances speed and memory
Early Stopping: Prevents unnecessary epochs (saves ~20-30% time)
In-Memory Data Loading: All data fits in RAM (40K samples × 5 features)
Lazy Visualization: Plots generated only when needed

Results

Performance Metrics

Expected results after training:

MSE_train: < 0.05 (on training set with seed #1)
MSE_test: ≈ MSE_train (on test set with seed #2)
Generalization: Good if |MSE_test - MSE_train| < 0.01

Visualizations

The project generates two main visualizations:

Single Frequency Comparison: Shows the extraction of one frequency (e.g., f2) with:
- Mixed noisy signal (background)
- Pure target (ground truth)
- LSTM predictions
All Frequencies: Four subplots showing extraction for each frequency

Key Implementation Highlights

1. State Reset Mechanism

for batch in dataloader:
    # CRITICAL: Reset state for each batch
    hidden = None  # PyTorch automatically initializes to zeros

    outputs, hidden = model(inputs, hidden)
    loss = criterion(outputs, targets)

This ensures the LSTM learns frequency patterns without relying on sequential dependencies between samples.

2. Noise Variation

# Generate random variations FOR EACH sample
for t in time_array:
    A_i[t] = np.random.uniform(0.8, 1.2)  # Amplitude varies
    phi_i[t] = np.random.uniform(0, 2*np.pi)  # Phase varies
    signal[t] = A_i[t] * np.sin(2*np.pi*f_i*t + phi_i[t])

3. Conditional Frequency Selection

# One-hot encoding for frequency selection
C = np.zeros(4)
C[frequency_index] = 1  # Select which frequency to extract
input = [S[t], *C]  # Concatenate signal and selection vector

Documentation

PRD.md: Complete product requirements and specifications
DESIGN.md: Technical architecture and design decisions
TASKS.md: Detailed implementation task breakdown

Testing

The project includes a comprehensive test suite with 121 tests covering all major functionality.

Running Tests

Run all tests:

uv run pytest tests/ -v

Run specific test file:

uv run pytest tests/test_model.py -v

Run with coverage report:

# Generate HTML coverage report
uv run pytest tests/ --cov=src --cov-report=html --cov-report=term

# View the report
open htmlcov/index.html  # macOS
xdg-open htmlcov/index.html  # Linux
start htmlcov/index.html  # Windows

Generate coverage badge:

# Terminal-only report with percentage
uv run pytest tests/ --cov=src --cov-report=term-missing

Expected Coverage: > 85% overall (targeting 90%+)

Test Coverage

Test File	Tests	Coverage
test_data_generator.py	28	Signal generation, dataset creation, file I/O, seed reproducibility
test_model.py	27	LSTM architecture, forward pass, hidden state management, dataset/dataloader
test_trainer.py	19	Training/validation epochs, model training, checkpoints, early stopping
test_evaluator.py	26	Model evaluation, metrics computation, predictions saving
test_visualizer.py	21	All plotting functions, file creation, visualization workflow

Key Test Areas

✅ Model Architecture: LSTM initialization, forward propagation, state reset
✅ Data Generation: Noisy signals, mixed signals, ground truth, one-hot encoding
✅ Training Loop: Weight updates, gradient handling, early stopping, checkpointing
✅ Evaluation: MSE calculation, frequency-separated results, generalization metrics
✅ Visualization: Training curves, frequency comparisons, metrics plots
✅ File I/O: Dataset saving/loading, model checkpoints, predictions export
✅ Edge Cases: Empty inputs, shape mismatches, zero division handling

Assignment Requirements Checklist

✅ Create 2 datasets (train with seed #1, test with seed #2)
✅ Implement LSTM network with proper architecture
✅ Reset internal state between samples (L=1)
✅ Train model and achieve low MSE
✅ Evaluate generalization on test set
✅ Generate required visualizations
✅ Demonstrate noise robustness
✅ Comprehensive test suite with 121 tests

Future Enhancements

Potential improvements and extensions:

Variable Sequence Length: Experiment with L > 1 (sliding window)
Comparison Study: Compare with other architectures (GRU, Transformer)
Additional Noise Models: Test with different noise distributions
Real-time Processing: Adapt for streaming data
Multiple Frequency Extraction: Extract multiple frequencies simultaneously
Hyperparameter Tuning: Systematic search for optimal parameters

Dependencies

Main dependencies (see pyproject.toml for complete list):

torch >= 2.0.0: PyTorch for LSTM implementation
numpy >= 1.24.0: Numerical computations
matplotlib >= 3.7.0: Visualization
jupyter >= 1.0.0: Interactive notebook environment

Development dependencies:

pytest >= 7.0.0: Testing framework
black >= 23.0.0: Code formatting
ruff >= 0.1.0: Linting

Author

Implementation for M.Sc. Assignment Original Assignment: Dr. Segal Yoram November 2025

License

References

Assignment PDF: L2-homework.pdf
PyTorch LSTM Documentation: https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html
Understanding LSTM Networks: http://colah.github.io/posts/2015-08-Understanding-LSTMs/

Troubleshooting

Common Issues

Issue: Dependencies fail to install

# Solution: Update UV
uv self update
uv sync

Issue: CUDA out of memory

# Solution: Reduce batch size in config
# Edit src/trainer.py: batch_size = 32

Issue: Training loss not decreasing

# Solution: Check state reset implementation
# Verify hidden=None in training loop

Issue: Poor generalization (MSE_test >> MSE_train)

# Solution:
# 1. Verify different seeds for train/test
# 2. Increase training epochs
# 3. Add regularization (dropout)

Contact

For questions or issues:

Open an issue on GitHub: https://github.com/TalHibner/llmcourse-hw2-lstm/issues
Check TASKS.md for implementation details
Review DESIGN.md for architecture explanations

Status: ✅ Implementation Complete Last Updated: November 2025 Test Coverage: 121 tests across 5 test files

Name		Name	Last commit message	Last commit date
data		data
models		models
notebooks		notebooks
prompts		prompts
results		results
src		src
tests		tests
.gitignore		.gitignore
DESIGN.md		DESIGN.md
L2-homework.pdf		L2-homework.pdf
PRD.md		PRD.md
README.md		README.md
TASKS.md		TASKS.md
pyproject.toml		pyproject.toml
run_experiment.py		run_experiment.py

Latest commit

History

Folders and files

Repository files navigation