pack-mm 0.0.21__py3-none-any.whl → 0.0.22__py3-none-any.whl

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/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,6 +139,52 @@ 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,
     molecule: str = "H2O",
@@ -161,6 +208,12 @@ def pack_molecules(
     cell_b: float = None,
     cell_c: float = None,
     out_path: str = ".",
+    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,
 ) -> float:
     """
     Pack molecules into a system based on the specified parameters.
@@ -185,6 +238,14 @@ 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"
+
     """
     kbt = temperature * kB
     validate_value("temperature", temperature)
@@ -198,51 +259,31 @@ 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 = ""
+    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 +291,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 +301,11 @@ 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
+                )
+
             tsys.calc = calc
             en = tsys.get_potential_energy()
             de = en - e
@@ -270,34 +317,55 @@ def pack_molecules(
             if u <= acc:
                 accept = True
                 break
+            if every > 0 and _itry / every == 0:
+                csys = save_the_day(
+                    Path(out_path) / f"{sysname}{i}{Path(molecule).stem}.cif",
+                    device,
+                    arch,
+                    model,
+                    fmax,
+                    out_path,
+                    md_temperature,
+                    md_steps,
+                    md_timestep,
+                    relax_strategy,
+                )
 
         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(
+                Path(out_path) / f"{sysname}{i}{Path(molecule).stem}.cif",
                 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",
+            Path(out_path) / f"{sysname}{nmols}{Path(molecule).stem}.cif",
             device,
             arch,
             model,
             fmax,
             out_path,
+            True,
         )
     return energy_final
 
@@ -329,7 +397,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,6 +409,72 @@ def rotate_molecule(mol):
     return mol
 
 
+def save_the_day(
+    struct_path: str = "",
+    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":
+        _ = optimize_geometry(
+            struct_path,
+            device,
+            arch,
+            model,
+            fmax,
+            out_path,
+        )
+        return read(Path(out_path) / f"{Path(struct_path).stem}-opt.cif")
+    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,
     device: str,
@@ -349,14 +482,16 @@ def optimize_geometry(
     model: str,
     fmax: float,
     out_path: str = ".",
+    opt_cell: bool = False,
 ) -> float:
     """Optimize the geometry of a structure."""
     geo = GeomOpt(
         struct_path=struct_path,
         device=device,
+        arch=arch,
         fmax=fmax,
         calc_kwargs={"model_paths": model},
-        filter_kwargs={"hydrostatic_strain": True},
+        filter_kwargs={"hydrostatic_strain": opt_cell},
     )
     geo.run()
     write(Path(out_path) / f"{Path(struct_path).stem}-opt.cif", geo.struct)

{pack_mm-0.0.21.dist-info → pack_mm-0.0.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pack-mm
-Version: 0.0.21
+Version: 0.0.22
 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
 
@@ -125,7 +125,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.22.dist-info/RECORD

@@ -0,0 +1,8 @@
+pack_mm-0.0.22.dist-info/METADATA,sha256=5XjBJexb_jY1kVAZdZSXFcXfdWvfHJFpnsYJ25tTQlc,13509
+pack_mm-0.0.22.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
+pack_mm-0.0.22.dist-info/entry_points.txt,sha256=ajKA2oehIa_LCVCP2XTRxV0VNgjGl9c2wYkwk0BasrQ,66
+pack_mm-0.0.22.dist-info/licenses/LICENSE,sha256=ZOYkPdn_vQ8wYJqZnjesow79F_grMbVlHcJ9V91G1pE,1100
+pack_mm/__init__.py,sha256=ct7qfCmTDwhLYip6JKYWRLasmmaGYt0ColbK0CpvYZk,150
+pack_mm/cli/packmm.py,sha256=E63MmBrxvn9g5GQMK5vlyq6CmHL8hgQ8CICeX2N9z1E,6398
+pack_mm/core/core.py,sha256=R0d-FTmuTj-5BSwG5S4Fa6f3BFdYrvBUsV2cSwcHi2c,15269
+pack_mm-0.0.22.dist-info/RECORD,,

pack_mm-0.0.21.dist-info/RECORD

@@ -1,8 +0,0 @@
-pack_mm-0.0.21.dist-info/METADATA,sha256=jW4IZ3HFn-e7AxJ8dAteajsDSs1dsybX3G9l080Gr5Y,13506
-pack_mm-0.0.21.dist-info/WHEEL,sha256=thaaA2w1JzcGC48WYufAs8nrYZjJm8LqNfnXFOFyCC4,90
-pack_mm-0.0.21.dist-info/entry_points.txt,sha256=ajKA2oehIa_LCVCP2XTRxV0VNgjGl9c2wYkwk0BasrQ,66
-pack_mm-0.0.21.dist-info/licenses/LICENSE,sha256=ZOYkPdn_vQ8wYJqZnjesow79F_grMbVlHcJ9V91G1pE,1100
-pack_mm/__init__.py,sha256=ct7qfCmTDwhLYip6JKYWRLasmmaGYt0ColbK0CpvYZk,150
-pack_mm/cli/packmm.py,sha256=VqumDT_f1Nf1LCZ1WsF5D6MoLmKEQPEimYenibQHIU4,4944
-pack_mm/core/core.py,sha256=vNMQOVDIXyRCUWosww9sSHUxwqbjPRk5JLsb9muwpVA,11369
-pack_mm-0.0.21.dist-info/RECORD,,