perf: vectorize metadata DataFrame operations

Replaces row-wise pandas apply() usage in metadata/cohort preparation with vectorized numpy/pandas operations to reduce Python-level iteration.

Made-with: Cursor

Author: khushthecoder

.gitignore

@@ -3,6 +3,7 @@
 __pycache__
 *.pyc
 dist
+.venv/
 .coverage
 coverage.xml
 .ipynb_checkpoints/

malariagen_data/anoph/base.py

@@ -607,9 +607,9 @@ def _read_sample_sets_manifest(self, *, single_release: str):
             # Get today's date in ISO format
             today_date_iso = date.today().isoformat()
             # Add an "unrestricted_use" column, set to True if terms-of-use expiry date <= today's date.
-            df["unrestricted_use"] = df[terms_of_use_expiry_date_column].apply(
-                lambda d: True if pd.isna(d) else (d <= today_date_iso)
-            )
+            # Vectorized operation: True if NaN, else (d <= today_date_iso)
+            s = df[terms_of_use_expiry_date_column]
+            df["unrestricted_use"] = (s.isna() | (s <= today_date_iso))
             # Make the "unrestricted_use" column a nullable boolean, to allow missing data.
             df["unrestricted_use"] = df["unrestricted_use"].astype(pd.BooleanDtype())
 

malariagen_data/anoph/frq_base.py

@@ -43,40 +43,39 @@ def _prep_samples_for_cohort_grouping(
 
     # Add period column.
 
-    # Map supported period_by values to functions that return either the relevant pd.Period or pd.NaT per row.
-    period_by_funcs = {
-        "year": _make_sample_period_year,
-        "quarter": _make_sample_period_quarter,
-        "month": _make_sample_period_month,
+    # Map supported period_by values to vectorized functions that create Period arrays.
+    period_by_funcs_vectorized = {
+        "year": _make_sample_periods_year_vectorized,
+        "quarter": _make_sample_periods_quarter_vectorized,
+        "month": _make_sample_periods_month_vectorized,
     }
 
     # Get the matching function for the specified period_by value, or None.
-    period_by_func = period_by_funcs.get(period_by)
+    period_by_func_vectorized = period_by_funcs_vectorized.get(period_by)
 
     # If there were no matching functions for the specified period_by value...
-    if period_by_func is None:
+    if period_by_func_vectorized is None:
         # Raise a ValueError if the specified period_by value is not a column in the DataFrame.
         if period_by not in df_samples.columns:
             raise ValueError(
                 f"Invalid value for `period_by`: {period_by!r}. Either specify the name of an existing column "
                 "or a supported period: 'year', 'quarter', or 'month'."
             )
 
-        # Raise a ValueError if the specified period_by column does not contain instances pd.Period.
-        if (
-            not df_samples[period_by]
-            .apply(lambda value: pd.isnull(value) or isinstance(value, pd.Period))
-            .all()
-        ):
-            raise TypeError(
-                f"Invalid values in {period_by!r} column. Must be either pandas.Period or null."
-            )
+        # Validate the specified period_by column contains pandas Periods (or nulls).
+        s_period_by = df_samples[period_by]
+        if not pd.api.types.is_period_dtype(s_period_by.dtype):
+            non_null = s_period_by.dropna()
+            if len(non_null) > 0 and not non_null.map(type).eq(pd.Period).all():
+                raise TypeError(
+                    f"Invalid values in {period_by!r} column. Must be either pandas.Period or null."
+                )
 
         # Copy the specified period_by column to a new "period" column.
         df_samples["period"] = df_samples[period_by]
     else:
-        # Apply the matching period_by function to create a new "period" column.
-        df_samples["period"] = df_samples.apply(period_by_func, axis="columns")
+        # Use the vectorized period creation function.
+        df_samples["period"] = period_by_func_vectorized(df_samples)
 
     # Copy the specified area_by column to a new "area" column.
     df_samples["area"] = df_samples[area_by]
@@ -101,22 +100,35 @@ def _build_cohorts_from_sample_grouping(
     df_cohorts = df_cohorts.reset_index()
 
     # Add cohort helper variables.
-    cohort_period_start = df_cohorts["period"].apply(lambda v: v.start_time)
-    cohort_period_end = df_cohorts["period"].apply(lambda v: v.end_time)
-    df_cohorts["period_start"] = cohort_period_start
-    df_cohorts["period_end"] = cohort_period_end
+    # Vectorized extraction of period start/end times.
+    period = df_cohorts["period"]
+    if pd.api.types.is_period_dtype(period.dtype):
+        df_cohorts["period_start"] = period.dt.start_time
+        df_cohorts["period_end"] = period.dt.end_time
+    else:
+        # Fallback for object dtype Period values.
+        df_cohorts["period_start"] = period.map(
+            lambda v: v.start_time if pd.notna(v) else pd.NaT
+        )
+        df_cohorts["period_end"] = period.map(
+            lambda v: v.end_time if pd.notna(v) else pd.NaT
+        )
+    
     # Create a label that is similar to the cohort metadata,
     # although this won't be perfect.
+    # Vectorized string operations
     if taxon_by == frq_params.taxon_by_default:
-        df_cohorts["label"] = df_cohorts.apply(
-            lambda v: f"{v.area}_{v[taxon_by][:4]}_{v.period}", axis="columns"
-        )
+        # Default case: area_taxon_short_period
+        area_str = df_cohorts["area"].astype(str)
+        taxon_short = df_cohorts[taxon_by].astype(str).str.slice(0, 4)
+        period_str = df_cohorts["period"].astype(str)
+        df_cohorts["label"] = area_str + "_" + taxon_short + "_" + period_str
     else:
-        # Replace non-alphanumeric characters in the taxon with underscores.
-        df_cohorts["label"] = df_cohorts.apply(
-            lambda v: f"{v.area}_{re.sub(r'[^A-Za-z0-9]+', '_', str(v[taxon_by]))}_{v.period}",
-            axis="columns",
-        )
+        # Non-default case: replace non-alphanumeric characters with underscores
+        area_str = df_cohorts["area"].astype(str)
+        taxon_clean = df_cohorts[taxon_by].astype(str).str.replace(r"[^A-Za-z0-9]+", "_", regex=True)
+        period_str = df_cohorts["period"].astype(str)
+        df_cohorts["label"] = area_str + "_" + taxon_clean + "_" + period_str
 
     # Apply minimum cohort size.
     df_cohorts = df_cohorts.query(f"size >= {min_cohort_size}").reset_index(drop=True)
@@ -173,6 +185,50 @@ def _make_sample_period_year(row):
         return pd.NaT
 
 
+def _make_sample_periods_month_vectorized(df_samples):
+    year = df_samples["year"]
+    month = df_samples["month"]
+    valid = (year > 0) & (month > 0)
+
+    out = pd.Series(pd.NaT, index=df_samples.index, dtype="period[M]")
+    if valid.any():
+        out.loc[valid] = pd.PeriodIndex.from_fields(
+            year=year.loc[valid].to_numpy(),
+            month=month.loc[valid].to_numpy(),
+            freq="M",
+        )
+    return out
+
+
+def _make_sample_periods_quarter_vectorized(df_samples):
+    year = df_samples["year"]
+    month = df_samples["month"]
+    valid = (year > 0) & (month > 0)
+
+    out = pd.Series(pd.NaT, index=df_samples.index, dtype="period[Q-DEC]")
+    if valid.any():
+        out.loc[valid] = pd.PeriodIndex.from_fields(
+            year=year.loc[valid].to_numpy(),
+            month=month.loc[valid].to_numpy(),
+            freq="Q-DEC",
+        )
+    return out
+
+
+def _make_sample_periods_year_vectorized(df_samples):
+    year = df_samples["year"]
+    valid = year > 0
+
+    out = pd.Series(pd.NaT, index=df_samples.index, dtype="period[Y-DEC]")
+    if valid.any():
+        out.loc[valid] = pd.PeriodIndex.from_fields(
+            year=year.loc[valid].to_numpy(),
+            month=np.full(int(valid.sum()), 12, dtype="int64"),
+            freq="Y-DEC",
+        )
+    return out
+
+
 class AnophelesFrequencyAnalysis(AnophelesBase):
     def __init__(
         self,
@@ -263,14 +319,10 @@ def plot_frequencies_heatmap(
             index = list(index_names_as_list)
         df = df.reset_index().copy()
         if isinstance(index, list):
-            index_col = (
-                df[index]
-                .astype(str)
-                .apply(
-                    lambda row: ", ".join([o for o in row if o is not None]),
-                    axis="columns",
-                )
-            )
+            idx_vals = df[index].astype(str).to_numpy()
+            index_col = pd.Series(idx_vals[:, 0], index=df.index)
+            for j in range(1, idx_vals.shape[1]):
+                index_col = index_col + ", " + idx_vals[:, j]
         else:
             assert isinstance(index, str)
             index_col = df[index].astype(str)

malariagen_data/anoph/genome_features.py

@@ -446,29 +446,25 @@ def plot_genes(
 
             # Put gene pointers (▲ or ▼) in a new column, depending on the strand.
             # Except if the gene_label is null or an empty string, which should not be shown.
-            data["gene_pointer"] = data.apply(
-                lambda row: ("▼" if row["strand"] == "+" else "▲")
-                if row["gene_label"]
-                else "",
-                axis=1,
-            )
+            # Vectorized operation: use np.where for conditional logic
+            has_label = data["gene_label"].astype(bool)
+            pointer_symbol = np.where(data["strand"] == "+", "▼", "▲")
+            data["gene_pointer"] = np.where(has_label, pointer_symbol, "")
 
             # Put the pointer above or below the gene rectangle, depending on + or - strand.
             neg_strand_pointer_y = orig_mid_y_range - 1.1
             pos_strand_pointer_y = orig_mid_y_range + 1.1
-            data["pointer_y"] = data["strand"].apply(
-                lambda strand: pos_strand_pointer_y
-                if strand == "+"
-                else neg_strand_pointer_y
+            # Vectorized operation: use np.where instead of Series.apply
+            data["pointer_y"] = np.where(
+                data["strand"] == "+", pos_strand_pointer_y, neg_strand_pointer_y
             )
 
             # Put the label above or below the gene rectangle, depending on + or - strand.
             neg_strand_label_y = orig_mid_y_range - 1.25
             pos_strand_label_y = orig_mid_y_range + 1.3
-            data["label_y"] = data["strand"].apply(
-                lambda strand: pos_strand_label_y
-                if strand == "+"
-                else neg_strand_label_y
+            # Vectorized operation: use np.where instead of Series.apply
+            data["label_y"] = np.where(
+                data["strand"] == "+", pos_strand_label_y, neg_strand_label_y
             )
 
             # Get the data as a ColumnDataSource.

malariagen_data/anoph/sample_metadata.py

@@ -181,9 +181,9 @@ def _parse_general_metadata(
             df["release"] = release
 
             # Derive a quarter column from month.
-            df["quarter"] = df.apply(
-                lambda row: ((row.month - 1) // 3) + 1 if row.month > 0 else -1,
-                axis="columns",
+            # Vectorized operation: quarter = ((month - 1) // 3) + 1 if month > 0 else -1
+            df["quarter"] = np.where(
+                df["month"] > 0, ((df["month"] - 1) // 3) + 1, -1
             )
 
             # Add study columns.