audit-labs/audit-tools

 1# Import packages
 2import math
 3
 4import pandas as pd
 5
 6# Load data
 7df = pd.read_csv("FILENAME_GOES_HERE.csv")
 8
 9# ALTERNATIVE: If you use Excel, use this instead. Supports xls, xlsx, xlsm,
10# xlsb, odf, ods and odt file extensions.
11# df = pd.read_excel("FILENAME_GOES_HERE.xlsx")
12
13# Print totals prior to sampling
14print("Dataframe size (rows, columns):", df.shape)
15
16# User-defined parameters
17SAMPLE_SIZE = 25
18STRATIFY_COLUMN = "Category"  # <- Change this to your column name
19
20# Define stratum proportions (as fractions)
21# Example: if you have categories A, B, and C
22stratum_proportions = {"A": 0.4, "B": 0.4, "C": 0.2}
23
24# Validate proportions sum to 1
25if not math.isclose(sum(stratum_proportions.values()), 1.0):
26    raise ValueError("Stratum proportions must sum to 1.")
27
28# Check that all strata exist in the data
29missing_strata = set(stratum_proportions.keys()) - set(df[STRATIFY_COLUMN].unique())
30if missing_strata:
31    raise ValueError(
32        f"Strata {missing_strata} not found in column '{STRATIFY_COLUMN}'."
33    )
34
35# Perform stratified sampling
36samples = []
37for stratum, proportion in stratum_proportions.items():
38    stratum_df = df[df[STRATIFY_COLUMN] == stratum]
39    n_samples = math.floor(SAMPLE_SIZE * proportion)
40    if n_samples > len(stratum_df):
41        raise ValueError(
42            f"Not enough data in stratum '{stratum}' to sample {n_samples} rows."
43        )
44    stratum_sample = stratum_df.sample(n=n_samples, random_state=42)
45    samples.append(stratum_sample)
46
47# Combine all stratum samples into one DataFrame
48final_sample = pd.concat(samples).reset_index()
49
50# If needed, randomly sample extra rows to fill any rounding gap
51current_sample_size = len(final_sample)
52if current_sample_size < SAMPLE_SIZE:
53    remaining = SAMPLE_SIZE - current_sample_size
54    remaining_sample = df.sample(n=remaining, random_state=42)
55    final_sample = pd.concat([final_sample, remaining_sample])
56
57# Print sample results
58print("Final sample size:", final_sample.shape[0])
59print("Sample breakdown by stratum:\n", final_sample[STRATIFY_COLUMN].value_counts())
60print("\nSample:\n", final_sample)
61
62# Optionally, save the sample to a new CSV
63# final_sample.to_csv("sample_output.csv", index=False)