max_atoms_to_try either not working as intended or misleading for memory estimation #276
Description
Hi everyone,
I am currently implementing the InFlightAutoBatcher and noticed a malfunction or at least a misleading definition of the max_atoms_to_try parameter when initializing the Autobatcher.
I am particularly looking at runtime optimization for cases where the number of structures that are optimized are so small, that the memory will not be filled anyway (I know the batcher doesn't really make sense then, but these cases might happen in dynamic environments from time to time). Let's look at an extreme case for clarity: If my SimState only contains 3 geometries with 20 atoms each, it doesnt really make sense to test the GPU memory to the limit or using the default max_atoms_to_try parameter of 500k. Instead, I would dynamically set the parameter "max_atoms_to_try" to the minimum min(simstate.n_atoms, 500_000).
So far so good, when I did that, I noticed that first of all, this parameter is not passed to the BinningAutobatcher in runner.py optimize function, where _chunked_apply is called to initialize the state. This should be easily fixed by adding the corresponding kwarg.
However, more importantly, in the autobatching.py file, determine_max_batch_size function, I noticed that the while loop that generates the list of system sizes to memory-test is effectively comparing the number of batches with the number of atoms. This causes the function not to work as intended, effectively rendering the max_atoms_to_test parameter useless. As a possible solution, maybe this should be harmonized, then it should be easy to fix, I guess?
Thank you!