add posibility to calculate descriptors

[alinelena]

the descriptors for a structure or for a set of structures shall be returned in dictionary for python or a file for cli...

some bad poc in below

# -*- coding: utf-8 -*-
# Author; alin m elena, alin@elena.re
# Contribs;
# Date: 01-03-2024
# ©alin m elena, GPL v3 https://www.gnu.org/licenses/gpl-3.0.en.html
from ase.io import read
import numpy as np
from mace.calculators import mace_mp


import argparse

cli=argparse.ArgumentParser()
cli.add_argument(
  "--device",
  type=str,
  default='cpu',
  choices = ['cuda','cpu','mps'],
  help = ' provide device to run. default %(default)s'
)

cli.add_argument(
  "--precision",
  type=str,
  default='float32',
  choices = ['float64','float32'],
  help = ' provide precision to run. default %(default)s'
)

cli.add_argument(
  "--invariant",
  action="store_true",
  help = 'enable invariant only? default %(default)s'
  )

cli.add_argument(
  "--reference",
  type=str,
  default=None,
  help = 'rederence structure, ase readable, default %(default)s'
)

cli.add_argument(
  "--trajectory",
  type=str,
  default=None,
  help = 'trajectory, default %(default)s'
)

cli.add_argument(
  "--model",
  type=str,
  default='medium',
  help = 'provide the mace model file. default: %(default)s'
)

cli.add_argument(
  "--out",
  type=str,
  default='descriptors.dat',
  help = 'output file. default: %(default)s'
)


def pme(t,el,mem,d,rmem=None,rd=None):
  m = " ".join(list(map(lambda x: f"{x:12.5E}" , [mem[e] for e in el ])))
  s = f"{t:10d} {m} {d:12.5E}"
  if rmem and rd:
     s += " ".join(list(map(lambda x: f"{x:12.5E}" , [mem[e]-rmem[e] for e in el ])))
     s += f"{d-rd:12.5E}"
  return s

def simil(conf,inv):
  species = sorted(set(conf.get_chemical_symbols()))
  patterns = {}
  for e in species:
    patterns[e]=[atom.index for atom in conf if atom.symbol == e]
  descriptors = calculator.get_descriptors(conf,invariants_only=inv)
  me = {}
  mem = {}
  for e in species:
    me[e] = descriptors[patterns[e],:].mean(axis=0)
    mem[e] = np.mean(me[e])
  return species, me, mem, np.mean(descriptors)

args = cli.parse_args()

model = args.model
device = args.device
precision = args.precision

reference = args.reference
traj = args.trajectory

inv = args.invariant
out = args.out

out = open(out,"w+")

# for descriptors single precision is ok
calculator = mace_mp(model=model, default_dtype=precision, device=device)
ref = read(reference)

el, rme, rmem, rd  = simil(ref, inv)
header = f"#Time     | {' | '.join([f'{x:10s}' for x in el])} | Descriptor |"
print("# ", header,file=out)
print("# ref ", pme(0,el, rmem,rd),file=out)
header += f"{' | '.join([f'*{x:10s}' for x in el])}| Descriptor*|"
print(header, file=out)

conf = read(traj, index=":")
for t,f in enumerate(conf):
  el, me, mem, d  = simil(f, inv)
  print(f"{pme(t,el,mem,d,rmem,rd)}",file=out)

out.close()