Skip to content

Commit

Permalink
np.newaxis -> None
Browse files Browse the repository at this point in the history
json->JSON, yaml->YAML in doc strings
  • Loading branch information
janosh committed Jun 8, 2024
1 parent 14988c0 commit 36cde41
Show file tree
Hide file tree
Showing 28 changed files with 92 additions and 99 deletions.
4 changes: 2 additions & 2 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ ci:

repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.4.7
rev: v0.4.8
hooks:
- id: ruff
args: [--fix, --unsafe-fixes]
Expand Down Expand Up @@ -64,6 +64,6 @@ repos:
args: [--drop-empty-cells, --keep-output]

- repo: https://github.com/RobertCraigie/pyright-python
rev: v1.1.365
rev: v1.1.366
hooks:
- id: pyright
11 changes: 4 additions & 7 deletions dev_scripts/regen_libxcfunc.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
"""
This script regenerates the enum values in pymatgen.core.libxc_func.py.
It requires in input the path of the `libxc_docs.txt` file contained in libxc/src
The script parses this file, creates a new json file inside pymatgen.core
The script parses this file, creates a new JSON file inside pymatgen.core
and update the enum values declared in LibxcFunc.
The script must be executed inside pymatgen/dev_scripts.
"""
Expand All @@ -16,10 +16,7 @@


def parse_libxc_docs(path):
"""
Parse libxc_docs.txt file, return dictionary with mapping:
libxc_id --> info_dict.
"""
"""Parse libxc_docs.txt file, return dictionary {libxc_id: info_dict}."""

def parse_section(section):
dct = {}
Expand All @@ -46,7 +43,7 @@ def parse_section(section):


def write_libxc_docs_json(xc_funcs, json_path):
"""Write json file with libxc metadata to path jpath."""
"""Write JSON file with libxc metadata to path jpath."""
xc_funcs = deepcopy(xc_funcs)

# Remove XC_FAMILY from Family and XC_ from Kind to make strings more human-readable.
Expand Down Expand Up @@ -85,7 +82,7 @@ def main():

xc_funcs = parse_libxc_docs(path)

# Generate new json file in pycore
# Generate new JSON file in pycore
pmg_core = os.path.abspath("../pymatgen/core/")
json_path = f"{pmg_core}/libxc_docs.json"
write_libxc_docs_json(xc_funcs, json_path)
Expand Down
2 changes: 1 addition & 1 deletion pymatgen/analysis/bond_valence.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@
for key, val in loadfn(f"{module_dir}/bvparam_1991.yaml").items():
BV_PARAMS[Element(key)] = val

# Read in yaml containing data-mined ICSD BV data.
# Read in YAML containing data-mined ICSD BV data.
all_data = loadfn(f"{module_dir}/icsd_bv.yaml")
ICSD_BV_DATA = {Species.from_str(sp): data for sp, data in all_data["bvsum"].items()}
PRIOR_PROB = {Species.from_str(sp): data for sp, data in all_data["occurrence"].items()}
Expand Down
4 changes: 2 additions & 2 deletions pymatgen/analysis/dimensionality.py
Original file line number Diff line number Diff line change
Expand Up @@ -403,9 +403,9 @@ def find_connected_atoms(struct, tolerance=0.45, ldict=None):

n_atoms = len(struct.species)
fc = np.array(struct.frac_coords)
fc_copy = np.repeat(fc[:, :, np.newaxis], 27, axis=2)
fc_copy = np.repeat(fc[:, :, None], 27, axis=2)
neighbors = np.array(list(itertools.product([0, 1, -1], [0, 1, -1], [0, 1, -1]))).T
neighbors = np.repeat(neighbors[np.newaxis, :, :], 1, axis=0)
neighbors = np.repeat(neighbors[None, :, :], 1, axis=0)
fc_diff = fc_copy - neighbors
species = list(map(str, struct.species))
# in case of charged species
Expand Down
12 changes: 7 additions & 5 deletions pymatgen/analysis/ferroelectricity/polarization.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,9 @@ class Polarization:
electron Angstroms along the three lattice directions (a,b,c).
"""

def __init__(self, p_elecs, p_ions, structures, p_elecs_in_cartesian=True, p_ions_in_cartesian=False):
def __init__(
self, p_elecs, p_ions, structures: Sequence[Structure], p_elecs_in_cartesian=True, p_ions_in_cartesian=False
):
"""
p_elecs: np.array of electronic contribution to the polarization with shape [N, 3]
p_ions: np.array of ionic contribution to the polarization with shape [N, 3]
Expand Down Expand Up @@ -260,8 +262,8 @@ def get_same_branch_polarization_data(self, convert_to_muC_per_cm2=True, all_in_
p_tot = p_elec + p_ion
p_tot = np.array(p_tot)

lattices = [s.lattice for s in self.structures]
volumes = np.array([s.lattice.volume for s in self.structures])
lattices = [struct.lattice for struct in self.structures]
volumes = np.array([latt.volume for latt in lattices])

n_elecs = len(p_elec)

Expand All @@ -273,7 +275,7 @@ def get_same_branch_polarization_data(self, convert_to_muC_per_cm2=True, all_in_
# convert polarizations and lattice lengths prior to adjustment
if convert_to_muC_per_cm2 and not all_in_polar:
# Convert the total polarization
p_tot = np.multiply(units.T[:, np.newaxis], p_tot)
p_tot = np.multiply(units.T[:, None], p_tot)
# adjust lattices
for idx in range(n_elecs):
lattice = lattices[idx]
Expand Down Expand Up @@ -316,7 +318,7 @@ def get_lattice_quanta(self, convert_to_muC_per_cm2=True, all_in_polar=True):
"""Get the dipole / polarization quanta along a, b, and c for
all structures.
"""
lattices = [s.lattice for s in self.structures]
lattices = [struct.lattice for struct in self.structures]
volumes = np.array([struct.volume for struct in self.structures])

n_structs = len(self.structures)
Expand Down
4 changes: 2 additions & 2 deletions pymatgen/analysis/graphs.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ def __init__(self, structure: Structure, graph_data: dict | None = None) -> None
self.structure = structure
self.graph = nx.readwrite.json_graph.adjacency_graph(graph_data)

# tidy up edge attr dicts, reading to/from json duplicates information
# tidy up edge attr dicts, reading to/from JSON duplicates information
for _, _, _, data in self.graph.edges(keys=True, data=True):
for key in ("id", "key"):
data.pop(key, None)
Expand Down Expand Up @@ -1585,7 +1585,7 @@ def __init__(self, molecule, graph_data=None):
self.molecule = molecule
self.graph = nx.readwrite.json_graph.adjacency_graph(graph_data)

# tidy up edge attr dicts, reading to/from json duplicates
# tidy up edge attr dicts, reading to/from JSON duplicates
# information
for _, _, _, data in self.graph.edges(keys=True, data=True):
for key in ("id", "key"):
Expand Down
56 changes: 27 additions & 29 deletions pymatgen/analysis/magnetism/heisenberg.py
Original file line number Diff line number Diff line change
Expand Up @@ -555,18 +555,18 @@ def get_interaction_graph(self, filename=None):
logging.warning(warning_msg)

# J_ij exchange interaction matrix
for i, _node in enumerate(sgraph.graph.nodes):
connections = sgraph.get_connected_sites(i)
for idx in range(len(sgraph.graph.nodes)):
connections = sgraph.get_connected_sites(idx)
for c in connections:
jimage = c[1] # relative integer coordinates of atom j
j = c[2] # index of neighbor
dist = c[-1] # i <-> j distance

j_exc = self._get_j_exc(i, j, dist)
j_exc = self._get_j_exc(idx, j, dist)

igraph.add_edge(i, j, to_jimage=jimage, weight=j_exc, warn_duplicates=False)
igraph.add_edge(idx, j, to_jimage=jimage, weight=j_exc, warn_duplicates=False)

# Save to a json file if desired
# Save to a JSON file if desired
if filename:
if not filename.endswith(".json"):
filename += ".json"
Expand Down Expand Up @@ -641,7 +641,7 @@ def get_heisenberg_model(self):
hm_nni = self.nn_interactions
hm_d = self.dists

# Exchange matrix DataFrame in json format
# Exchange matrix DataFrame in JSON format
hm_em = self.ex_mat.to_json()
hm_ep = self.get_exchange()
hm_javg = self.estimate_exchange()
Expand Down Expand Up @@ -860,29 +860,27 @@ def __init__(
self.igraph = igraph

def as_dict(self):
"""Because some dicts have tuple keys, some sanitization is required for json compatibility."""
dct = {}
dct["@module"] = type(self).__module__
dct["@class"] = type(self).__name__
dct["@version"] = __version__
dct["formula"] = self.formula
dct["structures"] = [s.as_dict() for s in self.structures]
dct["energies"] = self.energies
dct["cutoff"] = self.cutoff
dct["tol"] = self.tol
dct["sgraphs"] = [sgraph.as_dict() for sgraph in self.sgraphs]
dct["dists"] = self.dists
dct["ex_params"] = self.ex_params
dct["javg"] = self.javg
dct["igraph"] = self.igraph.as_dict()

# Sanitize tuple & int keys
dct["ex_mat"] = jsanitize(self.ex_mat)
dct["nn_interactions"] = jsanitize(self.nn_interactions)
dct["unique_site_ids"] = jsanitize(self.unique_site_ids)
dct["wyckoff_ids"] = jsanitize(self.wyckoff_ids)

return dct
"""Because some dicts have tuple keys, some sanitization is required for JSON compatibility."""
return {
"@module": type(self).__module__,
"@class": type(self).__name__,
"@version": __version__,
"formula": self.formula,
"structures": [struct.as_dict() for struct in self.structures],
"energies": self.energies,
"cutoff": self.cutoff,
"tol": self.tol,
"sgraphs": [sgraph.as_dict() for sgraph in self.sgraphs],
"dists": self.dists,
"ex_params": self.ex_params,
"javg": self.javg,
"igraph": self.igraph.as_dict(),
# Sanitize tuple & int keys
"ex_mat": jsanitize(self.ex_mat),
"nn_interactions": jsanitize(self.nn_interactions),
"unique_site_ids": jsanitize(self.unique_site_ids),
"wyckoff_ids": jsanitize(self.wyckoff_ids),
}

@classmethod
def from_dict(cls, dct: dict) -> Self:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -50,11 +50,9 @@ class SubstitutionProbability:
def __init__(self, lambda_table=None, alpha=-5):
"""
Args:
lambda_table:
json table of the weight functions lambda if None,
lambda_table: JSON table of the weight functions lambda if None,
will use the default lambda.json table
alpha:
weight function for never observed substitutions.
alpha (float): weight function for never observed substitutions.
"""
if lambda_table is not None:
self._lambda_table = lambda_table
Expand Down Expand Up @@ -94,8 +92,8 @@ def get_lambda(self, s1, s2):
Returns:
Lambda values
"""
k = frozenset([get_el_sp(s1), get_el_sp(s2)])
return self._l.get(k, self.alpha)
key = frozenset([get_el_sp(s1), get_el_sp(s2)])
return self._l.get(key, self.alpha)

def get_px(self, sp):
"""
Expand Down
2 changes: 1 addition & 1 deletion pymatgen/analysis/transition_state.py
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,7 @@ def as_dict(self):
"r": jsanitize(self.r),
"energies": jsanitize(self.energies),
"forces": jsanitize(self.forces),
"structures": [s.as_dict() for s in self.structures],
"structures": [struct.as_dict() for struct in self.structures],
}


Expand Down
2 changes: 1 addition & 1 deletion pymatgen/cli/pmg.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ def main():
metavar="dir_name",
nargs=2,
help="Initial directory where the CP2K data is located and the output directory where the "
"CP2K yaml data files will be written",
"CP2K YAML data files will be written",
)

parser_config.add_argument(
Expand Down
6 changes: 3 additions & 3 deletions pymatgen/command_line/critic2_caller.py
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,7 @@ def from_chgcar(
critical points to a file 'table.cml' in the working directory
useful for visualization
write_json (bool): Whether to write out critical points
and YT json. YT integration will be performed with this setting.
and YT JSON. YT integration will be performed with this setting.
zpsp (dict): Dict of element/symbol name to number of electrons
(ZVAL in VASP pseudopotential), with which to properly augment core regions
and calculate charge transfer. Optional.
Expand Down Expand Up @@ -453,8 +453,8 @@ class with bonding information. By default, this returns
structure: associated Structure
stdout: stdout from running critic2 in automatic mode
stderr: stderr from running critic2 in automatic mode
cpreport: json output from CPREPORT command
yt: json output from YT command
cpreport: JSON output from CPREPORT command
yt: JSON output from YT command
zpsp (dict): Dict of element/symbol name to number of electrons
(ZVAL in VASP pseudopotential), with which to calculate charge transfer.
Optional.
Expand Down
2 changes: 1 addition & 1 deletion pymatgen/command_line/vampire_caller.py
Original file line number Diff line number Diff line change
Expand Up @@ -400,7 +400,7 @@ class VampireOutput(MSONable):
def __init__(self, parsed_out=None, nmats=None, critical_temp=None):
"""
Args:
parsed_out (json): json rep of parsed stdout DataFrame.
parsed_out (str): JSON rep of parsed stdout DataFrame.
nmats (int): Number of distinct materials (1 for each specie and up/down spin).
critical_temp (float): Monte Carlo Tc result.
"""
Expand Down
2 changes: 1 addition & 1 deletion pymatgen/core/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,7 @@ def _load_pmg_settings() -> dict[str, Any]:
except Exception as exc:
# If there are any errors, default to using environment variables
# if present.
warnings.warn(f"Error loading {file_path}: {exc}.\nYou may need to reconfigure your yaml file.")
warnings.warn(f"Error loading {file_path}: {exc}.\nYou may need to reconfigure your YAML file.")

# Override .pmgrc.yaml with env vars (if present)
for key, val in os.environ.items():
Expand Down
4 changes: 2 additions & 2 deletions pymatgen/core/libxcfunc.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
__status__ = "Production"
__date__ = "May 16, 2016"

# Load libxc info from json file
# Load libxc info from JSON file
with open(os.path.join(os.path.dirname(__file__), "libxc_docs.json"), encoding="utf-8") as file:
_all_xcfuncs = {int(k): v for k, v in json.load(file).items()}

Expand Down Expand Up @@ -497,5 +497,5 @@ def from_dict(cls, dct: dict) -> Self:
return cls[dct["name"]]

def to_json(self) -> str:
"""Get a json string representation of the LibxcFunc."""
"""Get a JSON string representation of the LibxcFunc."""
return json.dumps(self.as_dict(), cls=MontyEncoder)
2 changes: 1 addition & 1 deletion pymatgen/core/periodic_table.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@

from pymatgen.util.typing import SpeciesLike

# Load element data from json file
# Load element data from JSON file
with open(Path(__file__).absolute().parent / "periodic_table.json", encoding="utf-8") as ptable_json:
_pt_data = json.load(ptable_json)

Expand Down
4 changes: 2 additions & 2 deletions pymatgen/core/surface.py
Original file line number Diff line number Diff line change
Expand Up @@ -1679,7 +1679,7 @@ def generate_all_slabs(
return all_slabs


# Load the reconstructions_archive json file
# Load the reconstructions_archive JSON file
module_dir = os.path.dirname(os.path.abspath(__file__))
with open(f"{module_dir}/reconstructions_archive.json", encoding="utf-8") as data_file:
RECONSTRUCTIONS_ARCHIVE = json.load(data_file)
Expand Down Expand Up @@ -1715,7 +1715,7 @@ class ReconstructionGenerator:
the reconstructed slab. Only the a and b lattice vectors are
actually changed while the c vector remains the same.
This matrix is what the Wood's notation is based on.
reconstruction_json (dict): The full json or dictionary containing
reconstruction_json (dict): The full JSON or dictionary containing
the instructions for building the slab.
Todo:
Expand Down
2 changes: 1 addition & 1 deletion pymatgen/electronic_structure/plotter.py
Original file line number Diff line number Diff line change
Expand Up @@ -4062,7 +4062,7 @@ def plot_fermi_surface(

polydata = cp.actor.actors[0].mapper.input
pts = np.array(polydata.points) # - 1
polydata.points = np.dot(pts, cell / np.array(data.shape)[:, np.newaxis])
polydata.points = np.dot(pts, cell / np.array(data.shape)[:, None])

cx, cy, cz = (np.mean(np.array(polydata.points)[:, i]) for i in range(3))

Expand Down
2 changes: 1 addition & 1 deletion pymatgen/ext/matproj_legacy.py
Original file line number Diff line number Diff line change
Expand Up @@ -364,7 +364,7 @@ def get_doc(self, materials_id):
materials_id (str): e.g. mp-1143 for Al2O3
Returns:
Dict of json document of all data that is displayed on a materials
Dict of JSON document of all data that is displayed on a materials
details page.
"""
return self._make_request(f"/materials/{materials_id}/doc", mp_decode=False)
Expand Down
2 changes: 1 addition & 1 deletion pymatgen/io/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -276,7 +276,7 @@ def to_hdf5(self, filename):
f["fcoords"]: Fractional coords
f["lattice"]: Lattice in the pymatgen.core.Lattice matrix
format
f.attrs["structure_json"]: String of json representation
f.attrs["structure_json"]: String of JSON representation
Args:
filename (str): Filename to output to.
Expand Down
4 changes: 2 additions & 2 deletions pymatgen/io/phonopy.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ def get_displaced_structures(pmg_structure, atom_disp=0.01, supercell_matrix=Non
atom_disp (float): Atomic displacement. Default is 0.01 $\\AA$.
supercell_matrix (3x3 array): Scaling matrix for supercell.
yaml_fname (str): If not None, it represents the full path to
the outputting displacement yaml file, e.g. disp.yaml.
the outputting displacement YAML file, e.g. disp.yaml.
**kwargs: Parameters used in Phonopy.generate_displacement method.
Returns:
Expand Down Expand Up @@ -411,7 +411,7 @@ def get_phonon_band_structure_symm_line_from_fc(
def get_gruneisenparameter(gruneisen_path, structure=None, structure_path=None) -> GruneisenParameter:
"""Get Gruneisen object from gruneisen.yaml file, as obtained from phonopy (Frequencies in THz!).
The order is structure > structure path > structure from gruneisen dict.
Newer versions of phonopy include the structure in the yaml file,
Newer versions of phonopy include the structure in the YAML file,
the structure/structure_path is kept for compatibility.
Args:
Expand Down
Loading

0 comments on commit 36cde41

Please sign in to comment.