DefaultRadii

Author: amirhajibabaei

theforce/calculator/active.py

@@ -2,7 +2,7 @@
 from theforce.regression.gppotential import PosteriorPotential, PosteriorPotentialFromFolder
 from theforce.descriptor.atoms import TorchAtoms, AtomsData, LocalsData
 from theforce.similarity.sesoap import SeSoapKernel
-from theforce.math.sesoap import SpecialRadii
+from theforce.math.sesoap import DefaultRadii
 from theforce.util.tensors import padded
 from theforce.util.util import date, timestamp
 from theforce.io.sgprio import SgprIO
@@ -18,7 +18,7 @@
 
 
 def default_kernel(cutoff=6.):
-    return SeSoapKernel(3, 3, 4, cutoff, radii=SpecialRadii({1: 0.5}))
+    return SeSoapKernel(3, 3, 4, cutoff, radii=DefaultRadii())
 
 
 class FilterDeltas(Filter):
@@ -215,6 +215,7 @@ def __init__(self, covariance=None, calculator=None, process_group=None, meta=No
         self.stdout = True
         self.step = 0
         self.log('active calculator says Hello!', mode='w')
+        self.log(f'kernel: {self.model.descriptors}')
         self.log_settings()
         self.log('model size: {} {}'.format(*self.size))
         self.pckl = pckl

theforce/cl/__init__.py

@@ -2,7 +2,7 @@
 from theforce.util.util import get_default_args
 from theforce.util.parallel import mpi_init
 from theforce.calculator.socketcalc import SocketCalculator
-from theforce.calculator.active import ActiveCalculator, kcal_mol, inf
+from theforce.calculator.active import ActiveCalculator, kcal_mol, inf, SeSoapKernel, DefaultRadii
 import os
 
 

theforce/math/sesoap.py

@@ -83,6 +83,25 @@ def state_args(self):
         return f'{dct}, {self.others}'
 
 
+class DefaultRadii(Radii):
+
+    def __init__(self, default=1., special={1: 0.5}):
+        self.default = default
+        self.special = special
+
+    def get(self, number):
+        try:
+            return self.special[number]
+        except KeyError:
+            return self.default
+
+    @property
+    def state_args(self):
+        default = float(self.default)
+        special = {z: float(r) for z, r in self.special.items()}
+        return f'{default}, {special}'
+
+
 class SeSoap(Module):
 
     def __init__(self, lmax, nmax, radial, radii=1., flatten=True, normalize=True):