Optimize cython find_points_in _spheres

pymatgen/cli/pmg_config.py

@@ -108,6 +108,95 @@ def setup_cp2k_data(cp2k_data_dirs: list[str]) -> None:
     print("\n Start a new terminal to ensure that your environment variables are properly set.")
 
 
+def setup_cp2k_data(cp2k_data_dirs: list[str]) -> None:
+    """Setup CP2K basis and potential data directory"""
+    data_dir, target_dir = (os.path.abspath(dir) for dir in cp2k_data_dirs)
+    try:
+        os.mkdir(target_dir)
+    except OSError:
+        reply = input("Destination directory exists. Continue (y/n)?")
+        if reply != "y":
+            print("Exiting ...")
+            raise SystemExit(0)
+    print("Generating pymatgen resource directory for CP2K...")
+
+    import glob
+
+    from monty.json import jsanitize
+    from ruamel import yaml
+
+    from pymatgen.core import Element
+    from pymatgen.io.cp2k.inputs import GaussianTypeOrbitalBasisSet, GthPotential
+    from pymatgen.io.cp2k.utils import chunk
+
+    basis_files = glob.glob(os.path.join(data_dir, "*BASIS*"))
+    potential_files = glob.glob(os.path.join(data_dir, "*POTENTIAL*"))
+
+    settings: dict[str, dict] = {str(el): {"potentials": {}, "basis_sets": {}} for el in Element}
+
+    for potential_file in potential_files:
+        print(f"Processing... {potential_file}")
+        with open(potential_file) as file:
+            try:
+                chunks = chunk(file.read())
+            except IndexError:
+                continue
+        for chk in chunks:
+            try:
+                potential = GthPotential.from_string(chk)
+                potential.filename = os.path.basename(potential_file)
+                potential.version = None
+                settings[potential.element.symbol]["potentials"][potential.get_hash()] = jsanitize(
+                    potential, strict=True
+                )
+            except ValueError:
+                # Chunk was readable, but the element is not pmg recognized
+                continue
+            except IndexError:
+                # Chunk was readable, but invalid. Mostly likely "N/A" for this potential
+                continue
+
+    for basis_file in basis_files:
+        print(f"Processing... {basis_file}")
+        with open(basis_file) as file:
+            try:
+                chunks = chunk(file.read())
+            except IndexError:
+                continue
+        for chk in chunks:
+            try:
+                basis = GaussianTypeOrbitalBasisSet.from_string(chk)
+                basis.filename = os.path.basename(basis_file)
+                settings[basis.element.symbol]["basis_sets"][basis.get_hash()] = jsanitize(  # type: ignore
+                    basis, strict=True
+                )
+            except ValueError:
+                # Chunk was readable, but the element is not pmg recognized
+                continue
+            except IndexError:
+                # Chunk was readable, but invalid. Mostly likely "N/A" for this potential
+                continue
+
+    print("Done processing cp2k data files")
+
+    for el in settings:
+        print(f"Writing {el} settings file")
+        with open(os.path.join(target_dir, el), "w") as file:
+            yaml.dump(settings.get(el), file, default_flow_style=False)
+
+    print(
+        "\n CP2K resource directory generated. It is recommended that you run:"
+        f"\n  'pmg config --add PMG_CP2K_DATA_DIR {os.path.abspath(target_dir)}' "
+    )
+    print(
+        "\n It is also recommended that you set the following (with example values):"
+        "\n  'pmg config --add PMG_DEFAULT_CP2K_FUNCTIONAL PBE' "
+        "\n  'pmg config --add PMG_DEFAULT_CP2K_BASIS_TYPE TZVP-MOLOPT' "
+        "\n  'pmg config --add PMG_DEFAULT_CP2K_AUX_BASIS_TYPE pFIT' "
+    )
+    print("\n Start a new terminal to ensure that your environment variables are properly set.")
+
+
 def setup_potcars(potcar_dirs: list[str]):
     """Setup POTCAR directories."""
     psp_dir, target_dir = (os.path.abspath(d) for d in potcar_dirs)

pymatgen/core/lattice.py

@@ -1482,16 +1482,16 @@ def get_points_in_sphere(
         except ImportError:
             return self.get_points_in_sphere_py(frac_points=frac_points, center=center, r=r, zip_results=zip_results)
         else:
-            frac_points = np.ascontiguousarray(frac_points, dtype=np.float_)
+            frac_points = np.ascontiguousarray(frac_points, dtype=float)
+            lattice_matrix = np.ascontiguousarray(self.matrix, dtype=float)
+            cart_coords = np.ascontiguousarray(self.get_cartesian_coords(frac_points), dtype=float)
+            pbc = np.ascontiguousarray(self.pbc, dtype=int)
             r = float(r)
-            lattice_matrix = np.array(self.matrix)
-            lattice_matrix = np.ascontiguousarray(lattice_matrix)
-            cart_coords = self.get_cartesian_coords(frac_points)
             _, indices, images, distances = find_points_in_spheres(
                 all_coords=cart_coords,
                 center_coords=np.ascontiguousarray([center], dtype=float),
                 r=r,
-                pbc=np.array(self.pbc, dtype=int),
+                pbc=pbc,
                 lattice=lattice_matrix,
                 tol=1e-8,
             )

pymatgen/core/structure.py

@@ -1380,12 +1380,13 @@ def get_neighbor_list(
             site_coords = np.array([site.coords for site in sites], dtype=float)
             cart_coords = np.ascontiguousarray(np.array(self.cart_coords), dtype=float)
             lattice_matrix = np.ascontiguousarray(np.array(self.lattice.matrix), dtype=float)
+            pbc = np.ascontiguousarray(self.pbc, dtype=int)
             r = float(r)
             center_indices, points_indices, images, distances = find_points_in_spheres(
                 cart_coords,
                 site_coords,
                 r=r,
-                pbc=np.array(self.pbc, dtype=int),
+                pbc=pbc,
                 lattice=lattice_matrix,
                 tol=numerical_tol,
             )

pymatgen/electronic_structure/tests/test_dos.py

@@ -292,6 +292,52 @@ def test_dos_fp_exceptions(self):
             "projections unavailable in input DOS or there's a typo in type."
         )
 
+    def test_get_dos_fp(self):
+        # normalize=True
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        bin_width = np.diff(dos_fp.energies)[0][0]
+        assert max(dos_fp.energies[0]) <= 0
+        assert min(dos_fp.energies[0]) >= -10
+        assert len(dos_fp.energies[0]) == 56
+        assert dos_fp.type == "s"
+        self.assertAlmostEqual(sum(dos_fp.densities * bin_width), 1, delta=0.001)
+        # normalize=False
+        dos_fp2 = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, n_bins=56, normalize=False)
+        bin_width2 = np.diff(dos_fp2.energies)[0][0]
+        self.assertAlmostEqual(sum(dos_fp2.densities * bin_width2), 7.279303571428509, delta=0.001)
+        self.assertAlmostEqual(dos_fp2.bin_width, bin_width2, delta=0.001)
+        # binning=False
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=None, max_e=None, n_bins=56, normalize=True, binning=False)
+        assert dos_fp.n_bins == len(self.dos.energies)
+
+    def test_get_dos_fp_similarity(self):
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        dos_fp2 = self.dos.get_dos_fp(type="tdos", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        similarity_index = self.dos.get_dos_fp_similarity(dos_fp, dos_fp2, col=1, tanimoto=True)
+        self.assertAlmostEqual(similarity_index, 0.3342481451042263, delta=0.0001)
+
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        dos_fp2 = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        similarity_index = self.dos.get_dos_fp_similarity(dos_fp, dos_fp2, col=1, tanimoto=True)
+        assert similarity_index == 1
+
+    def test_dos_fp_exceptions(self):
+        dos_fp = self.dos.get_dos_fp(type="s", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        dos_fp2 = self.dos.get_dos_fp(type="tdos", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        # test exceptions
+        with self.assertRaises(ValueError) as err:
+            self.dos.get_dos_fp_similarity(dos_fp, dos_fp2, col=1, tanimoto=True, normalize=True)
+        assert (
+            err.exception.__str__()
+            == "Cannot compute similarity index. Please set either normalize=True or tanimoto=True or both to False."
+        )
+        with self.assertRaises(ValueError) as err:
+            self.dos.get_dos_fp(type="k", min_e=-10, max_e=0, n_bins=56, normalize=True)
+        assert (
+            err.exception.__str__()
+            == "Please recheck type requested, either the orbital projections unavailable in input DOS or there's a typo in type."
+        )
+
 
 class DOSTest(PymatgenTest):
     def setUp(self):