Skip to content

Commit

Permalink
more var name and comment cleanups
Browse files Browse the repository at this point in the history
  • Loading branch information
@DanielYang59
DanielYang59 committed Jun 19, 2024
1 parent d08f0c2 commit 18499c4
Show file tree
Hide file tree
Showing 3 changed files with 281 additions and 241 deletions.
4 changes: 2 additions & 2 deletions pymatgen/io/lobster/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
"""
This package implements modules for input and output to and from Lobster. It
imports the key classes form both lobster.inputs and lobster_outputs to allow most
This package implements modules for input and output to and from LOBSTER. It
imports the key classes form both lobster.inputs and lobster.outputs to allow most
classes to be simply called as pymatgen.io.lobster.Lobsterin for example, to retain
backwards compatibility.
"""
Expand Down
93 changes: 53 additions & 40 deletions pymatgen/io/lobster/lobsterenv.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
"""This module provides classes to perform analyses of
the local environments (e.g., finding near neighbors)
of single sites in molecules and structures based on
bonding analysis with LOBSTER.
"""This module provides classes to perform analyses of the local
environments (e.g., finding near neighbors) of single sites in molecules
and structures based on bonding analysis with LOBSTER.
If you use this module, please cite:
J. George, G. Petretto, A. Naik, M. Esters, A. J. Jackson, R. Nelson, R. Dronskowski, G.-M. Rignanese, G. Hautier,
Expand Down Expand Up @@ -36,9 +35,10 @@
from typing import Any, Literal

import matplotlib as mpl
from numpy.typing import NDArray
from typing_extensions import Self

from pymatgen.core import Site, Structure
from pymatgen.core import PeriodicNeighbor, Site, Structure
from pymatgen.core.periodic_table import Element
from pymatgen.electronic_structure.cohp import IcohpCollection, IcohpValue
from pymatgen.util.typing import PathLike
Expand Down Expand Up @@ -92,7 +92,7 @@ def __init__(
filename_icohp (PathLike): Path to ICOHPLIST.lobster or
ICOOPLIST.lobster or ICOBILIST.lobster.
obj_icohp (Icohplist): Icohplist object.
structure (Structure): Typically constructed by Structure.from_file("POSCAR").
structure (Structure): Typically constructed by Structure.from_file("POSCAR").
are_coops (bool): Whether the file is a ICOOPLIST.lobster (True) or a
ICOHPLIST.lobster (False). Only tested for ICOHPLIST.lobster so far.
are_cobis (bool): Whether the file is a ICOBILIST.lobster (True) or
Expand Down Expand Up @@ -431,6 +431,7 @@ def get_info_icohps_to_neighbors(
labels = []
atoms = []
final_isites = []
assert self.Icohpcollection is not None
for idx, _site in enumerate(self.structure):
if idx in isites:
for keys, icohpsum in zip(self.list_keys[idx], self.list_icohps[idx]):
Expand Down Expand Up @@ -557,6 +558,7 @@ def get_info_cohps_to_neighbors(

# Check that the number of bonds in ICOHPLIST and COHPCAR are identical
# TODO: Further checks could be implemented
assert self.Icohpcollection is not None
if len(self.Icohpcollection._list_atom1) != len(self.completecohp.bonds):
raise ValueError("COHPCAR and ICOHPLIST do not fit together")

Expand Down Expand Up @@ -663,6 +665,7 @@ def get_info_icohps_between_neighbors(
number_bonds = 0
labels = []
atoms = []
assert self.Icohpcollection is not None
for isite in isites:
for in_site, n_site in enumerate(self.list_neighsite[isite]):
for in_site2, n_site2 in enumerate(self.list_neighsite[isite]):
Expand Down Expand Up @@ -756,6 +759,7 @@ def _evaluate_ce(
"""
# Get extremum
if lowerlimit is None and upperlimit is None:
assert self.Icohpcollection is not None
limits = self._get_limit_from_extremum(
self.Icohpcollection,
percentage=perc_strength_icohp,
Expand Down Expand Up @@ -790,8 +794,8 @@ def _evaluate_ce(
{
"site": neighbor,
"image": tuple(
int(round(i))
for i in (
int(round(idx))
for idx in (
neighbor.frac_coords
- self.structure[
next(
Expand Down Expand Up @@ -831,8 +835,8 @@ def _evaluate_ce(
{
"site": neighbor,
"image": tuple(
int(round(i))
for i in (
int(round(idx))
for idx in (
neighbor.frac_coords
- self.structure[
next(
Expand Down Expand Up @@ -864,25 +868,34 @@ def _find_environments(
lowerlimit: float,
upperlimit: float,
only_bonds_to: list[str] | None,
) -> tuple[list, list, list, list, list, list]: # TODO: DanielYang: more specific type
) -> tuple[
list[list[IcohpValue]],
list[list[str]],
list[list[float]],
list[list[int]],
list[list[PeriodicNeighbor]],
list[list[NDArray]],
]:
"""Find all relevant neighbors based on certain restrictions.
Args:
additional_condition (int): Additional condition (see above).
lowerlimit (float): Lower limit that tells you which ICOHPs are considered.
upperlimit (float): Upper limit that tells you which ICOHPs are considered.
additional_condition (int): Additional condition.
lowerlimit (float): Lower limit that ICOHPs are considered.
upperlimit (float): Upper limit that ICOHPs are considered.
only_bonds_to (list[str]): Only bonds to these elements will be considered.
Returns:
tuple of ICOHPs, keys, lengths, neighisite, neighsite, coords.
Tuple of ICOHPs, keys, lengths, neighisite, neighsite, coords.
"""
list_icohps: list[list[IcohpValue]] = []
list_keys: list[list[str]] = []
list_lengths: list[list[float]] = []
list_neighisite: list[list[int]] = []
list_neighsite: list[list[PeriodicNeighbor]] = []
list_coords: list[list[NDArray]] = []

# Run over structure
list_neighsite = []
list_neighisite = []
list_coords = []
list_icohps = []
list_lengths = []
list_keys = []
assert self.Icohpcollection is not None
for idx, site in enumerate(self.structure):
icohps = self._get_icohps(
icohpcollection=self.Icohpcollection,
Expand Down Expand Up @@ -979,23 +992,23 @@ def _find_environments(
def _find_relevant_atoms_additional_condition(
self,
isite: int,
icohps, # TODO: need type
icohps: dict[str, IcohpValue],
additional_condition: Literal[0, 1, 2, 3, 4, 5, 6],
) -> tuple[list, list, list, list]: # TODO: more specific type
) -> tuple[list[str], list[float], list[int], list[IcohpValue]]:
"""Find all relevant atoms that fulfill the additional condition.
Args:
isite (int): Index of site in structure (start from 0).
icohps: icohps
additional_condition (int): additional condition.
isite (int): Site index in structure (start from 0).
icohps (dict[str, IcohpValue]): ICOHP values.
additional_condition (int): Additional condition.
Returns:
tuple: keys, lengths, neighbors from selected ICOHPs and selected ICOHPs.
"""
neighbors_from_ICOHPs = []
lengths_from_ICOHPs = []
icohps_from_ICOHPs = []
keys_from_ICOHPs = []
keys_from_ICOHPs: list[str] = []
lengths_from_ICOHPs: list[float] = []
neighbors_from_ICOHPs: list[int] = []
icohps_from_ICOHPs: list[IcohpValue] = []

for key, icohp in icohps.items():
atomnr1 = self._get_atomnumber(icohp._atom1)
Expand All @@ -1008,8 +1021,8 @@ def _find_relevant_atoms_additional_condition(
val1 = self.valences[atomnr1]
val2 = self.valences[atomnr2]

# NO_ADDITIONAL_CONDITION
if additional_condition == 0:
# NO_ADDITIONAL_CONDITION
if atomnr1 == isite:
neighbors_from_ICOHPs.append(atomnr2)
lengths_from_ICOHPs.append(icohp._length)
Expand All @@ -1021,8 +1034,8 @@ def _find_relevant_atoms_additional_condition(
icohps_from_ICOHPs.append(icohp.summed_icohp)
keys_from_ICOHPs.append(key)

# ONLY_ANION_CATION_BONDS
elif additional_condition == 1:
# ONLY_ANION_CATION_BONDS
if (val1 < 0.0 < val2) or (val2 < 0.0 < val1): # type: ignore[operator]
if atomnr1 == isite:
neighbors_from_ICOHPs.append(atomnr2)
Expand All @@ -1036,8 +1049,8 @@ def _find_relevant_atoms_additional_condition(
icohps_from_ICOHPs.append(icohp.summed_icohp)
keys_from_ICOHPs.append(key)

# NO_ELEMENT_TO_SAME_ELEMENT_BONDS
elif additional_condition == 2:
# NO_ELEMENT_TO_SAME_ELEMENT_BONDS
if icohp._atom1.rstrip("0123456789") != icohp._atom2.rstrip("0123456789"):
if atomnr1 == isite:
neighbors_from_ICOHPs.append(atomnr2)
Expand All @@ -1051,8 +1064,8 @@ def _find_relevant_atoms_additional_condition(
icohps_from_ICOHPs.append(icohp.summed_icohp)
keys_from_ICOHPs.append(key)

# ONLY_ANION_CATION_BONDS_AND_NO_ELEMENT_TO_SAME_ELEMENT_BONDS
elif additional_condition == 3:
# ONLY_ANION_CATION_BONDS_AND_NO_ELEMENT_TO_SAME_ELEMENT_BONDS
if ((val1 < 0.0 < val2) or (val2 < 0.0 < val1)) and icohp._atom1.rstrip( # type: ignore[operator]
"0123456789"
) != icohp._atom2.rstrip("0123456789"):
Expand All @@ -1068,8 +1081,8 @@ def _find_relevant_atoms_additional_condition(
icohps_from_ICOHPs.append(icohp.summed_icohp)
keys_from_ICOHPs.append(key)

# ONLY_ELEMENT_TO_OXYGEN_BONDS
elif additional_condition == 4:
# ONLY_ELEMENT_TO_OXYGEN_BONDS
if icohp._atom1.rstrip("0123456789") == "O" or icohp._atom2.rstrip("0123456789") == "O":
if atomnr1 == isite:
neighbors_from_ICOHPs.append(atomnr2)
Expand All @@ -1083,8 +1096,8 @@ def _find_relevant_atoms_additional_condition(
icohps_from_ICOHPs.append(icohp.summed_icohp)
keys_from_ICOHPs.append(key)

# DO_NOT_CONSIDER_ANION_CATION_BONDS
elif additional_condition == 5:
# DO_NOT_CONSIDER_ANION_CATION_BONDS
if (val1 > 0.0 and val2 > 0.0) or (val1 < 0.0 and val2 < 0.0): # type: ignore[operator]
if atomnr1 == isite:
neighbors_from_ICOHPs.append(atomnr2)
Expand All @@ -1098,8 +1111,8 @@ def _find_relevant_atoms_additional_condition(
icohps_from_ICOHPs.append(icohp.summed_icohp)
keys_from_ICOHPs.append(key)

# ONLY_CATION_CATION_BONDS
elif additional_condition == 6 and val1 > 0.0 and val2 > 0.0: # type: ignore[operator]
# ONLY_CATION_CATION_BONDS
if atomnr1 == isite:
neighbors_from_ICOHPs.append(atomnr2)
lengths_from_ICOHPs.append(icohp._length)
Expand All @@ -1126,9 +1139,9 @@ def _get_icohps(
Args:
icohpcollection (IcohpCollection): IcohpCollection object.
isite (int): Index of a site.
lowerlimit (float): Lower limit that tells you which ICOHPs are considered.
upperlimit (float): Upper limit that tells you which ICOHPs are considered.
isite (int): Site index.
lowerlimit (float): Lower limit that ICOHPs are considered.
upperlimit (float): Upper limit that ICOHPs are considered.
only_bonds_to (list[str]): Only bonds to these elements will be considered, e.g. ["O"].
Returns:
Expand Down
Loading

0 comments on commit 18499c4

Please sign in to comment.