fixed converted and summary

Author: antonkulaga

src/cell2sentence4longevity/explore.py

@@ -86,7 +86,7 @@ def _build_summary_expressions(
         exprs.extend([
             pl.col('age_years').min().alias('age_years_min'),
             pl.col('age_years').max().alias('age_years_max'),
-            pl.col('age_years').mean().alias('age_years_mean'),
+            pl.col('age_years').mean().round(2).alias('age_years_mean'),
             pl.col('age_years').is_not_null().sum().alias('cells_with_age_years')
         ])
     else:
@@ -101,7 +101,7 @@ def _build_summary_expressions(
         exprs.extend([
             pl.col('age_months').min().alias('age_months_min'),
             pl.col('age_months').max().alias('age_months_max'),
-            pl.col('age_months').mean().alias('age_months_mean'),
+            pl.col('age_months').mean().round(2).alias('age_months_mean'),
             pl.col('age_months').is_not_null().sum().alias('cells_with_age_months')
         ])
     else:
@@ -305,6 +305,37 @@ def _add_categorical_summaries(
     return summary_df
 
 
+def _mask_age_months_for_non_mouse(summary_df: pl.DataFrame) -> pl.DataFrame:
+    """Ensure age-in-months summary columns are only populated for mouse datasets.
+    
+    For rows where organism is not 'Mus musculus', all age_months-related columns
+    are set to null while keeping the columns present in the schema.
+    """
+    if 'organism' not in summary_df.columns:
+        return summary_df
+
+    age_months_columns = [
+        'age_months_min',
+        'age_months_max',
+        'age_months_mean',
+        'cells_with_age_months',
+        'unique_ages_months',
+    ]
+    existing_age_months_columns = [
+        column for column in age_months_columns if column in summary_df.columns
+    ]
+    if not existing_age_months_columns:
+        return summary_df
+
+    return summary_df.with_columns([
+        pl.when(pl.col('organism') == 'Mus musculus')
+        .then(pl.col(column))
+        .otherwise(None)
+        .alias(column)
+        for column in existing_age_months_columns
+    ])
+
+
 def _coerce_all_null_object_columns(
     df: pl.DataFrame,
     obs_schema: dict[str, pl.DataType]
@@ -746,6 +777,7 @@ def extract_fields_from_h5ad(
 
                 # Add categorical summaries (tissues, cell types, etc.)
                 summary_df = _add_categorical_summaries(summary_df, summary_lazy, summary_schema)
+                summary_df = _mask_age_months_for_non_mouse(summary_df)
 
                 # Log statistics
                 if 'dataset_id' in summary_df.columns:
@@ -1129,6 +1161,7 @@ def batch(
 
                             # Add categorical summaries (tissues, cell types, etc.)
                             summary_df = _add_categorical_summaries(summary_df, meta_lazy, meta_schema)
+                            summary_df = _mask_age_months_for_non_mouse(summary_df)
 
                             # Write summary to temp file
                             summary_file = summary_temp_dir / f"summary_{idx:04d}.parquet"
@@ -1173,6 +1206,7 @@ def batch(
                         for organism in organisms:
                             if organism and organism != 'unknown':
                                 organism_summaries = combined.filter(pl.col('organism') == organism)
+                                organism_summaries = _mask_age_months_for_non_mouse(organism_summaries)
 
                                 # Create filename: replace spaces with underscores, lowercase
                                 safe_organism = organism.lower().replace(' ', '_').replace('-', '_')

src/cell2sentence4longevity/preprocessing/h5ad_converter.py

@@ -721,17 +721,18 @@ def convert_h5ad_to_parquet(
                 # Process a batch
                 batch_end = min(cell_idx + batch_size, n_cells)
 
-                # Create cell sentences for batch without materializing a large dense block
-                # This keeps memory bounded even for very wide matrices or large h5ad files.
+                # Create cell sentences for batch using a dense block slice of adata.X.
+                # This is significantly faster than per-row indexing while keeping
+                # memory bounded by the configured batch_size.
                 batch_sentence_results: list[dict[str, str]] = []
-                for row_idx in range(cell_idx, batch_end):
-                    cell_row = adata.X[row_idx]
-                    # For sparse matrices, getrow().toarray() returns a small dense vector
-                    if hasattr(cell_row, "toarray"):
-                        cell_expr = cell_row.toarray().ravel()
-                    else:
-                        # Dense backend: rely on numpy array view
-                        cell_expr = np.asarray(cell_row).ravel()
+                cell_block = adata.X[cell_idx:batch_end]
+                if hasattr(cell_block, "toarray"):
+                    cell_block = cell_block.toarray()
+                else:
+                    cell_block = np.asarray(cell_block)
+
+                for local_idx in range(cell_block.shape[0]):
+                    cell_expr = np.asarray(cell_block[local_idx]).ravel()
                     batch_sentence_results.append(
                         create_cell_sentence(cell_expr, gene_symbols, top_genes)
                     )
@@ -835,7 +836,7 @@ def convert_h5ad_to_parquet(
 
             n_chunks = chunk_idx
         else:
-            # Row-based chunking (original logic)
+            # Row-based chunking (original logic, optimized to use block slices of adata.X)
             n_chunks = (n_cells + chunk_size - 1) // chunk_size
             action.log(message_type="using_row_based_chunking", chunk_size=chunk_size, n_chunks=n_chunks)
 
@@ -844,15 +845,18 @@ def convert_h5ad_to_parquet(
                     start_idx = chunk_idx * chunk_size
                     end_idx = min(start_idx + chunk_size, n_cells)
 
-                    # Create cell sentences for each cell in chunk without creating a large dense block.
-                    # This keeps memory bounded even for very wide matrices or extremely large h5ad files.
+                    # Create cell sentences for each cell in chunk using a dense block
+                    # slice of adata.X instead of per-row indexing. This is much faster
+                    # while still keeping memory bounded by chunk_size.
                     chunk_sentence_results: list[dict[str, str]] = []
-                    for row_idx in range(start_idx, end_idx):
-                        cell_row = adata.X[row_idx]
-                        if hasattr(cell_row, "toarray"):
-                            cell_expr = cell_row.toarray().ravel()
-                        else:
-                            cell_expr = np.asarray(cell_row).ravel()
+                    cell_block = adata.X[start_idx:end_idx]
+                    if hasattr(cell_block, "toarray"):
+                        cell_block = cell_block.toarray()
+                    else:
+                        cell_block = np.asarray(cell_block)
+
+                    for local_idx in range(cell_block.shape[0]):
+                        cell_expr = np.asarray(cell_block[local_idx]).ravel()
                         chunk_sentence_results.append(
                             create_cell_sentence(cell_expr, gene_symbols, top_genes)
                         )
@@ -1216,15 +1220,18 @@ def convert_h5ad_to_train_test(
                 start_idx = chunk_idx * chunk_size
                 end_idx = min(start_idx + chunk_size, n_cells)
 
-                # Create cell sentences for each cell in chunk without materializing a large dense block.
-                # This keeps memory bounded even for very wide matrices or very large h5ad files.
+                # Create cell sentences for each cell in chunk using a dense block slice
+                # of adata.X instead of per-row indexing. This is much faster while
+                # keeping memory bounded by chunk_size.
                 sentence_results: list[dict[str, str]] = []
-                for row_idx in range(start_idx, end_idx):
-                    cell_row = adata.X[row_idx]
-                    if hasattr(cell_row, "toarray"):
-                        cell_expr = cell_row.toarray().ravel()
-                    else:
-                        cell_expr = np.asarray(cell_row).ravel()
+                cell_block = adata.X[start_idx:end_idx]
+                if hasattr(cell_block, "toarray"):
+                    cell_block = cell_block.toarray()
+                else:
+                    cell_block = np.asarray(cell_block)
+
+                for local_idx in range(cell_block.shape[0]):
+                    cell_expr = np.asarray(cell_block[local_idx]).ravel()
                     sentence_results.append(
                         create_cell_sentence(
                             cell_expr,