Enable user-specified spin flips to canonical MC

[juliayang]

Is your feature request related to a problem? Please describe.
There are some examples which require modification to the way spin flips are currently implemented:

1. Species which can exist on different sublattices should be flipped. For example, Li+ exists in both tetrahedral or octahedral sites. Therefore, Li+ on sublattice 1 (tetrahedral) or sublattice 2 (octahedral) should be allowed to flip between the two. Currently a Li+ on sublattice 1 will only be allowed to flip within sublattice 1.
2. Mn disproportionation are relevant in some systems. Therefore, Mn3+ and Mn3+ on sublattice 1 should be allowed to flip to Mn2+ and Mn4+ on sublattice 1 (or to another sublattice).

It has been suggested by Gerd to include a feature where a user can specify allowed perturbations since it seems that other perturbations, besides for the two examples above, may want to be enforced.

Describe the solution you'd like
Ideally, user specifies a table of allowed perturbations and a probability by which the perturbation can occur. In example 1, a user could specify: {'Li+': {'sublattice_1': 0.5}, {'sublattice_2': 0.5}} indicating that a Li+ can equally flip into sublattice 1 or sublattice 2. The code will then enforce that such possible flips are picked with equal probability.

Note that detailed balance should be obeyed at all times, which could be either up the user, or the code, to enforce.

Describe alternatives you've considered
@tchen0965 has code which can address the first issue of Li+ spin flips between different sublattices. She may want to submit a PR soon.

[tchen0965]

Just to clarify, my solution involves providing the CanonicalEnsemble object with a new set of sublattices, defined by the sites occupied by certain species for a given initial occupation, when the object is first instantiated. Maybe we could add this as an option during CanonicalEnsemble instantiation to specify a new or additional set of sublattices. The reason I mention the possibility of adding possible sublattices rather than replacing the current set of sublattices is because there is an option to choose which sublattices you want to perform the MC on in the run function, so you could potentially run an annealing over the all-cation sublattice then a normal MC run within the Li-Va sublattice with a single CanonicalEnsemble instance. Of course, the load to just create separate CanonicalEnsemble instances for the all-cation sublattice and the Li-Va sublattice isn't high, so I'm not sure if there is a benefit either way.

[lbluque]

Thanks @juliayang and @tchen0965 this seems to be a pretty important option for us if using a canonical ensemble.

I think as @juliayang mentions this are two (3?) issues.

1. The first is handling overlapping species in multisublattice canonical ensemble and allowing those species to be able to "swap" to another sublattice. This is a general enough issue, that I think only plays a role the canonical case, and so we should look to implement this as a general (default?) option in the CanonicalEnsemble class.
2. My first thought on general spin flips is to just create a new derived class for whatever ensemble we are want to. We do need to be pretty careful we do not break balance (or stick with detailed balance). If we stick with well known and published flipping rules we should be good, but if we start adding our own recipes then we should have proof that recipe indeed reaches the target probability.
3. General settings for the sublattices are allowed. This is what you enabled right @tchen0965 ?
   

   smol/smol/moca/ensembles/base.py

   Lines 21 to 45 in f55cb96

   	def __init__(self, processor, sample_interval, initial_occupancy,
   	sublattices=None, seed=None):
   	"""Initialize class instance.
   	Args:
   	processor (Processor):
   	A processor that can compute the change in a property given
   	a set of flips.
   	sample_interval (int):
   	interval of steps to save the current occupancy and property
   	initial_occupancy (ndarray or list):
   	Initial occupancy vector. The occupancy can be encoded
   	according to the processor or the species names directly.
   	sublattices (dict): optional
   	dictionary with keys identifying the active sublattices
   	(i.e. "anion" or the allowed species in that sublattice
   	"Li+/Vacancy". The values should be a dictionary
   	with two items {'sites': array with the site indices for all
   	sites corresponding to that sublattice in the occupancy vector,
   	'site_space': OrderedDict of allowed species in sublattice}
   	All sites in a sublattice need to have the same set of allowed
   	species.
   	seed (int): optional
   	seed for random number generator
   	"""

   
   But when having these as input, users need to be pretty careful that they are not creating situations where a species ends up in a site where it should never be.

[tchen0965]

@lbluque Yes, that's what I used. As a note, in a previous pull request, I added the option to set the sublattices in the CanonicalEnsemble instantiation too.

I agree that some care has to be taken that species aren't swapped onto sites on which they are not in the domain (site space). I think this takes some knowledge from the user that I'm not sure we should assume they have, which is partially why I'm not quite sure how to handle it. For example, I was thinking of adding the option to add sublattices by species (e.g. user can provide ['Li+', 'Vacancy'] as another sublattice), but I don't think you should provide an option which can break the MC (e.g. if you provide ['Li+', 'Va', 'Mn2+', 'Mn3+'] as possible sublattice but Mn3+ can only be in oct site), so I don't think that's a reasonable solution.

[lbluque]

Oooh. I like your idea! Something along the lines of set intersections of site spaces?

[tchen0965]

I think so? The set interactions would be only those sites which contain the given species (for example, Li and Vacancy). As long as both Li and Vacancy are all allowed on the same set of sites, even if their list of possible sites has different site spaces, they shouldn't end up in a site that they shouldn't be on. It may require a bit of editing on the flip function too though, if the bit of a species changes from one site space to another. If this is an ok solution, I can work on a version for the pull request, and we can decide when the request is submitted if we want to make changes in how it's implemented.

[lbluque]

I think this can be a very good solution, although fully generalizing in a robust manner will require a bit of thinking, but I actually think its a very fun problem.

It should only involve adapting the _get_flips function which should augment the swap_options list to include allowed sites to flip with on other sublattices.

smol/smol/moca/ensembles/canonical.py

Lines 208 to 231 in f55cb96

	def _get_flips(self, sublattices=None):
	"""Get a possible canonical flip. A swap between two sites.
	Args:
	sublattices (list of str): optional
	If only considering one sublattice.
	Returns:
	tuple
	"""
	if sublattices is None:
	sublattices = self.sublattices
	sublattice_name = random.choice(sublattices)
	sites = self._active_sublatts[sublattice_name]['sites']
	site1 = random.choice(sites)
	swap_options = [i for i in sites
	if self._occupancy[i] != self._occupancy[site1]]
	if swap_options:
	site2 = random.choice(swap_options)
	return ((site1, self._occupancy[site2]),
	(site2, self._occupancy[site1]))
	else:
	# inefficient, maybe re-call method? infinite recursion problem
	return tuple()

Two not so trivial things are to define a general, robust and fast way to find all the sites that are needed to augment the swap_options; and then make sure that the Monte Carlo is still all valid.

I actually have some ideas on how to do the first part, so if you want I may just create a quick feature branch for this so we can work on it and then merge it into master and/or experimental once we are happy.

[lbluque]

Hi @tchen0965, have a look at the overlap_canonical branch I just created.

I added the following things as a sketch to implement the possible swapping of species among sublattices. I haven't even checked it, so this still needs quite a bit of work. But feel free to checkout this branch and add/edit to it.

This is an attempt to get the info necessary of overlaps. Its a start, but still not sure what the best way to store it is.

smol/smol/moca/ensembles/canonical.py

Lines 68 to 79 in d7382c7

	self._sublattice_overlap = None
	# get sublattice species intersection sets, fill overlap dict
	if site_space_overlap:
	self._sublattice_overlap = defaultdict(list)
	for sublatt1, sublatt2 in combinations(self._sublattices, 2):
	species1 = sublatt1['site_space'].keys()
	species2 = sublatt2['site_space'].keys()
	overlap = set(species1).intersection(species2)
	if len(overlap) >= 2:
	for sp in overlap:
	overlap_info = ((sublatt1, sublatt2), overlap - {sp})
	self._sublattice_overlap[sp].append(overlap_info)

Here is where we augment the list of sites to choose a flip from. This is ok, the problem is there is a lot of calls and stuff going on so not sure how it affects the MC. I think the best is to get as much possible information in the step above (only done in construction) so that this part can be as efficient as possible.

smol/smol/moca/ensembles/canonical.py

Lines 247 to 260 in d7382c7

	if self._sublattice_overlap:
	sspace = self._active_sublatts[sublatt_name]['site_space']
	sp = list(sspace.keys()).index(occu1)
	for (sublatt1, sublatt2), sp_compliment in self._sublattice_overlap[sp]:
	if sublatt1['site_space'] == sspace:
	swap_sublatt = sublatt2
	elif sublatt2['site_space'] == sspace:
	swap_sublatt = sublatt1
	else:
	raise RuntimeError('Somthing has gone real off!!!')
	swap_species = list(swap_sublatt['site_space'].keys())
	allowed_swaps = (swap_species.index(s) for s in sp_compliment)
	swap_options += [i for i in swap_sublatt['sites']
	if i in allowed_swaps]

[lbluque]

I'm closing this since it seems to have largely been addressed in #271