agox 3.7.0rc1__tar.gz → 3.8.0__tar.gz

{agox-3.7.0rc1/agox.egg-info → agox-3.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: agox
-Version: 3.7.0rc1
+Version: 3.8.0
 Summary: Atomistic Global Optimziation X is a framework for structure optimization in materials science.
 Author: AGOX Developers
 Author-email: "Mads-Peter V. Christiansen" <machri@phys.au.dk>, Nikolaj Rønne <nronne@phys.au.dk>, Bjørk Hammer <hammer@phys.au.dk>
@@ -29,6 +29,8 @@ Requires-Dist: importlib-resources>1.3; python_version < "3.9"
 Requires-Dist: dscribe>=2.0.0
 Requires-Dist: nbformat
 Requires-Dist: typing_extensions
+Requires-Dist: jax
+Requires-Dist: jaxlib
 Requires-Dist: ray
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
@@ -43,6 +45,9 @@ Requires-Dist: agox[docs,test]; extra == "all"
 
 # Atomistic Global Optimization X (AGOX)
 
+![logo](docs/source/logos/logo-carbon.gif)
+
+
 AGOX is a package for global optimization of atomic system using e.g. the energy 
 calculated from density functional theory as the objective function. AGOX interfaces 
 with the Atomistic Simulation Environment (ASE) and as such supports any of 

{agox-3.7.0rc1 → agox-3.8.0}/README.md

@@ -1,5 +1,8 @@
 # Atomistic Global Optimization X (AGOX)
 
+![logo](docs/source/logos/logo-carbon.gif)
+
+
 AGOX is a package for global optimization of atomic system using e.g. the energy 
 calculated from density functional theory as the objective function. AGOX interfaces 
 with the Atomistic Simulation Environment (ASE) and as such supports any of 

{agox-3.7.0rc1 → agox-3.8.0}/agox/__init__.py

@@ -1,7 +1,7 @@
 import os
 
 # Versioning
-__version_info__ = (3, 7, 0)
+__version_info__ = (3, 8, 0)
 __version__ = '{}.{}.{}'.format(*__version_info__)
 
 # Extra versioning - Mainly for managing Pypi releases.

{agox-3.7.0rc1 → agox-3.8.0}/agox/acquisitors/ABC_acquisitor.py

@@ -30,6 +30,7 @@ class AcquisitorBaseClass(ABC, Observer, Writer):
         use_counter=True,
         prefix="",
         surname="",
+        skip_chance=0.0,
     ):
         Observer.__init__(self, gets=gets, sets=sets, order=order, surname=surname)
         Writer.__init__(self, verbose=verbose, use_counter=use_counter, prefix=prefix)
@@ -42,6 +43,8 @@ class AcquisitorBaseClass(ABC, Observer, Writer):
             handler_identifier="AGOX",
         )
 
+        self.skip_chance = skip_chance
+
     @abstractmethod
     def calculate_acquisition_function(self, candidates):
         """
@@ -79,6 +82,10 @@ class AcquisitorBaseClass(ABC, Observer, Writer):
             State object that contains the iteration data.
         """
 
+        if self.skip_chance > 0 and np.random.rand() < self.skip_chance:
+            state.add_to_cache(self, self.set_key, [], mode="a")
+            return
+
         # Get data from the iteration data dict.
         candidate_list = state.get_from_cache(self, self.get_key)
         candidate_list = list(filter(None, candidate_list))

agox-3.8.0/agox/acquisitors/LCB_penalty.py

@@ -0,0 +1,135 @@
+from collections import Counter
+from collections.abc import Hashable
+from typing import List
+
+import numpy as np
+from numpy.typing import NDArray
+
+from agox.acquisitors.LCB import LowerConfidenceBoundAcquisitor
+from agox.candidates.standard import StandardCandidate
+from agox.databases.ABC_database import DatabaseBaseClass
+from agox.models.ABC_model import ModelBaseClass
+from agox.models.descriptors import SpectralGraphDescriptor, Voronoi
+from agox.models.descriptors.ABC_descriptor import DescriptorBaseClass
+from agox.observer import Observer
+from agox.writer import agox_writer
+
+
+class LCBPenaltyAcquisitor(LowerConfidenceBoundAcquisitor):
+
+    """Lower confidence bound acquisitor with a penalty term for repeated
+    descriptor values.
+
+    The acquisition function for selecting candidates is given by
+
+    .. math::
+
+        F(x) = E(x) - \kappa \sigma(x) + \alpha n_D(x),
+
+    where :math:`n_D` is the number of times the descriptor value already
+    exists for the candidates in the database, and :math:`\alpha` is the
+    penalty scale value.
+
+    The acquisition function for relaxing candidates is equivalent to that of
+    regular LCB, as the :math:`\alpha n_D(x)` term is discontinuous.
+
+    Parameters
+    ----------
+    model : ModelBaseClass
+        Model to obtain energy and uncertainty predictions from.
+    descriptor : DescriptorBaseClass
+        Descriptor to evaluate descriptor values from.
+    penalty_scale : float, optional
+        Scale value for penalty (in eV), by default 1.0
+    """
+
+    def __init__(self,
+                 model: ModelBaseClass,
+                 descriptor: DescriptorBaseClass,
+                 penalty_scale: float = 1.0,
+                 **kwargs):
+        super().__init__(model, **kwargs)
+
+        if not isinstance(descriptor, (SpectralGraphDescriptor, Voronoi)):
+            raise ValueError('Descriptor can currently only be graph-type descriptor')
+
+        self.descriptor = descriptor
+        self.penalty_scale = penalty_scale
+
+        self.descriptor_value_counts = Counter()
+
+        self.add_observer_method(self.update_descriptor_values,
+                                 gets={}, sets={}, order=self.order[0],
+                                 handler_identifier='database')
+
+    def calculate_acquisition_function(self, candidates: List[StandardCandidate]) -> NDArray:
+        fitness = np.zeros(len(candidates))
+
+        for i, candidate in enumerate(candidates):
+            E, sigma = self.model.predict_energy_and_uncertainty(candidate)
+
+            descriptor_value = self._get_descriptor_value(candidate)
+            count = self.descriptor_value_counts[descriptor_value]
+            penalty = count * self.penalty_scale
+
+            fitness[i] = self.acquisition_function(E, sigma) + penalty
+
+            candidate.add_meta_information('model_energy', E)
+            candidate.add_meta_information('uncertainty', sigma)
+            candidate.add_meta_information('penalty', penalty)
+
+        return fitness
+
+    def print_information(self, candidates: List[StandardCandidate], acquisition_values: List[float]) -> None:
+        if self.model.ready_state:
+            for i, candidate in enumerate(candidates):
+                fitness = acquisition_values[i]
+                Emodel = candidate.get_meta_information('model_energy')
+                sigma = candidate.get_meta_information('uncertainty')
+                penalty = candidate.get_meta_information('penalty')
+                self.writer('Candidate: E={:8.3f}, s={:8.3f}, p={:7.2f}, F={:8.3f}'.format(
+                    Emodel, sigma, penalty, fitness))
+
+    def attach_to_database(self, database: DatabaseBaseClass) -> None:
+        if not isinstance(database, DatabaseBaseClass):
+            raise TypeError(f'{database} is not a DatabaseBaseClass object')
+
+        print(f'{self.name}: Attaching to database: {database}')
+        self.attach(database)
+
+    @agox_writer
+    @Observer.observer_method
+    def update_descriptor_values(self, database: DatabaseBaseClass, _) -> None:
+        """Update the internal counter of descriptor values present in the
+        database.
+        """
+        candidates = database.get_all_candidates()
+        new_candidates = candidates[sum(self.descriptor_value_counts):]
+
+        for candidate in new_candidates:
+            descriptor_value = self._get_descriptor_value(candidate)
+            self.descriptor_value_counts[descriptor_value] += 1
+
+        s = 's' if len(new_candidates) != 1 else ''
+        self.writer(f'Added {len(new_candidates)} structure{s} to the descriptor list')
+
+    def _get_descriptor_value(self, candidate: StandardCandidate) -> Hashable:
+        """Obtain a hashable descriptor value to store in the internal counter.
+
+        Parameters
+        ----------
+        candidate : StandardCandidate
+            Candidate to evaluate the descriptor for.
+
+        Returns
+        -------
+        Hashable
+            Hashable descriptor value.
+        """
+        descriptor_value = self.descriptor.get_features(candidate)
+
+        # convert to hashable object
+        if isinstance(descriptor_value, np.ndarray):
+            descriptor_value = descriptor_value.data.tobytes()
+
+        return descriptor_value

{agox-3.7.0rc1 → agox-3.8.0}/agox/acquisitors/__init__.py

@@ -1,28 +1,31 @@
 """
 Acquisitors are used to select the next candidate to evaluate among a number of
-candidates, usually by sorting according to some metric. 
-The concept of acquisitors comes from Bayesian search algorithms. 
+candidates, usually by sorting according to some metric.
+The concept of acquisitors comes from Bayesian search algorithms.
 
-A common acquisitor is the Lower Confidence Bound (LCB) acquisitor, which 
+A common acquisitor is the Lower Confidence Bound (LCB) acquisitor, which
 selects according to the expression
 
 .. math::
     LCB(x) = E(x) - \kappa \sigma(x).
 
 Where E(x) is a prediction of the objective function at x, and :math:`\sigma(x)`
-is the uncertainty of the prediction. :math:`\kappa` is a parameter that 
-controls the trade-off between exploration and exploitation. 
-This lets an algorithm choose which candidates to evaluate, but 
+is the uncertainty of the prediction. :math:`\kappa` is a parameter that
+controls the trade-off between exploration and exploitation.
+This lets an algorithm choose which candidates to evaluate, but
 does require the generation of multiple candidates to choose from.
 """
 
-from .ABC_acquisitor import AcquisitorBaseClass, AcquisitonCalculatorBaseClass
+from .ABC_acquisitor import AcquisitonCalculatorBaseClass, AcquisitorBaseClass
 from .LCB import LowerConfidenceBoundAcquisitor
+from .LCB_penalty import LCBPenaltyAcquisitor
 from .LCB_power import PowerLowerConfidenceBoundAcquisitor
+from .meta_acquisitor import MetaInformationAcquisitor
 
-__all__ = ['AcquisitorBaseClass', 
-           'LowerConfidenceBoundAcquisitor', 
+__all__ = ['AcquisitorBaseClass',
+           'LowerConfidenceBoundAcquisitor',
+           'LCBPenaltyAcquisitor',
            'PowerLowerConfidenceBoundAcquisitor'
            'AcquisitonCalculatorBaseClass',
+           'MetaInformationAcquisitor',
            ]
-

agox-3.8.0/agox/acquisitors/meta_acquisitor.py

@@ -0,0 +1,28 @@
+import numpy as np
+from agox.acquisitors.ABC_acquisitor import AcquisitorBaseClass
+
+class MetaInformationAcquisitor(AcquisitorBaseClass):
+
+    name = 'MetaInformationAcquisitor'
+
+    def __init__(self, meta_key, mode='min', **kwargs):
+        super().__init__(**kwargs)
+        self.meta_key = meta_key
+        self.mode = mode
+
+        assert mode in ['max', 'min'] # Either maximize or minimize according to the meta key.
+
+    def calculate_acquisition_function(self, candidates):
+        sign = 1 if self.mode == 'min' else -1
+        fitness = np.zeros(len(candidates))
+        
+        for i, candidate in enumerate(candidates):
+            fitness[i] = sign * candidate.get_meta_information(self.meta_key)
+
+        return fitness
+    
+    
+
+
+
+

{agox-3.7.0rc1 → agox-3.8.0}/agox/collectors/__init__.py

@@ -5,7 +5,6 @@ and this is done using parallelization from Ray.
 """
 from agox.collectors.standard import StandardCollector
 from agox.collectors.ray_collector import ParallelCollector
+from agox.collectors.parallel_tempering import ParallelTemperingCollector
 
-# __all__ = ['StandardCollector', 'ParallelCollector']
-
-__all__ = ['StandardCollector', 'ParallelCollector']
+__all__ = ['StandardCollector', 'ParallelCollector', 'ParallelTemperingCollector']

agox-3.8.0/agox/collectors/parallel_tempering.py

@@ -0,0 +1,165 @@
+import numpy as np
+import ray
+
+from agox.collectors.ABC_collector import CollectorBaseClass
+from agox.samplers.fixed_sampler import FixedSampler
+from agox.utils.ray import RayPoolUser
+
+
+def remote_generate(generator, sampler, environment):
+    return generator(sampler, environment)
+
+
+class ParallelTemperingCollector(CollectorBaseClass, RayPoolUser):
+    name = "ParallelTemperingCollector"
+
+    """
+    Parameters
+    ----------
+    num_candidates : dict or list
+        If dict, it must have the form {0: [], 500: []} where the keys are the iteration numbers
+        and the values are the number of candidates to generate for that iteration.
+        If a list it must provide the number of candidates for each generator used. 
+    """
+
+    def __init__(self, generator_dict, dynamic=True, **kwargs):
+        # Get the kwargs that are relevant for RayPoolUser
+        RayPoolUser.__init__(self)
+        self.dynamic = dynamic
+
+        # Get all the generators from the generator_dict
+        unique_generators = []
+        for generator in generator_dict.values():
+            if generator not in unique_generators:
+                unique_generators.append(generator)
+
+        # Figure out which unique generator each generator in the generator_dict corresponds to.
+        self.mapping = {}
+        for key, generator in generator_dict.items():
+            self.mapping[key] = unique_generators.index(generator)
+
+        # Give the rest to the CollectorBaseClass.
+        CollectorBaseClass.__init__(self, unique_generators, **kwargs)
+
+        self.generator_keys = []
+        for generator in self.generators:
+            key = self.pool_add_module(generator)
+            self.generator_keys.append(key)
+
+    def make_candidates(self):
+
+        if self.dynamic:
+            self.update_generators()
+        # We need to build the args and kwargs sent to the actors.
+
+        # If sampler is empty and initialized, then use fallback generator
+        if len(self.sampler) == 0 and self.sampler.initialized:
+            number_of_candidates = self.get_number_of_candidates()
+            fallback_generator = self.get_fallback_generator()
+            generator_id = self.pool_add_module(fallback_generator)
+            modules = [[generator_id]] * np.sum(number_of_candidates)
+        else: 
+            if self.get_iteration_counter() == 1 or self.get_iteration_counter() is None:
+                random_key = self.mapping['random']
+                modules = [[self.generator_keys[random_key]]] * self.sampler.sample_size
+            else:
+                # This specifies which module each actor of the pool will use.
+                modules = []
+                for temperature in self.sampler.temperatures:
+                    modules += [[self.get_generator_key(temperature)]]
+        
+        environment_id = ray.put(self.environment)
+        fixed_samplers = self.get_dummy_samplers()
+        if len(fixed_samplers) == 0:
+            sampler_ids = [ray.put(self.sampler)] * self.sampler.sample_size
+        else:
+            sampler_ids = [ray.put(sampler) for sampler in fixed_samplers]
+            
+        # The args and kwargs passed to the function - in this case the remote_generate
+        # function defined above.
+        args = [[sampler_id, environment_id] for sampler_id in sampler_ids]
+        # kwargs = [{}] * np.sum(number_of_candidates)
+        kwargs = [{} for _ in range(np.sum(self.sampler.sample_size))]
+
+        # Generate in parallel using the pool.
+        candidates = self.pool_map(remote_generate, modules, args, kwargs)
+
+        # Flatten the output which is a list of lists.
+        flat_candidates = []
+        for cand_list in candidates:
+            for cand in cand_list:
+                flat_candidates.append(cand.copy())
+
+        for i, cand in enumerate(flat_candidates):
+            cand.add_meta_information("walker_index", i)
+
+        return flat_candidates
+
+    def get_dummy_samplers(self):
+        sample = self.sampler.get_all_members()
+        samplers = []
+        for member in sample:
+            sampler = FixedSampler(member)
+            samplers.append(sampler)
+        return samplers
+
+    def get_number_of_candidates(self):
+        return [self.sampler.sample_size]
+
+    @classmethod
+    def from_sampler(cls, sampler, environment, amplitudes, random_generator_kwargs=None, **kwargs):
+        from agox.generators import RandomGenerator, RattleGenerator
+
+        # Generators to produce candidates structures
+        if random_generator_kwargs is None:
+            random_generator_kwargs = {}
+        random_generator = RandomGenerator(**environment.get_confinement(), **random_generator_kwargs)
+
+        if not len(amplitudes) == len(sampler.temperatures):
+            raise ValueError(
+                "The number of amplitudes must match the number of temperatures in the sampler."
+            )
+
+        # Dict specificies how many candidates are created with and the dict-keys are iterations.
+        generator_dict = {}
+        generator_dict["random"] = random_generator
+
+        N_atoms = len(environment.get_missing_indices())
+
+        for i, t in enumerate(sampler.temperatures):
+            generator_dict[t] = RattleGenerator(
+                **environment.get_confinement(),
+                rattle_amplitude=amplitudes[i],
+                n_rattle=N_atoms,
+            )
+            generator_dict[t].add_dynamic_attribute("rattle_amplitude")  # The rattle amplitude is dynamic.
+
+        # Collector creates a number of structures in each iteration.
+        return cls(
+            generator_dict,
+            sampler=sampler,
+            environment=environment,
+            order=1,
+            **kwargs
+        )
+
+    def update_generators(self):
+        for i, temperature in enumerate(self.sampler.temperatures):
+            generator = self.get_generator(temperature)
+            acceptance = self.sampler.get_acceptance_rate(i, start=-10) # Acceptance rate over the last 10 iterations.
+
+            if acceptance.size == 0:
+                continue
+            
+            # Rattle update amplitude
+            if hasattr(generator, 'rattle_amplitude'):
+                if acceptance < 0.5:
+                    generator.rattle_amplitude *= 0.99
+                else:
+                    generator.rattle_amplitude = min(5, generator.rattle_amplitude * 1.01)
+
+    def get_generator(self, temperature):
+        return self.generators[self.mapping[temperature]]
+    
+    def get_generator_key(self, temperature):
+        return self.generator_keys[self.mapping[temperature]]

{agox-3.7.0rc1 → agox-3.8.0}/agox/models/GPR/GPR.py

@@ -48,7 +48,6 @@ class GPR(ModelBaseClass, RayPoolUser):
 
     supported_descriptor_types = ["global"]
     implemented_properties = ["energy", "forces", "uncertainty", "forces_uncertainty"]
-    dynamic_attributes = ["alpha", "K_inv", "X", "K", "kernel", "Y", "mean_energy", "_ready_state"]
 
     def __init__(
         self,
@@ -85,6 +84,9 @@ class GPR(ModelBaseClass, RayPoolUser):
             descriptor.descriptor_type in self.supported_descriptor_types
         ), "Only global descriptors are supported, use SparseGPR for local descriptors instead."
 
+        for attr in ["alpha", "K_inv", "X", "K", "kernel", "Y", "mean_energy", "_ready_state"]:
+            self.add_dynamic_attribute(attr)
+
         self.descriptor = descriptor
         self.kernel = kernel
         self.prior = prior