From 8f71b1bc7d5882a9d26753a383c81ceda0a4e44c Mon Sep 17 00:00:00 2001 From: Callow Date: Fri, 7 Jun 2024 11:17:26 +0200 Subject: [PATCH] Add ace potential to finish generation of coupling coefficients Co-authored-by: James Goff --- mala/common/parameters.py | 2 + mala/descriptors/__init__.py | 1 + mala/descriptors/ace.py | 21 ++ mala/descriptors/ace_potential.py | 398 ++++++++++++++++++++++++++++++ 4 files changed, 422 insertions(+) create mode 100644 mala/descriptors/ace_potential.py diff --git a/mala/common/parameters.py b/mala/common/parameters.py index c1829b6d2..411c22418 100644 --- a/mala/common/parameters.py +++ b/mala/common/parameters.py @@ -371,7 +371,9 @@ def __init__(self): self.ace_ranks = [1, 2, 3] self.ace_nmax = [6, 2, 2] self.ace_lmax = [0, 2, 2] + self.ace_nradbase = max(self.ace_nmax) self.ace_lmin = [0, 0, 0] + self.ace_reference_ens = [0.0, 0.0] self.ace_nshell = 2.0 self.ace_apply_shift = False diff --git a/mala/descriptors/__init__.py b/mala/descriptors/__init__.py index 49f58a8df..41852d868 100644 --- a/mala/descriptors/__init__.py +++ b/mala/descriptors/__init__.py @@ -5,3 +5,4 @@ from .descriptor import Descriptor from .minterpy_descriptors import MinterpyDescriptors from .ace import ACE +from .ace_potential import AcePot diff --git a/mala/descriptors/ace.py b/mala/descriptors/ace.py index d6f9ec8b0..4b9823498 100755 --- a/mala/descriptors/ace.py +++ b/mala/descriptors/ace.py @@ -15,6 +15,7 @@ from mala.common.parallelizer import printout from mala.descriptors.lammps_utils import extract_compute_np from mala.descriptors.descriptor import Descriptor +from mala.descriptors.ace_potential import AcePot import mala.descriptors.ace_coupling_utils as acu import mala.descriptors.wigner_coupling as wigner_coupling import mala.descriptors.cg_coupling as cg_coupling @@ -522,6 +523,26 @@ def calculate_coupling_coeffs(self): limit_nus = self.calc_limit_nus() + # permutation symmetry adapted ACE labels + Apot = AcePot( + self.parameters.ace_elements, + self.parameters.ace_reference_ens, + self.parameters.ace_ranks, + self.parameters.ace_nmax, + self.parameters.ace_lmax, + self.parameters.ace_nradbase, + rcutfac, + lmbda, + rcinner, + drcinner, + lmin=self.parameters.ace_lmin, + **{"input_nus": limit_nus, "ccs": ccs[self.parameters.ace_M_R]} + ) + Apot.write_pot("coupling_coefficients_fullbasis") + + Apot.set_funcs(nulst=limit_nus, muflg=True, print_0s=True) + Apot.write_pot("coupling_coefficients") + def calc_limit_nus(self): ranked_chem_nus = [] for ind, rank in enumerate(self.parameters.ace_ranks): diff --git a/mala/descriptors/ace_potential.py b/mala/descriptors/ace_potential.py new file mode 100644 index 000000000..dba402c00 --- /dev/null +++ b/mala/descriptors/ace_potential.py @@ -0,0 +1,398 @@ +import itertools +import json + +import numpy as np + +import mala.descriptors.ace_coupling_utils as acu + +class AcePot: + def __init__( + self, + elements, + reference_ens, + ranks, + nmax, + lmax, + nradbase, + rcut, + lmbda, + rcutinner=0.0, + drcutinner=0.01, + lmin=1, + b_basis="pa_tabulated", #'pa_tabulated', 'minsub', 'ysg_x_so3' + manuallabs=None, + **kwargs + ): + self.ccs = None + if kwargs != None: + self.__dict__.update(kwargs) + + if self.ccs == None: + self.global_ccs = generate_ccs() + else: + self.global_ccs = self.ccs + self.global_ccs[1] = {"0": {tuple([]): {"0": 1.0}}} + self.E0 = reference_ens + self.ranks = ranks + self.elements = elements + self.betas = None + self.nus = None + self.deltaSplineBins = 0.001 + self.global_ndensity = 1 + self.global_FSparams = [1.0, 1.0] + self.global_rhocut = 100000 + self.global_drhocut = 250 + # assert the same nmax,lmax,nradbase (e.g. same basis) for each bond type + self.radbasetype = "ChebExpCos" + self.global_nmax = nmax + self.global_lmax = lmax + self.b_basis = b_basis + assert len(nmax) == len(lmax), "nmax and lmax arrays must be same size" + self.global_nradbase = nradbase + if type(rcut) != dict and type(rcut) != list: + self.global_rcut = rcut + self.global_lmbda = lmbda + self.global_rcutinner = rcutinner + self.global_drcutinner = drcutinner + self.global_lmin = lmin + else: + self.rcut = rcut + self.lmbda = lmbda + self.rcutinner = rcutinner + self.drcutinner = drcutinner + self.lmin = lmin + self.manuallabs = manuallabs + self.set_embeddings() + self.set_bonds() + self.set_bond_base() + + lmax_dict = { + rank: lv for rank, lv in zip(self.ranks, self.global_lmax) + } + try: + lmin_dict = {rank: lv for rank, lv in zip(self.ranks, self.lmin)} + except AttributeError: + lmin_dict = { + rank: lv + for rank, lv in zip( + self.ranks, self.global_lmin * len(self.ranks) + ) + } + nradmax_dict = { + rank: nv for rank, nv in zip(self.ranks, self.global_nmax) + } + mumax_dict = {rank: len(self.elements) for rank in self.ranks} + if self.manuallabs != None: + with open(self.manuallabs, "r") as readjson: + labdata = json.load(readjson) + nulst_1 = [list(ik) for ik in list(labdata.values())] + elif self.manuallabs == None and self.b_basis == "minsub": + from fitsnap3lib.lib.sym_ACE.rpi_lib import descriptor_labels_YSG + + nulst_1 = [ + descriptor_labels_YSG( + rank, + nradmax_dict[rank], + lmax_dict[rank], + mumax_dict[rank], + lmin_dict[rank], + ) + for rank in self.ranks + ] + elif self.manuallabs == None and self.b_basis == "ysg_x_so3": + # TODO until all are tabulated + nulst_1 = [] + else: + nulst_1 = [] + for rank in self.ranks: + PA_lammps, not_compat = acu.pa_labels_raw( + rank, + nradmax_dict[rank], + lmax_dict[rank], + mumax_dict[rank], + lmin_dict[rank], + ) + nulst_1.append(PA_lammps) + nus_unsort = [item for sublist in nulst_1 for item in sublist] + nus = nus_unsort.copy() + mu0s = [] + mus = [] + ns = [] + ls = [] + for nu in nus_unsort: + mu0ii, muii, nii, lii = acu.get_mu_n_l(nu) + mu0s.append(mu0ii) + mus.append(tuple(muii)) + ns.append(tuple(nii)) + ls.append(tuple(lii)) + nus.sort(key=lambda x: mus[nus_unsort.index(x)], reverse=False) + nus.sort(key=lambda x: ns[nus_unsort.index(x)], reverse=False) + nus.sort(key=lambda x: ls[nus_unsort.index(x)], reverse=False) + nus.sort(key=lambda x: mu0s[nus_unsort.index(x)], reverse=False) + nus.sort(key=lambda x: len(x), reverse=False) + nus.sort(key=lambda x: mu0s[nus_unsort.index(x)], reverse=False) + self.nus = nus + self.set_funcs(nus) + + return None + + def set_embeddings( + self, npoti="FinnisSinclair", FSparams=[1.0, 1.0] + ): # default for linear models in lammps PACE + # embeddings =dict()#OrderedDict() #{ind:None for ind in range(len(self.elements))} + embeddings = {ind: None for ind in range(len(self.elements))} + for elemind in range(len(self.elements)): + embeddings[elemind] = { + "ndensity": self.global_ndensity, + "FS_parameters": FSparams, + "npoti": npoti, + "rho_core_cutoff": self.global_rhocut, + "drho_core_cutoff": self.global_drhocut, + } + self.embeddings = embeddings + + def set_bonds(self): + bondinds = range(len(self.elements)) + bond_lsts = [list(b) for b in itertools.product(bondinds, bondinds)] + self.bondlsts = bond_lsts + + def set_bond_base(self): + bondstrs = ["[%d, %d]" % (b[0], b[1]) for b in self.bondlsts] + bonds = {bondstr: None for bondstr in bondstrs} + + # radial basis function expansion coefficients + # saved in n,l,k shape + # defaults to orthogonal delta function [g(n,k)] basis of drautz 2019 + try: + nradmax = max(self.global_nmax[:]) + except ValueError: + nradmax = max(self.global_nmax) + lmax = max(self.global_lmax) + nradbase = self.global_nradbase + crad = np.zeros((nradmax, lmax + 1, nradbase), dtype=int) + for n in range(nradmax): + for l in range(lmax + 1): + crad[n][l] = np.array( + [1 if k == n else 0 for k in range(nradbase)] + ) + + cnew = np.zeros((nradbase, nradmax, lmax + 1)) + for n in range(1, nradmax + 1): + for l in range(lmax + 1): + for k in range(1, nradbase + 1): + cnew[k - 1][n - 1][l] = crad[n - 1][l][k - 1] + + for bondind, bondlst in enumerate(self.bondlsts): + bstr = "[%d, %d]" % (bondlst[0], bondlst[1]) + try: + bonds[bstr] = { + "nradmax": nradmax, + "lmax": max(self.global_lmax), + "nradbasemax": self.global_nradbase, + "radbasename": self.radbasetype, + "radparameters": [self.global_lmbda], + "radcoefficients": crad.tolist(), + "prehc": 0, + "lambdahc": self.global_lmbda, + "rcut": self.global_rcut, + "dcut": 0.01, + "rcut_in": self.global_rcutinner, + "dcut_in": self.global_drcutinner, + "inner_cutoff_type": "distance", + } + except AttributeError: + if type(self.rcut) == dict: + bonds[bstr] = { + "nradmax": nradmax, + "lmax": max(self.global_lmax), + "nradbasemax": self.global_nradbase, + "radbasename": self.radbasetype, + "radparameters": [self.lmbda[bstr]], + "radcoefficients": crad.tolist(), + "prehc": 0, + "lambdahc": self.lmbda[bstr], + "rcut": self.rcut[bstr], + "dcut": 0.01, + "rcut_in": self.rcutinner[bstr], + "dcut_in": self.drcutinner[bstr], + "inner_cutoff_type": "distance", + } + elif type(self.rcut) == list: + bonds[bstr] = { + "nradmax": nradmax, + "lmax": max(self.global_lmax), + "nradbasemax": self.global_nradbase, + "radbasename": self.radbasetype, + "radparameters": [self.lmbda[bondind]], + "radcoefficients": crad.tolist(), + "prehc": 0, + "lambdahc": self.lmbda[bondind], + "rcut": self.rcut[bondind], + "dcut": 0.01, + "rcut_in": self.rcutinner[bondind], + "dcut_in": self.drcutinner[bondind], + "inner_cutoff_type": "distance", + } + + self.bonds = bonds + + def set_funcs(self, nulst=None, muflg=True, print_0s=True): + if nulst == None: + if self.nus != None: + nulst = self.nus.copy() + else: + raise AttributeError("No list of descriptors found/specified") + nus_per_rank = {} + permu0 = {b: [] for b in range(len(self.elements))} + permunu = {b: [] for b in range(len(self.elements))} + if self.betas != None: + betas = self.betas + else: + # betas = {ind:{nu:1.0 for nu in nulst} for ind in range(len(self.elements))} + betas = {ind: {} for ind in range(len(self.elements))} + for nu in nulst: + mu0, mu, n, l, L = acu.get_mu_n_l(nu, return_L=True) + betas[mu0][nu] = 1.0 + for nu in nulst: + mu0, mu, n, l, L = acu.get_mu_n_l(nu, return_L=True) + rank = acu.get_mu_nu_rank(nu) + try: + nus_per_rank[rank].append(nu) + except KeyError: + nus_per_rank[rank] = [nu] + bkeys = list(betas[mu0].keys()) + llst = ["%d"] * rank + # print (nu,l,oldfmt,muflg) + lstr = ",".join(b for b in llst) % tuple(l) + if L != None: + ccs = self.global_ccs[rank][lstr][tuple(L)] + elif L == None: + try: + ccs = self.global_ccs[rank][lstr][()] + except KeyError: + ccs = self.global_ccs[rank][lstr] + ms = list(ccs.keys()) + mslsts = [[int(k) for k in m.split(",")] for m in ms] + msflat = [item for sublist in mslsts for item in sublist] + if print_0s or betas[mu0][nu] != 0.0: + ccoeffs = list(np.array(list(ccs.values())) * betas[mu0][nu]) + permu0[mu0].append( + { + "mu0": mu0, + "rank": rank, + "ndensity": self.global_ndensity, + "num_ms_combs": len(ms), + "mus": mu, + "ns": n, + "ls": l, + "ms_combs": msflat, + "ctildes": ccoeffs, + } + ) + permunu[mu0].append(nu) + elif betas[mu0][nu] == 0.0 and not print_0s: + print("Not printing descriptor: %s, coefficient is 0" % nu) + self.nus_per_rank = nus_per_rank + + # for b in range(len(self.elements)): + # for i in permunu[b]: + # print (b,i) + # for b in range(len(self.elements)): + # print (b,len(permu0[b])) + self.funcs = permu0 + self.permunu = permunu + + def set_betas(self, betas, has_zeros=False): + if type(betas) != dict: + if not has_zeros: + assert len(betas) == len( + self.nus + ), "list of betas must be the same size as list of descriptors (0th order coefficient should NOT be included in this list" + elif has_zeros: + with_nu_inds = len(self.nus) + len(self.elements) + base_N_nu_per_ind = int(len(self.nus) / len(self.elements)) + e0inds = [] + for i in range(len(self.elements)): + e0ind = (i * base_N_nu_per_ind) + i + e0inds.append(e0ind) + e0s = [betas[e0ind] for e0ind in e0inds] + self.E0 = e0s + betas = [b for i, b in enumerate(betas) if i not in e0inds] + + betas_dict = {ind: {} for ind in range(len(self.elements))} + for nu, beta in zip(self.nus, betas): + mu0, mu, n, l, L = acu.get_mu_n_l(nu, return_L=True) + betas_dict[mu0][nu] = beta + self.betas = betas_dict + elif type(betas) == dict: + self.betas = betas + + def read_acecoeff(self, name, remove_0s=False): + f = "%s.acecoeff" % name + coeff_dict = process_acepot(f, self.elements) + e0s = [] + beta_dict = {mu0: {} for mu0 in range(len(self.elements))} + remove_keys = {mu0: [] for mu0 in range(len(self.elements))} + for ind, element in enumerate(self.elements): + for key in coeff_dict.keys(): + if key == "%d_0" % ind: + e0s.append(coeff_dict[key]) + remove_keys[ind].append(key) + else: + mu0, mu, n, l = acu.get_mu_n_l(key) + beta_dict[mu0][key] = coeff_dict[key] + + for ind, element in enumerate(self.elements): + for key in remove_keys[ind]: + try: + del beta_dict[ind][key] + except KeyError: + pass + + self.E0 = e0s + return beta_dict + + def write_pot(self, name): + srt_nus = [] + with open("%s.yace" % name, "w") as writeout: + e0lst = ["%f"] * len(self.elements) + e0str = ", ".join(b for b in e0lst) % tuple(self.E0) + elemlst = ["%s"] * len(self.elements) + elemstr = ", ".join(b for b in elemlst) % tuple(self.elements) + writeout.write("elements: [%s] \n" % elemstr) + writeout.write("E0: [%s] \n" % e0str) + writeout.write("deltaSplineBins: %f \n" % self.deltaSplineBins) + writeout.write("embeddings:\n") + for mu0, embed in self.embeddings.items(): + writeout.write(" %d: " % mu0) + ystr = json.dumps(embed) + "\n" + ystr = ystr.replace('"', "") + writeout.write(ystr) + writeout.write("bonds:\n") + bondstrs = ["[%d, %d]" % (b[0], b[1]) for b in self.bondlsts] + for bondstr in bondstrs: + writeout.write(" %s: " % bondstr) + bstr = json.dumps(self.bonds[bondstr]) + "\n" + bstr = bstr.replace('"', "") + writeout.write(bstr) + writeout.write("functions:\n") + for mu0 in range(len(self.elements)): + writeout.write(" %d:\n" % (mu0)) + mufuncs = self.funcs[mu0] + for mufunc in mufuncs: + mufuncstr = ( + " - " + json.dumps(mufunc, cls=NpEncoder) + "\n" + ) + mufuncstr = mufuncstr.replace('"', "") + writeout.write(mufuncstr) + + +class NpEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, np.integer): + return int(obj) + if isinstance(obj, np.floating): + return float(obj) + if isinstance(obj, np.ndarray): + return obj.tolist() + return super(NpEncoder, self).default(obj)