[Python] How to resample a dataset to achieve a uniform distribution

I have a dataset with a schema like:

df = pl.DataFrame(
{
"target": [
[1.0, 1.0, 0.0],
[1.0, 1.0, 0.1],
[1.0, 1.0, 0.2],
[1.0, 1.0, 0.8],
[1.0, 1.0, 0.9],
[1.0, 1.0, 1.0],
],
"feature": ["a", "b", "c", "d", "e", "f"],
},
schema={
"target": pl.Array(pl.Float32, 3),
"feature": pl.String,
},
)


If I make a histogram of the target-z values it looks like:

I want to resample the data so its flat along z.

I managed to do it in a hacky-many-steps way (also very slow). I was wondering if people could suggest a cleaner (and more efficient) way?

What I am doing is:

1. Find the bin edges of said histogram:

bins = 2 # Use e.g. 100 or larger in reality
z = df.select(z=pl.col("target").arr.get(2))
z_min = z.min()
z_max = z.max()
breaks = np.linspace(z_min, z_max, num=bins+1)

1. Find how many counts are in the bin with the fewest counts:

counts = (
df.with_columns(bin=pl.col("target").arr.get(2).cut(breaks))
.with_columns(counter=pl.int_range(pl.len()).over("bin"))
.group_by("bin")
.agg(pl.col("counter").max())
.filter(pl.col("counter") > 0) # <- Nasty way of filtering the (-inf, min] bin
.select(pl.col("counter").min())
).item()

1. Choose only "count" elements on each bin:

df = (
df.with_columns(bin=pl.col("target").arr.get(2).cut(breaks))
.with_columns(counter=pl.int_range(pl.len()).over("bin"))
.filter(pl.col("counter") <= counts)
.select("target", "feature")
)


This gives me:

Do people have any suggestions?

Continue reading...