pack-mm 0.0.21__tar.gz → 0.0.24__tar.gz

{pack_mm-0.0.21 → pack_mm-0.0.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pack-mm
-Version: 0.0.21
+Version: 0.0.24
 Summary: packing materials and molecules in boxes using for machine learnt interatomic potentials
 Author: Alin M. Elena
 Classifier: Programming Language :: Python
@@ -42,7 +42,7 @@ It provides both a cli and a python api, with some examples below.
    uv pip install pack-mm
 
 ```
-or install the lates
+or install the latest
 
 ```bash
 
@@ -50,6 +50,11 @@ or install the lates
 
 ```
 
+## Jupyter notebook examples
+
+
+- [Basics](docs/source/tutorials/basics.ipynb) [![badge](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ddmms/pack-mm/blob/main/docs/source/tutorials/basic.ipynb)
+
 ## CLI examples
 
 
@@ -125,7 +130,7 @@ packmm --system Pd-super.cif --molecule H2 --nmols 50  --where anywhere   --mode
 
 before optimisation
 
-![](examples/pics/Pd-H2-noopt.webp)
+![](examples/pics/Pd-super+50H2.webp)
 
 
 after optimisation

{pack_mm-0.0.21 → pack_mm-0.0.24}/README.md

@@ -21,7 +21,7 @@ It provides both a cli and a python api, with some examples below.
    uv pip install pack-mm
 
 ```
-or install the lates
+or install the latest
 
 ```bash
 
@@ -29,6 +29,11 @@ or install the lates
 
 ```
 
+## Jupyter notebook examples
+
+
+- [Basics](docs/source/tutorials/basics.ipynb) [![badge](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ddmms/pack-mm/blob/main/docs/source/tutorials/basic.ipynb)
+
 ## CLI examples
 
 
@@ -104,7 +109,7 @@ packmm --system Pd-super.cif --molecule H2 --nmols 50  --where anywhere   --mode
 
 before optimisation
 
-![](examples/pics/Pd-H2-noopt.webp)
+![](examples/pics/Pd-super+50H2.webp)
 
 
 after optimisation

{pack_mm-0.0.21 → pack_mm-0.0.24}/pack_mm/cli/packmm.py

@@ -25,6 +25,22 @@ class InsertionMethod(str, Enum):
     ELLIPSOID = "ellipsoid"
 
 
+class InsertionStrategy(str, Enum):
+    """Insertion options."""
+
+    # propose randomly a point
+    MC = "mc"
+    # hybrid monte carlo
+    HMC = "hmc"
+
+
+class RelaxStrategy(str, Enum):
+    """Relaxation options."""
+
+    GEOMETRY_OPTIMISATION = "geometry_optimisation"
+    MD = "md"
+
+
 app = typer.Typer(no_args_is_help=True)
 
 
@@ -44,12 +60,26 @@ def packmm(
     ntries: int = typer.Option(
         50, help="Maximum number of attempts to insert each molecule."
     ),
+    every: int = typer.Option(
+        -1, help="Run MD-NVE or Geometry optimisation everyth insertion."
+    ),
     seed: int = typer.Option(2025, help="Random seed for reproducibility."),
+    md_steps: int = typer.Option(10, help="Number of steps to run MD."),
+    md_timestep: float = typer.Option(1.0, help="Timestep for MD integration, in fs."),
     where: InsertionMethod = typer.Option(
         InsertionMethod.ANYWHERE,
         help="""Where to insert the molecule. Choices: 'anywhere', 'sphere',
         'box', 'cylinderZ', 'cylinderY', 'cylinderX', 'ellipsoid'.""",
     ),
+    insert_strategy: InsertionStrategy = typer.Option(
+        InsertionStrategy.MC,
+        help="""How to insert a new molecule. Choices: 'mc', 'hmc',""",
+    ),
+    relax_strategy: RelaxStrategy = typer.Option(
+        RelaxStrategy.GEOMETRY_OPTIMISATION,
+        help="""How to relax the system to get more favourable structures.
+            Choices: 'geometry_optimisation', 'md',""",
+    ),
     centre: str | None = typer.Option(
         None,
         help="""Centre of the insertion zone, coordinates in Å,
@@ -84,6 +114,9 @@ def packmm(
     temperature: float = typer.Option(
         300.0, help="Temperature for the Monte Carlo acceptance rule."
     ),
+    md_temperature: float = typer.Option(
+        100.0, help="Temperature for the Molecular dynamics relaxation."
+    ),
     cell_a: float = typer.Option(
         20.0, help="Side of the empty box along the x-axis in Å."
     ),
@@ -125,6 +158,12 @@ def packmm(
     print(f"{fmax=}")
     print(f"{geometry=}")
     print(f"{out_path=}")
+    print(f"{every=}")
+    print(f"insert_strategy={insert_strategy.value}")
+    print(f"relax_strategy={relax_strategy.value}")
+    print(f"{md_steps=}")
+    print(f"{md_timestep=}")
+    print(f"{md_temperature=}")
     if nmols == -1:
         print("nothing to do, no molecule to insert")
         raise typer.Exit(0)
@@ -161,8 +200,10 @@ def packmm(
         cell_b=cell_b,
         cell_c=cell_c,
         out_path=out_path,
+        every=every,
+        relax_strategy=relax_strategy,
+        insert_strategy=insert_strategy,
+        md_steps=md_steps,
+        md_timestep=md_timestep,
+        md_temperature=md_temperature,
     )
-
-
-if __name__ == "__main__":
-    app()

{pack_mm-0.0.21 → pack_mm-0.0.24}/pack_mm/core/core.py

@@ -13,6 +13,7 @@ from ase.build import molecule as build_molecule
 from ase.io import read, write
 from ase.units import kB
 from janus_core.calculations.geom_opt import GeomOpt
+from janus_core.calculations.md import NVE
 from janus_core.helpers.mlip_calculators import choose_calculator
 from numpy import cos, exp, pi, random, sin, sqrt
 
@@ -130,7 +131,7 @@ def random_point_in_cylinder(
     return (x, y, z)
 
 
-def validate_value(label, x):
+def validate_value(label: str, x: float | int) -> None:
     """Validate input value, and raise an exception."""
     if x is not None and x < 0.0:
         err = f"Invalid {label}, needs to be positive"
@@ -138,8 +139,54 @@ def validate_value(label, x):
         raise Exception(err)
 
 
+def set_random_seed(seed: int) -> None:
+    """Set random seed."""
+    random.seed(seed)
+
+
+def set_defaults(
+    cell: (float, float, float),
+    centre: (float, float, float) | None = None,
+    where: str | None = None,
+    a: float | None = None,
+    b: float | None = None,
+    c: float | None = None,
+    radius: float | None = None,
+    height: float | None = None,
+) -> tuple(
+    (float, float, float),
+    float | None,
+    float | None,
+    float | None,
+    float | None,
+    float | None,
+):
+    """Set defaults for insertion areas."""
+    if centre is None:
+        centre = (cell[0] * 0.5, cell[1] * 0.5, cell[2] * 0.5)
+
+    if where == "anywhere":
+        a, b, c = cell[0], cell[1], cell[2]
+    elif where == "sphere":
+        radius = radius or min(cell) * 0.5
+    elif where == "cylinderZ":
+        radius = radius or min(cell[0], cell[1]) * 0.5
+        height = height or 0.5 * cell[2]
+    elif where == "cylinderY":
+        radius = radius or min(cell[0], cell[2]) * 0.5
+        height = height or 0.5 * cell[1]
+    elif where == "cylinderX":
+        radius = radius or min(cell[2], cell[1]) * 0.5
+        height = height or 0.5 * cell[0]
+    elif where == "box":
+        a, b, c = a or cell[0], b or cell[1], c or cell[2]
+    elif where == "ellipsoid":
+        a, b, c = a or cell[0] * 0.5, b or cell[1] * 0.5, c or cell[2] * 0.5
+    return (centre, a, b, c, radius, height)
+
+
 def pack_molecules(
-    system: str = None,
+    system: str | Atoms = None,
     molecule: str = "H2O",
     nmols: int = -1,
     arch: str = "cpu",
@@ -161,13 +208,19 @@ def pack_molecules(
     cell_b: float = None,
     cell_c: float = None,
     out_path: str = ".",
-) -> float:
+    every: int = -1,
+    relax_strategy: str = "geometry_optimisation",
+    insert_strategy: str = "random",
+    md_steps: int = 10,
+    md_timestep: float = 1.0,
+    md_temperature: float = 100.0,
+) -> tuple(float, Atoms):
     """
     Pack molecules into a system based on the specified parameters.
 
     Parameters
     ----------
-        system (str): Path to the system file or name of the system.
+        system (str|Atoms): Path to the system file or name of the system.
         molecule (str): Path to the molecule file or name of the molecule.
         nmols (int): Number of molecules to insert.
         arch (str): Architecture for the calculator.
@@ -185,6 +238,19 @@ def pack_molecules(
         geometry (bool): Whether to perform geometry optimization after insertion.
         cell_a, cell_b, cell_c (float): Cell dimensions if system is empty.
         out_path (str): path to save various outputs
+        every (int): After how many instertions to do a relaxation,
+                     default -1 means none..
+        md_temperature (float): Temperature in Kelvin for MD.
+        md_steps (int): Number of steps for MD.
+        md_timestep (float): Timestep in fs for MD.
+        insert_strategy (str): Insert strategy, "random" or "md"
+        relax_strategy (str): Relax strategy, "geometry_optimisation" or "md"
+
+    Returns
+    -------
+        tuple: A tuple energy and Atoms object containing
+               original system and added molecules..
+
     """
     kbt = temperature * kB
     validate_value("temperature", temperature)
@@ -198,51 +264,34 @@ def pack_molecules(
     validate_value("ntries", ntries)
     validate_value("cell box cell a", cell_a)
     validate_value("cell box cell b", cell_b)
-    validate_value("nmols", nmols)
     validate_value("cell box cell c", cell_c)
+    validate_value("nmols", nmols)
+    validate_value("MD steps", md_steps)
+    validate_value("MD timestep", md_timestep)
+    validate_value("MD temperature", md_temperature)
 
-    random.seed(seed)
+    set_random_seed(seed)
 
-    try:
-        sys = read(system)
-        sysname = Path(system).stem
-    except Exception:
+    if system is None:
         sys = Atoms(cell=[cell_a, cell_b, cell_c], pbc=[True, True, True])
-        sysname = "empty"
-
-    cell = sys.cell.lengths()
+        sysname = ""
+    elif isinstance(system, Atoms):
+        sys = system.copy()
+        sysname = sys.get_chemical_formula()
+    else:
+        sys = read(system)
+        sysname = Path(system).stem + "+"
 
     # Print summary
     print(f"Inserting {nmols} {molecule} molecules in {sysname}.")
     print(f"Using {arch} model {model} on {device}.")
     print(f"Insert in {where}.")
 
-    if center is None:
-        center = (cell[0] * 0.5, cell[1] * 0.5, cell[2] * 0.5)
+    cell = sys.cell.lengths()
 
-    if where == "anywhere":
-        a, b, c = cell[0], cell[1], cell[2]
-    elif where == "sphere":
-        if radius is None:
-            radius = min(cell) * 0.5
-    elif where in ["cylinderZ", "cylinderY", "cylinderX"]:
-        if radius is None:
-            if where == "cylinderZ":
-                radius = min(cell[0], cell[1]) * 0.5
-                if height is None:
-                    height = 0.5 * cell[2]
-            elif where == "cylinderY":
-                radius = min(cell[0], cell[2]) * 0.5
-                if height is None:
-                    height = 0.5 * cell[1]
-            elif where == "cylinderX":
-                radius = min(cell[2], cell[1]) * 0.5
-                if height is None:
-                    height = 0.5 * cell[0]
-    elif where == "box":
-        a, b, c = a or cell[0], b or cell[1], c or cell[2]
-    elif where == "ellipsoid":
-        a, b, c = a or cell[0], b or cell[1], c or cell[2]
+    center, a, b, c, radius, height = set_defaults(
+        cell, center, where, a, b, c, radius, height
+    )
 
     calc = choose_calculator(arch=arch, model_path=model, device=device)
     sys.calc = calc
@@ -250,7 +299,8 @@ def pack_molecules(
     e = sys.get_potential_energy() if len(sys) > 0 else 0.0
 
     csys = sys.copy()
-    for i in range(nmols):
+    i = 0
+    while i < nmols:
         accept = False
         for _itry in range(ntries):
             mol = load_molecule(molecule)
@@ -259,6 +309,25 @@ def pack_molecules(
             mol.translate(tv)
 
             tsys = csys.copy() + mol.copy()
+            if insert_strategy == "hmc":
+                tsys = run_md_nve(
+                    tsys, md_temperature, md_steps, md_timestep, arch, model, device
+                )
+
+            if every > 0 and _itry / every == 0:
+                tsys = save_the_day(
+                    struct_path=tsys,
+                    device=device,
+                    arch=arch,
+                    model=model,
+                    fmax=fmax,
+                    out_path=out_path,
+                    md_temperature=md_temperature,
+                    md_steps=md_steps,
+                    md_timestep=md_timestep,
+                    relax_strategy=relax_strategy,
+                )
+
             tsys.calc = calc
             en = tsys.get_potential_energy()
             de = en - e
@@ -270,36 +339,43 @@ def pack_molecules(
             if u <= acc:
                 accept = True
                 break
-
         if accept:
             csys = tsys.copy()
             e = en
-            print(f"Inserted particle {i + 1}")
-            write(Path(out_path) / f"{sysname}+{i + 1}{Path(molecule).stem}.cif", csys)
+            i += 1
+            print(f"Inserted particle {i}")
+            write(Path(out_path) / f"{sysname}{i}{Path(molecule).stem}.cif", csys)
         else:
             # Things are bad, maybe geomatry optimisation saves us
+            # once you hit here is bad, this can keep looping
             print(f"Failed to insert particle {i + 1} after {ntries} tries")
-            _ = optimize_geometry(
-                f"{sysname}+{i + 1}{Path(molecule).stem}.cif",
+            csys = save_the_day(
+                csys,
                 device,
                 arch,
                 model,
                 fmax,
                 out_path,
+                md_temperature,
+                md_steps,
+                md_timestep,
+                relax_strategy,
             )
+
     energy_final = e
 
     # Perform final geometry optimization if requested
     if geometry:
-        energy_final = optimize_geometry(
-            f"{sysname}+{nmols}{Path(molecule).stem}.cif",
+        energy_final, csys = optimize_geometry(
+            Path(out_path) / f"{sysname}{nmols}{Path(molecule).stem}.cif",
             device,
             arch,
             model,
             fmax,
             out_path,
+            True,
         )
-    return energy_final
+    return (energy_final, csys)
 
 
 def load_molecule(molecule: str):
@@ -329,7 +405,6 @@ def get_insertion_position(
     if where in ["cylinderZ", "cylinderY", "cylinderX"]:
         axis = where[-1].lower()
         return random_point_in_cylinder(center, radius, height, axis)
-    # now is anywhere
     return random.random(3) * [a, b, c]
 
 
@@ -342,22 +417,101 @@ def rotate_molecule(mol):
     return mol
 
 
+def save_the_day(
+    struct_path: str | Atoms,
+    device: str = "",
+    arch: str = "",
+    model: str = "",
+    fmax: float = 0.01,
+    out_path: str = ".",
+    md_temperature: float = 100.0,
+    md_steps: int = 10,
+    md_timestep: float = 1.0,
+    relax_strategy: str = "geometry_optimisation",
+) -> Atoms:
+    """Geometry optimisation or MD to get a better structure."""
+    if relax_strategy == "geometry_optimisation":
+        _, a = optimize_geometry(
+            struct_path,
+            device,
+            arch,
+            model,
+            fmax,
+            out_path,
+        )
+        return a
+    if relax_strategy == "md":
+        return run_md_nve(
+            struct_path, md_temperature, md_steps, md_timestep, arch, model, device
+        )
+    return None
+
+
+def run_md_nve(
+    struct_path: str | Atoms,
+    temp: float = 100.0,
+    steps: int = 10,
+    timestep: float = 1.0,
+    arch: str = "",
+    model: str = "",
+    device: str = "",
+) -> Atoms:
+    """Run nve simulation."""
+    if isinstance(struct_path, Atoms):
+        md = NVE(
+            struct=struct_path,
+            temp=temp,
+            device=device,
+            arch=arch,
+            calc_kwargs={"model_paths": model},
+            stats_every=1,
+            steps=steps,
+            timestep=timestep,
+        )
+    else:
+        md = NVE(
+            struct_path=struct_path,
+            temp=temp,
+            device=device,
+            arch=arch,
+            calc_kwargs={"model_paths": model},
+            stats_every=1,
+            steps=steps,
+            timestep=timestep,
+        )
+    md.run()
+    return md.struct
+
+
 def optimize_geometry(
-    struct_path: str,
+    struct_path: str | Atoms,
     device: str,
     arch: str,
     model: str,
     fmax: float,
     out_path: str = ".",
-) -> float:
+    opt_cell: bool = False,
+) -> tuple(float, Atoms):
     """Optimize the geometry of a structure."""
-    geo = GeomOpt(
-        struct_path=struct_path,
-        device=device,
-        fmax=fmax,
-        calc_kwargs={"model_paths": model},
-        filter_kwargs={"hydrostatic_strain": True},
-    )
-    geo.run()
-    write(Path(out_path) / f"{Path(struct_path).stem}-opt.cif", geo.struct)
-    return geo.struct.get_potential_energy()
+    if isinstance(struct_path, Atoms):
+        geo = GeomOpt(
+            struct=struct_path,
+            device=device,
+            arch=arch,
+            fmax=fmax,
+            calc_kwargs={"model_paths": model},
+            filter_kwargs={"hydrostatic_strain": opt_cell},
+        )
+        geo.run()
+    else:
+        geo = GeomOpt(
+            struct_path=struct_path,
+            device=device,
+            arch=arch,
+            fmax=fmax,
+            calc_kwargs={"model_paths": model},
+            filter_kwargs={"hydrostatic_strain": opt_cell},
+        )
+        geo.run()
+        write(Path(out_path) / f"{Path(struct_path).stem}-opt.cif", geo.struct)
+    return (geo.struct.get_potential_energy(), geo.struct)

{pack_mm-0.0.21 → pack_mm-0.0.24}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pack-mm"
-version = "0.0.21"
+version = "0.0.24"
 description = "packing materials and molecules in boxes using for machine learnt interatomic potentials"
 authors = [
     { name = "Alin M. Elena" },
@@ -53,6 +53,8 @@ docs = [
     "sphinx-autodoc-typehints<3.0.0,>=2.5.0",
     "sphinx-collapse>=0.1.3",
     "sphinx-copybutton<1.0.0,>=0.5.2",
+    "data-tutorials",
+    "weas-widget",
 ]
 pre-commit = [
     "pre-commit<4.0.0,>=3.6.0",

pack_mm-0.0.24/tests/test_advanced.py

@@ -0,0 +1,85 @@
+"""Advanced tests for packmm."""
+
+# Author; alin m elena, alin@elena.re
+# Contribs;
+# Date: 23-02-2025
+# ©alin m elena,
+from __future__ import annotations
+
+from ase.build import bulk
+from ase.io import read, write
+import pytest
+from typer.testing import CliRunner
+
+from pack_mm.cli.packmm import app
+
+runner = CliRunner()
+
+err = 1.0e-8
+
+
+def test_packmm_hmc(tmp_path):
+    """Check values."""
+    result = runner.invoke(
+        app,
+        [
+            "--nmols",
+            "2",
+            "--model",
+            "small-0b2",
+            "--insert-strategy",
+            "hmc",
+            "--seed",
+            "2042",
+            "--cell-a",
+            "10",
+            "--cell-b",
+            "10",
+            "--cell-c",
+            "10",
+            "--out-path",
+            tmp_path,
+        ],
+    )
+    assert result.exit_code == 0
+    assert (tmp_path / "1H2O.cif").exists()
+    assert (tmp_path / "2H2O.cif").exists()
+    f = read(tmp_path / "2H2O.cif")
+    assert f[0].position == pytest.approx([7.83420049, 6.21661184, 5.21814887], abs=err)
+    assert (tmp_path / "2H2O-opt.cif").exists()
+    f = read(tmp_path / "2H2O-opt.cif")
+    assert f[0].position == pytest.approx([7.83120709, 6.21742874, 5.22597213], abs=err)
+
+
+def test_packmm_every(tmp_path):
+    """Check values."""
+    na = bulk("NaCl", "rocksalt", a=5.61, cubic=True)
+    write(tmp_path / "system.cif", na)
+    write("system.cif", na)
+
+    result = runner.invoke(
+        app,
+        [
+            "--system",
+            tmp_path / "system.cif",
+            "--molecule",
+            "H2",
+            "--every",
+            "1",
+            "--nmols",
+            "2",
+            "--model",
+            "small-0b2",
+            "--seed",
+            "2042",
+            "--out-path",
+            tmp_path,
+            "--no-geometry",
+        ],
+    )
+    assert result.exit_code == 0
+    assert (tmp_path / "system.cif").exists()
+    assert (tmp_path / "system+1H2.cif").exists()
+    assert (tmp_path / "system+2H2.cif").exists()
+    f = read(tmp_path / "system+2H2.cif")
+    assert f[0].position == pytest.approx([1.24701529, 0.024216, 0.11632905], abs=err)

pack_mm-0.0.21/tests/test_packmm.py → pack_mm-0.0.24/tests/test_cli.py

@@ -25,14 +25,14 @@ def test_packmm_custom_molecule():
     """Check molecule."""
     result = runner.invoke(app, ["--molecule", "CO2"])
     assert result.exit_code == 0
-    assert "CO2" in strip_ansi_codes(result.output)
+    assert "molecule='CO2'" in strip_ansi_codes(result.output)
 
 
 def test_packmm_custom_nmols():
     """Check nmols."""
-    result = runner.invoke(app, ["--nmols", "1"])
-    assert result.exit_code == 0
-    assert "nmols=1" in strip_ansi_codes(result.output)
+    result = runner.invoke(app, ["--nmols", "-2"])
+    assert result.exit_code == 1
+    assert "nmols=-2" in strip_ansi_codes(result.output)
 
 
 def test_packmm_custom_ntries():
@@ -49,6 +49,13 @@ def test_packmm_custom_seed():
     assert "seed=1234" in strip_ansi_codes(result.output)
 
 
+def test_packmm_custom_every():
+    """Check seed."""
+    result = runner.invoke(app, ["--every", "10"])
+    assert result.exit_code == 0
+    assert "every=10" in strip_ansi_codes(result.output)
+
+
 def test_packmm_custom_insertion_method():
     """Check insertion."""
     result = runner.invoke(app, ["--where", "sphere"])
@@ -56,6 +63,20 @@ def test_packmm_custom_insertion_method():
     assert "where=sphere" in strip_ansi_codes(result.output)
 
 
+def test_packmm_custom_insert_strategy():
+    """Check insertion."""
+    result = runner.invoke(app, ["--insert-strategy", "hmc"])
+    assert result.exit_code == 0
+    assert "insert_strategy=hmc" in strip_ansi_codes(result.output)
+
+
+def test_packmm_custom_relax_strategy():
+    """Check relax."""
+    result = runner.invoke(app, ["--relax-strategy", "md"])
+    assert result.exit_code == 0
+    assert "relax_strategy=md" in strip_ansi_codes(result.output)
+
+
 def test_packmm_custom_center():
     """Check centre."""
     result = runner.invoke(app, ["--centre", "0.5,0.5,0.5"])
@@ -122,6 +143,27 @@ def test_packmm_custom_temperature():
     assert "temperature=400.0" in strip_ansi_codes(result.output)
 
 
+def test_packmm_md_temperature():
+    """Check md temperature."""
+    result = runner.invoke(app, ["--md-temperature", "300.0"])
+    assert result.exit_code == 0
+    assert "md_temperature=300.0" in strip_ansi_codes(result.output)
+
+
+def test_packmm_md_timestep():
+    """Check md temperature."""
+    result = runner.invoke(app, ["--md-timestep", "1.0"])
+    assert result.exit_code == 0
+    assert "md_timestep=1.0" in strip_ansi_codes(result.output)
+
+
+def test_packmm_md_steps():
+    """Check md steps."""
+    result = runner.invoke(app, ["--md-steps", "10"])
+    assert result.exit_code == 0
+    assert "md_steps=10" in strip_ansi_codes(result.output)
+
+
 def test_packmm_custom_fmax():
     """Check fmax."""
     result = runner.invoke(app, ["--fmax", "0.05"])
@@ -143,16 +185,41 @@ def test_packmm_invalid_insertion_method():
     assert "Invalid value for '--where'" in strip_ansi_codes(result.output)
 
 
-def test_packmm_invalid_centre_format():
-    """Check centre."""
-    result = runner.invoke(app, ["--nmols", "1", "--centre", "0.5,0.5"])
+def test_packmm_invalid_insertion_strategy():
+    """Check insertion strategt."""
+    result = runner.invoke(app, ["--insert-strategy", "invalid_method"])
     assert result.exit_code != 0
-    assert "Invalid centre" in strip_ansi_codes(result.output)
+    assert "Invalid value for '--insert-strategy'" in strip_ansi_codes(result.output)
+
+
+def test_packmm_invalid_relax_strategy():
+    """Check insertion strategt."""
+    result = runner.invoke(app, ["--relax-strategy", "invalid_method"])
+    assert result.exit_code != 0
+    assert "Invalid value for '--relax-strategy'" in strip_ansi_codes(result.output)
+
+
+def test_packmm_invalid_md_steps():
+    """Check md steps."""
+    result = runner.invoke(app, ["--nmols", "1", "--md-steps", "-10"])
+    assert "Invalid MD steps" in strip_ansi_codes(result.output)
+
+
+def test_packmm_invalid_md_temperature():
+    """Check md steps."""
+    result = runner.invoke(app, ["--nmols", "1", "--md-temperature", "-10.0"])
+    assert "Invalid MD temperature" in strip_ansi_codes(result.output)
 
 
-def test_packmm_invalid_centre_value():
+def test_packmm_invalid_md_timestep():
+    """Check md timestep."""
+    result = runner.invoke(app, ["--nmols", "1", "--md-timestep", "-10.0"])
+    assert "Invalid MD timestep" in strip_ansi_codes(result.output)
+
+
+def test_packmm_invalid_centre_format():
     """Check centre."""
-    result = runner.invoke(app, ["--nmols", "1", "--centre", "-0.6,0.5,0.5"])
+    result = runner.invoke(app, ["--nmols", "1", "--centre", "0.5,0.5"])
     assert result.exit_code != 0
     assert "Invalid centre" in strip_ansi_codes(result.output)
 

pack_mm-0.0.24/tests/test_core.py

@@ -0,0 +1,416 @@
+"""Tests for core."""
+
+# -*- coding: utf-8 -*-
+# Author; alin m elena, alin@elena.re
+# Contribs;
+# Date: 22-02-2025
+# ©alin m elena,
+from __future__ import annotations
+
+from ase import Atoms
+from ase.build import molecule as build_molecule
+from ase.io import write
+from numpy import random
+import pytest
+
+from pack_mm.core.core import (
+    get_insertion_position,
+    load_molecule,
+    optimize_geometry,
+    pack_molecules,
+    random_point_in_box,
+    random_point_in_cylinder,
+    random_point_in_ellipsoid,
+    random_point_in_sphere,
+    rotate_molecule,
+    run_md_nve,
+    save_the_day,
+    set_defaults,
+    validate_value,
+)
+
+err = 1.0e-8
+
+
+# Set a fixed seed for reproducibility in tests
+@pytest.fixture(autouse=True)
+def set_random_seed():
+    """Set random seed."""
+    random.seed(2042)
+
+
+def test_random_point_in_sphere():
+    """Test point in sphere."""
+    center = (2, 2, 2)
+    radius = 8.0
+    x, y, z = random_point_in_sphere(center, radius)
+    assert x == pytest.approx(-3.299696236298196, abs=err)
+    assert y == pytest.approx(-3.046619861327052, abs=err)
+    assert z == pytest.approx(1.1884891239565165, abs=err)
+
+
+def test_random_point_in_ellipsoid():
+    """Test point in ellipsoid."""
+    center = (0, 0, 0)
+    a, b, c = 2.0, 2.0, 4.0
+    x, y, z = random_point_in_ellipsoid(center, a, b, c)
+    assert x == pytest.approx(0.27914659851849705, abs=err)
+    assert y == pytest.approx(-1.4815802529721946, abs=err)
+    assert z == pytest.approx(-1.2059956538672925, abs=err)
+
+
+def test_random_point_in_box():
+    """Test point in box."""
+    center = (0, 0, 0)
+    a, b, c = 1.0, 2.0, 3.0
+    x, y, z = random_point_in_box(center, a, b, c)
+    assert x == pytest.approx(0.2796391455704136, abs=err)
+    assert y == pytest.approx(0.24221552377157196, abs=err)
+    assert z == pytest.approx(0.10546160764197499, abs=err)
+
+
+def test_random_point_in_cylinder():
+    """Test point in cylinder."""
+    center = (0, 0, 0)
+    radius = 2.0
+    height = 5.0
+
+    direction = "z"
+    x, y, z = random_point_in_cylinder(center, radius, height, direction)
+    assert x == pytest.approx(0.29184067570623795, abs=err)
+    assert y == pytest.approx(-1.5489545079008842, abs=err)
+    assert z == pytest.approx(0.1757693460699583, abs=err)
+
+    direction = "y"
+    x, y, z = random_point_in_cylinder(center, radius, height, direction)
+    assert x == pytest.approx(-1.417921218944154, abs=err)
+    assert y == pytest.approx(-0.8295636754899149, abs=err)
+    assert z == pytest.approx(-1.2828806053300128, abs=err)
+
+    direction = "x"
+    x, y, z = random_point_in_cylinder(center, radius, height, direction)
+    assert x == pytest.approx(0.5063181926884202, abs=err)
+    assert y == pytest.approx(-0.5839696769917422, abs=err)
+    assert z == pytest.approx(1.10139156645954, abs=err)
+
+
+def test_validate_value_positive():
+    """Test point in test value."""
+    validate_value("test_value", 1.0)  # Should not raise an exception
+
+
+def test_validate_value_negative():
+    """Test point in test value."""
+    with pytest.raises(Exception, match="Invalid test_value, needs to be positive"):
+        validate_value("test_value", -1.0)
+
+
+def test_load_molecule_from_file(tmp_path):
+    """Test point in load molecule."""
+    molecule = build_molecule("H2O")
+    molecule_file = tmp_path / "water.xyz"
+    write(molecule_file, molecule)
+    loaded_molecule = load_molecule(str(molecule_file))
+    assert len(loaded_molecule) == 3
+
+
+def test_load_molecule_from_name():
+    """Test point in load molecule."""
+    molecule = load_molecule("H2O")
+    assert len(molecule) == 3
+
+
+def test_get_insertion_position_sphere():
+    """Test point in sphere."""
+    center = (5, 5, 5)
+    radius = 10.0
+    x, y, z = get_insertion_position("sphere", center, radius=radius)
+    assert x == pytest.approx(-1.624620295372745, abs=err)
+    assert y == pytest.approx(-1.3082748266588151, abs=err)
+    assert z == pytest.approx(3.9856114049456455, abs=err)
+
+
+def test_get_insertion_position_box():
+    """Test point in box."""
+    center = (5, 5, 5)
+    a = 10.0
+    x, y, z = get_insertion_position("box", center, a=a, b=a, c=a)
+    assert x == pytest.approx(7.796391455704136, abs=err)
+    assert y == pytest.approx(6.21107761885786, abs=err)
+    assert z == pytest.approx(5.351538692139917, abs=err)
+
+
+def test_get_insertion_position_ellipsoid():
+    """Test point in ellipsoid."""
+    center = (5, 5, 5)
+    x, y, z = get_insertion_position("ellipsoid", center, a=2.0, b=2.0, c=4.0)
+    assert x == pytest.approx(5.279146598518497, abs=err)
+    assert y == pytest.approx(3.5184197470278056, abs=err)
+    assert z == pytest.approx(3.7940043461327075, abs=err)
+
+
+def test_get_insertion_position_cylinder():
+    """Test point in cylinder."""
+    center = (5, 5, 5)
+    x, y, z = get_insertion_position("cylinderZ", center, radius=3.0, height=10.0)
+    assert x == pytest.approx(5.437761013559357, abs=err)
+    assert y == pytest.approx(2.6765682381486737, abs=err)
+    assert z == pytest.approx(5.351538692139917, abs=err)
+
+
+def test_set_defaults_centre():
+    """Test centre."""
+    cell = [10, 10, 10]
+    centre, _, _, _, _, _ = set_defaults(cell, centre=None)
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert centre[1] == pytest.approx(5.0, abs=err)
+    assert centre[2] == pytest.approx(5.0, abs=err)
+
+
+def test_set_defaults_box():
+    """Test box."""
+    cell = [10, 10, 10]
+    centre, a, b, c, _, _ = set_defaults(cell, where="box", b=5.0, centre=None)
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert a == pytest.approx(10.0, abs=err)
+    assert b == pytest.approx(5.0, abs=err)
+    assert c == pytest.approx(10.0, abs=err)
+
+
+def test_set_defaults_anywhere():
+    """Test box."""
+    cell = [10, 10, 10]
+    centre, a, b, c, _, _ = set_defaults(cell, where="anywhere", centre=None)
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert a == pytest.approx(10.0, abs=err)
+    assert b == pytest.approx(10.0, abs=err)
+    assert c == pytest.approx(10.0, abs=err)
+
+
+def test_set_defaults_ellipsoid():
+    """Test ellipsoid."""
+    cell = [10, 10, 10]
+    centre, a, b, c, _, _ = set_defaults(
+        cell, where="ellipsoid", b=3.0, a=3.0, centre=None
+    )
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert a == pytest.approx(3.0, abs=err)
+    assert b == pytest.approx(3.0, abs=err)
+    assert c == pytest.approx(5.0, abs=err)
+
+
+def test_set_defaults_cylinder():
+    """Test ellipsoid."""
+    cell = [10, 12, 14]
+    centre, _, _, _, radius, height = set_defaults(cell, where="cylinderZ", centre=None)
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert radius == pytest.approx(5.0, abs=err)
+    assert height == pytest.approx(7.0, abs=err)
+
+    centre, _, _, _, radius, height = set_defaults(cell, where="cylinderX", centre=None)
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert radius == pytest.approx(6.0, abs=err)
+    assert height == pytest.approx(5.0, abs=err)
+
+    centre, _, _, _, radius, height = set_defaults(cell, where="cylinderY", centre=None)
+    assert centre[0] == pytest.approx(5.0, abs=err)
+    assert radius == pytest.approx(5.0, abs=err)
+    assert height == pytest.approx(6.0, abs=err)
+
+
+def test_rotate_molecule():
+    """Test rotate molecule."""
+    molecule = build_molecule("H2O")
+    a1 = molecule.get_angle(2, 0, 1)
+    a2 = molecule.get_angle(2, 1, 0)
+
+    rotated_molecule = rotate_molecule(molecule)
+    assert a1 == pytest.approx(rotated_molecule.get_angle(2, 0, 1))
+    assert a2 == pytest.approx(rotated_molecule.get_angle(2, 1, 0))
+
+
+def test_optimize_geometry(tmp_path):
+    """Test go."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    structure_file = tmp_path / "water.cif"
+    write(structure_file, molecule)
+    optimized_energy, _ = optimize_geometry(
+        str(structure_file),
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        fmax=0.01,
+        out_path=tmp_path,
+        opt_cell=True,
+    )
+    assert optimized_energy == pytest.approx(-14.17098106193308, abs=err)
+
+
+def test_pack_molecules(tmp_path):
+    """Test pack molecule."""
+    system = Atoms(
+        "Ca", positions=[(5.0, 5.0, 5.0)], cell=[10, 10, 10], pbc=[True, True, True]
+    )
+    system_file = tmp_path / "system.cif"
+    write(system_file, system)
+
+    e, _ = pack_molecules(
+        system=str(system_file),
+        molecule="H2O",
+        nmols=2,
+        arch="mace_mp",
+        model="small-0b2",
+        device="cpu",
+        where="sphere",
+        center=(5.0, 5.0, 5.0),
+        radius=5.0,
+        seed=2042,
+        temperature=300,
+        ntries=10,
+        geometry=False,
+        fmax=0.1,
+        out_path=tmp_path,
+    )
+
+    assert (tmp_path / "system+1H2O.cif").exists()
+    assert (tmp_path / "system+2H2O.cif").exists()
+    assert e == pytest.approx(-29.21589570470306, abs=err)
+
+
+def test_pack_molecules_atoms(tmp_path):
+    """Test pack molecule."""
+    system = Atoms(
+        "Ca", positions=[(5.0, 5.0, 5.0)], cell=[10, 10, 10], pbc=[True, True, True]
+    )
+
+    e, _ = pack_molecules(
+        system=system,
+        molecule="H2O",
+        nmols=2,
+        arch="mace_mp",
+        model="small-0b2",
+        device="cpu",
+        where="sphere",
+        center=(5.0, 5.0, 5.0),
+        radius=5.0,
+        seed=2042,
+        temperature=300,
+        ntries=10,
+        geometry=False,
+        fmax=0.1,
+        out_path=tmp_path,
+    )
+
+    assert e == pytest.approx(-29.21589570470306, abs=err)
+
+
+def test_pack_molecules_2(tmp_path, capsys):
+    """Test pack molecule."""
+    system = Atoms(
+        "Ca", positions=[(2.5, 2.5, 2.5)], cell=[5, 5, 5], pbc=[True, True, True]
+    )
+    system_file = tmp_path / "system.cif"
+    write(system_file, system)
+
+    e, _ = pack_molecules(
+        system=str(system_file),
+        molecule="H2O",
+        nmols=3,
+        arch="mace_mp",
+        model="small-0b2",
+        device="cpu",
+        where="sphere",
+        center=(2.5, 2.5, 2.5),
+        radius=2.5,
+        seed=2042,
+        temperature=300,
+        ntries=2,
+        geometry=False,
+        fmax=0.1,
+        out_path=tmp_path,
+    )
+    captured = capsys.readouterr()
+
+    assert "Failed to insert particle 3 after 2 tries" in captured.out
+    assert e == pytest.approx(-47.194755808249454, abs=err)
+
+
+def test_save_the_day(tmp_path):
+    """Test save the day."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    molecule.center()
+    structure_file = tmp_path / "water.cif"
+    write(structure_file, molecule)
+    s = save_the_day(
+        str(structure_file),
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        fmax=0.01,
+        out_path=tmp_path,
+    )
+    assert s[0].position == pytest.approx([5.0, 4.99619815, 5.30704738], abs=err)
+
+
+def test_save_the_day_md(tmp_path):
+    """Test save the day."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    molecule.center()
+    structure_file = tmp_path / "water.cif"
+    write(structure_file, molecule)
+    s = save_the_day(
+        str(structure_file),
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        relax_strategy="md",
+        md_timestep=1.0,
+        md_steps=10.0,
+        md_temperature=100.0,
+    )
+    assert s[0].position == pytest.approx([4.99684244, 5.00440785, 5.2987255], abs=err)
+
+
+def test_save_the_day_invalid(tmp_path):
+    """Test save the day."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    molecule.center()
+    structure_file = tmp_path / "water.cif"
+    write(structure_file, molecule)
+    s = save_the_day(
+        str(structure_file),
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        relax_strategy="invalid_stratregy",
+        md_timestep=1.0,
+        md_steps=10.0,
+        md_temperature=100.0,
+    )
+    assert s is None
+
+
+def test_run_md(tmp_path):
+    """Test md."""
+    molecule = build_molecule("H2O")
+    molecule.set_cell([10, 10, 10])
+    molecule.set_pbc([True, True, True])
+    molecule.center()
+    s = run_md_nve(
+        molecule,
+        device="cpu",
+        arch="mace_mp",
+        model="small-0b2",
+        timestep=1.0,
+        steps=10.0,
+        temp=100.0,
+    )
+    assert s[0].position == pytest.approx([4.99684244, 5.00440785, 5.2987255], abs=err)

pack_mm-0.0.21/tests/test_core.py

@@ -1,178 +0,0 @@
-"""Test cli for core."""
-
-# -*- coding: utf-8 -*-
-# Author; alin m elena, alin@elena.re
-# Contribs;
-# Date: 22-02-2025
-# ©alin m elena,
-from __future__ import annotations
-
-from ase import Atoms
-from ase.build import molecule as build_molecule
-from ase.io import write
-import numpy as np
-from numpy import random
-import pytest
-
-from pack_mm.core.core import (
-    get_insertion_position,
-    load_molecule,
-    optimize_geometry,
-    pack_molecules,
-    random_point_in_box,
-    random_point_in_cylinder,
-    random_point_in_ellipsoid,
-    random_point_in_sphere,
-    rotate_molecule,
-    validate_value,
-)
-
-
-# Set a fixed seed for reproducibility in tests
-@pytest.fixture(autouse=True)
-def set_random_seed():
-    """Set random seed."""
-    random.seed(2042)
-
-
-def test_random_point_in_sphere():
-    """Test point in sphere."""
-    center = (0, 0, 0)
-    radius = 10.0
-    point = random_point_in_sphere(center, radius)
-    assert len(point) == 3
-    distance = np.linalg.norm(np.array(point) - np.array(center))
-    assert distance <= radius
-
-
-def test_random_point_in_ellipsoid():
-    """Test point in ellipsoid."""
-    center = (0, 0, 0)
-    a, b, c = 1.0, 2.0, 3.0
-    point = random_point_in_ellipsoid(center, a, b, c)
-    assert len(point) == 3
-    x, y, z = point
-    assert (x**2 / a**2) + (y**2 / b**2) + (z**2 / c**2) <= 1.0
-
-
-def test_random_point_in_box():
-    """Test point in box."""
-    center = (0, 0, 0)
-    a, b, c = 1.0, 2.0, 3.0
-    point = random_point_in_box(center, a, b, c)
-    assert len(point) == 3
-    x, y, z = point
-    assert center[0] - a * 0.5 <= x <= center[0] + a * 0.5
-    assert center[1] - b * 0.5 <= y <= center[1] + b * 0.5
-    assert center[2] - c * 0.5 <= z <= center[2] + c * 0.5
-
-
-def test_random_point_in_cylinder():
-    """Test point in cylinder."""
-    center = (0, 0, 0)
-    radius = 1.0
-    height = 2.0
-    direction = "z"
-    point = random_point_in_cylinder(center, radius, height, direction)
-    assert len(point) == 3
-    x, y, z = point
-    assert x**2 + y**2 <= radius**2
-    assert center[2] - height * 0.5 <= z <= center[2] + height * 0.5
-
-
-def test_validate_value_positive():
-    """Test point in test value."""
-    validate_value("test_value", 1.0)  # Should not raise an exception
-
-
-def test_validate_value_negative():
-    """Test point in test value."""
-    with pytest.raises(Exception, match="Invalid test_value, needs to be positive"):
-        validate_value("test_value", -1.0)
-
-
-def test_load_molecule_from_file(tmp_path):
-    """Test point in load molecule."""
-    molecule = build_molecule("H2O")
-    molecule_file = tmp_path / "water.xyz"
-    molecule.write(molecule_file)
-    loaded_molecule = load_molecule(str(molecule_file))
-    assert isinstance(loaded_molecule, Atoms)
-    assert len(loaded_molecule) == 3  # H2O has 3 atoms
-
-
-def test_load_molecule_from_name():
-    """Test point in load molecule."""
-    molecule = load_molecule("H2O")
-    assert isinstance(molecule, Atoms)
-    assert len(molecule) == 3  # H2O has 3 atoms
-
-
-def test_get_insertion_position_sphere():
-    """Test point in sphere."""
-    center = (0, 0, 0)
-    radius = 10.0
-    point = get_insertion_position("sphere", center, radius=radius)
-    assert len(point) == 3
-    distance = np.linalg.norm(np.array(point) - np.array(center))
-    assert distance <= radius
-
-
-def test_rotate_molecule():
-    """Test rotate molecule."""
-    molecule = build_molecule("H2O")
-    rotated_molecule = rotate_molecule(molecule)
-    assert isinstance(rotated_molecule, Atoms)
-    assert len(rotated_molecule) == 3  # H2O has 3 atoms
-
-
-def test_optimize_geometry(tmp_path):
-    """Test go."""
-    # Create a temporary structure file
-    molecule = build_molecule("H2O")
-    molecule.set_cell([10, 10, 10])
-    molecule.set_pbc([True, True, True])
-    structure_file = tmp_path / "water.cif"
-    write(structure_file, molecule)
-    optimized_energy = optimize_geometry(
-        str(structure_file),
-        device="cpu",
-        arch="mace_mp",
-        model="medium-omat-0",
-        fmax=0.01,
-    )
-    assert optimized_energy == pytest.approx(-13.759273983276572, abs=1.0e-8)
-
-
-# Test pack_molecules with a simple case
-def test_pack_molecules(tmp_path):
-    """Test pack molecule."""
-    # Create a temporary system file
-    system = Atoms(
-        "Ca", positions=[(5.0, 5.0, 5.0)], cell=[10, 10, 10], pbc=[True, True, True]
-    )
-    system_file = tmp_path / "system.cif"
-    write(system_file, system)
-
-    # Test packing molecules
-    e = pack_molecules(
-        system=str(system_file),
-        molecule="H2O",
-        nmols=2,
-        arch="mace_mp",
-        model="medium-omat-0",
-        device="cpu",
-        where="sphere",
-        center=(5.0, 5.0, 5.0),
-        radius=5.0,
-        seed=2042,
-        temperature=300,
-        ntries=10,
-        geometry=False,
-        fmax=0.1,
-        out_path=tmp_path,
-    )
-
-    assert (tmp_path / "system+1H2O.cif").exists()
-    assert (tmp_path / "system+2H2O.cif").exists()
-    assert e == pytest.approx(-28.251229837533085, abs=1.0e-6)