Change this to something like "Default is None, in which case a linear schedule from 0 to 1 will be used"

I don't think this is the right way to change this function to work for small molecules. Specifying "not water" will include solvent ions, which we definitely don't want to include as part of the REST region.

We first need to determine what the desired behavior is. The current behavior of this function for small molecules is to automatically include all small molecule atoms in the REST region, and also include protein atoms within X angstroms (where X is user specified) of the small molecule. I think we discussed that you would run some experiments testing this approach, though this will require finding transformations where REST is known to improve sampling / speed up convergence.

One worry I have with this approach is that including all small molecule atoms in the REST region is not the best practices approach (Schrodinger papers typically choose the alchemically-changing (aka unique old and new) atoms and some nearby protein atoms) and there may be risk of small molecule unbinding. Therefore, I think the desired approach should be to modify these lines:

        # Retrieve the residue index of the residue that is being alchemified based on the first unique old atom
        atom_index = list(self._atom_classes['unique_old_atoms'])[0]
        hybrid_topology_atoms = list(self._hybrid_topology.atoms)
        alchemical_residue_index = hybrid_topology_atoms[atom_index].residue.index
        
       # Retrieve indices of all atoms in the alchemical residue
        hybrid_topology_residues = list(self._hybrid_topology.residues)
        mutated_res = hybrid_topology_residues[alchemical_residue_index]
        query_indices = [atom.index for atom in mutated_res.atoms]
        _logger.info(f"Generating rest_region with query indices: {query_indices}")

such that for small molecule transformations only, query_indices is all unique old/new atoms/core atoms. This may requiring introducing an argument to the function like is_small_molecule that tells the function whether its dealing with a small molecule transformation.

One of the things that we noticed is that the solute_atoms after the selection is empty for solvent phase. But we probably want to have the option to do REST scaling in solvent phase as well, such that the neighboring part is not needed, but the molecule atoms are still scaled according to the REST protocol.

Notes from discussion:

• As noted above, there should be three code paths: 1) protein mutation, 2) small molecule complex phase, 3) small molecule solvent phase
• Double check that "is protein" mdtraj selection only contains protein atoms (and does not interpret small molecule atoms as protein atoms). If indeed "is protein" only corresponds to protein atoms, the way that the function is written will fail to include small molecule atoms in the rest_atoms_all. Will need to make sure the relevant small molecule atoms are included in this.