Surface maker and atoms describer update.

jarvis/__init__.py

@@ -1,6 +1,6 @@
 """Version number."""
 
-__version__ = "2024.4.10"
+__version__ = "2024.4.20"
 
 import os
 

jarvis/analysis/defects/surface.py

@@ -35,10 +35,11 @@ def __init__(
         atoms=None,
         indices=[0, 0, 1],
         layers=3,
-        thickness=25,
+        thickness=None,
         vacuum=18.0,
         tol=1e-10,
         from_conventional_structure=True,
+        use_thickness_c=True,
     ):
         """Initialize the class.
 
@@ -67,6 +68,7 @@ def __init__(
         self.vacuum = vacuum
         self.layers = layers
         self.thickness = thickness
+        self.use_thickness_c = use_thickness_c
         # Note thickness overwrites layers
 
     def to_dict(self):
@@ -147,7 +149,26 @@ def make_surface(self):
             cartesian=True,
         )
         if self.thickness is not None and (self.thickness) > 0:
-            self.layers = int(self.thickness / new_atoms.lattice.c) + 1
+            if not self.use_thickness_c:
+                new_lat = new_atoms.lattice_mat  # lat_lengths()
+                a1 = new_lat[0]
+                a2 = new_lat[1]
+                a3 = new_lat[2]
+                new_lat = np.array(
+                    [
+                        a1,
+                        a2,
+                        np.cross(a1, a2)
+                        * np.dot(a3, np.cross(a1, a2))
+                        / norm(np.cross(a1, a2)) ** 2,
+                    ]
+                )
+
+                a3 = new_lat[2]
+                self.layers = int(self.thickness / np.linalg.norm(a3)) + 1
+            else:
+                self.layers = int(self.thickness / new_atoms.lattice.c) + 1
+            # print("self.layers", self.layers)
             # dims=get_supercell_dims(new_atoms,enforce_c_size=self.thickness)
             # print ('dims=',dims,self.layers)
             # surf_atoms = new_atoms.make_supercell_matrix([1, 1, dims[2]])

jarvis/core/atoms.py

@@ -20,10 +20,26 @@
 import string
 import datetime
 from collections import defaultdict
+from sklearn.metrics import mean_absolute_error
+import zipfile
+import json
 
 amu_gm = 1.66054e-24
 ang_cm = 1e-8
 
+with zipfile.ZipFile(
+    str(
+        os.path.join(
+            os.path.dirname(__file__), "mineral_name_prototype.json.zip"
+        )
+    ),
+    "r",
+) as z:
+    # Open the specific JSON file within the zip
+    with z.open("mineral_name_prototype.json") as file:
+        # Load the JSON data from the file
+        mineral_json_file = json.load(file)
+
 
 class Atoms(object):
     """
@@ -649,10 +665,10 @@ def check_polar(self):
                 dn = dn + Specie(element).Z
         polar = False
         if up != dn:
-            print("Seems like a polar materials.")
+            # print("Seems like a polar materials.")
             polar = True
         if up == dn:
-            print("Non-polar")
+            # print("Non-polar")
             polar = False
         return polar
 
@@ -1039,6 +1055,106 @@ def density(self):
         )
         return den
 
+    def plot_atoms(
+        self=None,
+        colors=[],
+        sizes=[],
+        cutoff=1.9,
+        opacity=0.5,
+        bond_width=2,
+        filename=None,
+    ):
+        """Plot atoms using plotly."""
+        import plotly.graph_objects as go
+
+        fig = go.Figure()
+        if not colors:
+            colors = ["blue", "green", "red"]
+
+        unique_elements = self.uniq_species
+        if len(unique_elements) > len(colors):
+            raise ValueError("Provide more colors.")
+        color_map = {}
+        size_map = {}
+        for ii, i in enumerate(unique_elements):
+            color_map[i] = colors[ii]
+            size_map[i] = Specie(i).Z * 2
+        cart_coords = self.cart_coords
+        elements = self.elements
+        atoms_arr = []
+
+        for ii, i in enumerate(cart_coords):
+            atoms_arr.append(
+                [
+                    i[0],
+                    i[1],
+                    i[2],
+                    color_map[elements[ii]],
+                    size_map[elements[ii]],
+                ]
+            )
+        # atoms = [
+        #     (0, 0, 0, 'red'),   # Atom 1
+        #     (1, 1, 1, 'blue'),  # Atom 2
+        #     (2, 0, 1, 'green')  # Atom 3
+        # ]
+
+        # Create a scatter plot for the 3D points
+        trace1 = go.Scatter3d(
+            x=[atom[0] for atom in atoms_arr],
+            y=[atom[1] for atom in atoms_arr],
+            z=[atom[2] for atom in atoms_arr],
+            mode="markers",
+            marker=dict(
+                size=[atom[4] for atom in atoms_arr],  # Marker size
+                color=[atom[3] for atom in atoms_arr],  # Marker color
+                opacity=opacity,
+            ),
+        )
+        fig.add_trace(trace1)
+
+        # Update plot layout
+        fig.update_layout(
+            title="3D Atom Coordinates",
+            scene=dict(
+                xaxis_title="X Coordinates",
+                yaxis_title="Y Coordinates",
+                zaxis_title="Z Coordinates",
+            ),
+            margin=dict(l=0, r=0, b=0, t=0),  # Tight layout
+        )
+        if bond_width is not None:
+            nbs = self.get_all_neighbors(r=5)
+            bonds = []
+            for i in nbs:
+                for j in i:
+                    if j[2] <= cutoff:
+                        bonds.append([j[0], j[1]])
+            for bond in bonds:
+                # print(bond)
+                # Extract coordinates of the first and second atom in each bond
+                x_coords = [atoms_arr[bond[0]][0], atoms_arr[bond[1]][0]]
+                y_coords = [atoms_arr[bond[0]][1], atoms_arr[bond[1]][1]]
+                z_coords = [atoms_arr[bond[0]][2], atoms_arr[bond[1]][2]]
+
+                fig.add_trace(
+                    go.Scatter3d(
+                        x=x_coords,
+                        y=y_coords,
+                        z=z_coords,
+                        mode="lines",
+                        line=dict(color="grey", width=bond_width),
+                        marker=dict(
+                            size=0.1
+                        ),  # Small marker size to make lines prominent
+                    )
+                )
+        # Show the plot
+        if filename is not None:
+            fig.write_image(filename)
+        else:
+            fig.show()
+
     @property
     def atomic_numbers(self):
         """Get list of atomic numbers of atoms in the atoms object."""
@@ -1168,6 +1284,105 @@ def packing_fraction(self):
         pf = np.array([4 * np.pi * total_rad / (3 * self.volume)])
         return round(pf[0], 5)
 
+    def get_alignn_feats(
+        self,
+        model_name="jv_formation_energy_peratom_alignn",
+        max_neighbors=12,
+        neighbor_strategy="k-nearest",
+        use_canonize=True,
+        atom_features="cgcnn",
+        line_graph=True,
+        cutoff=8,
+        model="",
+    ):
+        """Get ALIGNN features."""
+
+        def get_val(model, g, lg):
+            activation = {}
+
+            def getActivation(name):
+                # the hook signature
+                def hook(model, input, output):
+                    activation[name] = output.detach()
+
+                return hook
+
+            h = model.readout.register_forward_hook(getActivation("readout"))
+            out = model([g, lg])
+            del out
+            h.remove()
+            return activation["readout"][0]
+
+        from alignn.graphs import Graph
+        from alignn.pretrained import get_figshare_model
+
+        g, lg = Graph.atom_dgl_multigraph(
+            self,
+            cutoff=cutoff,
+            atom_features=atom_features,
+            max_neighbors=max_neighbors,
+            neighbor_strategy=neighbor_strategy,
+            compute_line_graph=line_graph,
+            use_canonize=use_canonize,
+        )
+        if model == "":
+            model = get_figshare_model(
+                model_name="jv_formation_energy_peratom_alignn"
+            )
+        h = get_val(model, g, lg)
+        return h
+
+    def get_mineral_prototype_name(
+        self, prim=True, include_c_over_a=False, digits=3
+    ):
+        from jarvis.analysis.structure.spacegroup import Spacegroup3D
+
+        spg = Spacegroup3D(self)
+        number = spg.space_group_number
+        cnv_atoms = spg.conventional_standard_structure
+        n_conv = cnv_atoms.num_atoms
+        # hall_number=str(spg._dataset['hall_number'])
+        wyc = "".join(list((sorted(set(spg._dataset["wyckoffs"])))))
+        name = (
+            (self.composition.prototype_new)
+            + "_"
+            + str(number)
+            + "_"
+            + str(wyc)
+            + "_"
+            + str(n_conv)
+        )
+        # if include_com:
+        if include_c_over_a:
+            # print(cnv_atoms)
+            abc = cnv_atoms.lattice.abc
+            ca = round(abc[2] / abc[0], digits)
+            # com_positions = "_".join(map(str,self.get_center_of_mass()))
+            # round(np.sum(spg._dataset['std_positions']),round_digit)
+            name += "_" + str(ca)
+            # name+="_"+str(com_positions)
+        return name
+
+    def get_minaral_name(self, model=""):
+        """Get mineral prototype."""
+        mae = np.inf
+        feats = self.get_alignn_feats(model=model)
+        nm = self.get_mineral_prototype_name()
+        if nm in mineral_json_file:
+            for i in mineral_json_file[nm]:
+                maem = mean_absolute_error(i[1], feats)
+                if maem < mae:
+                    mae = maem
+                    name = i[0]
+        else:
+            for i, j in mineral_json_file.items():
+                for k in j:
+                    maem = mean_absolute_error(k[1], feats)
+                    if maem < mae:
+                        mae = maem
+                        name = k[0]
+        return name
+
     def lattice_points_in_supercell(self, supercell_matrix):
         """
         Adapted from Pymatgen.
@@ -1236,6 +1451,7 @@ def describe(
         """Describe for NLP applications."""
         from jarvis.analysis.diffraction.xrd import XRD
 
+        min_name = self.get_minaral_name()
         if include_spg:
             from jarvis.analysis.structure.spacegroup import Spacegroup3D
 
@@ -1264,6 +1480,7 @@ def describe(
             fracs[i] = round(j, 3)
         info = {}
         chem_info = {
+            "mineral_name": min_name,
             "atomic_formula": self.composition.reduced_formula,
             "prototype": self.composition.prototype,
             "molecular_weight": round(self.composition.weight / 2, 2),
@@ -1372,8 +1589,30 @@ def describe(
             + struct_info["wyckoff"]
             + "."
         )
+        if min_name is not None:
+            line3 = (
+                chem_info["atomic_formula"]
+                + " is "
+                + min_name
+                + "-derived structured and"
+                + " crystallizes in the "
+                + struct_info["crystal_system"]
+                + " "
+                + str(struct_info["spg_symbol"])
+                + " spacegroup."
+            )
+        else:
+            line3 = (
+                chem_info["atomic_formula"]
+                + " crystallizes in the "
+                + struct_info["crystal_system"]
+                + " "
+                + str(struct_info["spg_symbol"])
+                + " spacegroup."
+            )
         info["desc_1"] = line1
         info["desc_2"] = line2
+        info["desc_3"] = line3
         return info
 
     def make_supercell_matrix(self, scaling_matrix):