diff --git a/NOTEBOOK_TUTORIALS.md b/NOTEBOOK_TUTORIALS.md index b5d2d60..025979d 100644 --- a/NOTEBOOK_TUTORIALS.md +++ b/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. diff --git a/README.md b/README.md index faf5171..980b0fc 100644 --- a/README.md +++ b/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. diff --git a/client.ts b/client.ts index 84498b2..9a74e51 100644 --- a/client.ts +++ b/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 }); } diff --git a/package.json b/package.json index e4ea3ec..5174958 100644 --- a/package.json +++ b/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" ], diff --git a/python/README.md b/python/README.md index 44e43e7..3b1e614 100644 --- a/python/README.md +++ b/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)`). diff --git a/python/lammps/__init__.py b/python/lammps/__init__.py index 5485e20..c17f668 100644 --- a/python/lammps/__init__.py +++ b/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. diff --git a/python/pyproject.toml b/python/pyproject.toml index a855ff3..06b1faa 100644 --- a/python/pyproject.toml +++ b/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" diff --git a/tests/atomify.spec.ts b/tests/atomify.spec.ts index f1bd58b..cf565a5 100644 --- a/tests/atomify.spec.ts +++ b/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(() => diff --git a/tests/client.unit.spec.ts b/tests/client.unit.spec.ts index 371cb8b..5a52b3e 100644 --- a/tests/client.unit.spec.ts +++ b/tests/client.unit.spec.ts @@ -199,6 +199,48 @@ describe("LammpsClient lifecycle and delegation", () => { expect(mocks.startWithArgs).toHaveBeenCalledWith(["-k", "on", "t", "4", "-sf", "kk"]); }); + it('variant: "kokkos" starts with Kokkos args like kokkos: true', () => { + const { moduleMock } = createModuleMock(); + const { instanceMock, mocks } = createInstanceMock(); + const client = createClient(moduleMock, instanceMock, { variant: "kokkos" }); + + client.start(); + + expect(mocks.start).not.toHaveBeenCalled(); + expect(mocks.startWithArgs).toHaveBeenCalledTimes(1); + const args = mocks.startWithArgs.mock.calls[0][0] as string[]; + expect(args.slice(0, 3)).toEqual(["-k", "on", "t"]); + expect(args.slice(4)).toEqual(["-sf", "kk"]); + }); + + it('variant: "atomify" is a threaded build: starts with Kokkos args, kokkos option tunes them', () => { + const { moduleMock } = createModuleMock(); + const { instanceMock, mocks } = createInstanceMock(); + const client = createClient(moduleMock, instanceMock, { + variant: "atomify", + kokkos: { threads: 4 } + }); + + client.start(); + + expect(mocks.start).not.toHaveBeenCalled(); + expect(mocks.startWithArgs).toHaveBeenCalledWith(["-k", "on", "t", "4", "-sf", "kk"]); + }); + + it('explicit variant: "serial" wins over kokkos: true (plain start)', () => { + const { moduleMock } = createModuleMock(); + const { instanceMock, mocks } = createInstanceMock(); + const client = createClient(moduleMock, instanceMock, { + variant: "serial", + kokkos: true + }); + + client.start(); + + expect(mocks.startWithArgs).not.toHaveBeenCalled(); + expect(mocks.start).toHaveBeenCalledTimes(1); + }); + it("clamps kokkos threads to the pthread pool size and honors suffix: false", () => { const { moduleMock } = createModuleMock(); const { instanceMock, mocks } = createInstanceMock(); diff --git a/tests/worker-host.spec.ts b/tests/worker-host.spec.ts index 7fe83a7..be73914 100644 --- a/tests/worker-host.spec.ts +++ b/tests/worker-host.spec.ts @@ -117,6 +117,27 @@ describe("installLammpsWorker", () => { }); }); + it("forwards the variant from init to the client factory", async () => { + const { scope, posted, dispatch } = createScope(); + const mocks = createClientMocks(); + const createClient = vi.fn(async () => { + const clientMock: Partial = mocks; + return clientMock as unknown as LammpsClient; + }); + + installLammpsWorker(scope, { createClient }); + dispatch({ id: 1, type: "init", variant: "atomify" }); + + await vi.waitFor(() => { + expect(posted).toContainEqual({ type: "response", id: 1, ok: true, result: undefined }); + }); + expect(createClient).toHaveBeenCalledWith(expect.anything(), { + workdir: undefined, + kokkos: undefined, + variant: "atomify" + }); + }); + it("rejects requests before initialization", async () => { const { scope, posted, dispatch } = createScope(); installLammpsWorker(scope, { diff --git a/types/client.d.ts b/types/client.d.ts index 651feb3..2651928 100644 --- a/types/client.d.ts +++ b/types/client.d.ts @@ -46,8 +46,34 @@ export interface SyncBoxOptions { copy?: boolean; } +export interface KokkosOptions { + /** Number of Kokkos threads (capped at the built pthread pool size, 8). */ + threads?: number; + /** Apply the kk accelerator suffix to all styles (`-sf kk`). Default true. */ + suffix?: boolean; +} + +/** + * Which wasm build the client loads: "serial" (default, dist/cpp/lammps.js), + * "kokkos" (multi-threaded, dist/cpp/lammps-kokkos.js) or "atomify" (full + * package set, also multi-threaded KOKKOS, dist/cpp/lammps-atomify.js). + */ +export type LammpsVariant = "serial" | "kokkos" | "atomify"; + export interface LammpsClientOptions { workdir?: string; + /** + * Use the multi-threaded KOKKOS wasm build. Requires a cross-origin + * isolated context in browsers. Shorthand for `variant: "kokkos"`; also + * tunes threads/suffix for `variant: "atomify"`. + */ + kokkos?: boolean | KokkosOptions; + /** + * Select the wasm build explicitly; takes precedence over `kokkos`. + * "kokkos" and "atomify" are multi-threaded builds and require a + * cross-origin isolated context in browsers (SharedArrayBuffer). + */ + variant?: LammpsVariant; } export declare class LammpsClient { diff --git a/worker-client.ts b/worker-client.ts index e4e2bf1..8e9008e 100644 --- a/worker-client.ts +++ b/worker-client.ts @@ -1,4 +1,4 @@ -import type { KokkosOptions } from "./client.js"; +import type { KokkosOptions, LammpsVariant } from "./client.js"; import type { LammpsWorkerRequest, LammpsWorkerRequestBody, @@ -25,6 +25,8 @@ export interface LammpsWorkerClientOptions { onError?: (error: Error) => void; /** Use the multi-threaded KOKKOS wasm build inside the worker. */ kokkos?: boolean | KokkosOptions; + /** Which wasm build to load inside the worker; see LammpsClientOptions.variant. */ + variant?: LammpsVariant; } interface PendingRequest { @@ -68,7 +70,12 @@ export class LammpsWorkerClient { ownsWorker = false ): Promise { const client = new LammpsWorkerClient(worker, options, ownsWorker); - await client.#request({ type: "init", workdir: options.workdir, kokkos: options.kokkos }); + await client.#request({ + type: "init", + workdir: options.workdir, + kokkos: options.kokkos, + variant: options.variant + }); return client; } diff --git a/worker-host.ts b/worker-host.ts index dc00cfe..c729e61 100644 --- a/worker-host.ts +++ b/worker-host.ts @@ -1,5 +1,5 @@ import { LammpsClient } from "./client.js"; -import type { KokkosOptions } from "./client.js"; +import type { KokkosOptions, LammpsVariant } from "./client.js"; import { serializeStepData } from "./worker-protocol.js"; import type { LammpsWorkerRequest, @@ -16,7 +16,7 @@ export interface InstallLammpsWorkerOptions { /** Injectable factory, used by tests to supply a mock client. */ createClient?: ( moduleOptions: Record, - clientOptions: { workdir?: string; kokkos?: boolean | KokkosOptions } + clientOptions: { workdir?: string; kokkos?: boolean | KokkosOptions; variant?: LammpsVariant } ) => Promise; } @@ -35,8 +35,10 @@ class RunAbortedError extends Error { export function installLammpsWorker(scope: WorkerScope, options: InstallLammpsWorkerOptions = {}): void { const createClient = options.createClient ?? - ((moduleOptions: Record, clientOptions: { workdir?: string; kokkos?: boolean | KokkosOptions }) => - LammpsClient.create(moduleOptions, clientOptions)); + (( + moduleOptions: Record, + clientOptions: { workdir?: string; kokkos?: boolean | KokkosOptions; variant?: LammpsVariant } + ) => LammpsClient.create(moduleOptions, clientOptions)); let client: LammpsClient | null = null; let abortRequested = false; @@ -127,7 +129,7 @@ export function installLammpsWorker(scope: WorkerScope, options: InstallLammpsWo print: (text: unknown) => post({ type: "output", stream: "stdout", text: String(text) }), printErr: (text: unknown) => post({ type: "output", stream: "stderr", text: String(text) }) }, - { workdir: request.workdir, kokkos: request.kokkos } + { workdir: request.workdir, kokkos: request.kokkos, variant: request.variant } ); client.start(); respond(request.id); diff --git a/worker-protocol.ts b/worker-protocol.ts index 642dc40..98905ed 100644 --- a/worker-protocol.ts +++ b/worker-protocol.ts @@ -1,4 +1,4 @@ -import type { AsyncStepData, KokkosOptions } from "./client.js"; +import type { AsyncStepData, KokkosOptions, LammpsVariant } from "./client.js"; export interface WorkerParticleData { count: number; @@ -42,7 +42,13 @@ export interface WorkerRunResult { } export type LammpsWorkerRequest = - | { id: number; type: "init"; workdir?: string; kokkos?: boolean | KokkosOptions } + | { + id: number; + type: "init"; + workdir?: string; + kokkos?: boolean | KokkosOptions; + variant?: LammpsVariant; + } | { id: number; type: "runCommand"; command: string } | { id: number; type: "runScript"; script: string } | { id: number; type: "runScriptAsync"; script: string; options: WorkerRunOptions }