diff --git a/docs/source/_get_started/ips.rst b/docs/source/_get_started/ips.rst index 50034569..abfeeb9e 100644 --- a/docs/source/_get_started/ips.rst +++ b/docs/source/_get_started/ips.rst @@ -26,4 +26,4 @@ Sort of like a Jigsaw multiple nodes can be connected together in order to creat IPSuite provides various Nodes in order to meet the needs of MPL creation and training. .. figure:: ../images/mlp.png - :alt: MLP Develepment Process. + :alt: MLP Development Process. diff --git a/docs/source/index.rst b/docs/source/index.rst index 76bb3bf8..e50f7196 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -17,7 +17,7 @@ Welcome to IPSuite's documentation! What is IPS? <_get_started/ips> Getting Started <_get_started/quickstart> IPSuite <_nodes/ipsuite> - Modules <_nodes/modlules> + Modules <_nodes/modules> Indices and tables diff --git a/ipsuite/__init__.py b/ipsuite/__init__.py index 05a6c80a..a69f2493 100644 --- a/ipsuite/__init__.py +++ b/ipsuite/__init__.py @@ -1,11 +1,58 @@ """The ipsuite package.""" -import lazy_loader as lazy +# +# IPSuite uses `lazy-loader` to only load the necessary subpackages and functions. To make sure that all the +# functions and nodes are exposed to the API and static type checkers, `mkinit` can automatically generate +# the `__init__.py` files. Before committing changes, please run (while in the `ipsuite` source directory) +# +# mkinit --relative --lazy_loader_typed -w +# +# and add the changed `__init__.py` and `__init__.pyi` files to the commit. from ipsuite.utils.helpers import setup_ase from ipsuite.utils.logs import setup_logging -__getattr__, __dir__, __all__ = lazy.attach_stub(__name__, __file__) - setup_logging(__name__) setup_ase() + +# fmt: off +# +import lazy_loader + + +__getattr__, __dir__, __all__ = lazy_loader.attach_stub(__name__, __file__) + +__all__ = ['ASEMD', 'ASEMDSafeSampling', 'ATOMS_LST', 'AddData', 'AddDataH5MD', + 'AllowedStructuresFilter', 'AnalyseAtoms', 'AnalyseDensity', + 'AnalyseGlobalForceSensitivity', 'AnalyseSingleForceSensitivity', + 'AnalyseStructureMeanForce', 'ApplyCalculator', 'Atoms', + 'BarycenterMapping', 'Berendsen', 'BondStretchAnalyses', + 'BoxHeatUp', 'BoxOscillatingRampModifier', 'BoxScale', + 'CalibrationMetrics', 'Check', 'CollectMDSteps', + 'ComparePredictions', 'ConfigurationSelection', 'ConnectivityCheck', + 'DebugCheck', 'DensityCheck', 'DipoleHistogram', 'EnergyHistogram', + 'EnergySpikeCheck', 'EnergyUncertaintyHistogram', 'FilterOutlier', + 'FixedBondLengthConstraint', 'FixedLayerConstraint', + 'FixedSphereConstraint', 'Flatten', 'ForceAngles', + 'ForceDecomposition', 'ForceUncertaintyDecomposition', + 'ForcesHistogram', 'ForcesUncertaintyHistogram', 'HasAtoms', + 'HasOrIsAtoms', 'HasSelectedConfigurations', 'HookeanConstraint', + 'IPSNode', 'IndexSelection', 'LangevinThermostat', 'MD22Dataset', + 'MDStability', 'MoveSingleParticle', 'NPTThermostat', 'NaNCheck', + 'NodeWithCalculator', 'NodeWithThermostat', 'Prediction', + 'PredictionMetrics', 'PressureRampModifier', 'ProcessAtoms', + 'ProcessSingleAtom', 'Project', 'RandomSelection', 'RattleAnalysis', + 'RattleAtoms', 'RescaleBoxModifier', 'RotateMolecules', + 'STATIC_PATH', 'SVCRBarostat', 'SplitSelection', 'StressHistogram', + 'SurfaceRasterMetrics', 'SurfaceRasterScan', 'TemperatureCheck', + 'TemperatureOscillatingRampModifier', 'TemperatureRampModifier', + 'ThresholdCheck', 'ThresholdSelection', 'TranslateMolecules', + 'UNION_ATOMS_OR_ATOMS_LST', 'UniformArangeSelection', + 'UniformEnergeticSelection', 'UniformTemporalSelection', + 'VelocityVerletDynamic', 'WrapModifier', 'analysis', 'ase_sim', + 'base', 'bootstrap', 'calc', 'calculators', 'combine', + 'configuration_generation', 'configuration_selection', 'conftest', + 'data_loading', 'datasets', 'docs', 'doctest_namespace', 'dynamics', + 'fields', 'geometry', 'helpers', 'interfaces', 'log', 'md', + 'metrics', 'models', 'nodes', 'project', 'static_data', 'utils', + 'version'] diff --git a/ipsuite/__init__.pyi b/ipsuite/__init__.pyi index a63c4cdf..510c0f13 100644 --- a/ipsuite/__init__.pyi +++ b/ipsuite/__init__.pyi @@ -125,107 +125,104 @@ from .project import Project from .version import __version__ # Update __all__ for lazy loading -__all__ = [ - "__version__", - # Base - "Flatten", - "base", - # Models - "EnsembleModel", - "CP2KModel", - "TBLiteModel", - "ORCAModel", - "MACEMPModel", - "GenericASEModel", - "TorchDFTD3", - # Configuration Selection - "IndexSelection", - "RandomSelection", - "SplitSelection", - "UniformArangeSelection", - "UniformEnergeticSelection", - "UniformTemporalSelection", - "ThresholdSelection", - "FilterOutlier", - # Configuration Generation - "Packmol", - "MultiPackmol", - "Smiles2Atoms", - "Smiles2Conformers", - "Smiles2Gromacs", - # Data - "AddData", - "AddDataH5MD", - # Datasets - "MD22Dataset", - # Bootstrap - "RattleAtoms", - "TranslateMolecules", - "RotateMolecules", - "SurfaceRasterScan", - "SurfaceRasterMetrics", - # Analysis - "DipoleHistogram", - "EnergyHistogram", - "ForcesHistogram", - "StressHistogram", - "ForcesUncertaintyHistogram", - "EnergyUncertaintyHistogram", - "PredictionMetrics", - "ForceAngles", - "RattleAnalysis", - "Prediction", - "CalibrationMetrics", - "BoxScale", - "BoxHeatUp", - "DebugCheck", - "NaNCheck", - "ConnectivityCheck", - "EnergySpikeCheck", - "MDStability", - "MoveSingleParticle", - "AnalyseGlobalForceSensitivity", - "AnalyseSingleForceSensitivity", - "ForceUncertaintyDecomposition", - "ForceDecomposition", - "ThresholdCheck", - "TemperatureCheck", - "AnalyseDensity", - "DensityCheck", - "CollectMDSteps", - "AllowedStructuresFilter", - "AnalyseStructureMeanForce", - # Calculators - "CP2KSinglePoint", - "ASEGeoOpt", - "ASEMD", - "ASEMDSafeSampling", - "xTBSinglePoint", - "LJSinglePoint", - "EMTSinglePoint", - "OrcaSinglePoint", - "LammpsSimulator", - "MixCalculator", - "LangevinThermostat", - "VelocityVerletDynamic", - "Berendsen", - "NPTThermostat", - "SVCRBarostat", - "RescaleBoxModifier", - "BoxOscillatingRampModifier", - "TemperatureRampModifier", - "TemperatureOscillatingRampModifier", - "FixedSphereConstraint", - "FixedLayerConstraint", - "FixedBondLengthConstraint", - "HookeanConstraint", - "PressureRampModifier", - "PlumedModel", - # Geometry - "BarycenterMapping", - # Project - "Project", - # Calc - "ApplyCalculator", - "WrapModifier", -] +from . import analysis +from . import base +from . import bootstrap +from . import calc +from . import calculators +from . import configuration_generation +from . import configuration_selection +from . import conftest +from . import data_loading +from . import datasets +from . import dynamics +from . import fields +from . import geometry +from . import interfaces +from . import models +from . import nodes +from . import project +from . import static_data +from . import utils +from . import version + +from .analysis import (AllowedStructuresFilter, AnalyseDensity, + AnalyseGlobalForceSensitivity, + AnalyseSingleForceSensitivity, + AnalyseStructureMeanForce, BondStretchAnalyses, + BoxHeatUp, BoxScale, CalibrationMetrics, CollectMDSteps, + DipoleHistogram, EnergyHistogram, + EnergyUncertaintyHistogram, ForceAngles, + ForceDecomposition, ForceUncertaintyDecomposition, + ForcesHistogram, ForcesUncertaintyHistogram, + MDStability, MoveSingleParticle, Prediction, + PredictionMetrics, RattleAnalysis, StressHistogram,) +from .base import (AnalyseAtoms, Check, ComparePredictions, Flatten, IPSNode, + ProcessAtoms, ProcessSingleAtom, interfaces,) +from .bootstrap import (RattleAtoms, RotateMolecules, SurfaceRasterMetrics, + SurfaceRasterScan, TranslateMolecules,) +from .calc import (ApplyCalculator,) +from .configuration_selection import (ConfigurationSelection, FilterOutlier, + IndexSelection, RandomSelection, + SplitSelection, ThresholdSelection, + UniformArangeSelection, + UniformEnergeticSelection, + UniformTemporalSelection,) +from .conftest import (doctest_namespace, project,) +from .data_loading import (AddData, AddDataH5MD,) +from .datasets import (MD22Dataset,) +from .dynamics import (ASEMD, ASEMDSafeSampling, Berendsen, + BoxOscillatingRampModifier, ConnectivityCheck, + DebugCheck, DensityCheck, EnergySpikeCheck, + FixedBondLengthConstraint, FixedLayerConstraint, + FixedSphereConstraint, HookeanConstraint, + LangevinThermostat, NPTThermostat, NaNCheck, + PressureRampModifier, RescaleBoxModifier, SVCRBarostat, + TemperatureCheck, TemperatureOscillatingRampModifier, + TemperatureRampModifier, ThresholdCheck, + VelocityVerletDynamic, WrapModifier,) +from .fields import (Atoms,) +from .geometry import (BarycenterMapping,) +from .interfaces import (ATOMS_LST, HasAtoms, HasOrIsAtoms, + HasSelectedConfigurations, NodeWithCalculator, + NodeWithThermostat, ProcessAtoms, + UNION_ATOMS_OR_ATOMS_LST,) +from .nodes import (nodes,) +from .project import (Project, log,) +from .static_data import (STATIC_PATH,) +from .utils import (ase_sim, combine, docs, helpers, md, metrics,) + +__all__ = ['ASEMD', 'ASEMDSafeSampling', 'ATOMS_LST', 'AddData', 'AddDataH5MD', + 'AllowedStructuresFilter', 'AnalyseAtoms', 'AnalyseDensity', + 'AnalyseGlobalForceSensitivity', 'AnalyseSingleForceSensitivity', + 'AnalyseStructureMeanForce', 'ApplyCalculator', 'Atoms', + 'BarycenterMapping', 'Berendsen', 'BondStretchAnalyses', + 'BoxHeatUp', 'BoxOscillatingRampModifier', 'BoxScale', + 'CalibrationMetrics', 'Check', 'CollectMDSteps', + 'ComparePredictions', 'ConfigurationSelection', 'ConnectivityCheck', + 'DebugCheck', 'DensityCheck', 'DipoleHistogram', 'EnergyHistogram', + 'EnergySpikeCheck', 'EnergyUncertaintyHistogram', 'FilterOutlier', + 'FixedBondLengthConstraint', 'FixedLayerConstraint', + 'FixedSphereConstraint', 'Flatten', 'ForceAngles', + 'ForceDecomposition', 'ForceUncertaintyDecomposition', + 'ForcesHistogram', 'ForcesUncertaintyHistogram', 'HasAtoms', + 'HasOrIsAtoms', 'HasSelectedConfigurations', 'HookeanConstraint', + 'IPSNode', 'IndexSelection', 'LangevinThermostat', 'MD22Dataset', + 'MDStability', 'MoveSingleParticle', 'NPTThermostat', 'NaNCheck', + 'NodeWithCalculator', 'NodeWithThermostat', 'Prediction', + 'PredictionMetrics', 'PressureRampModifier', 'ProcessAtoms', + 'ProcessSingleAtom', 'Project', 'RandomSelection', 'RattleAnalysis', + 'RattleAtoms', 'RescaleBoxModifier', 'RotateMolecules', + 'STATIC_PATH', 'SVCRBarostat', 'SplitSelection', 'StressHistogram', + 'SurfaceRasterMetrics', 'SurfaceRasterScan', 'TemperatureCheck', + 'TemperatureOscillatingRampModifier', 'TemperatureRampModifier', + 'ThresholdCheck', 'ThresholdSelection', 'TranslateMolecules', + 'UNION_ATOMS_OR_ATOMS_LST', 'UniformArangeSelection', + 'UniformEnergeticSelection', 'UniformTemporalSelection', + 'VelocityVerletDynamic', 'WrapModifier', 'analysis', 'ase_sim', + 'base', 'bootstrap', 'calc', 'calculators', 'combine', + 'configuration_generation', 'configuration_selection', 'conftest', + 'data_loading', 'datasets', 'docs', 'doctest_namespace', 'dynamics', + 'fields', 'geometry', 'helpers', 'interfaces', 'log', 'md', + 'metrics', 'models', 'nodes', 'project', 'static_data', 'utils', + 'version'] diff --git a/ipsuite/analysis/bond_stretch.py b/ipsuite/analysis/bond_stretch.py index a75d2fc2..31b2a3b5 100644 --- a/ipsuite/analysis/bond_stretch.py +++ b/ipsuite/analysis/bond_stretch.py @@ -104,6 +104,8 @@ def run(self): e_fig.savefig(self.plots_dir / f"energy_{chem_symbols[0]}_{chem_symbols[1]}.png") f_fig.savefig(self.plots_dir / f"force_{chem_symbols[0]}_{chem_symbols[1]}.png") + plt.close(fig=e_fig) + plt.close(fig=f_fig) def get_plots( self, diff --git a/ipsuite/analysis/forces.py b/ipsuite/analysis/forces.py index 3e3dcc12..c92236c2 100644 --- a/ipsuite/analysis/forces.py +++ b/ipsuite/analysis/forces.py @@ -96,3 +96,4 @@ def run(self): plt.tight_layout() plt.savefig(self.figure_path) + plt.close() diff --git a/ipsuite/analysis/md.py b/ipsuite/analysis/md.py index c4639e71..617138d6 100644 --- a/ipsuite/analysis/md.py +++ b/ipsuite/analysis/md.py @@ -30,6 +30,7 @@ def run(self): ax.set_xlabel("Step") fig.tight_layout() fig.savefig(self.figure) + plt.close() self.density = { "density": np.mean(densities[self.start : self.end]), diff --git a/ipsuite/analysis/molecules.py b/ipsuite/analysis/molecules.py index ecfe5de6..4df947ef 100644 --- a/ipsuite/analysis/molecules.py +++ b/ipsuite/analysis/molecules.py @@ -21,18 +21,22 @@ class AllowedStructuresFilter(base.IPSNode): The molecules that are allowed. smiles : list[str], optional The SMILES strings of the allowed molecules. + cutoffs : dict[str, float] | None, optional + The cutoffs for each element. + If None, use the `ase.data.covalent_radii`. Default: None """ data: list[ase.Atoms] = zntrack.deps() molecules: list[ase.Atoms] = zntrack.deps(default_factory=list) smiles: list[str] = zntrack.params(default_factory=list) + cutoffs: dict[str, float] | None = zntrack.params(None) fail: bool = zntrack.params(False) outliers: list[int] = zntrack.outs() def run(self): molecules = self.molecules + [rdkit2ase.smiles2atoms(s) for s in self.smiles] - mapping = BarycenterMapping() + mapping = BarycenterMapping(cutoffs=self.cutoffs) outliers_set = set() for idx, atoms in enumerate(tqdm.tqdm(self.data)): _, mols = mapping.forward_mapping(atoms) diff --git a/ipsuite/analysis/sensitivity.py b/ipsuite/analysis/sensitivity.py index d949ad0b..e11a90f1 100644 --- a/ipsuite/analysis/sensitivity.py +++ b/ipsuite/analysis/sensitivity.py @@ -83,6 +83,7 @@ def run(self): fig, ax = plt.subplots() nonuniform_imshow(ax, r_ij[0, :, 0], r_ij[0, :, 1], mean_forces) fig.savefig(self.plots / "2d_forces.png") + plt.close() fig, ax = plt.subplots() ax.scatter(d_ij[0], np.sum(std_forces, axis=1)) @@ -90,6 +91,7 @@ def run(self): ax.set_xlabel(r"distance $d ~ / ~ \AA$") ax.set_ylabel(r"standard deviation $\sigma ~ / ~ a.u.$") fig.savefig(self.plots / "std_forces.png") + plt.close() def _compute_std_leave_one_out(data): # Leave-One-Out Cross-Validation @@ -159,3 +161,4 @@ def run(self): ax.set_xlabel(r"Distance $r ~ / ~ \AA$") ax.set_yscale("log") fig.savefig(self.sensitivity_plot, bbox_inches="tight") + plt.close() diff --git a/ipsuite/bootstrap/surface_mods.py b/ipsuite/bootstrap/surface_mods.py index 342c7e5c..3d966f63 100644 --- a/ipsuite/bootstrap/surface_mods.py +++ b/ipsuite/bootstrap/surface_mods.py @@ -196,3 +196,4 @@ def plot_ture_vs_pred(x, y, z, name, height, plots_dir): fig.suptitle(rf"Additive {height} $\AA$ over the surface") fig.savefig(plots_dir / f"{name}-{height}-heat.png") + plt.close() diff --git a/ipsuite/calculators/ase_geoopt.py b/ipsuite/calculators/ase_geoopt.py index 48d7997b..7f59da6b 100644 --- a/ipsuite/calculators/ase_geoopt.py +++ b/ipsuite/calculators/ase_geoopt.py @@ -79,6 +79,7 @@ class ASEGeoOpt(base.IPSNode): sampling_rate: int = zntrack.params(1) maxstep: int = zntrack.params(None) + log_file: pathlib.Path = zntrack.outs_path(zntrack.nwd / "opt.log") traj_file: pathlib.Path = zntrack.outs_path(zntrack.nwd / "structures.h5") def run(self): # noqa: C901 @@ -105,7 +106,7 @@ def run(self): # noqa: C901 db = znh5md.IO(self.traj_file) optimizer = getattr(ase.optimize, self.optimizer) - dyn = optimizer(atoms, **self.init_kwargs) + dyn = optimizer(atoms, logfile=self.log_file, **self.init_kwargs) for step, _ in enumerate(dyn.irun(**self.run_kwargs)): stop = [] diff --git a/ipsuite/configuration_selection/base.py b/ipsuite/configuration_selection/base.py index db736103..0daf2c26 100644 --- a/ipsuite/configuration_selection/base.py +++ b/ipsuite/configuration_selection/base.py @@ -88,6 +88,7 @@ def _get_plot(self, atoms_lst: typing.List[ase.Atoms], indices: typing.List[int] ax.scatter(indices, line_data[indices], c="r") ax.set_xlabel("Configuration") fig.savefig(self.img_selection, bbox_inches="tight") + plt.close() class BatchConfigurationSelection(ConfigurationSelection): diff --git a/ipsuite/configuration_selection/filter.py b/ipsuite/configuration_selection/filter.py index 59469bca..047104a2 100644 --- a/ipsuite/configuration_selection/filter.py +++ b/ipsuite/configuration_selection/filter.py @@ -60,6 +60,7 @@ def run(self): ax[2].hist([values[i] for i in self.filtered_indices], bins=100) ax[2].set_title("Excluded") fig.savefig(self.histogram, bbox_inches="tight") + plt.close() @property def frames(self) -> list[ase.Atoms]: diff --git a/ipsuite/configuration_selection/threshold.py b/ipsuite/configuration_selection/threshold.py index 077c3d4f..3c71e5cb 100644 --- a/ipsuite/configuration_selection/threshold.py +++ b/ipsuite/configuration_selection/threshold.py @@ -140,3 +140,4 @@ def _get_plot(self, atoms_lst: typing.List[ase.Atoms], indices: typing.List[int] ax.set_xlabel("configuration") fig.savefig(self.img_selection, bbox_inches="tight") + plt.close() diff --git a/ipsuite/geometry/graphs.py b/ipsuite/geometry/graphs.py index 42030e33..f98f3367 100644 --- a/ipsuite/geometry/graphs.py +++ b/ipsuite/geometry/graphs.py @@ -1,23 +1,46 @@ import ase import networkx as nx import numpy as np -from ase.neighborlist import build_neighbor_list +from ase.neighborlist import build_neighbor_list, natural_cutoffs -def atoms_to_graph(atoms: ase.Atoms) -> nx.Graph: +def atoms_to_graph(atoms: ase.Atoms, cutoffs: dict[str, float] | None = None) -> nx.Graph: """Converts ASE Atoms into a Graph based on their bond connectivity. + + Args: + atoms (ase.Atoms): Atoms instance to convert + cutoffs (dict[str, float] | None): cutoffs of each atom. + Dictionary with keys for the symbols and values of the cutoff radii. + If None, use the `ase.data.covalent_radii`. Default: None + + Returns: + G (nx.Graph): Connectivity graph """ # This can be optimized by reusing the NL! - nl = build_neighbor_list(atoms, self_interaction=False) + if cutoffs is not None: + cutoffs = natural_cutoffs(atoms, **cutoffs) + nl = build_neighbor_list(atoms, self_interaction=False, cutoffs=cutoffs) cm = nl.get_connectivity_matrix(sparse=False) G = nx.from_numpy_array(cm) return G -def identify_molecules(atoms: ase.Atoms) -> list[np.ndarray]: - """Identifies molecules in a structure based on the connected subgraphs.""" - G = atoms_to_graph(atoms) +def identify_molecules( + atoms: ase.Atoms, cutoffs: dict[str, float] | None = None +) -> list[np.ndarray]: + """Identifies molecules in a structure based on the connected subgraphs. + + Args: + atoms (ase.Atoms): Atoms instance to identify molecules in + cutoffs (dict[str, float] | None): cutoffs of each element. + Dictionary with keys for the symbols and values of the cutoff radii. + If None, use the `ase.data.covalent_radii`. Default: None + + Returns: + c_list (np.ndarray): Array of lists of connected atom indices + """ + G = atoms_to_graph(atoms, cutoffs=cutoffs) components = nx.connected_components(G) c_list = [np.array(list(c)) for c in components] return c_list diff --git a/ipsuite/geometry/mapping.py b/ipsuite/geometry/mapping.py index 579afa9b..d064efb6 100644 --- a/ipsuite/geometry/mapping.py +++ b/ipsuite/geometry/mapping.py @@ -25,15 +25,20 @@ class BarycenterMapping: frozen: bool If True, the neighbor list is only constructed for the first configuration. The indices of the molecules will be frozen for all configurations. + cutoffs: dict[str, float] | None + cutoffs of each element. + Dictionary with keys for the symbols and values of the cutoff radii. + If None, use the `ase.data.covalent_radii`. Default: None """ frozen: bool = False + cutoffs: dict[str, float] | None = None _components: t.Any | None = None def forward_mapping(self, atoms: ase.Atoms) -> tuple[ase.Atoms, list[ase.Atoms]]: if self._components is None: - components = graphs.identify_molecules(atoms) + components = graphs.identify_molecules(atoms, cutoffs=self.cutoffs) else: components = self._components diff --git a/pyproject.toml b/pyproject.toml index 088c6e8e..ed15281b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -84,6 +84,7 @@ exclude = [ "node_modules", "venv", "tests", + "ipsuite/__init__.pyi", ] # Same as Black. @@ -98,6 +99,10 @@ ignore = [ "B905", # not supported in older python versions ] +[tool.ruff.lint.per-file-ignores] +"ipsuite/__init__.py*" = ["E402", "F811", "I001"] + + [tool.codespell] ignore-words-list = "dscribe" skip = "poetry.lock,ipsuite/static_data/*,docs/source/examples/06_Bootstrapping_Datasets.ipynb,*.out"