Add variant option (atomify wasm in the kernel) and ship python/ in the npm tarball

NOTEBOOK_TUTORIALS.md

@@ -49,11 +49,13 @@ the kernel is the browser.
 
 ### Constraints to keep in mind
 
-- **Package set**: the deployed wasm is the default `PACKAGES=MOLECULE` build.
+- **Package set**: by default the kernel loads the `PACKAGES=MOLECULE` build.
   LJ + molecular force fields work; EAM (`MANYBODY`), `KSPACE`, `REAXFF`,
-  `GRANULAR` do not. Tutorials that need them are marked below; shipping the
-  atomify-flavor wasm (`PACKAGES=atomify`) to the notebook site later unlocks
-  nearly all of them.
+  `GRANULAR` do not. The kernel can request the full-package atomify build
+  with `await lammps(variant="atomify")` — it needs the notebook site to
+  serve `dist/cpp/lammps-atomify.js` and, being a multithreaded KOKKOS
+  build, a cross-origin isolated page. Tutorials that need the extra
+  packages are marked below; the atomify variant unlocks nearly all of them.
 - **Potential/data files** are not bundled in the wasm FS. Notebooks fetch
   them (`pyfetch(lammps.site_url("files/data/…"))`, or any URL) and hand them
   to `lmp.file(path, contents=…)` / `lmp.write_file(...)` before running.

README.md

@@ -26,6 +26,24 @@ needs a [cross-origin-isolated](#required-headers-browsers-only) context
 (`SharedArrayBuffer`) to run in the browser. It is larger (~41 MB module,
 ~10 MB gzipped) than the default (~11 MB):
 
+The easiest way to use it is the `variant` client option — the client loads
+the atomify module and starts LAMMPS with Kokkos threads enabled (the kk
+accelerator suffix is applied by default; threads/suffix are tunable via the
+`kokkos` option):
+
+```ts
+import { LammpsClient } from "lammps.js/client";
+
+const lammps = await LammpsClient.create({}, { variant: "atomify" });
+lammps.start();
+lammps.runScript("pair_style vashishta\n"); // MANYBODY — not in the default build
+```
+
+`variant` accepts `"serial"` (default), `"kokkos"`, or `"atomify"`; an
+explicit `variant` takes precedence over the `kokkos` flag (`kokkos: true`
+is shorthand for `variant: "kokkos"`). The Python bindings mirror this as
+`await lammps(variant="atomify")`. The wasm module is also usable directly:
+
 ```ts
 import createModule from "lammps.js/wasm-atomify";
 
@@ -34,13 +52,9 @@ const lammps = new wasm.LAMMPSWeb();
 // Enable Kokkos threads and the kk accelerator suffix (styles with a /kk
 // variant run multithreaded; the rest fall back to serial).
 lammps.startWithArgs(["-k", "on", "t", "4", "-sf", "kk"]);
-lammps.runScript("pair_style vashishta\n"); // MANYBODY — not in the default build
+lammps.runScript("pair_style vashishta\n");
 ```
 
-There is no separate `LammpsClient` entry point for this variant yet; use the
-wasm module directly as above, or the equivalent low-level calls shown in
-[Usage (manual stepping, optional)](#usage-manual-stepping-optional).
-
 ## Usage (main flow: `runScriptAsync`)
 
 `runScriptAsync()` is the main API.

client.ts

@@ -61,6 +61,15 @@ export interface KokkosOptions {
   suffix?: boolean;
 }
 
+/**
+ * Which wasm build the client loads:
+ * - "serial" (default): the single-threaded MOLECULE build (dist/cpp/lammps.js)
+ * - "kokkos": the multi-threaded KOKKOS build (dist/cpp/lammps-kokkos.js)
+ * - "atomify": the full-package build (dist/cpp/lammps-atomify.js), which is
+ *   also a multi-threaded KOKKOS/pthreads build
+ */
+export type LammpsVariant = "serial" | "kokkos" | "atomify";
+
 export interface LammpsClientOptions {
   workdir?: string;
   /**
@@ -78,8 +87,19 @@ export interface LammpsClientOptions {
    * Use the multi-threaded KOKKOS wasm build (dist/cpp/lammps-kokkos.js).
    * Requires a cross-origin isolated context in browsers
    * (SharedArrayBuffer). `true` uses default options.
+   * Equivalent to `variant: "kokkos"`; with `variant: "atomify"` it still
+   * configures the Kokkos threads/suffix options.
    */
   kokkos?: boolean | KokkosOptions;
+  /**
+   * Select the wasm build explicitly. An explicit variant takes precedence
+   * over the `kokkos` flag (`kokkos: true` is shorthand for
+   * `variant: "kokkos"`). Both "kokkos" and "atomify" are multi-threaded
+   * KOKKOS/pthreads builds: they start LAMMPS with `-k on t N [-sf kk]`
+   * (tunable via the `kokkos` option) and require a cross-origin isolated
+   * context in browsers (SharedArrayBuffer).
+   */
+  variant?: LammpsVariant;
 }
 
 const KOKKOS_MAX_THREADS = 8;
@@ -91,6 +111,14 @@ function resolveKokkosOptions(value: boolean | KokkosOptions | undefined): Kokko
   return value === true ? {} : value;
 }
 
+/** Explicit `variant` wins; otherwise the `kokkos` flag selects "kokkos". */
+function resolveVariant(options: LammpsClientOptions): LammpsVariant {
+  if (options.variant) {
+    return options.variant;
+  }
+  return options.kokkos ? "kokkos" : "serial";
+}
+
 function defaultKokkosThreads(): number {
   const hardware =
     typeof navigator !== "undefined" && Number.isFinite(navigator.hardwareConcurrency)
@@ -219,7 +247,11 @@ export class LammpsClient {
     this.instance = instance;
     this.workdir = options.workdir ?? DEFAULT_WORKDIR;
     this._shifts = buildShiftMap(module);
-    this.#kokkos = resolveKokkosOptions(options.kokkos);
+    // "atomify" is itself a KOKKOS/pthreads build, so every non-serial
+    // variant starts with Kokkos arguments; the `kokkos` option only tunes
+    // threads/suffix when a threaded variant is selected.
+    const variant = resolveVariant(options);
+    this.#kokkos = variant === "serial" ? null : resolveKokkosOptions(options.kokkos) ?? {};
 
     try {
       module.FS.mkdir(this.workdir);
@@ -244,11 +276,15 @@ export class LammpsClient {
     if (clientOptions.worker) {
       return createWorkerBackedClient(moduleOptions, clientOptions, clientOptions.worker);
     }
-    // Both wasm modules are imported lazily so that consumers (and CI
+    // The wasm modules are imported lazily so that consumers (and CI
     // variants) only need the artifact they actually use.
-    const factory = clientOptions.kokkos
-      ? (await import("./cpp/lammps-kokkos.js")).default
-      : (await import("./cpp/lammps.js")).default;
+    const variant = resolveVariant(clientOptions);
+    const factory =
+      variant === "atomify"
+        ? (await import("./cpp/lammps-atomify.js")).default
+        : variant === "kokkos"
+          ? (await import("./cpp/lammps-kokkos.js")).default
+          : (await import("./cpp/lammps.js")).default;
     const module = await factory(moduleOptions);
     const instance = new module.LAMMPSWeb();
     return new LammpsClient(module, instance, clientOptions);
@@ -488,7 +524,8 @@ async function createWorkerBackedClient(
   const options: LammpsWorkerClientOptions = {
     workdir: clientOptions.workdir,
     onError: clientOptions.onError,
-    kokkos: clientOptions.kokkos
+    kokkos: clientOptions.kokkos,
+    variant: clientOptions.variant
   };
   if (print || printErr) {
     options.onOutput = (stream, text) => {
@@ -521,7 +558,8 @@ export async function createLammps(
   return LammpsClient.create(moduleOptions, {
     workdir: clientOptions.workdir,
     onError: clientOptions.onError,
-    kokkos: clientOptions.kokkos
+    kokkos: clientOptions.kokkos,
+    variant: clientOptions.variant
   });
 }
 

package.json

@@ -1,6 +1,6 @@
 {
   "name": "lammps.js",
-  "version": "1.5.1",
+  "version": "1.6.0",
   "description": "Run LAMMPS directly in the browser — a WebAssembly build with TypeScript-ready bindings.",
   "license": "MIT",
   "repository": {
@@ -33,6 +33,9 @@
   "files": [
     "dist",
     "cpp/build.py",
+    "python/lammps",
+    "python/pyproject.toml",
+    "python/README.md",
     "README.md",
     "LICENSE"
   ],

python/README.md

@@ -35,6 +35,13 @@ cross-origin isolated page:
 lmp = await lammps(cmdargs=["-k", "on", "t", "4", "-sf", "kk"])
 ```
 
+The full-package Atomify build (MANYBODY, KSPACE, REAXFF, GRANULAR, … — also
+a multithreaded KOKKOS build, so it needs a cross-origin isolated page too):
+
+```python
+lmp = await lammps(variant="atomify")
+```
+
 Differences from the native module are listed in the module docstring
 (`import lammps; help(lammps)`).
 

python/lammps/__init__.py

@@ -22,6 +22,11 @@
 
     lmp = await lammps(cmdargs=["-k", "on", "t", "4", "-sf", "kk"])
 
+The full-package Atomify build (also a multithreaded KOKKOS build, with e.g.
+MANYBODY, KSPACE, REAXFF, GRANULAR enabled) is selected with:
+
+    lmp = await lammps(variant="atomify")
+
 Differences from the native module, by necessity or convenience:
 
 - Creation is `await lammps(...)` instead of `lammps(...)`.
@@ -392,11 +397,18 @@ class lammps:  # noqa: N801 - mirrors the official class name
 
     - ``kokkos``: True or ``{"threads": n}`` to use the multithreaded
       KOKKOS wasm build (equivalently pass native-style ``cmdargs``).
+    - ``variant``: "serial", "kokkos" or "atomify" to select the wasm build
+      explicitly. "atomify" is the full-package build (MANYBODY, KSPACE,
+      REAXFF, GRANULAR, …) and, like "kokkos", is multithreaded. An explicit
+      ``variant`` takes precedence over ``kokkos``; ``kokkos=True`` is
+      shorthand for ``variant="kokkos"``.
     - ``output``: callable that receives each LAMMPS output line
       (default: ``print``). Pass ``None`` to silence output.
     - ``client_url``: URL of client.js, overriding auto-detection.
     """
 
+    _VARIANTS = ("serial", "kokkos", "atomify")
+
     def __init__(
         self,
         name: str = "",
@@ -405,10 +417,16 @@ def __init__(
         comm=None,
         *,
         kokkos=None,
+        variant: str | None = None,
         output=print,
         client_url: str | None = None,
     ):
         del name, ptr, comm  # accepted for API compatibility
+        if variant is not None and variant not in self._VARIANTS:
+            raise ValueError(
+                f"variant must be one of {self._VARIANTS}, got {variant!r}"
+            )
+        self._variant = variant
         self._cmdargs = [str(a) for a in cmdargs] if cmdargs else None
         kk_from_args = _parse_kokkos_cmdargs(self._cmdargs)
         if kokkos is None:
@@ -438,13 +456,19 @@ async def _astart(self):
             return self
         js, _run_js, create_proxy, to_js = self._ffi = _require_pyodide()
 
-        if self._kokkos is not None and not getattr(js, "crossOriginIsolated", False):
+        # Explicit variant wins; otherwise kokkos/-k arguments select the
+        # KOKKOS build (mirrors resolveVariant in client.ts).
+        variant = self._variant or (
+            "kokkos" if self._kokkos is not None else "serial"
+        )
+
+        if variant != "serial" and not getattr(js, "crossOriginIsolated", False):
             raise LammpsError(
-                "The KOKKOS multithreaded build needs SharedArrayBuffer, "
-                "which requires a cross-origin isolated page (COOP/COEP "
-                "headers). This page is not cross-origin isolated — create "
-                "the instance without kokkos/-k arguments to use the "
-                "single-threaded build."
+                f"The {variant} build is multithreaded and needs "
+                "SharedArrayBuffer, which requires a cross-origin isolated "
+                "page (COOP/COEP headers). This page is not cross-origin "
+                "isolated — create the instance without variant=/kokkos/-k "
+                "arguments to use the single-threaded build."
             )
 
         url = self._client_url or _default_client_url(js)
@@ -456,25 +480,29 @@ async def _astart(self):
 
         # Reuse the wasm module across lammps() calls — this shares the
         # in-memory filesystem so files written by one session are visible
-        # to the next without manual copying.
-        cache_key = f"{url}|{'kk' if self._kokkos is not None else 'serial'}"
+        # to the next without manual copying. Each variant is a distinct
+        # wasm module, so the cache is keyed per (client URL, variant).
+        cache_key = f"{url}|{variant}"
+        client_opts = to_js(
+            {
+                "kokkos": self._kokkos if self._kokkos is not None else False,
+                "variant": variant,
+            },
+            dict_converter=js.Object.fromEntries,
+        )
         wasm_module = _wasm_modules.get(cache_key)
         if wasm_module is None:
             module_opts = to_js(
                 {"print": print_proxy, "printErr": printerr_proxy},
                 dict_converter=js.Object.fromEntries,
             )
-            client_opts = to_js(
-                {"kokkos": self._kokkos if self._kokkos is not None else False},
-                dict_converter=js.Object.fromEntries,
-            )
             self._client = await mod.LammpsClient.create(module_opts, client_opts)
             _wasm_modules[cache_key] = self._client.module
         else:
             wasm_module.print = print_proxy
             wasm_module.printErr = printerr_proxy
             instance = wasm_module.LAMMPSWeb.new()
-            self._client = mod.LammpsClient.new(wasm_module, instance)
+            self._client = mod.LammpsClient.new(wasm_module, instance, client_opts)
 
         # Mount JupyterLite's DriveFS on /work so LAMMPS reads/writes go
         # directly to the notebook filesystem — no syncing needed.

python/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "lammps-js"
 # Tracks the lammps.js npm package version this binding is developed against.
-version = "1.5.0"
+version = "1.6.0"
 description = "LAMMPS Python bindings for the browser: drives the lammps.js WebAssembly engine from a Pyodide kernel (JupyterLite), mirroring the official `lammps` Python module API (imports as `lammps`)."
 readme = "README.md"
 requires-python = ">=3.10"

tests/atomify.spec.ts

@@ -7,12 +7,13 @@
 // moltemplate pair styles). Runs under a plain node environment because
 // emscripten pthreads are backed by worker_threads + SharedArrayBuffer, which
 // jsdom does not provide. It ships in this same package under the
-// ./wasm-atomify export, loaded directly here (LammpsClient's kokkos option
-// targets the separate lammps-kokkos.js).
+// ./wasm-atomify export, loaded both directly and through LammpsClient's
+// `variant: "atomify"` option.
 import { afterAll, describe, expect, it } from "vitest";
 import { existsSync, readFileSync } from "node:fs";
 import { join } from "node:path";
 
+import { LammpsClient } from "../dist/client.js";
 import type { LAMMPSWeb, LammpsModule } from "../types";
 
 const atomifyModulePath = join(process.cwd(), "dist", "cpp", "lammps-atomify.js");
@@ -117,6 +118,22 @@ pair_style vashishta
     ).not.toThrow();
   });
 
+  it('is selected by LammpsClient\'s variant: "atomify" option (Kokkos start args)', async () => {
+    const client = await LammpsClient.create(
+      { print: () => undefined, printErr: () => undefined },
+      { variant: "atomify", kokkos: { threads: 2 } }
+    );
+    instances.push(client.instance);
+    client.start();
+
+    // The client resolved the atomify module (full package set) and started
+    // LAMMPS with the Kokkos runtime enabled (kk suffix on by default).
+    expect(client.instance.hasPackage("KOKKOS")).toBe(true);
+    expect(client.instance.hasPackage("MANYBODY")).toBe(true);
+    client.runScript(readFileSync(join(process.cwd(), "tests", "fixtures", "lj.mini.in"), "utf8"));
+    expect(client.getCurrentStep()).toBe(5);
+  });
+
   it("accepts the vendored moltemplate pair style", async () => {
     const lmp = await createInstance(2);
     expect(() =>