Merge pull request #152 from lbluque/master

Active and inactive sublattices

CHANGES.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 Use this section to keep track of changes in the works.
 ### Added
+* `SamplerContainer.to_hdf5` to save MC sample containers
+[\#151](https://github.com/CederGroupHub/smol/pull/151)
+  ([lbluque](https://github.com/lbluque))
 * `PottsSubspace` class to generate redundant frame expansions.
 [\#146](https://github.com/CederGroupHub/smol/pull/146)
   ([lbluque](https://github.com/lbluque))

smol/moca/ensemble/base.py

@@ -4,7 +4,6 @@
 
 from abc import ABC, abstractmethod
 from smol.moca import CompositeProcessor, CEProcessor, EwaldProcessor
-from .sublattice import get_sublattices
 
 
 class Ensemble(ABC):
@@ -25,7 +24,7 @@ class Ensemble(ABC):
 
     valid_mcmc_steps = None  # add this in derived classes
 
-    def __init__(self, processor, sublattices=None):
+    def __init__(self, processor, sublattices=None, inactive_sublattices=None):
         """Initialize class instance.
 
         Args:
@@ -35,13 +34,19 @@ def __init__(self, processor, sublattices=None):
             sublattices (list of Sublattice): optional
                 list of Lattice objects representing sites in the processor
                 supercell with same site spaces.
+            inactive_sublattices (list of InactiveSublattice): optional
+                list of Lattice objects representing sites in the processor
+                supercell with same site spaces.
         """
         if sublattices is None:
-            sublattices = get_sublattices(processor)
+            sublattices = processor.get_sublattices()
+        if inactive_sublattices is None:
+            inactive_sublattices = processor.get_inactive_sublattices()
         self.num_energy_coefs = len(processor.coefs)
         self.thermo_boundaries = {}  # not pretty way to save general info
         self._processor = processor
         self._sublattices = sublattices
+        self._inact_sublattices = inactive_sublattices
 
     @classmethod
     def from_cluster_expansion(cls, cluster_expansion, supercell_matrix,
@@ -104,13 +109,18 @@ def processor(self):
         """Get the ensemble processor."""
         return self._processor
 
-    # TODO make a setter for this that checks sublattices are correct and
+    # TODO make a setter for these that checks sublattices are correct and
     #  all sites are included.
     @property
     def sublattices(self):
         """Get list of sublattices included in ensemble."""
         return self._sublattices
 
+    @property
+    def inactive_sublattices(self):
+        """Get list of sublattices included in ensemble."""
+        return self._inact_sublattices
+
     @property
     def restricted_sites(self):
         """Get indices of all restricted sites."""

smol/moca/ensemble/canonical.py

@@ -12,7 +12,7 @@
 
 from smol.moca.ensemble.base import Ensemble
 from smol.moca.processor.base import Processor
-from .sublattice import Sublattice
+from smol.moca.sublattice import Sublattice
 
 
 class CanonicalEnsemble(Ensemble, MSONable):

smol/moca/ensemble/semigrand.py

@@ -20,7 +20,7 @@
 from smol.cofe.space.domain import get_species, Vacancy
 from smol.moca.processor.base import Processor
 from .base import Ensemble
-from .sublattice import Sublattice
+from smol.moca.sublattice import Sublattice
 
 
 class BaseSemiGrandEnsemble(Ensemble):

smol/moca/processor/base.py

@@ -17,6 +17,7 @@
 from pymatgen.core import Structure, PeriodicSite
 from monty.json import MSONable
 from smol.utils import get_subclasses
+from smol.moca.sublattice import Sublattice, InactiveSublattice
 from smol.cofe.space import (get_allowed_species, get_site_spaces,
                              Vacancy)
 
@@ -56,8 +57,12 @@ def __init__(self, cluster_subspace, supercell_matrix, coefficients):
 
         self.coefs = np.array(coefficients)
         # this can be used (maybe should) to check if a flip is valid
-        site_spaces = get_site_spaces(self._subspace.expansion_structure)
-        self.unique_site_spaces = tuple(set(site_spaces))
+        active_site_spaces = set(
+            get_site_spaces(self._subspace.expansion_structure))
+        self.unique_site_spaces = tuple(active_site_spaces)
+        # and keep a record of sites with no DOFs
+        all_site_spaces = set(get_site_spaces(self._subspace.structure))
+        self.inactive_site_spaces = tuple(all_site_spaces - active_site_spaces)
 
         self.allowed_species = get_allowed_species(self.structure)
         self.size = self._subspace.num_prims_from_matrix(supercell_matrix)
@@ -190,6 +195,30 @@ def decode_occupancy(self, encoded_occupancy):
         return [species[i] for i, species in
                 zip(encoded_occupancy, self.allowed_species)]
 
+    def get_sublattices(self):
+        """Get a list of sublattices from a processor.
+
+        Returns:
+            list of Sublattice
+        """
+        return [Sublattice(site_space,
+                           np.array([i for i, sp in
+                                     enumerate(self.allowed_species)
+                                     if sp == list(site_space.keys())]))
+                for site_space in self.unique_site_spaces]
+
+    def get_inactive_sublattices(self):
+        """Get a list of inactive sublattices from a processor.
+
+        Returns:
+            list of InactiveSublattice
+        """
+        return [InactiveSublattice(
+            site_space, np.array([i for i, sp in
+                                 enumerate(self.allowed_species)
+                                 if sp == list(site_space.keys())]))
+                for site_space in self.inactive_site_spaces]
+
     def compute_average_drift(self, iterations=1000):
         """Compute average forward and reverse drift for the given property.
 

smol/moca/sampler/container.py

@@ -14,7 +14,7 @@
 import numpy as np
 
 from monty.json import MSONable
-from smol.moca.ensemble.sublattice import Sublattice
+from smol.moca.sublattice import Sublattice
 
 try:
     import h5py

smol/moca/sampler/kernel.py

@@ -51,6 +51,7 @@ def __init__(self, ensemble, step_type, *args, **kwargs):
         try:
             self._usher = mcusher_factory(self.valid_mcushers[step_type],
                                           ensemble.sublattices,
+                                          ensemble.inactive_sublattices,
                                           *args, **kwargs)
         except KeyError:
             raise ValueError(f"Step type {step_type} is not valid for a "

smol/moca/sampler/mcusher.py

@@ -12,25 +12,36 @@
 
 from abc import ABC, abstractmethod
 import random
-
+import numpy as np
 from smol.utils import derived_class_factory
 
 
 class MCUsher(ABC):
     """Abstract base class for MC usher classes."""
 
-    def __init__(self, sublattices, sublattice_probabilities=None):
+    def __init__(self, sublattices, inactive_sublattices,
+                 sublattice_probabilities=None):
         """Initialize MCMCStep.
 
         Args:
             sublattices (list of Sublattice):
-                list of Sublattices to propose steps for.
+                list of active Sublattices to propose steps for. Active
+                sublattices are those that include sites with configuration
+                DOFs.
+            inactive_sublattices (list of Sublattice):
+                list of inactive Sublattices, i.e. those with no configuration
+                DOFs. These can be used to obtain auxiliary information for MC
+                step proposals for the active sublattices
             sublattice_probabilities (list of float): optional
-                list of probability to pick a site from a specific sublattice.
+                list of probability to pick a site from a specific active
+                sublattice.
         """
         self.sublattices = sublattices
+        self.inactive_sublattices = inactive_sublattices
+
         if sublattice_probabilities is None:
-            self._sublatt_probs = len(self.sublattices) * [1/len(self.sublattices), ]  # noqa
+            self._sublatt_probs = np.array(
+                len(self.sublattices) * [1/len(self.sublattices), ])
         elif len(sublattice_probabilities) != len(self.sublattices):
             raise AttributeError('Sublattice probabilites needs to be the '
                                  'same length as sublattices.')
@@ -48,12 +59,13 @@ def sublattice_probabilities(self):
     def sublattice_probabilities(self, value):
         """Set the sublattice probabilities."""
         if len(value) != len(self.sublattices):
-            raise AttributeError('Can not set sublattice probabilities. '
-                                 'Length must be the the same as the number '
-                                 f'of sublattices {len(self.sublattices)}')
+            raise AttributeError(
+                f'Can not set sublattice probabilities.\n Length must be the'
+                f' same as the number of sublattices {len(self.sublattices)}')
         elif sum(value) != 1:
-            raise ValueError('Can not set sublattice probabilities. '
-                             'Sublattice probabilites must sum to one.')
+            raise ValueError(
+                'Can not set sublattice probabilities.\n'
+                'Sublattice probabilites must sum to one.')
         self._sublatt_probs = value
 
     @abstractmethod
@@ -137,14 +149,18 @@ def propose_step(self, occupancy):
         return swap
 
 
-def mcusher_factory(usher_type, sublattices, *args, **kwargs):
+def mcusher_factory(usher_type, sublattices, inactive_sublattices, *args,
+                    **kwargs):
     """Get a MC Usher from string name.
 
     Args:
         usher_type (str):
             string specifying step to instantiate.
         sublattices (list of Sublattice):
                 list of Sublattices to propose steps for.
+        inactive_sublattices (list of InactiveSublattice):
+                list of InactiveSublattices for sites with no configuration
+                DOFs.
         *args:
             positional arguments passed to class constructor
         **kwargs:
@@ -153,5 +169,6 @@ def mcusher_factory(usher_type, sublattices, *args, **kwargs):
     Returns:
         MCUsher: instance of derived class.
     """
-    return derived_class_factory(usher_type.capitalize(), MCUsher,
-                                 sublattices, *args, **kwargs)
+    return derived_class_factory(
+        usher_type.capitalize(), MCUsher, sublattices, inactive_sublattices,
+        *args, **kwargs)

smol/moca/ensemble/sublattice.py → smol/moca/sublattice.py

@@ -7,28 +7,13 @@
 
 __author__ = "Luis Barroso-Luque"
 
+from dataclasses import dataclass, field
 import numpy as np
 from monty.json import MSONable
 from smol.cofe.space.domain import SiteSpace
 
 
-def get_sublattices(processor):
-    """Get a list of sublattices from a processor.
-
-    Args:
-        processor (Processor):
-            A processor object to extract sublattices from
-    Returns:
-        list of Sublattice
-    """
-    return [Sublattice(site_space,
-                       np.array([i for i, sp in
-                                 enumerate(processor.allowed_species)
-                                if sp == list(site_space.keys())]))
-            for site_space in processor.unique_site_spaces]
-
-
-# TODO consider adding the inactive sublattices?
+@dataclass
 class Sublattice(MSONable):
     """Sublattice class.
 
@@ -46,18 +31,13 @@ class Sublattice(MSONable):
         array of site indices for all unrestricted sites in the sublattice.
     """
 
-    def __init__(self, site_space, sites):
-        """Initialize Sublattice.
+    site_space: SiteSpace
+    sites: np.array
+    active_sites: np.array = field(init=False)
 
-        Args:
-            site_space (SiteSpace):
-                A site space object representing the sites in the sublattice
-            sites (ndarray):
-                array with the site indices
-        """
-        self.sites = sites
-        self.site_space = site_space
-        self.active_sites = sites.copy()
+    def __post_init__(self):
+        """Copy sites into active_sites."""
+        self.active_sites = self.sites.copy()
 
     @property
     def species(self):
@@ -88,18 +68,6 @@ def reset_restricted_sites(self):
         """Reset all restricted sites to active."""
         self.active_sites = self.sites.copy()
 
-    def __str__(self):
-        """Pretty print the sublattice species."""
-        string = f'Sublattice\n Site space: {dict(self.site_space)}\n'
-        string += f' Number of sites: {len(self.sites)}\n'
-        return string
-
-    def __repr__(self):
-        """Repr for nice viewing."""
-        rep = f'Sublattice Summary \n\n   site_space: {self.site_space}\n\n'
-        rep += f'   sites: {self.sites}\n\n active_sites: {self.active_sites}'
-        return rep
-
     def as_dict(self):
         """Get Json-serialization dict representation.
 
@@ -122,3 +90,38 @@ def from_dict(cls, d):
                          sites=np.array(d['sites']))
         sublattice.active_sites = np.array(d['active_sites'])
         return sublattice
+
+
+@dataclass
+class InactiveSublattice(MSONable):
+    """Same as above but for sublattices with no configuration DOFs.
+
+    Attributes:
+     site_space (SiteSpace):
+        SiteSpace with the allowed species and their random
+        state composition.
+     sites (ndarray):
+        array of site indices for all sites in sublattice
+    """
+
+    site_space: SiteSpace
+    sites: np.array
+
+    def as_dict(self):
+        """Get Json-serialization dict representation.
+
+        Returns:
+            MSONable dict
+        """
+        d = {'site_space': self.site_space.as_dict(),
+             'sites': self.sites.tolist()}
+        return d
+
+    @classmethod
+    def from_dict(cls, d):
+        """Instantiate a sublattice from dict representation.
+
+        Returns:
+            Sublattice
+        """
+        return cls(SiteSpace.from_dict(d['site_space']), np.array(d['sites']))