CategoricalSampler support

[selmanozleyen]

Description of feature

Hi @ilan-gold ,

After having the Sampler API we can finally support CategoricalSampler. I have an idea how this is going to look like.

Here are some of the decisions I made going forward, I'd appreciate any comments on them before I start:

• I don't want the sampler to take the obs dataframe at this level so I will just assume the datasets are sorted and the boundaries for the categories are known, higher level API can do these and pass in the boundaries, this way we also don't deal with pandas categoricals. The core algorithm doesn't also deal with it so it makes sense to me
• For my project w cellflow I need an config that ensures that
  • each batch is from same group, but one load request can have different groups (loadrequest1 can have = batch from grp A and batch from grp B but not mixed)
  • or a more natural option for the implementation would be each load request would have one group only but the above option should exist

Since we first sample the sequence of groups to sample from each batch like a,c,b,a,..~cat(a,b,c) it is going to be difficult to mix load requests with multiple workers without interprocess communitcation. I have one approach to tackle this: if each load request needs two batches then we actually sample pairs of categories per worker a,c~cat(a,b) then that worker pops the chunkids it has per it's chunk share from group a and b. Then it prepares it and yields the load request. But this will lead in a signal coming from the sampling scheme which isn't desired. Because for example within the same loadrequest only some pairs of rows can be seen. So there is this order signal of "when I see this row in this batch then that means I saw either this row before or after." Couldn't formulate it yet but it's there. This signal I think can only be gone when load request = batch without interproccess communication.

[ilan-gold]

with multiple workers without interprocess communitcation

Do you need workers though? I think we should benchmark first a non-worker implementation against what you have and hten we can worry about this problem. It should be easy to disallow workers in this sampler so people don't try it.

I don't want the sampler to take the obs dataframe at this level so I will just assume the datasets are sorted and the boundaries for the categories are known, higher level API can do these and pass in the boundaries, this way we also don't deal with pandas categoricals. The core algorithm doesn't also deal with it so it makes sense to me

Is there a use-case where you wouldn't have a pandas.Categorical? I suppose, but if most people will be using pandas.Categorical and then converting to slices, I think it makes sense to have at least a from_pandas method or something if it's a hugely common op. I definitely am not married to making it necessary though!

[selmanozleyen]

Do you need workers though? I think we should benchmark first a non-worker implementation against what you have and hten we can worry about this problem. It should be easy to disallow workers in this sampler so people don't try

No, I don't need it to run it. Depends on the performance difference but we might want to utilize the unused cores if there is no compression. We can have many solutions to this like gradient accumulation or sampling with replacement (this way each worker can be independent of each other when sampling). Just wanted to mention it because it is something that requires some thinking. For now, it makes sense to just go without workers because I feel like some of these problems will depend on how the pan-cellflow project goes.

Is there a use-case where you wouldn't have a pandas.Categorical? I suppose, but if most people will be using pandas.Categorical and then converting to slices, I think it makes sense to have at least a from_pandas method or something if it's a hugely common op.

Depends on which level and how you envision the API would look like. What would the users want to pass? like the CategoricalSampler.from_pandas(obs['categories']))?. Then it makes sense except the obs should be ready before add_dataset calls. Which is fine for me if it's for you.

[ilan-gold]

Depends on which level and how you envision the API would look like. What would the users want to pass? like the CategoricalSampler.from_pandas(obs['categories']))?. Then it makes sense except the obs should be ready before add_dataset calls. Which is fine for me if it's for you.

Right I think the user would have to concatenate all the on-disk obs in memory first so CategoricalSampler.from_pandas(obs_categorical_concatenated)) but yea other than that, this was what I was thinking. If everyone has to write a function to_slices that does to_slices(obs_categorical_concatenated), then I think we should do that.

[selmanozleyen]

if you see https://github.com/scverse/annbatch/blob/6e20323ed0ebfc6a4d2aaf37499c52a16af90a44/src/annbatch/samplers/_categorical_sampler.py I made an initial attempt. However this is now what pan-cellflow will use. What they want is basically stratified sampling which requires sampling with replacement. This way the concept of epoch won't also exist and we would have fixed number of iterations. But this will have very little in common with the ChunkSampler. So should I create a ChunkSamplerWithReplacement to have similar pattern or should I have the stratified sampler potentially not use ChunkSampler? Would like to discuss about this a bit before going further

[ilan-gold]

This way the concept of epoch won't also exist and we would have fixed number of iterations. But this will have very little in common with the ChunkSampler. So should I create a ChunkSamplerWithReplacement to have similar pattern or should I have the stratified sampler potentially not use ChunkSampler? Would like to discuss about this a bit before going furthe

Confused how to square this comment with

if you see https://github.com/scverse/annbatch/blob/6e20323ed0ebfc6a4d2aaf37499c52a16af90a44/src/annbatch/samplers/_categorical_sampler.py I made an initial attempt

Does this not work as intended?

What they want is basically stratified sampling which requires sampling with replacement.

They want sampling with replacement within a given category? Is there a difference between that and calling ChunkSampler._sample again? They want observations to be repeated across batches even?

[selmanozleyen]

They want sampling with replacement within a given category? Is there a difference between that and calling ChunkSampler._sample again? They want observations to be repeated across batches even?

Yes that's why, technically it should be possible to have repeated chunks observed. It is each batch c~Cat, x~c.

You know what never mind. I can of course ask for _sample again even though it's not mathematically the same at least the code will be cleaner I guess. Honestly I had sampling with replacement in mind to allow multi worker sampling. However I haven't had enough experience with annbatch and realistic use cases yet to say workers will be useful. Surely there are parameters and datasets they will be but I am not sure how realistic they would be. I really am not sure if a researcher will one day come to us and say my results are more accurate when chunk_size is very small but it's super slow! Or if they will one day need to write a compute heavy sampler. Should I create a separate issue: "When do we need torch workers"?

Does this not work as intended?

It does when each __iter__ needs to go through the whole dataset once.

[ilan-gold]

Should I create a separate issue: "When do we need torch workers"?

Yes. For now I wouldn't design anything around the use-case specifically.

I can of course ask for _sample again even though it's not mathematically the same at least the code will be cleaner I guess.

hmm yeah, this isn't ideal, but also I'm not clear how tight the "with replacement" has to be. Without something concrete it's tough to say. It's also possible that parametrizing ChunkSampler with an n_iters argument overrides the incoming n_obs from _sample. And it's also possible "sample with replacement" is easier done than said (i.e., just requires a change to the way we create chunk_indices and otherwise the logic is identical). Not being prescriptive at all! I think a "with replacement chunked sampler" could be cool if we get the API right!