Fix __init__.py and __init__.pyi

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.

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

ipsuite/__init__.py

@@ -1,11 +1,58 @@
 """The ipsuite package."""
 
-import lazy_loader as lazy
+# <AUTOGEN_INIT>
+# 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
+# <AUTOGEN_INIT>
+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']

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']

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,

ipsuite/analysis/forces.py

@@ -96,3 +96,4 @@ def run(self):
 
         plt.tight_layout()
         plt.savefig(self.figure_path)
+        plt.close()

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]),

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)

ipsuite/analysis/sensitivity.py

@@ -83,13 +83,15 @@ 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))
         ax.set_yscale("log")
         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()

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()

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 = []

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):

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]:

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()