src.data

  1import pandas as pd
  2import numpy as np
  3from sklearn.preprocessing import StandardScaler, LabelEncoder
  4from sklearn.model_selection import train_test_split
  5from sklearn.impute import SimpleImputer
  6import os
  7import urllib.request
  8import torch
  9from torch.utils.data import Dataset
 10import tempfile
 11
 12class HepatitisDataset(Dataset):
 13    """
 14    Custom PyTorch Dataset for Hepatitis data.
 15    
 16    This dataset can be reused with different models and training approaches.
 17
 18    Parameters
 19    -----------
 20    X : np.ndarray or pd.DataFrame
 21        Feature matrix.
 22    y : np.ndarray or pd.Series    
 23        Target vector.
 24
 25    Attributes
 26    -----------
 27    X : torch.FloatTensor
 28        Feature matrix as a FloatTensor.
 29    y : torch.LongTensor
 30        Target vector as a LongTensor.
 31        
 32    Examples
 33    ---------
 34    >>> from src.data import HepatitisDataset
 35    >>> dataset = HepatitisDataset(X_train, y_train)
 36    >>> loader = DataLoader(dataset, batch_size=32, shuffle=True)
 37    """
 38
 39    def __init__(self, X: np.ndarray, y: np.ndarray):
 40        self.X = torch.FloatTensor(X)
 41        self.y = torch.LongTensor(y.values if hasattr(y, 'values') else y)
 42    
 43    def __len__(self):
 44        return len(self.X)
 45    
 46    def __getitem__(self, idx: int):
 47        return self.X[idx], self.y[idx]
 48
 49def download_dataset(target_path: str = 'data/raw/hepatitis_data.csv', demo: bool = False) -> bool:
 50    """
 51    Download the Hepatitis C dataset from a public source if not already present.
 52    
 53    Parameters
 54    ------------
 55    target_path : str
 56        Path where the dataset should be saved.
 57    demo : bool
 58        If True, makes tempdirs and tempfiles deletable after use.
 59        
 60    Returns
 61    ------------
 62    bool
 63        True if download was successful or file already exists, False otherwise.
 64        
 65    Examples
 66    ---------
 67    >>> download_dataset()
 68    True
 69    """
 70    if demo:
 71        target_path = os.path.join(tempfile.gettempdir(), 'hepatitis_data.csv')
 72        print(f"Using temporary dataset path: {target_path}")
 73
 74    else:
 75        if os.path.exists(target_path):
 76            print(f"Dataset already exists at: {target_path}")
 77            return True
 78
 79        # Create directory if it doesn't exist
 80        os.makedirs(os.path.dirname(target_path), exist_ok=True)
 81
 82    # URL to a reliable source - using a direct CSV link
 83    # This is the UCI ML Repository version of the dataset
 84    url = "https://archive.ics.uci.edu/ml/machine-learning-databases/00571/hcvdat0.csv"
 85    
 86    try:
 87        print(f"Downloading dataset from {url}...")
 88        urllib.request.urlretrieve(url, target_path)
 89        print(f"✅ Dataset downloaded successfully to: {target_path}")
 90        if demo:
 91            return target_path
 92        return True
 93    except Exception as e:
 94        print(f"❌ Error downloading dataset: {e}")
 95        print("\n📌 Alternative: Download manually from:")
 96        print("   https://www.kaggle.com/datasets/fedesoriano/hepatitis-c-dataset")
 97        print(f"   and place it in: {target_path}")
 98        return False
 99
100def load_raw_data(filepath: str ='data/raw/hepatitis_data.csv', demo: bool = False) -> pd.DataFrame:
101    '''
102    Load raw data from a CSV file. If the file doesn't exist, attempts to download it automatically.
103    
104    Parameters
105    ------------
106    filepath : str
107        Path to the CSV file to be loaded.
108    demo : bool
109        If True, uses a temporary file path for demo purposes.
110
111    Returns
112    ------------
113    pd.DataFrame
114        Loaded dataset as a pandas DataFrame.   
115
116    Examples
117    ---------
118    >>> df = load_raw_data()
119    >>> df.head()
120    '''
121
122    try:
123        df = pd.read_csv(filepath)
124        # Rename the unnamed index column to Patient ID
125        if 'Unnamed: 0' in df.columns:
126            df = df.rename(columns={'Unnamed: 0': 'Patient ID'})
127        print(f"Dataset loaded successfully: {df.shape}")
128        return df
129    except FileNotFoundError:
130        print(f"File not found: {filepath}")
131        print("Attempting to download dataset automatically...")
132        
133        if download_dataset(filepath):
134            # Try loading again after download
135            try:
136                df = pd.read_csv(filepath)
137                if 'Unnamed: 0' in df.columns:
138                    df = df.rename(columns={'Unnamed: 0': 'Patient ID'})
139                print(f"Dataset loaded successfully: {df.shape}")
140                return df
141            except Exception as e:
142                print(f"Error loading downloaded dataset: {e}")
143                return None
144        else:
145            print("Please download the dataset manually from Kaggle and place it in data/raw/")
146            return None
147
148def get_data_info(df):
149    if df is None:  
150        return None
151    
152    info = {
153        'shape': df.shape,
154        'columns': list(df.columns),
155        'missing_values': df.isnull().sum(),
156        'target_distribution': df['Category'].value_counts() if 'Category' in df.columns else None,
157        'data_types': df.dtypes
158    }
159    
160    return info
161
162def clean_data(df: pd.DataFrame) -> pd.DataFrame:
163    """
164    Clean and preprocess the dataset.
165    
166    Parameters
167    ------------
168    df : pd.DataFrame
169        Raw dataset to be cleaned.
170        
171    Returns
172    ------------
173    pd.DataFrame
174        Cleaned dataset with necessary transformations applied.
175
176    Examples
177    ---------
178    >>> cleaned_df = clean_data(df)
179    >>> cleaned_df.head()
180    """
181
182    if df is None:
183        return None
184    data = df.copy()
185
186    # Keep Patient ID column for identification
187    # Remove it only if needed for modeling
188    if 'Unnamed: 0' in data.columns:
189        data = data.drop('Unnamed: 0', axis=1)
190    
191    def simplify_category(category: str) -> int:
192        """
193        Simplify the category labels.
194
195        Parameters
196        ------------
197        category : str
198            Original category label.
199
200        Returns
201        ------------
202        int
203            Simplified category label. 0 for healthy, 1 for hepatitis C.
204
205        Examples
206        ---------
207        >>> simplify_category('0=Blood Donor')
208        0
209        >>> simplify_category('1=Hepatitis C')
210        1
211        """
212
213        if category in ['0=Blood Donor', '0s=suspect Blood Donor']:
214            return 0
215        else:
216            return 1
217    
218    data['target'] = data['Category'].apply(simplify_category)
219    
220    sex_encoder = LabelEncoder()
221    data['sex_encoded'] = sex_encoder.fit_transform(data['Sex'])
222    
223    print(f"Data cleaned successfully")
224    print(f"Healthy: {sum(data['target'] == 0)} samples")
225    print(f"Hepatitis C: {sum(data['target'] == 1)} samples")
226    
227    return data, sex_encoder
228
229def prepare_features(data: pd.DataFrame) -> pd.DataFrame:
230    """
231    Prepare features for modeling.
232
233    Parameters
234    ------------
235    data : pd.DataFrame
236        Cleaned dataset with necessary transformations applied.
237
238    Returns
239    ------------
240    pd.DataFrame
241        Feature matrix ready for modeling.
242
243    Examples
244    ---------
245    >>> prepared_features = prepare_features(cleaned_df)
246    >>> prepared_features.head()
247    """
248
249    if data is None:
250        return None
251    feature_columns = ['Age', 'ALB', 'ALP', 'ALT', 'AST', 'BIL', 'CHE', 'CHOL', 'CREA', 'GGT', 'PROT', 'sex_encoded']
252    
253    X = data[feature_columns]
254    y = data['target']
255
256    imputer = SimpleImputer(strategy='median')
257    X_imputed = pd.DataFrame(
258        imputer.fit_transform(X), 
259        columns=X.columns, 
260        index=X.index
261    )
262    
263    print(f"Features prepared: {X_imputed.shape}")
264    print(f"Missing values after imputation: {X_imputed.isnull().sum().sum()}")
265    
266    return X_imputed, y, imputer
267
268def split_and_scale_data(X: pd.DataFrame, y: pd.Series, test_size: float = 0.2, random_state: int = 42) -> tuple[np.ndarray, np.ndarray, pd.Series, pd.Series, StandardScaler]:
269    """
270    Split and scale the dataset.
271
272    Parameters
273    ------------
274    X : pd.DataFrame
275        Feature matrix.
276    y : pd.Series
277        Target vector.
278    test_size : float
279        Proportion of the dataset to include in the test split.
280    random_state : int
281        Random seed for reproducibility.
282
283    Returns
284    ------------
285    tuple
286        (X_train_scaled, X_test_scaled, y_train, y_test, scaler)
287    """
288
289    X_train, X_test, y_train, y_test = train_test_split(
290        X, y, test_size=test_size, random_state=random_state, stratify=y
291    )
292    scaler = StandardScaler()
293    X_train_scaled = scaler.fit_transform(X_train)
294    X_test_scaled = scaler.transform(X_test)
295    
296    print(f"✅ Data split and scaled:")
297    print(f"   Training set: {X_train_scaled.shape}")
298    print(f"   Test set: {X_test_scaled.shape}")
299    
300    return X_train_scaled, X_test_scaled, y_train, y_test, scaler