Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Self return type on Lattice methods #3707

Merged
merged 12 commits into from
Mar 25, 2024
Merged

Self return type on Lattice methods #3707

merged 12 commits into from
Mar 25, 2024

Conversation

janosh
Copy link
Member

@janosh janosh commented Mar 25, 2024

define new type PbcLike = tuple[bool, bool, bool]. improve tests for typing module

@janosh janosh added the types Type all the things label Mar 25, 2024
@janosh janosh enabled auto-merge (squash) March 25, 2024 10:33

@staticmethod
Copy link
Contributor

@DanielYang59 DanielYang59 Mar 25, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch! I notice some other methods should not be staticmethods as well (sorry I got too many mypy errors to fix at this moment and could not help much on this.).

Maybe we could discuss and fix some in a separate PR?

For example:

@staticmethod
def from_filenames(poscar_filenames, transformations=None, extend_collection=False):
"""Convenient constructor to generates a POSCAR transmuter from a list of
POSCAR filenames.
Args:
poscar_filenames: List of POSCAR filenames
transformations: New transformations to be applied to all
structures.
extend_collection:
Same meaning as in __init__.
"""
trafo_structs = []
for filename in poscar_filenames:
with open(filename) as file:
trafo_structs.append(TransformedStructure.from_poscar_str(file.read(), []))
return StandardTransmuter(trafo_structs, transformations, extend_collection=extend_collection)

@staticmethod
def with_edges(structure: Structure, edges: dict) -> StructureGraph:
"""
Constructor for MoleculeGraph, using pre-existing or pre-defined edges
with optional edge parameters.
Args:
structure: Structure object
edges: dict representing the bonds of the functional
group (format: {(from_index, to_index, from_image, to_image): props},
where props is a dictionary of properties, including weight.
Props should be None if no additional properties are to be
specified.
Returns:
sg, a StructureGraph
"""
struct_graph = StructureGraph.with_empty_graph(
structure, name="bonds", edge_weight_name="weight", edge_weight_units=""
)
for edge, props in edges.items():
try:
from_index = edge[0]
to_index = edge[1]
from_image = edge[2]
to_image = edge[3]
except TypeError:
raise ValueError("Edges must be given as (from_index, to_index, from_image, to_image) tuples")
if props is not None:
weight = props.pop("weight", None)
if len(props.items()) == 0:
props = None
else:
weight = None
nodes = struct_graph.graph.nodes
if not (from_index in nodes and to_index in nodes):
raise ValueError(
"Edges cannot be added if nodes are not present in the graph. Please check your indices."
)
struct_graph.add_edge(
from_index,
to_index,
from_jimage=from_image,
to_jimage=to_image,
weight=weight,
edge_properties=props,
)
struct_graph.set_node_attributes()
return struct_graph
@staticmethod
def with_local_env_strategy(
structure: Structure, strategy: NearNeighbors, weights: bool = False, edge_properties: bool = False
) -> StructureGraph:
"""
Constructor for StructureGraph, using a strategy
from pymatgen.analysis.local_env.
Args:
structure: Structure object
strategy: an instance of a pymatgen.analysis.local_env.NearNeighbors object
weights(bool): if True, use weights from local_env class (consult relevant class for their meaning)
edge_properties(bool): if True, edge_properties from neighbors will be used
"""
if not strategy.structures_allowed:
raise ValueError("Chosen strategy is not designed for use with structures! Please choose another strategy.")
struct_graph = StructureGraph.with_empty_graph(structure, name="bonds")
for idx, neighbors in enumerate(strategy.get_all_nn_info(structure)):
for neighbor in neighbors:
# local_env will always try to add two edges
# for any one bond, one from site u to site v
# and another form site v to site u: this is
# harmless, so warn_duplicates=False
struct_graph.add_edge(
from_index=idx,
from_jimage=(0, 0, 0),
to_index=neighbor["site_index"],
to_jimage=neighbor["image"],
weight=neighbor["weight"] if weights else None,
edge_properties=neighbor["edge_properties"] if edge_properties else None,
warn_duplicates=False,
)
return struct_graph

@staticmethod
def with_edges(molecule: Molecule, edges: dict[tuple[int, int], None | dict]) -> MoleculeGraph:
"""
Constructor for MoleculeGraph, using pre-existing or pre-defined edges
with optional edge parameters.
Args:
molecule: Molecule object
edges: dict representing the bonds of the functional
group (format: {(u, v): props}, where props is a dictionary of
properties, including weight. Props should be None if no
additional properties are to be specified.
Returns:
A MoleculeGraph
"""
mg = MoleculeGraph.with_empty_graph(molecule, name="bonds", edge_weight_name="weight", edge_weight_units="")
for edge, props in edges.items():
try:
from_index = edge[0]
to_index = edge[1]
except TypeError:
raise ValueError("Edges must be given as (from_index, to_index) tuples")
if props is None:
weight = None
else:
weight = props.pop("weight", None)
if len(props.items()) == 0:
props = None
nodes = mg.graph.nodes
if not (from_index in nodes and to_index in nodes):
raise ValueError(
"Edges cannot be added if nodes are not present in the graph. Please check your indices."
)
mg.add_edge(from_index, to_index, weight=weight, edge_properties=props)
mg.set_node_attributes()
return mg
@staticmethod
def with_local_env_strategy(molecule, strategy) -> MoleculeGraph:
"""
Constructor for MoleculeGraph, using a strategy
from pymatgen.analysis.local_env.
molecule: Molecule object
strategy: an instance of a
pymatgen.analysis.local_env.NearNeighbors object
Returns:
mg, a MoleculeGraph
"""
if not strategy.molecules_allowed:
raise ValueError(f"{strategy=} is not designed for use with molecules! Choose another strategy.")
extend_structure = strategy.extend_structure_molecules
mg = MoleculeGraph.with_empty_graph(molecule, name="bonds", edge_weight_name="weight", edge_weight_units="")
# NearNeighbor classes only (generally) work with structures
# molecules have to be boxed first
coords = molecule.cart_coords
if extend_structure:
a = max(coords[:, 0]) - min(coords[:, 0]) + 100
b = max(coords[:, 1]) - min(coords[:, 1]) + 100
c = max(coords[:, 2]) - min(coords[:, 2]) + 100
structure = molecule.get_boxed_structure(a, b, c, no_cross=True, reorder=False)
else:
structure = None
for n in range(len(molecule)):
neighbors = strategy.get_nn_info(molecule, n) if structure is None else strategy.get_nn_info(structure, n)
for neighbor in neighbors:
# all bonds in molecules should not cross
# (artificial) periodic boundaries
if not np.array_equal(neighbor["image"], [0, 0, 0]):
continue
if n > neighbor["site_index"]:
from_index = neighbor["site_index"]
to_index = n
else:
from_index = n
to_index = neighbor["site_index"]
mg.add_edge(
from_index=from_index,
to_index=to_index,
weight=neighbor["weight"],
warn_duplicates=False,
)
duplicates = []
for edge in mg.graph.edges:
if edge[2] != 0:
duplicates.append(edge)
for duplicate in duplicates:
mg.graph.remove_edge(duplicate[0], duplicate[1], key=duplicate[2])
mg.set_node_attributes()
return mg

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I might open an issue to keep things organized. Instead of having issues scattered in comments. Sorry for any possible inconvenience.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no worries, there's no rush with #3705. i'll take care of any merge conflicts

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

frankly, mypy can be very annoying and if there are some errors that you can't fix, i'm happy to take a look but also feel free to add type: ignores

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no worries, there's no rush with #3705. i'll take care of any merge conflicts

Deeply appreciate that.

frankly, mypy can be very annoying and if there are some errors that you can't fix, i'm happy to take a look but also feel free to add type: ignores

Yes I would try my best to fix as many as I could, and ping you whenever necessary (grateful for this). I would try to avoid suppress error unless we have a good reason (to make the code as robust as we can).

@janosh janosh merged commit 2bd3906 into master Mar 25, 2024
22 checks passed
@janosh janosh deleted the lattice-methods-ret-self branch March 25, 2024 10:59
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
types Type all the things
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants