Skip to content

Commit

Permalink
Clean up and improved docstrings
Browse files Browse the repository at this point in the history
  • Loading branch information
oashour committed Sep 23, 2024
1 parent 0f27dc1 commit 82bd520
Show file tree
Hide file tree
Showing 2 changed files with 150 additions and 65 deletions.
16 changes: 13 additions & 3 deletions pymatgen/io/espresso/outputs/projwfc.py
Original file line number Diff line number Diff line change
Expand Up @@ -167,7 +167,9 @@ def __add__(self, other: "Projwfc"):
if not isinstance(other, Projwfc):
raise ValueError("Can only add Projwfc objects to other Projwfc objects.")
if self != other:
raise ValueError("Can only add Projwfc objects from the same calculation.")
raise InconsistentProjwfcDataError(
"Can only add Projwfc objects from the same calculation."
)

# Check that one is spin up and the other is spin down
# Get all the spins of each object. These are the keys of the projection
Expand All @@ -179,15 +181,17 @@ def __add__(self, other: "Projwfc"):
spin for state in other.atomic_states for spin in state.projections.keys()
}
if len(spin1) != 1 or len(spin2) != 1:
raise ValueError(
raise InconsistentProjwfcDataError(
(
"You are trying to add two Projwfc objects with multiple spins. "
"This should only be used to add objects with one spin each."
)
)
spin1, spin2 = spin1.pop(), spin2.pop()
if spin1 == spin2:
raise ValueError("Can only add Projwfc objects with opposite spins.")
raise InconsistentProjwfcDataError(
"Can only add Projwfc objects with opposite spins."
)

result = deepcopy(self)
for s1, s2 in zip(result.atomic_states, other.atomic_states, strict=True):
Expand Down Expand Up @@ -973,3 +977,9 @@ class ProjwfcParserError(Exception):
"""
Exception class for Projwfc parsing.
"""


class InconsistentProjwfcDataError(Exception):
"""
Exception class for Projwfc addition.
"""
Loading

0 comments on commit 82bd520

Please sign in to comment.