TreeGrid - Disable collapse if child is selected

Any idea how to disable parent/child collapsing if a child is selected? If a child is selected, I want the parent node to be Collapse Allowed false.
I’ve tied to setItemCollapseAllowedProvider. however it doesn’t seem to be triggered after grid selection change event.

// Disable parent collapse if child selected
portfolioTree.setItemCollapseAllowedProvider(item->{
// Is root item and child is not selected
return portfolioTree.getTreeData().getRootItems().contains(item) &&
Collections.disjoint(portfolioTree.getSelectedItems(), portfolioTree.getTreeData().getChildren(item));
});