diff --git a/README.rst b/README.rst index c14fb6636..784f9c85b 100644 --- a/README.rst +++ b/README.rst @@ -81,7 +81,7 @@ in the form of pre-compiled packages that can be easily installed with e.g.:: Create a new conda_ environment (let's call it ``abienv``) with:: - conda create --name abienv + conda create --name abienv python=3.11 and activate it with:: @@ -106,7 +106,10 @@ in the `anaconda howto `_ with:: git clone https://github.com/abinit/abipy @@ -118,7 +121,7 @@ For pip, use:: If you are using conda_ (see `Installing conda`_ to install conda itself), create a new environment (``abienv``) with:: - conda create -n abienv + conda create -n abienv python=3.11 source activate abienv Add ``conda-forge``, and ``abinit`` to your channels with:: @@ -143,22 +146,6 @@ or alternately:: to install the package in developmental mode. This is the recommended approach, especially if you are planning to implement new features. -Note, however, that the developmental version of AbiPy is kept in sync with the -developmental version of pymatgen thus ```python setup.py develop``` may -try to download new versions from the PyPi portal and then fail with e.g. the error message:: - - ... - processing dependencies for abipy==0.6.0.dev0 - error: scipy 1.0.0 is installed but scipy>=1.0.1 is required by {'pymatgen'} - -due to inconsistent dependencies. -To solve the problem, use conda to update scipy to a version >= 1.0.1 with:: - - conda install "scipy>=1.0.1" - -then issue again python setup.py develop. If this fails, supposing you were upgrading abipy inside -an already existing conda environment, try to restart by creating from scratch a fresh conda environment, see above. - Also note that the BLAS/Lapack libraries provided by conda have multithreading support activated by default. Each process will try to use all of the cores on your machine, which quickly overloads things if there are multiple processes running. diff --git a/abipy/abio/robots.py b/abipy/abio/robots.py index 07bf958c3..77522c154 100644 --- a/abipy/abio/robots.py +++ b/abipy/abio/robots.py @@ -817,36 +817,6 @@ def close(self) -> None: print("Exception while closing: ", abifile.filepath) print(exc) - #@classmethod - #def open(cls, obj, nids=None, **kwargs): - # """ - # Flexible constructor. obj can be a :class:`Flow` or a string with the directory containing the Flow. - # `nids` is an optional list of :class:`Node` identifiers used to filter the set of :class:`Task` in the Flow. - # """ - # has_dirpath = False - # if is_string(obj): - # try: - # from abipy.flowtk import Flow - # obj = Flow.pickle_load(obj) - # except: - # has_dirpath = True - - # if not has_dirpath: - # # We have a Flow. smeth is the name of the Task method used to open the file. - # items = [] - # smeth = "open_" + cls.EXT.lower() - # for task in obj.iflat_tasks(nids=nids): #, status=obj.S_OK): - # open_method = getattr(task, smeth, None) - # if open_method is None: continue - # abifile = open_method() - # if abifile is not None: items.append((task.pos_str, abifile)) - # return cls(*items) - - # else: - # # directory --> search for files with the appropriate extension and open it with abiopen. - # if nids is not None: raise ValueError("nids cannot be used when obj is a directory.") - # return cls.from_dir(obj) - #def get_attributes(self, attr_name, obj=None, retdict=False): # od = OrderedDict() # for label, abifile in self.items(): @@ -1241,10 +1211,14 @@ def get_baserobot_code_cells(self, title=None) -> list: @staticmethod def get_yvals_item_abifiles(item: Any, abifiles: list) -> np.ndarray: """Extract values for a list of Abinit files.""" + def _float(obj): + if obj is None: return obj + return float(obj) + if callable(item): - return np.array([float(item(a)) for a in abifiles]) + return np.array([_float(item(a)) for a in abifiles]) else: - return np.array([float(duck.getattrd(a, item)) for a in abifiles]) + return np.array([_float(duck.getattrd(a, item)) for a in abifiles]) @staticmethod def plot_xvals_or_xstr_ax(ax, xs, yvals, fontsize, **kwargs) -> list: diff --git a/abipy/electrons/ebands.py b/abipy/electrons/ebands.py index 04c50e6b3..fecf41f27 100644 --- a/abipy/electrons/ebands.py +++ b/abipy/electrons/ebands.py @@ -3933,7 +3933,7 @@ def boxplot(self, e0="fermie", brange=None, swarm=False, fontsize=8, **kwargs) - ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, sharex=False, sharey=True, squeeze=False) ax_list = ax_list.ravel() - # don't show the last ax if numeb is odd. + # don't show the last ax if num_plots is odd. if num_plots % ncols != 0: ax_list[-1].axis("off") for (label, ebands), ax in zip(self.ebands_dict.items(), ax_list): diff --git a/abipy/electrons/gwr.py b/abipy/electrons/gwr.py index 6a4ff5978..9d4d99258 100644 --- a/abipy/electrons/gwr.py +++ b/abipy/electrons/gwr.py @@ -38,6 +38,7 @@ class _MyQpkindsList(list): """Returned by find_qpkinds.""" + @dataclasses.dataclass class MinimaxMesh: """ diff --git a/abipy/electrons/lruj.py b/abipy/electrons/lruj.py new file mode 100644 index 000000000..6c1f10c25 --- /dev/null +++ b/abipy/electrons/lruj.py @@ -0,0 +1,399 @@ +# coding: utf-8 +"""Classes to analyse LRUJ results.""" +from __future__ import annotations + +import dataclasses +import numpy as np +import pandas as pd +from pathlib import Path +#from typing import Any +#from monty.string import is_string, list_strings, marquee +from abipy.core.mixins import AbinitNcFile, Has_Header, Has_Structure, Has_ElectronBands, NotebookWriter +from abipy.iotools import ETSF_Reader +from abipy.tools.iotools import yaml_safe_load +from abipy.core.structure import Structure +from abipy.tools.typing import Figure, PathLike +from abipy.tools.plotting import (set_axlims, add_fig_kwargs, get_ax_fig_plt, get_axarray_fig_plt, + get_ax3d_fig_plt, rotate_ticklabels, set_visible, plot_unit_cell, set_ax_xylabels, get_figs_plotly) + + + +#class LrujFile(AbinitNcFile, Has_Header, Has_Structure): #, Has_ElectronBands, NotebookWriter): +# """ +# File containing the results of a ground-state calculation. +# +# Usage example: +# +# .. code-block:: python +# +# with GsrFile("foo_GSR.nc") as gsr: +# print("energy: ", gsr.energy) +# gsr.ebands.plot() +# +# .. rubric:: Inheritance Diagram +# .. inheritance-diagram:: LrujFile +# """ +# +# @classmethod +# def from_file(cls, filepath: str) -> GsrFile: +# """Initialize the object from a netcdf_ file.""" +# return cls(filepath) +# +# def __init__(self, filepath: str): +# super().__init__(filepath) +# self.r = r = EtsfReader(filepath) + + +#=============================================================================================================== +#=============================================================================================================== +@dataclasses.dataclass(kw_only=True) +class LrujResults: + """ + This object stores the results produced by lruj. + """ + natom: int + npert: int + ndata: int + pawujat: int + macro_uj: int + diem_token: str + diem: float + alphas: np.ndarray + occ_unscr: np.ndarray + occ_scr: np.ndarray + chi0_coefficients: dict + chi_coefficients: dict + maxdeg: int + dmatpuopt: int + pert_name: str + parname: str + metric: str + fit_df: pd.DataFrame + + @classmethod + def from_file(cls, filepath: PathLike): + """ + Extract results from the main ouput file produced by lruj. + """ + with open(filepath, "rt") as fh: + lines = [line.lstrip() for line in fh] + + # Extract the Yaml document with the chi/chi0 coefficients + in_doc = False + yaml_lines = [] + for i, line in enumerate(lines): + + if line.startswith("--- !LRUJ_Abipy_Plots"): + in_doc = True + continue + + if in_doc and line.startswith("..."): + data = yaml_safe_load("".join(yaml_lines)) + break + + if in_doc: + yaml_lines.append(line) + + natom = data['natom'] + ndata = data['ndata'] + pawujat = data['pawujat'] + macro_uj = data['macro_uj'] + diem_token = data['diem_token'] + diem = data['diem'] + npert = ndata - 1 + if macro_uj==4: + pert_name = r"$\beta$" + metric = r"M $(n^{\uparrow} - n^{\downarrow})$" + parname = 'J' + else: + pert_name = r"$\alpha$" + metric = r"N $(n^{\uparrow} + n^{\downarrow})$" + parname = 'U' + + chi0_coefficients = {} + chi_coefficients = {} + for k, v in data.items(): + magic = "chi0_coefficients_degree" + if k.startswith(magic): + degree = int(k.replace(magic, "")) + chi0_coefficients[degree] = v + magic = "chi_coefficients_degree" + if k.startswith(magic): + degree = int(k.replace(magic, "")) + chi_coefficients[degree] = v + + #print(f"{chi_coefficients=}") + + def find(header, dtype=None): + for i, line in enumerate(lines): + if line.startswith(header): + after = line.replace(header, "", 1).strip() + if dtype: after = dtype(after) + return i, after + raise ValueError(f"Cannot find {header=} in {filepath=}") + + _, maxdeg = find("Maximum degree of polynomials analyzed:", dtype=int) + _, dmatpuopt = find("Value of dmatpuopt:", dtype=int) + + # Parse the section with perturbations and occupations. + """ + Perturbations Occupations + -------------- ----------------------------- + alpha [eV] Unscreened Screened + -------------- ----------------------------- + 0.0000000000 8.6380182458 8.6380182458 + -0.1500000676 8.6964981922 8.6520722003 + + -OR- + + Perturbations Magnetizations + --------------- ----------------------------- + beta [eV] Unscreened Screened + --------------- ----------------------------- + + """ + i, _ = find("Perturbations",dtype=None) + i += 4 + vals = [] + for ipert in range(ndata): + vals.append([float(t) for t in lines[i+ipert].split()]) + vals = np.reshape(vals, (ndata, 3)) + alphas, occ_unscr, occ_scr = vals[:,0], vals[:,1], vals[:,2] + """ + RMS Errors + --------------------------------------- + Regression Chi0 [eV^-1] Chi [eV^-1] J [eV] | Chi0 [eV^-1] Chi [eV^-1] J [eV] + --------------------------------------------------------|--------------------------------------- + Linear: -0.8594082 -0.0949434 9.3689952 | 0.0023925 0.0000878 0.1139297 + Quadratic: -0.8574665 -0.0955791 9.2963099 | 0.0023777 0.0000129 0.0681722 + Cubic: -0.8007858 -0.0952726 9.2474220 | 0.0001546 0.0000015 0.0200543 + ************************************************************************************************* + """ + i, _ = find("Regression Chi0 [eV^-1]") + i += 2 + keys = ["Chi0", "Chi", "HP", "rms_Chi0", "rms_Chi", "rms_HP"] + dict_list = [] + for irow in range(maxdeg): + l = lines[i+irow].replace("|", " ") + tokens = l.split() + d = dict(zip(keys, [float(t) for t in tokens[-6:]])) + d["degree"] = irow + 1 + dict_list.append(d) + + fit_df = pd.DataFrame(dict_list) + + # Build instance from locals dict. + _data = locals() + return cls(**{k: _data[k] for k in [field.name for field in dataclasses.fields(cls)]}) + +#=============================================================================================================== +#=============================================================================================================== + @add_fig_kwargs + def plot(self, ax=None, degrees="all", inset=True, insetdegree=1, insetlocale="lower left", ptcolor0='k', ptcolor='k', gradcolor1='#3575D5',gradcolor2='#FDAE7B', ptitle="default", fontsize=12, **kwargs) -> Figure: + """ + Plot + + Args: + ax: |matplotlib-Axes| or None if a new figure should be created. + degrees: List of degrees to plot. If None, no polynomials are plotted. + ptcolor0: Color of unscreened response data points (default: black) + ptcolor: Color of screened response data points (default: black) + gradcolor1: Hex code of linear regression color (default: Blue #3575D5) + gradcolor2: Hex code of color of regression of maximum degree in list (default: Salmon #FDAE7B) + ptitle: Plot title (default: "Linear Response for atom ) + inset: Plots inset with LR parameter for polynomial fit of degree (default: True) + insetdegree: Degree of polynomial fit information printed in the inset (default: 1) + insetlocale: Position of inset in the plot. Standard matplotlob locations. (default: "lower left") + """ + import seaborn as sns + sns.set() + ax, fig, plt = get_ax_fig_plt(ax=ax) + + # Plot data + yshift = self.occ_unscr[np.where(self.alphas == 0.0000)] * (self.diem - 1.0) + data0 = 1.0/self.diem * (self.occ_unscr + yshift) + ax.scatter(self.alphas, data0, s=70, color=ptcolor0, facecolors='none', linewidths=2, label="Unscreened") + ax.scatter(self.alphas, self.occ_scr, s=70, color=ptcolor, label="Screened") + ax.axvline(x=0.0, color="white", linestyle="--", lw=0.5) + + # Generate mesh for polynomial functions + xstart, xstop = 1.1 * self.alphas.min(), 1.1 * self.alphas.max() + xs = np.arange(xstart, xstop, step=0.01) + + #Prepare colors and coefficients for polynomials the use wants to plot + if degrees == "all": + degrees = self.chi0_coefficients.keys() + + def hex_to_RGB(hex_str): + return [int(hex_str[i:i+2], 16) for i in range(1,6,2)] + + def get_color_gradient(c1, c2, n): + assert n > 1 + c1_rgb = np.array(hex_to_RGB(c1))/255 + c2_rgb = np.array(hex_to_RGB(c2))/255 + mix_pcts = [x/(n-1) for x in range(n)] + rgb_colors = [((1-mix)*c1_rgb + (mix*c2_rgb)) for mix in mix_pcts] + return ["#" + "".join([format(int(round(val*255)), "02x") for val in item]) for item in rgb_colors] + + linecolors=get_color_gradient(gradcolor1,gradcolor2,len(degrees)) + + # Plot polynomial functions + def poly0(coeffs): + return lambda x: 1.0/self.diem * (sum((coeff*x**i for i,coeff in enumerate(coeffs))) + yshift) + + def poly(coeffs): + return lambda x: sum((coeff*x**i for i,coeff in enumerate(coeffs))) + + for degree in degrees: + polynomial0=poly0(self.chi0_coefficients[degree]) + polynomial=poly(self.chi_coefficients[degree]) + if degree == 1: + Labelstring='Linear' + elif degree == 2: + Labelstring='Quadratic' + elif degree == 3: + Labelstring='Cubic' + else: + Labelstring=' '.join(['Degree',str(degree)]) + + if insetdegree==degree: + deginfo = ' '.join(['Parameters for',Labelstring,'fit']) + insetcolor = linecolors[degree-1] + + ax.plot(xs,polynomial0(xs),color=linecolors[degree-1],linewidth=2.0,linestyle='dashed') + ax.plot(xs,polynomial(xs),color=linecolors[degree-1],linewidth=2.0,label=Labelstring) + + ax.legend(loc="best", fontsize=fontsize, shadow=True) + if ptitle=="default": + ptitle=' '.join(["Linear Response",self.parname,"on atom",str(self.pawujat)]) + plt.title(ptitle) + ax.grid(True) + ax.set_xlabel(' '.join([self.pert_name,"(eV)"])) + ax.set_ylabel(self.metric) + + # Generate inset with numerical information on LR + if inset: + from matplotlib.offsetbox import AnchoredText + select = self.fit_df["degree"] == insetdegree + def dfvalue(keywo): + tablevalue = self.fit_df[select][keywo] + return "%.4f" % tablevalue.values[0] + + X0str = ' '.join([r"$\chi_0$","=",dfvalue("Chi0"),r"$\pm$",dfvalue("rms_Chi0"),r"[eV]$^{-1}$"]) + Xstr = ' '.join([r"$\chi$","=",dfvalue("Chi"),r"$\pm$",dfvalue("rms_Chi"),r"[eV]$^{-1}$"]) + HPstr = ' '.join([self.parname,"=",dfvalue("HP"),r"$\pm$",dfvalue("rms_HP"),r"[eV]"]) + insettxt = '\n'.join([deginfo,X0str,Xstr,HPstr]) + parambox = AnchoredText(insettxt,loc=insetlocale) + parambox.patch.set_linewidth(2) + parambox.patch.set_edgecolor(insetcolor) + parambox.patch.set_facecolor('white') + ax.add_artist(parambox) + return fig + + + +#=============================================================================================================== +#=============================================================================================================== +class LrujAnalyzer: + """ + Analyzes multiple sets of LRUJ files. + """ + def __init__(self, manager=None, verbose=0): + self.ncfiles_of_key = {} + self.results_of_key = {} + self.manager = manager + self.verbose = verbose + + #def __str__(self): + # return self.to_string() + + #def to_string(self, verbose: int = 0) -> str: + # lines = [] + # app = lines.append + # return "\n".join(lines) + + #@classmethod + #def from_dir(cls, key: str, directory: PathLike) -> None: + # new = cls() + # new.scan_dir(key, directory) + # return new + + #def walk_dir(self, key: str, top: PathLike) -> None: + # nc_paths = None + # self.add_ncpaths(key, nc_paths) + + def add_ncpaths(self, key: str, nc_paths: list[str]) -> None: + self.ncfiles_of_key[key] = nc_paths + self.results_of_key[key] = None + self.run() + + def run(self, workdir=None) -> None: + """ + Invoke lruj + """ + from abipy.flowtk.wrappers import Lruj + from abipy.core.globals import get_workdir + workdir = Path(get_workdir(workdir)) + for key, nc_paths in self.ncfiles_of_key.items(): + if self.results_of_key[key] is not None: continue + wdir = workdir / f"run_{key}" + wdir.mkdir() + Lruj(manager=self.manager, verbose=self.verbose).run(nc_paths, workdir=wdir) + self.results_of_key[key] = LrujResults.from_file(wdir / "lruj.stdout") + + @add_fig_kwargs + def plot(self, **kwargs) -> Figure: + """ + Plot + + Args: + """ + keys = self.results_of_key.keys() + + # Build grid of plots. + num_plots, ncols, nrows = len(keys), 1, 1 + if num_plots > 1: + ncols = 2 + nrows = (num_plots // ncols) + (num_plots % ncols) + + ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, + sharex=True, sharey=True, squeeze=False) + ax_list = ax_list.ravel() + # don't show the last ax if num_plots is odd. + if num_plots % ncols != 0: ax_list[-1].axis("off") + + for ix, (key, ax) in enumerate(zip(keys, ax_list)): + res = self.results_of_key[key] + res.plot(ax=ax, show=False, title=key) + + return fig + + +#class LrujInputGenerator +# +# def __init__(self, gs_input, elements, site_inds=None, dmatpuopt=3): +# self.dmatpuopt = dmatpuopt +# self.gs_input = gs_input.copy() +# self.gs_input.set_vars( +# chkprim=0, # Will complain otherwise with AFII magnetic state +# chksymbreak=0, # Don't check for symmetry breaking +# # DFT+U +# usepawu=1, # Alert Abinit to use of DFT+U +# lpawu=[2, 2, 1], # Subspaces of interest: Ni 3d, O 2p +# upawu="*0.0", # Raw (non-corrected) XC functional required to establish U(J) +# jpawu="*0.0", +# dmatpuopt=self.dmatpuopt, # PAW projector for density matrix +# ) +# +# def gen_inputs(self, macro_uj, alphas, supercells): +# #supercells = np.respahe(supercells, (-1, +# for supercell in supercells: +# for alpha in alphas: +# if alpha == 0.0: continue +# gs_input.new_with_vars( +# macro_uj=macro_uj, +# pawujat=1, # Label of perturbed atom +# pawujv1=f"{alpha} eV", # List of perturbation strengths +# prtwf=0, +# prtebands 0 # Don't print ebands +# ) diff --git a/abipy/flowtk/abiinspect.py b/abipy/flowtk/abiinspect.py index bf8879546..a928fafc6 100644 --- a/abipy/flowtk/abiinspect.py +++ b/abipy/flowtk/abiinspect.py @@ -11,7 +11,6 @@ from typing import Union import numpy as np -#import ruamel.yaml as yaml from monty.collections import AttrDict from monty.functools import lazy_property from tabulate import tabulate diff --git a/abipy/flowtk/qutils.py b/abipy/flowtk/qutils.py index 9034fca20..f050f4d74 100644 --- a/abipy/flowtk/qutils.py +++ b/abipy/flowtk/qutils.py @@ -181,8 +181,8 @@ class SlurmJobArray: #SBATCH --job-name=abiml_md #SBATCH --time=0-16:0:0 #SBATCH --partition=batch -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 # 1 node has 128 cores +#SBATCH --nodes=1 # 1 node has 128 cores +#SBATCH --ntasks-per-node=1 #SBATCH --cpus-per-task=1 conda activate env3.10 @@ -194,7 +194,7 @@ class SlurmJobArray: arr_options = ["--help", "--version"] job_array = SlurmJobArray(header, command, arr_options) print(job_array) -queue_id = job_array.sbatch("job.sh") +queue_id = job_array.sbatch("job_array.sh") """ def __init__(self, header: str, command: str, arr_options: list[str]): @@ -212,6 +212,7 @@ def __str__(self): break else: raise ValueError("Cannot find line starting with #SBATCH") + lines.insert(il, f"#SBATCH --array=0-{len(self.arr_options)-1}") header = "\n".join(lines) @@ -234,28 +235,46 @@ def __str__(self): env """ % (self.arr_options_str) - end = f"{self.command} ${{OPTS}} > job_${{index}}.log 2> job_${{index}}.err" + end = f""" +{self.command} ${{OPTS}} > job_${{index}}.log 2> job_${{index}}.err + +# Remove the file with the Slurm job id +me=$(basename "$0") +rm ${{me}}.qid +""" return header + select_opts + end def sbatch(self, slurm_filepath: PathLike) -> int: + """ + Write slurm submission script to slurm_filepath and submits it. + Return Slurm JOB id. + """ + # Make sure no slurm job is already running by checking for a .qid file. + path_qid = slurm_filepath + ".qid" + if os.path.exists(path_qid): + with open(path_qid, "rt") as fh: + queue_id = int(fh.read().split("#")) + err_msg = f"Found slurm job ID {queue_id} in {path_qid}" + \ + "This usually indicates that a similar array job is already running\n" + \ + f"If this not the case, please remove {path_qid} and rerun the script." + raise RuntimeError(err_msg) + with open(slurm_filepath, "wt") as fh: fh.write(str(self)) queue_id = slurm_sbatch(slurm_filepath) - save_qid = slurm_filepath + ".qid", - print("Saving Slurm job ID to file:", save_qid) - with open(save_qid, "wt") as fh: - fh.write("# Slurm job id\n") - fh.write(str(queue_id)) + print("Saving slurm job ID in:", path_qid) + with open(path_qid, "wt") as fh: + fh.write(str(queue_id) + " # Slurm job id") return queue_id def slurm_sbatch(script_file) -> int: """ - Submit a job script to the queue with sbatch. Return JOB ID. + Submit a job script to the queue with sbatch. Return Slurm JOB ID. """ from subprocess import Popen, PIPE # need string not bytes so must use universal_newlines @@ -267,7 +286,7 @@ def slurm_sbatch(script_file) -> int: try: # output should of the form '2561553.sdb' or '352353.jessup' - just grab the first part for job id queue_id = int(out.split()[3]) - print('Job submission was successful and queue_id is {}'.format(queue_id)) + print(f"Job submission was successful and {queue_id=}") return queue_id except Exception as exc: # probably error parsing job code diff --git a/abipy/flowtk/wrappers.py b/abipy/flowtk/wrappers.py index 12ad7c7ba..a60c75d2a 100644 --- a/abipy/flowtk/wrappers.py +++ b/abipy/flowtk/wrappers.py @@ -99,7 +99,6 @@ def _execute(self, workdir, with_mpirun=False, exec_args=None) -> int: qjob, process = qadapter.submit_to_queue(script_file) self.stdout_data, self.stderr_data = process.communicate() self.returncode = process.returncode - #raise self.Error("%s returned %s\n cmd_str: %s" % (self, self.returncode, self.cmd_str)) return self.returncode @@ -296,7 +295,7 @@ def cut3d(self, cut3d_input, workdir) -> tuple[str, str]: retcode = self._execute(workdir, with_mpirun=False) if retcode != 0: - raise RuntimeError("Error while running cut3d in %s." % workdir) + raise RuntimeError("Error while running cut3d in %s" % workdir) output_filepath = cut3d_input.output_filepath @@ -345,3 +344,33 @@ def unfold(self, wfkpath, folds, workdir=None) -> str: raise RuntimeError("Cannot find *_FOLD2BLOCH.nc file in: %s" % str(os.listdir(workdir))) return os.path.join(workdir, filepaths[0]) + + +class Lruj(ExecWrapper): + """ + Wraps the lruj Fortran executable. + """ + _name = "lruj" + + def run(self, nc_paths: list[str], workdir=None) -> int: + """ + Execute lruj inside directory `workdir` to analyze `nc_paths`. + """ + workdir = get_workdir(workdir) + + self.stdin_fname = None + self.stdout_fname, self.stderr_fname = \ + map(os.path.join, 2 * [workdir], ["lruj.stdout", "lruj.stderr"]) + + # We work with absolute paths. + nc_paths = [os.path.abspath(s) for s in list_strings(nc_paths)] + + retcode = self.execute(workdir, exec_args=nc_paths) + if retcode != 0: + print("stdout:") + print(self.stdout_data) + print("stderr:") + print(self.stderr_data) + raise RuntimeError("Error while running lruj in %s" % workdir) + + return retcode diff --git a/abipy/ml/aseml.py b/abipy/ml/aseml.py index 0a25a0698..dd9d4560a 100644 --- a/abipy/ml/aseml.py +++ b/abipy/ml/aseml.py @@ -201,12 +201,13 @@ def diff_two_structures(label1, structure1, label2, structure2, fmt, file=sys.st @dataclasses.dataclass class AseResults: """ - Container with the results produced by the ASE calculator. + Container with the results produced by an ASE calculator. """ atoms: Atoms ene: float - stress: np.ndarray # 3x3 matrix with stress + stress: np.ndarray # 3x3 matrix with stress tensor. forces: np.ndarray + magmoms: np.ndarray # None if calculator does not provide magmoms. @classmethod def from_traj_inds(cls, trajectory: Trajectory, *inds) -> AseResults: @@ -221,13 +222,20 @@ def from_atoms(cls, atoms: Atoms, calc=None) -> AseResults: from ase.stress import voigt_6_to_full_3x3_strain stress_voigt = atoms.get_stress() - #print(stress_voigt) stress = voigt_6_to_full_3x3_strain(stress_voigt) + from ase.calculators.calculator import PropertyNotImplementedError + try: + magmoms = atoms.get_magnetic_moments() + except PropertyNotImplementedError: + magmoms = None + results = cls(atoms=atoms.copy(), ene=float(atoms.get_potential_energy()), stress=stress, - forces=atoms.get_forces()) + forces=atoms.get_forces(), + magmoms=magmoms) + if calc is not None: atoms.calc = None return results @@ -1187,8 +1195,8 @@ def install_nn_names(nn_names="all", update=False, verbose=0) -> None: Install NN potentials in the environment using pip. Args: - nn_names: - update: + nn_names: List of NN potentisl to install. + update: True if packages should be updated. verbose: Verbosity level. """ def pip_install(nn_name) -> int: @@ -1248,9 +1256,9 @@ class CalcBuilder: "chgnet", "alignn", "mace", - #"pyace", - #"quip", - #"nequip", + "pyace", + "nequip", + "metatensor", ] def __init__(self, name: str, **kwargs): @@ -1268,11 +1276,11 @@ def __init__(self, name: str, **kwargs): raise ValueError(f"Invalid {name=}, it should be in {self.ALL_NN_TYPES=}") self._model = None - #print(self) def __str__(self): if self.model_name is not None: return f"{self.__class__.__name__} nn_type: {self.nn_type}, model_name: {self.model_name}" + return f"{self.__class__.__name__} nn_type: {self.nn_type}" # pickle support. @@ -1399,11 +1407,14 @@ class MyPyACECalculator(_MyCalculator, PyACECalculator): class MyMACECalculator(_MyCalculator, MACECalculator): """Add abi_forces and abi_stress""" + self.model_path = os.path.expanduser("~/NN_MODELS/2023-08-14-mace-universal.model") + print("Using MACE model_path:", self.model_path) + if self.model_path is None: raise RuntimeError("MACECalculator requires model_path e.g. nn_name='mace@FILEPATH'") cls = MyMACECalculator if with_delta else MACECalculator - return cls(model_path=self.model_path, device="cpu") #, default_dtype='float32') + return cls(model_paths=self.model_path, device="cpu") #, default_dtype='float32') if self.nn_type == "nequip": try: @@ -2778,10 +2789,9 @@ def run(self, steps: int): class MlCompareNNs(MlBase): """ - Compare energies, forces and stresses obtaiend with different ML potentials. + Compare energies, forces and stresses obtained with different ML potentials. Also profile the time required. """ - def __init__(self, atoms, nn_names, num_tests, rattle, stdev_rvol, verbose, workdir, prefix=None): """ Args: @@ -2856,3 +2866,207 @@ def run(self, print_dataframes=True) -> AseResultsComparator: self._finalize() return comp + + + +class MlCwfEos(MlBase): + """ + Compute EOS with different ML potentials. + """ + def __init__(self, elements, nn_names, verbose, workdir, prefix=None): + """ + Args: + atoms: ASE atoms + elements: String or List of strings with the chemical symbols to consider. + nn_names: String or list of strings defining the NN potential. See also CalcBuilder. + verbose: Verbosity level. + workdir: Working directory. + """ + super().__init__(workdir, prefix) + self.elements = list_strings(elements) + self.nn_names = list_strings(nn_names) + self.verbose = verbose + + self.configurations_set_name = { + "unaries-verification-PBE-v1": ["X/SC", "X/FCC", "X/BCC", "X/Diamond"], + "oxides-verification-PBE-v1": ["XO", "XO2", "XO3", "X2O", "X2O3", "X2O5"], + } + + root = Path("/Users/giantomassi/git_repos/acwf-verification-scripts/0-preliminary-do-not-run") + self.dirpath_set_name = { + "unaries-verification-PBE-v1": root / "unaries" / "xsfs-unaries-verification-PBE-v1", + "oxides-verification-PBE-v1": root / "oxides" / "xsfs-oxides-verification-PBE-v1", + } + + def to_string(self, verbose=0) -> str: + """String representation with verbosity level `verbose`.""" + return f"""\ + +{self.__class__.__name__} parameters: + + elements = {self.elements} + nn_names = {self.nn_names} + workdir = {self.workdir} + verbose = {self.verbose} +""" + + def run(self): + for nn_name in self.nn_names: + for set_name in self.configurations_set_name: + self.run_nn_name_set_name(nn_name, set_name) + + def run_nn_name_set_name(self, nn_name: str, set_name: str) -> dict: + print(f"Computing ML EOS with {nn_name=}, {set_name=} ...") + # This piece is taken from get_results.py + try: + from acwf_paper_plots.eosfit_31_adapted import BM, echarge + except ImportError as exc: + raise ImportError("ase not installed. Try `pip install ase`.") from exc + + calc = as_calculator(nn_name) + warning_lines = [] + + uuid_mapping = {} + all_missing_outputs = {} + completely_off = [] + failed_wfs = [] + all_eos_data = {} + all_stress_data = {} + all_BM_fit_data = {} + num_atoms_in_sim_cell = {} + + import uuid + for element in self.elements: + for configuration in self.configurations_set_name[set_name]: + my_uuid = uuid.uuid4().hex + uuid_mapping[f'{element}-{configuration}'] = { + 'structure': my_uuid, + 'eos_workflow': my_uuid + } + + #count = 7 + #increment = 0.02 + #return tuple(float(1 + i * increment - (count - 1) * increment / 2) for i in range(count)) + + conf = configuration.replace("X/", "") + filepath = self.dirpath_set_name[set_name] / f"{element}-{conf}.xsf" + v0_atoms = read(filepath, format='xsf') + v0 = v0_atoms.get_volume() + volumes = (v0 * np.array([0.94, 0.96, 0.98, 1.00, 1.02, 1.04, 1.06])).tolist() + energies, stresses = [], [] + + # Initialize to None if the outputs are not there + eos_data = None + stress_data = None + BM_fit_data = None + num_atoms = len(v0_atoms) + + try: + for volume in volumes: + #ase = structure.get_ase().copy() + #ase.set_cell(ase.get_cell() * float(scale_factor)**(1 / 3), scale_atoms=True) + atoms = to_ase_atoms(Structure.as_structure(v0_atoms).scale_lattice(volume)) + r = AseResults.from_atoms(atoms, calc=calc) + energies.append(r.ene) + stresses.append(r.stress.tolist()) + + eos_data = list(zip(volumes, energies)) + stress_data = list(zip(volumes, stresses)) + # This line disables the visualization of stress + stress_data = None + #for v, e in eos_data: print(v, e) + #except IndexError: + + # Check if the central point was completely off (i.e. the minimum of the energies is + # on the very left or very right of the volume range) + min_loc = np.array(energies).argmin() + if min_loc == 0: + # Side is whether the minimum occurs on the left side (small volumes) or right side (large volumes) + completely_off.append({'element': element, 'configuration': configuration, 'side': 'left'}) + elif min_loc == len(energies) - 1: + completely_off.append({'element': element, 'configuration': configuration, 'side': 'right'}) + + try: + min_volume, E0, bulk_modulus_internal, bulk_deriv, residuals = BM(np.array(eos_data)) + bulk_modulus_GPa = bulk_modulus_internal * echarge * 1.0e21 + #1 eV/Angstrom3 = 160.21766208 GPa + bulk_modulus_ev_ang3 = bulk_modulus_GPa / 160.21766208 + BM_fit_data = { + 'min_volume': min_volume, + 'E0': E0, + 'bulk_modulus_ev_ang3': bulk_modulus_ev_ang3, + 'bulk_deriv': bulk_deriv, + 'residuals': residuals[0] + } + if residuals[0] > 1.e-3: + warning_lines.append(f"WARNING! High fit residuals: {residuals[0]} for {element} {configuration}") + except ValueError as exc: + # If we cannot find a minimum + # Note that BM_fit_data was already set to None at the top + warning_lines.append(f"WARNING! Unable to fit for {element=} {configuration=}") + #print(str(exc)) + + except Exception as exc: + warning_lines.append(f"WARNING! Unable to compute E(V) for {element=} {configuration=}") + #print(str(exc)) + + all_eos_data[f'{element}-{configuration}'] = eos_data + num_atoms_in_sim_cell[f'{element}-{configuration}'] = num_atoms + if stress_data is None: + stress_data = list(zip(volumes, [None for _ in range(len(volumes))])) + all_stress_data[f'{element}-{configuration}'] = stress_data + all_BM_fit_data[f'{element}-{configuration}'] = BM_fit_data + + data = { + 'script_version': "0.0.4", + 'set_name': set_name, + # Mapping from strings like "He-X2O" to a dictionary with the UUIDs of the structure and the EOS workflow + 'uuid_mapping': uuid_mapping, + # A list of dictionaries with information on the workchains that did not finish with a 0 exit code + 'failed_wfs': failed_wfs, + # A dictionary that indicate for which elements and configurations there are missing outputs, + # (only for the workchains that still had enough volumes to be considered for a fit) + 'missing_outputs': all_missing_outputs, + # A list of dictionaries that indicate which elements and configurations have been computed completely + # off-centre (meaning that the minimum of all computed energies is on either of the two edges, i.e. for + # the smallest or largest volume) + 'completely_off': completely_off, + # Dictionary with the EOS data (volumes and energies datapoints). The keys are the same as the `uuid_mapping`. + # Values can be None. + 'eos_data': all_eos_data, + 'stress_data': all_stress_data, + # Birch-Murnaghan fit data. See above for the keys. Can be None. + 'BM_fit_data': all_BM_fit_data, + 'num_atoms_in_sim_cell': num_atoms_in_sim_cell + } + + # Print some statistics on the results + warning_lines.append("") + #warning_lines.append("Counter of states: " + str(Counter(states))) + good_cnt = len([eos_data for eos_data in data['eos_data'] if eos_data is not None]) + warning_lines.append("") + warning_lines.append(f"Minimum completely off for {len(completely_off)}/{good_cnt}") + warning_lines.append("Completely off systems (symbol indicates if the minimum is on the very left or right):") + for system in data['completely_off']: + warning_lines.append( + f"- {system['element']} {system['configuration']} " + f"({'<' if system['side'] == 'left' else '>'})" + ) + + SET_NAME = set_name + PLUGIN_NAME = nn_name + + fname = self.workdir / f"warnings-{SET_NAME}-{PLUGIN_NAME}.txt" + with open(fname, 'w') as fhandle: + for line in warning_lines: + fhandle.write(f"{line}\n") + print(line) + print(f"Warning log written to: '{fname}'.") + + # Output results to file + fname = self.workdir / f"results-{SET_NAME}-{PLUGIN_NAME}.json" + with open(fname, 'w') as fhandle: + json.dump(data, fhandle, indent=2, sort_keys=True) + print(f"Output results written to: '{fname}'.") + + return data diff --git a/abipy/ml/ml_phonopy.py b/abipy/ml/ml_phonopy.py index 81cfcaadc..2f4b6376f 100644 --- a/abipy/ml/ml_phonopy.py +++ b/abipy/ml/ml_phonopy.py @@ -26,7 +26,6 @@ Phonopy = None - def cprint_traceback(color="red") -> None: """Print traceback.""" import traceback @@ -110,9 +109,6 @@ def __init__(self, ddb_filepath, prefix: Prefix for workdir supercell: with supercell dimensions. None to use the supercell from the DDB file. """ - # Store args for reconstruction - #self.init_kwargs = {k: v for k, v in locals().items() if k not in ["self", "__class__", "kwargs"]} - super().__init__(workdir, prefix) self.distance = float(distance) @@ -270,7 +266,6 @@ def _run_nn_name(self, nn_name: str) -> None: show = False from abipy.dfpt.phonons import PhononBands, PhononBandsPlotter - with Timer(header="Starting phonopy ph-bands computation...", footer=""): phonon.run_band_structure(self.py_qpoints, with_eigenvectors=True) @@ -285,7 +280,6 @@ def _run_nn_name(self, nn_name: str) -> None: for q_list, w_list, eig_list in zip(bands_dict['qpoints'], bands_dict['frequencies'], bands_dict['eigenvectors']): nqpt += len(q_list) py_phfreqs.extend(w_list) - #print(eig_list) py_displ_cart.extend(eig_list) py_phfreqs = np.reshape(py_phfreqs, (nqpt, 3*natom)) / abu.eV_to_THz @@ -347,11 +341,7 @@ def __init__(self, structure, supercell, verbose: Verbosity level. workdir: Working directory, None to generate temporary directory automatically. prefix: Prefix for workdir. - """ - # Store args for reconstruction - #self.init_kwargs = {k: v for k, v in locals().items() if k not in ["self", "__class__", "kwargs"]} - super().__init__(workdir, prefix) self.initial_atoms = structure.to_ase_atoms() @@ -465,7 +455,7 @@ def _run_nn_name(self, nn_name: str) -> None: with_group_velocities=False, plot=True, write_yaml=False, - filename=workdir / f"{nn_name}_band.yml") + filename=workdir / f"{nn_name}_band.yml", ) plt.savefig(workdir / f"phonopy_{nn_name}_phbands.png") if show: plt.show() diff --git a/abipy/ml/relax_scanner.py b/abipy/ml/relax_scanner.py index 371939f2d..5792502cf 100644 --- a/abipy/ml/relax_scanner.py +++ b/abipy/ml/relax_scanner.py @@ -28,7 +28,27 @@ from abipy.tools.serialization import HasPickleIO from abipy.tools.printing import print_dataframe from abipy.ml.aseml import (relax_atoms, get_atoms, as_calculator, ase_optimizer_cls, RX_MODE, fix_atoms, - MlNeb, MlGsList, CalcBuilder, make_ase_neb, nprocs_for_ntasks) + MlNeb, MlGsList, CalcBuilder, make_ase_neb) + + + +def nprocs_for_ntasks(nprocs, ntasks, title=None) -> int: + """ + Return the number of procs to be used in a multiprocessing Pool. + If negative or None, use all procs in the system. + """ + import os + if nprocs is None or nprocs <= 0: + nprocs = max(1, os.cpu_count()) + else: + nprocs = int(nprocs) + + nprocs = min(nprocs, ntasks) + if title is not None: + print(title) + print(f"Using multiprocessing pool with {nprocs=} for {ntasks=} ...") + return nprocs + @dataclasses.dataclass diff --git a/abipy/scripts/abicomp.py b/abipy/scripts/abicomp.py index 468f7664d..181f62c62 100755 --- a/abipy/scripts/abicomp.py +++ b/abipy/scripts/abicomp.py @@ -16,7 +16,7 @@ from monty.functools import prof_main from monty.termcolor import cprint from abipy import abilab -from abipy.tools.plotting import get_ax_fig_plt, GenericDataFilesPlotter +from abipy.tools.plotting import get_ax_fig_plt, GenericDataFilesPlotter, FilesPlotter def remove_disordered(structures, paths): @@ -257,6 +257,15 @@ def abicomp_data(options): return 0 +def abicomp_png(options): + """ + Use matplotlib to plot multiple png files on a grid. + """ + plotter = FilesPlotter(options.filepaths) + plotter.plot() + return 0 + + def abicomp_ebands(options): """ Plot electron bands on a grid. @@ -1055,6 +1064,9 @@ def get_parser(with_epilog=False): p_data.add_argument("-i", "--use-index", default=False, action="store_true", help="Use the row index as x-value in the plot. By default the plotter uses the first column as x-values") + # Subparser for png command. + p_png = subparsers.add_parser('png', parents=[copts_parser], help=abicomp_png.__doc__) + # Subparser for ebands command. p_ebands = subparsers.add_parser('ebands', parents=[copts_parser, ipy_parser, pandas_parser], help=abicomp_ebands.__doc__) diff --git a/abipy/scripts/abiml.py b/abipy/scripts/abiml.py index 1cdf6fc9f..f178a538f 100755 --- a/abipy/scripts/abiml.py +++ b/abipy/scripts/abiml.py @@ -146,8 +146,7 @@ def add_workdir_verbose_opts(f): def add_nn_name_opt(f): """Add CLI options to select the NN potential.""" f = click.option("--nn-name", "-nn", default=DEFAULT_NN, show_default=True, - #type=click.Choice(aseml.CalcBuilder.ALL_NN_TYPES), - help='ML potential to be used')(f) + help=f"ML potential to be used. Supported values are: {aseml.CalcBuilder.ALL_NN_TYPES}")(f) return f @@ -672,7 +671,7 @@ def install(ctx, nn_names, update, verbose): @click.option("-srv", "--stdev-rvol", default=0.1, type=float, show_default=True, help="Scale volumes randomly around input v0 with stdev: v0*value") @add_workdir_verbose_opts -@click.option('--config', default='abiml_time.yml', type=click.Path(), callback=set_default, is_eager=True, expose_value=False) +@click.option('--config', default='abiml_compare.yml', type=click.Path(), callback=set_default, is_eager=True, expose_value=False) def compare(ctx, filepath, nn_names, num_tests, rattle, stdev_rvol, workdir, verbose @@ -681,7 +680,6 @@ def compare(ctx, filepath, nn_names, Compare different neural networks. """ atoms = _get_atoms_from_filepath(filepath) - nn_names = _get_nn_names(nn_names) ml_comp = aseml.MlCompareNNs(atoms, nn_names, num_tests, rattle, stdev_rvol, verbose, workdir, prefix="_abiml_comp_") print(ml_comp.to_string(verbose=verbose)) @@ -689,6 +687,54 @@ def compare(ctx, filepath, nn_names, return 0 +@main.command() +@herald +@click.pass_context +@click.argument("filepath", type=str) +@add_nn_name_opt +@click.option('--config', default='abiml_gs.yml', type=click.Path(), callback=set_default, is_eager=True, expose_value=False) +def gs(ctx, filepath, nn_name): + """ + Compute grounde-state properties and magnetic moments with ML potential. + """ + atoms = _get_atoms_from_filepath(filepath) + calc = aseml.CalcBuilder(nn_name).get_calculator() + #magmoms = atoms.get_magnetic_moments() + from abipy.ml.aseml import AseResults + res = AseResults.from_atoms(atoms, calc=calc) + + for ia, (atom, magmoms) in enumerate(zip(res.atoms, res.magmoms)): + print(atom, magmoms) + + return 0 + + +@main.command() +@herald +@click.pass_context +@click.argument("elements", nargs=-1, type=str) +@add_nn_names_opt +@add_workdir_verbose_opts +@click.option('--config', default='abiml_cwf_eos.yml', type=click.Path(), callback=set_default, is_eager=True, expose_value=False) +def cwf_eos(ctx, elements, nn_names, + workdir, verbose + ): + """ + Compute CWF EOS with ML potentials. + """ + nn_names = _get_nn_names(nn_names) + if "all" in elements: + if len(elements) != 1: + raise ValueError(f"When all is used for elements len(elements) should be 1 while {elements=}") + from ase.data import chemical_symbols + elements = [chemical_symbols[Z] for Z in range(1, 96+1)] + + ml_cwf_eos = aseml.MlCwfEos(elements, nn_names, verbose, workdir, prefix="_abiml_cwf_eos_") + print(ml_cwf_eos.to_string(verbose=verbose)) + ml_cwf_eos.run() + return 0 + + #@main.command() #@herald #@click.pass_context diff --git a/abipy/scripts/abiopen.py b/abipy/scripts/abiopen.py index 850a1c2de..eb2196282 100755 --- a/abipy/scripts/abiopen.py +++ b/abipy/scripts/abiopen.py @@ -239,6 +239,9 @@ def show_examples_and_exit(err_msg=None, error_code=1): if options.filepath.endswith(".json"): return handle_json(options) + if options.filepath.endswith(".traj"): + return handle_ase_traj(options) + if os.path.basename(options.filepath) == "flows.db": from abipy.flowtk.launcher import print_flowsdb_file return print_flowsdb_file(options.filepath) @@ -328,6 +331,15 @@ def show_examples_and_exit(err_msg=None, error_code=1): return 0 +def handle_ase_traj(options): + """Handle ASE trajectory file.""" + from ase.io import read + traj = read(options.filepath, index=":") + print(f"ASE trajectory with {len(traj)} configurations") + from abipy.ml.aseml import AseResults + first, last = AseResults(traj[0]), AseResults(traj[-1]) + + def handle_json(options): """Handle JSON file.""" diff --git a/abipy/tools/iotools.py b/abipy/tools/iotools.py index 49dde128e..0ec6c96f3 100644 --- a/abipy/tools/iotools.py +++ b/abipy/tools/iotools.py @@ -342,6 +342,29 @@ def workdir_with_prefix(workdir, prefix, exist_ok=False) -> Path: return Path(workdir).absolute() +def change_ext_from_top(top: PathLike, old_ext: str, new_ext: str) -> int: + """ + Change the extension of all the files with extension old_ext with new_ext. + + Args: + top: Walk the file system starting from top. + old_ext: Old file extension. + new_ext: New file extension. + + Return: Number of files whose extension has been changed. + """ + count = 0 + for root, dirs, files in os.walk(str(top)): + root = Path(root) + for filepath in files: + filepath = root / Path(filepath) + if not filepath.name.endswith(old_ext): continue + new_name = filepath.name[:-len(old_ext)] + new_ext + filepath.rename(root / new_name) + count += 1 + + return count + class _Script: """ diff --git a/abipy/tools/plotting.py b/abipy/tools/plotting.py index 3c2d7fde0..b52f46375 100644 --- a/abipy/tools/plotting.py +++ b/abipy/tools/plotting.py @@ -17,6 +17,7 @@ from collections import namedtuple, OrderedDict from typing import Any, Callable, Iterator +from monty.string import list_strings from pymatgen.util.plotting import add_fig_kwargs from abipy.tools import duck from abipy.tools.iotools import dataframe_from_filepath @@ -66,11 +67,43 @@ ) +class FilesPlotter: + """ + Use matplotlib to plot multiple png files on a grid. + + Example: + + FilesPlotter(["file1.png", file2.png"]).plot() + """ + def __init__(self, filepaths: list[str]): + self.filepaths = list_strings(filepaths) + + @add_fig_kwargs + def plot(self, **kwargs) -> Figure: + """Loop through the PNG files and display them in subplots.""" + # Build grid of plots. + num_plots, ncols, nrows = len(self.filepaths), 1, 1 + if num_plots > 1: + ncols = 2 + nrows = (num_plots // ncols) + (num_plots % ncols) + + ax_list, fig, plt = get_axarray_fig_plt(None, nrows=nrows, ncols=ncols, + sharex=False, sharey=False, squeeze=False) + ax_list = ax_list.ravel() + # don't show the last ax if num_plots is odd. + if num_plots % ncols != 0: ax_list[-1].axis("off") + + for i, (filepath, ax) in enumerate(zip(self.filepaths, ax_list)): + ax.axis('off') + ax.imshow(plt.imread(filepath)) + + return fig + @functools.cache def get_color_symbol(style: str="VESTA") -> dict: """ - Dictionary mapping chemical symbol to RGB color. + Dictionary mapping chemical symbols to RGB color. Args: style: "VESTA" or "Jmol". @@ -84,15 +117,13 @@ def get_color_symbol(style: str="VESTA") -> dict: return color_symbol - - ################### # Matplotlib tools ################### - def get_ax_fig_plt(ax=None, **kwargs): - """Helper function used in plot functions supporting an optional Axes argument. + """ + Helper function used in plot functions supporting an optional Axes argument. If ax is None, we build the `matplotlib` figure and create the Axes else we return the current active figure. @@ -116,7 +147,8 @@ def get_ax_fig_plt(ax=None, **kwargs): def get_ax3d_fig_plt(ax=None, **kwargs): - """Helper function used in plot functions supporting an optional Axes3D + """ + Helper function used in plot functions supporting an optional Axes3D argument. If ax is None, we build the `matplotlib` figure and create the Axes3D else we return the current active figure. @@ -140,7 +172,8 @@ def get_ax3d_fig_plt(ax=None, **kwargs): def get_axarray_fig_plt( ax_array, nrows=1, ncols=1, sharex=False, sharey=False, squeeze=True, subplot_kw=None, gridspec_kw=None, **fig_kw ): - """Helper function used in plot functions that accept an optional array of Axes + """ + Helper function used in plot functions that accept an optional array of Axes as argument. If ax_array is None, we build the `matplotlib` figure and create the array of Axes by calling plt.subplots else we return the current active figure. @@ -298,7 +331,6 @@ def set_ticks_fontsize(ax_or_axlist, fontsize: int, xy_string="xy", **kwargs) -> ax.tick_params(axis='y', labelsize=fontsize, **kwargs) - def set_grid_legend(ax_or_axlist, fontsize: int, xlabel=None, ylabel=None, grid=True, legend=True, direction=None, title=None, legend_loc="best") -> None: """ @@ -370,8 +402,6 @@ def rotate_ticklabels(ax, rotation: float, axname: str ="x") -> None: tick.set_rotation(rotation) - - def hspan_ax_line(ax, line, abs_conv, hatch, alpha=0.2, with_label=True) -> None: """ Add hspan to ax showing the convergence region of width `abs_conv`. @@ -2646,12 +2676,12 @@ def parse_latex(label): new_label = new_label.replace("{", "") if not latex else new_label new_label = new_label.replace("}", "") if not latex else new_label # plotly latex needs an extra \ for parsing python strings - # new_label = new_label.replace(" ", "\\ ") if latex else new_label + # new_label = new_label.replace(" ", "\\ ") if latex else new_label # Wrap the label in dollar signs for LaTeX, if needed unless empty`` new_label = f"${new_label}$" if latex and len(new_label) > 0 else new_label - + return new_label - + for ax in fig.get_axes(): # TODO improve below logic to add new scatter plots? # Loop backwards through the collections to avoid modifying the list as we iterate @@ -2659,7 +2689,7 @@ def parse_latex(label): if isinstance(coll, mcoll.PathCollection): # Use the remove() method to remove the scatter plot collection from the axes coll.remove() - + # Process the axis title, x-label, and y-label for label in [ax.get_title(), ax.get_xlabel(), ax.get_ylabel()]: # Few differences in how mpl and ply parse/encode symbols @@ -2671,7 +2701,7 @@ def parse_latex(label): ax.set_xlabel(new_label) elif label == ax.get_ylabel(): ax.set_ylabel(new_label) - + # Check if the axis has a legend if ax.get_legend(): legend = ax.get_legend() @@ -2685,7 +2715,7 @@ def parse_latex(label): # Convert to plotly figure plotly_fig = mpl_to_plotly(fig) - + plotly_fig.update_layout(template = "plotly_white", title = { "xanchor": "center", "yanchor": "top", @@ -2694,7 +2724,7 @@ def parse_latex(label): "size": 14 }, }) - + # Iterate over the axes in the figure to retrieve the custom line attributes for ax in fig.get_axes(): if hasattr(ax, '_custom_rc_lines'): @@ -2711,8 +2741,8 @@ def parse_latex(label): for trace in plotly_fig.data: # Retrieve the current label and remove any $ signs new_label = trace.name.replace("$", "") - + # Update the trace's name (which is used for the legend label) trace.name = new_label - - return plotly_fig \ No newline at end of file + + return plotly_fig diff --git a/dev_scripts/abiml_md.py b/dev_scripts/abiml_md.py new file mode 100755 index 000000000..f5ec57d2c --- /dev/null +++ b/dev_scripts/abiml_md.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python + +import os +from abipy.flowtk.qutils import SlurmJobArray + +def main(): + # IMPORTANT: You need customize the slurm options according to your machine. + conda_env = os.environ['CONDA_DEFAULT_ENV'] + print(f"Slurm script will be executed in {conda_env=}") + + header = f"""\ +#!/bin/bash + +#SBATCH --account=battab +#SBATCH --job-name=abiml_md +#SBATCH --time=0-1:0:00 +#SBATCH --partition=batch +#SBATCH --nodes=1 # 1 node has 128 cores +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=1 +#SBATCH --mem-per-cpu=4000 + +export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK +ulimit -s unlimited + +# Important: This is the conda environment providing the NN potential(s) nn_names +conda activate {conda_env} +""" + filepath = "LLZO_cubic.vasp" + nn_names = ["m3gnet", "chgnet"] + temp_list = [600, 800, 1000, 1200] + steps = 4000 + + print(f"""\ +Performing MD calculations with the following parameters: + +{filepath=} +{nn_names=} +{temp_list=} +{steps=} +""") + + arr_options = [] + for nn_name in nn_names: + for temperature in temp_list: + workdir = f"nn-{nn_name}_T-{temperature}" + opts = f"{filepath} --nn-name {nn_name} --temperature {temperature} --timestep 1 \ + --loginterval 100 --steps {steps} --ensemble nvt -w {workdir}" + arr_options.append(opts) + + job_array = SlurmJobArray(header, "abiml.py md", arr_options) + job_array.sbatch("job_array.sh") + + +if __name__ == "__main__": + main() diff --git a/dev_scripts/abiml_neb.py b/dev_scripts/abiml_neb.py new file mode 100755 index 000000000..8fdfe4421 --- /dev/null +++ b/dev_scripts/abiml_neb.py @@ -0,0 +1,139 @@ +cat run_neb.py +#!/usr/bin/env python + +from pathlib import Path +import sys +import os +import json +import numpy as np +import pandas as pd + +from monty.os.path import find_exts +from abipy.tools.iotools import make_executable + +from abipy.flowtk.qutils import SlurmJobArray + +NN_NAMES = ["chgnet", "m3gnet", "matgl"] + +TOP = Path(os.path.dirname(os.path.realpath(__file__))) + +VERBOSE = 0 + + +def run(): + # IMPORTANT: Please customize the slurm options according to your machine. + conda_env = os.environ['CONDA_DEFAULT_ENV'] + print(f"Slurm script will be executed in {conda_env=}") + + header = f"""\ +#!/bin/bash + +#SBATCH --account=battab +#SBATCH --job-name=abiml_neb +#SBATCH --time=0-1:0:00 +#SBATCH --partition=batch +#SBATCH --nodes=1 # 1 node has 128 cores +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=1 +#SBATCH --mem-per-cpu=4000 + +export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK +ulimit -s unlimited + +# Important: This is the conda environment providing the NN potential(s) +conda activate {conda_env} +""" + npaths = 3 + + print(f"""\ +Performing NEB calculations with the following parameters: + +{npaths=} +{NN_NAMES=} +""") + + arr_options = [] + for nn_name in NN_NAMES: + for i in range(npaths): + dirpath = Path(f"path{i+1}") + workdir = dirpath / nn_name + ini_vasp = str(dirpath / "ini.vasp") + fin_vasp = str(dirpath / "fin.vasp") + opts = f"{ini_vasp} {fin_vasp} --nn-name {nn_name} -w {str(workdir)}" + arr_options.append(opts) + + command = "abiml.py neb" + job_array = SlurmJobArray(header, command, arr_options) + job_array.sbatch("job_array.sh") + + +def process(): + """Post-process the results.""" + json_paths = find_exts(TOP, "neb_data.json") + neb_data_list = [] + for path in json_paths: + if VERBOSE: print("About to read json data from", path) + parts = path.split(os.sep) + path_index, nn_name = parts[-3], parts[-2] + path_index = int(path_index.replace("path", "")) + if VERBOSE: print(parts, "\n", path_index, nn_name) + with open(path, "rt") as fh: + d = json.load(fh) + d["path"] = path_index + d["nn_name"] = nn_name + neb_data_list.append(d) + + # Sort dict by path. + neb_data_list = sorted(neb_data_list, key=lambda d: d["path"]) + + keys = [ + "nn_name", + "path", + "barrier_with_fit", + "barrier_without_fit", + "energy_change_with_fit", + "energy_change_without_fit", + ] + + d_list = [] + for d in neb_data_list: + d_list.append({k: d[k] for k in keys}) + + df = pd.DataFrame(d_list) + #df.to_csv("XXX_ML_barriers.csv") + print(df) + + from abipy.tools.plotting import get_axarray_fig_plt + ax_list, fig, plt = get_axarray_fig_plt( + None, nrows=1, ncols=3, sharex=True, sharey=True, squeeze=False) + ax_list = ax_list.ravel() + cmap = plt.get_cmap("jet") + fontsize = 8 + + for ix, (ax, nn_name) in enumerate(zip(ax_list, NN_NAMES)): + my_data_list = [d for d in neb_data_list if d["nn_name"] == nn_name] + my_data_list = sorted(my_data_list, key=lambda d: d["path"]) + + for i, data in enumerate(my_data_list): + enes = np.array(data["energies_images"]) + ax.plot(enes - enes[0], label=f"path{i+1}", color=cmap(i/len(my_data_list))) + + ax.set_title(nn_name) + ax.set_xlabel('Image index', fontsize=fontsize) + ax.set_ylabel('$\Delta$ energy [eV]', fontsize=fontsize) + ax.legend(loc="best", shadow=True, fontsize=fontsize) + + fig.suptitle("K2Cu3Ge5O14") + plt.show() + + +if __name__ == "__main__": + cmd = sys.argv[1] + if cmd == "md_gen": + md_generate(sys.argv[2]) + elif cmd == "run": + run() + elif cmd == "process": + process() + else: + raise ValueError(f"Invalid {cmd=}") diff --git a/docs/api/electrons_api.rst b/docs/api/electrons_api.rst index f07923fbd..a23ed3039 100644 --- a/docs/api/electrons_api.rst +++ b/docs/api/electrons_api.rst @@ -93,6 +93,15 @@ electrons Package :undoc-members: :show-inheritance: + +:mod:`lruj` Module +------------------ + +.. automodule:: abipy.electrons.lruj + :members: + :undoc-members: + :show-inheritance: + :mod:`lobster` Module ---------------------