Skip to content

Observable and visitor plugins

Register a factory with polarism.results.register_observable(name) and select it with cfg.result.observables = (name,). Factories receive the initialized SimulationController and return one ResultNode or a list of nodes. Unknown names fail configuration validation and list registered observables. There are no built-in names in this registry.

ResultNode.compute_fn(**context) receives state, grid, P_total, and scalar_groups, the context keys passed by SimulationController._record_results. Use compute_engine.xp for NumPy or CuPy operations. For is_field=False, storage records the value returned by compute_fn; reduce_dim_fn is applied only for CPU visitors such as live view. Return a 0-D backend array for a scalar to avoid a GPU sync during result computation. Set save and expose to control storage and live view.

"""Condensate x centre of mass observable."""
from polarism.compute_engine import compute_engine
from polarism.results import ResultNode, register_observable


@register_observable("example-center-x")
def center_of_mass_x(controller):
    """Create a device-computed condensate centre of mass scalar."""
    def compute(**context):
        xp = compute_engine.xp
        density = xp.abs(context["state"].psi) ** 2
        return (density * context["grid"].X).sum() / xp.maximum(density.sum(), 1e-30)

    return ResultNode(
        name="center_x", label="Condensate centre x", units="µm",
        compute_fn=compute, reduce_dim_fn=lambda value: value,
        cmap=None, scaling=None, clim=None, expose=True, save=True,
        cut=None, is_field=False,
    )

The example computes the condensate x centre of mass in µm. It reads the current wavefunction and the grid's x coordinates without transferring arrays to the host.

For a custom consumer, subclass ResultVisitor and call controller.add_visitor(visitor) before controller.run(). Nodes are built when the first visitor is added, even when no configured output is enabled. needs(nodes) returns the set of node names to compute; visit_all(nodes, t=..., cached=..., **context) receives each sampled frame. cached maps node names to (field, reduced) pairs. Set data_location to "cpu" or "device" to choose host or backend arrays. fatal_on_error controls whether a visit error stops the run. An optional finalize() is called at the end of the run, including after an exception. Sampling uses result.save_interval and always includes the initial state.